miri/intrinsics/x86/
avx512.rs

1use rustc_span::Symbol;
2
3use super::{
4    packssdw, packsswb, packusdw, packuswb, permute, permute2, pmaddbw, pmaddwd, psadbw, pshufb,
5};
6use crate::*;
7
8impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
9pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
10    fn emulate_x86_avx512_intrinsic(
11        &mut self,
12        link_name: Symbol,
13        args: &[OpTy<'tcx>],
14        dest: &MPlaceTy<'tcx>,
15    ) -> InterpResult<'tcx, EmulateItemResult> {
16        let this = self.eval_context_mut();
17        // Prefix should have already been checked.
18        let unprefixed_name = link_name.as_str().strip_prefix("llvm.x86.avx512.").unwrap();
19
20        match unprefixed_name {
21            // Used by the ternarylogic functions.
22            "pternlog.d.128" | "pternlog.d.256" | "pternlog.d.512" => {
23                this.expect_target_feature_for_intrinsic(link_name, "avx512f")?;
24                if matches!(unprefixed_name, "pternlog.d.128" | "pternlog.d.256") {
25                    this.expect_target_feature_for_intrinsic(link_name, "avx512vl")?;
26                }
27
28                let [a, b, c, imm8] = this.check_shim_sig_unadjusted(link_name, args)?;
29
30                assert_eq!(dest.layout, a.layout);
31                assert_eq!(dest.layout, b.layout);
32                assert_eq!(dest.layout, c.layout);
33
34                // The signatures of these operations are:
35                //
36                // ```
37                // fn vpternlogd(a: i32x16, b: i32x16, c: i32x16, imm8: i32) -> i32x16;
38                // fn vpternlogd256(a: i32x8, b: i32x8, c: i32x8, imm8: i32) -> i32x8;
39                // fn vpternlogd128(a: i32x4, b: i32x4, c: i32x4, imm8: i32) -> i32x4;
40                // ```
41                //
42                // The element type is always a 32-bit integer, the width varies.
43
44                let (a, _a_len) = this.project_to_simd(a)?;
45                let (b, _b_len) = this.project_to_simd(b)?;
46                let (c, _c_len) = this.project_to_simd(c)?;
47                let (dest, dest_len) = this.project_to_simd(dest)?;
48
49                // Compute one lane with ternary table.
50                let tern = |xa: u32, xb: u32, xc: u32, imm: u32| -> u32 {
51                    let mut out = 0u32;
52                    // At each bit position, select bit from imm8 at index = (a << 2) | (b << 1) | c
53                    for bit in 0..32 {
54                        let ia = (xa >> bit) & 1;
55                        let ib = (xb >> bit) & 1;
56                        let ic = (xc >> bit) & 1;
57                        let idx = (ia << 2) | (ib << 1) | ic;
58                        let v = (imm >> idx) & 1;
59                        out |= v << bit;
60                    }
61                    out
62                };
63
64                let imm8 = this.read_scalar(imm8)?.to_u32()? & 0xFF;
65                for i in 0..dest_len {
66                    let a_lane = this.project_index(&a, i)?;
67                    let b_lane = this.project_index(&b, i)?;
68                    let c_lane = this.project_index(&c, i)?;
69                    let d_lane = this.project_index(&dest, i)?;
70
71                    let va = this.read_scalar(&a_lane)?.to_u32()?;
72                    let vb = this.read_scalar(&b_lane)?.to_u32()?;
73                    let vc = this.read_scalar(&c_lane)?.to_u32()?;
74
75                    let r = tern(va, vb, vc, imm8);
76                    this.write_scalar(Scalar::from_u32(r), &d_lane)?;
77                }
78            }
79            // Used to implement the _mm512_sad_epu8 function.
80            "psad.bw.512" => {
81                this.expect_target_feature_for_intrinsic(link_name, "avx512bw")?;
82
83                let [left, right] = this.check_shim_sig_unadjusted(link_name, args)?;
84
85                psadbw(this, left, right, dest)?
86            }
87            // Used to implement the _mm512_madd_epi16 function.
88            "pmaddw.d.512" => {
89                this.expect_target_feature_for_intrinsic(link_name, "avx512bw")?;
90
91                let [left, right] = this.check_shim_sig_unadjusted(link_name, args)?;
92
93                pmaddwd(this, left, right, dest)?;
94            }
95            // Used to implement the _mm512_maddubs_epi16 function.
96            "pmaddubs.w.512" => {
97                let [left, right] = this.check_shim_sig_unadjusted(link_name, args)?;
98
99                pmaddbw(this, left, right, dest)?;
100            }
101            // Used to implement the _mm512_permutexvar_epi32/_mm512_permutexvar_epi64 functions.
102            "permvar.si.512" | "permvar.di.512" => {
103                let [left, right] = this.check_shim_sig_unadjusted(link_name, args)?;
104
105                permute(this, left, right, dest)?;
106            }
107            "permvar.qi.512" | "permvar.qi.256" | "permvar.qi.128" => {
108                this.expect_target_feature_for_intrinsic(link_name, "avx512vbmi")?;
109                if !unprefixed_name.ends_with("512") {
110                    this.expect_target_feature_for_intrinsic(link_name, "avx512vl")?;
111                }
112
113                let [left, right] = this.check_shim_sig_unadjusted(link_name, args)?;
114
115                permute(this, left, right, dest)?;
116            }
117            // Used to implement the _mm512_permutex2var_epi64 intrinsic.
118            "vpermi2var.q.512" => {
119                let [left, indices, right] = this.check_shim_sig_unadjusted(link_name, args)?;
120
121                permute2(this, left, indices, right, dest)?;
122            }
123            // Used to implement the _mm512_permutex2var_epi8 intrinsic.
124            "vpermi2var.qi.512" => {
125                this.expect_target_feature_for_intrinsic(link_name, "avx512vbmi")?;
126
127                let [left, indices, right] = this.check_shim_sig_unadjusted(link_name, args)?;
128
129                permute2(this, left, indices, right, dest)?;
130            }
131            // Used to implement the _mm512_shuffle_epi8 intrinsic.
132            "pshuf.b.512" => {
133                let [left, right] = this.check_shim_sig_unadjusted(link_name, args)?;
134
135                pshufb(this, left, right, dest)?;
136            }
137
138            // Used to implement the _mm512_dpbusd_epi32 function.
139            "vpdpbusd.512" | "vpdpbusd.256" | "vpdpbusd.128" => {
140                this.expect_target_feature_for_intrinsic(link_name, "avx512vnni")?;
141                if matches!(unprefixed_name, "vpdpbusd.128" | "vpdpbusd.256") {
142                    this.expect_target_feature_for_intrinsic(link_name, "avx512vl")?;
143                }
144
145                let [src, a, b] = this.check_shim_sig_unadjusted(link_name, args)?;
146
147                vpdpbusd(this, src, a, b, dest)?;
148            }
149            // Used to implement the _mm512_packs_epi16 function
150            "packsswb.512" => {
151                this.expect_target_feature_for_intrinsic(link_name, "avx512bw")?;
152
153                let [a, b] = this.check_shim_sig_unadjusted(link_name, args)?;
154
155                packsswb(this, a, b, dest)?;
156            }
157            // Used to implement the _mm512_packus_epi16 function
158            "packuswb.512" => {
159                this.expect_target_feature_for_intrinsic(link_name, "avx512bw")?;
160
161                let [a, b] = this.check_shim_sig_unadjusted(link_name, args)?;
162
163                packuswb(this, a, b, dest)?;
164            }
165            // Used to implement the _mm512_packs_epi32 function
166            "packssdw.512" => {
167                this.expect_target_feature_for_intrinsic(link_name, "avx512bw")?;
168
169                let [a, b] = this.check_shim_sig_unadjusted(link_name, args)?;
170
171                packssdw(this, a, b, dest)?;
172            }
173            // Used to implement the _mm512_packus_epi32 function
174            "packusdw.512" => {
175                this.expect_target_feature_for_intrinsic(link_name, "avx512bw")?;
176
177                let [a, b] = this.check_shim_sig_unadjusted(link_name, args)?;
178
179                packusdw(this, a, b, dest)?;
180            }
181            _ => return interp_ok(EmulateItemResult::NotSupported),
182        }
183        interp_ok(EmulateItemResult::NeedsReturn)
184    }
185}
186
187/// Multiply groups of 4 adjacent pairs of unsigned 8-bit integers in `a` with corresponding signed
188/// 8-bit integers in `b`, producing 4 intermediate signed 16-bit results. Sum these 4 results with
189/// the corresponding 32-bit integer in `src` (using wrapping arighmetic), and store the packed
190/// 32-bit results in `dst`.
191///
192/// <https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_dpbusd_epi32>
193/// <https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_dpbusd_epi32>
194/// <https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm512_dpbusd_epi32>
195fn vpdpbusd<'tcx>(
196    ecx: &mut crate::MiriInterpCx<'tcx>,
197    src: &OpTy<'tcx>,
198    a: &OpTy<'tcx>,
199    b: &OpTy<'tcx>,
200    dest: &MPlaceTy<'tcx>,
201) -> InterpResult<'tcx, ()> {
202    let (src, src_len) = ecx.project_to_simd(src)?;
203    let (a, a_len) = ecx.project_to_simd(a)?;
204    let (b, b_len) = ecx.project_to_simd(b)?;
205    let (dest, dest_len) = ecx.project_to_simd(dest)?;
206
207    // fn vpdpbusd(src: i32x16, a: u8x64, b: i8x64) -> i32x16;
208    // fn vpdpbusd256(src: i32x8, a: u8x32, b: i8x32) -> i32x8;
209    // fn vpdpbusd128(src: i32x4, a: u8x16, b: i8x16) -> i32x4;
210    assert_eq!(src_len, dest_len);
211    assert_eq!(a_len, dest_len.strict_mul(4));
212    assert_eq!(b_len, a_len);
213
214    for i in 0..dest_len {
215        let src = ecx.read_scalar(&ecx.project_index(&src, i)?)?.to_i32()?;
216        let dest = ecx.project_index(&dest, i)?;
217
218        let mut intermediate_sum: i32 = 0;
219        for j in 0..4 {
220            let idx = i.strict_mul(4).strict_add(j);
221            let a = ecx.read_scalar(&ecx.project_index(&a, idx)?)?.to_u8()?;
222            let b = ecx.read_scalar(&ecx.project_index(&b, idx)?)?.to_i8()?;
223
224            let product = i32::from(a).strict_mul(i32::from(b));
225            intermediate_sum = intermediate_sum.strict_add(product);
226        }
227
228        // Use `wrapping_add` because `src` is an arbitrary i32 and the addition can overflow.
229        let res = Scalar::from_i32(intermediate_sum.wrapping_add(src));
230        ecx.write_scalar(res, &dest)?;
231    }
232
233    interp_ok(())
234}
miri/intrinsics/x86/avx512.rs

miri/intrinsics/x86/
avx512.rs
