miri/shims/x86/
avx2.rs

1use rustc_abi::CanonAbi;
2use rustc_middle::mir;
3use rustc_middle::ty::Ty;
4use rustc_middle::ty::layout::LayoutOf as _;
5use rustc_span::Symbol;
6use rustc_target::callconv::FnAbi;
7
8use super::{
9    ShiftOp, horizontal_bin_op, int_abs, mask_load, mask_store, mpsadbw, packssdw, packsswb,
10    packusdw, packuswb, pmulhrsw, psign, shift_simd_by_scalar, shift_simd_by_simd,
11};
12use crate::*;
13
14impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
15pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
16    fn emulate_x86_avx2_intrinsic(
17        &mut self,
18        link_name: Symbol,
19        abi: &FnAbi<'tcx, Ty<'tcx>>,
20        args: &[OpTy<'tcx>],
21        dest: &MPlaceTy<'tcx>,
22    ) -> InterpResult<'tcx, EmulateItemResult> {
23        let this = self.eval_context_mut();
24        this.expect_target_feature_for_intrinsic(link_name, "avx2")?;
25        // Prefix should have already been checked.
26        let unprefixed_name = link_name.as_str().strip_prefix("llvm.x86.avx2.").unwrap();
27
28        match unprefixed_name {
29            // Used to implement the _mm256_abs_epi{8,16,32} functions.
30            // Calculates the absolute value of packed 8/16/32-bit integers.
31            "pabs.b" | "pabs.w" | "pabs.d" => {
32                let [op] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
33
34                int_abs(this, op, dest)?;
35            }
36            // Used to implement the _mm256_h{add,adds,sub}_epi{16,32} functions.
37            // Horizontally add / add with saturation / subtract adjacent 16/32-bit
38            // integer values in `left` and `right`.
39            "phadd.w" | "phadd.sw" | "phadd.d" | "phsub.w" | "phsub.sw" | "phsub.d" => {
40                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
41
42                let (which, saturating) = match unprefixed_name {
43                    "phadd.w" | "phadd.d" => (mir::BinOp::Add, false),
44                    "phadd.sw" => (mir::BinOp::Add, true),
45                    "phsub.w" | "phsub.d" => (mir::BinOp::Sub, false),
46                    "phsub.sw" => (mir::BinOp::Sub, true),
47                    _ => unreachable!(),
48                };
49
50                horizontal_bin_op(this, which, saturating, left, right, dest)?;
51            }
52            // Used to implement `_mm{,_mask}_{i32,i64}gather_{epi32,epi64,pd,ps}` functions
53            // Gathers elements from `slice` using `offsets * scale` as indices.
54            // When the highest bit of the corresponding element of `mask` is 0,
55            // the value is copied from `src` instead.
56            "gather.d.d" | "gather.d.d.256" | "gather.d.q" | "gather.d.q.256" | "gather.q.d"
57            | "gather.q.d.256" | "gather.q.q" | "gather.q.q.256" | "gather.d.pd"
58            | "gather.d.pd.256" | "gather.q.pd" | "gather.q.pd.256" | "gather.d.ps"
59            | "gather.d.ps.256" | "gather.q.ps" | "gather.q.ps.256" => {
60                let [src, slice, offsets, mask, scale] =
61                    this.check_shim(abi, CanonAbi::C, link_name, args)?;
62
63                assert_eq!(dest.layout, src.layout);
64
65                let (src, _) = this.project_to_simd(src)?;
66                let (offsets, offsets_len) = this.project_to_simd(offsets)?;
67                let (mask, mask_len) = this.project_to_simd(mask)?;
68                let (dest, dest_len) = this.project_to_simd(dest)?;
69
70                // There are cases like dest: i32x4, offsets: i64x2
71                // If dest has more elements than offset, extra dest elements are filled with zero.
72                // If offsets has more elements than dest, extra offsets are ignored.
73                let actual_len = dest_len.min(offsets_len);
74
75                assert_eq!(dest_len, mask_len);
76
77                let mask_item_size = mask.layout.field(this, 0).size;
78                let high_bit_offset = mask_item_size.bits().strict_sub(1);
79
80                let scale = this.read_scalar(scale)?.to_i8()?;
81                if !matches!(scale, 1 | 2 | 4 | 8) {
82                    panic!("invalid gather scale {scale}");
83                }
84                let scale = i64::from(scale);
85
86                let slice = this.read_pointer(slice)?;
87                for i in 0..actual_len {
88                    let mask = this.project_index(&mask, i)?;
89                    let dest = this.project_index(&dest, i)?;
90
91                    if this.read_scalar(&mask)?.to_uint(mask_item_size)? >> high_bit_offset != 0 {
92                        let offset = this.project_index(&offsets, i)?;
93                        let offset =
94                            i64::try_from(this.read_scalar(&offset)?.to_int(offset.layout.size)?)
95                                .unwrap();
96                        let ptr = slice.wrapping_signed_offset(offset.strict_mul(scale), &this.tcx);
97                        // Unaligned copy, which is what we want.
98                        this.mem_copy(
99                            ptr,
100                            dest.ptr(),
101                            dest.layout.size,
102                            /*nonoverlapping*/ true,
103                        )?;
104                    } else {
105                        this.copy_op(&this.project_index(&src, i)?, &dest)?;
106                    }
107                }
108                for i in actual_len..dest_len {
109                    let dest = this.project_index(&dest, i)?;
110                    this.write_scalar(Scalar::from_int(0, dest.layout.size), &dest)?;
111                }
112            }
113            // Used to implement the _mm256_madd_epi16 function.
114            // Multiplies packed signed 16-bit integers in `left` and `right`, producing
115            // intermediate signed 32-bit integers. Horizontally add adjacent pairs of
116            // intermediate 32-bit integers, and pack the results in `dest`.
117            "pmadd.wd" => {
118                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
119
120                let (left, left_len) = this.project_to_simd(left)?;
121                let (right, right_len) = this.project_to_simd(right)?;
122                let (dest, dest_len) = this.project_to_simd(dest)?;
123
124                assert_eq!(left_len, right_len);
125                assert_eq!(dest_len.strict_mul(2), left_len);
126
127                for i in 0..dest_len {
128                    let j1 = i.strict_mul(2);
129                    let left1 = this.read_scalar(&this.project_index(&left, j1)?)?.to_i16()?;
130                    let right1 = this.read_scalar(&this.project_index(&right, j1)?)?.to_i16()?;
131
132                    let j2 = j1.strict_add(1);
133                    let left2 = this.read_scalar(&this.project_index(&left, j2)?)?.to_i16()?;
134                    let right2 = this.read_scalar(&this.project_index(&right, j2)?)?.to_i16()?;
135
136                    let dest = this.project_index(&dest, i)?;
137
138                    // Multiplications are i16*i16->i32, which will not overflow.
139                    let mul1 = i32::from(left1).strict_mul(right1.into());
140                    let mul2 = i32::from(left2).strict_mul(right2.into());
141                    // However, this addition can overflow in the most extreme case
142                    // (-0x8000)*(-0x8000)+(-0x8000)*(-0x8000) = 0x80000000
143                    let res = mul1.wrapping_add(mul2);
144
145                    this.write_scalar(Scalar::from_i32(res), &dest)?;
146                }
147            }
148            // Used to implement the _mm256_maddubs_epi16 function.
149            // Multiplies packed 8-bit unsigned integers from `left` and packed
150            // signed 8-bit integers from `right` into 16-bit signed integers. Then,
151            // the saturating sum of the products with indices `2*i` and `2*i+1`
152            // produces the output at index `i`.
153            "pmadd.ub.sw" => {
154                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
155
156                let (left, left_len) = this.project_to_simd(left)?;
157                let (right, right_len) = this.project_to_simd(right)?;
158                let (dest, dest_len) = this.project_to_simd(dest)?;
159
160                assert_eq!(left_len, right_len);
161                assert_eq!(dest_len.strict_mul(2), left_len);
162
163                for i in 0..dest_len {
164                    let j1 = i.strict_mul(2);
165                    let left1 = this.read_scalar(&this.project_index(&left, j1)?)?.to_u8()?;
166                    let right1 = this.read_scalar(&this.project_index(&right, j1)?)?.to_i8()?;
167
168                    let j2 = j1.strict_add(1);
169                    let left2 = this.read_scalar(&this.project_index(&left, j2)?)?.to_u8()?;
170                    let right2 = this.read_scalar(&this.project_index(&right, j2)?)?.to_i8()?;
171
172                    let dest = this.project_index(&dest, i)?;
173
174                    // Multiplication of a u8 and an i8 into an i16 cannot overflow.
175                    let mul1 = i16::from(left1).strict_mul(right1.into());
176                    let mul2 = i16::from(left2).strict_mul(right2.into());
177                    let res = mul1.saturating_add(mul2);
178
179                    this.write_scalar(Scalar::from_i16(res), &dest)?;
180                }
181            }
182            // Used to implement the _mm_maskload_epi32, _mm_maskload_epi64,
183            // _mm256_maskload_epi32 and _mm256_maskload_epi64 functions.
184            // For the element `i`, if the high bit of the `i`-th element of `mask`
185            // is one, it is loaded from `ptr.wrapping_add(i)`, otherwise zero is
186            // loaded.
187            "maskload.d" | "maskload.q" | "maskload.d.256" | "maskload.q.256" => {
188                let [ptr, mask] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
189
190                mask_load(this, ptr, mask, dest)?;
191            }
192            // Used to implement the _mm_maskstore_epi32, _mm_maskstore_epi64,
193            // _mm256_maskstore_epi32 and _mm256_maskstore_epi64 functions.
194            // For the element `i`, if the high bit of the element `i`-th of `mask`
195            // is one, it is stored into `ptr.wapping_add(i)`.
196            // Unlike SSE2's _mm_maskmoveu_si128, these are not non-temporal stores.
197            "maskstore.d" | "maskstore.q" | "maskstore.d.256" | "maskstore.q.256" => {
198                let [ptr, mask, value] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
199
200                mask_store(this, ptr, mask, value)?;
201            }
202            // Used to implement the _mm256_mpsadbw_epu8 function.
203            // Compute the sum of absolute differences of quadruplets of unsigned
204            // 8-bit integers in `left` and `right`, and store the 16-bit results
205            // in `right`. Quadruplets are selected from `left` and `right` with
206            // offsets specified in `imm`.
207            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_mpsadbw_epu8
208            "mpsadbw" => {
209                let [left, right, imm] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
210
211                mpsadbw(this, left, right, imm, dest)?;
212            }
213            // Used to implement the _mm256_mulhrs_epi16 function.
214            // Multiplies packed 16-bit signed integer values, truncates the 32-bit
215            // product to the 18 most significant bits by right-shifting, and then
216            // divides the 18-bit value by 2 (rounding to nearest) by first adding
217            // 1 and then taking the bits `1..=16`.
218            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_mulhrs_epi16
219            "pmul.hr.sw" => {
220                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
221
222                pmulhrsw(this, left, right, dest)?;
223            }
224            // Used to implement the _mm256_packs_epi16 function.
225            // Converts two 16-bit integer vectors to a single 8-bit integer
226            // vector with signed saturation.
227            "packsswb" => {
228                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
229
230                packsswb(this, left, right, dest)?;
231            }
232            // Used to implement the _mm256_packs_epi32 function.
233            // Converts two 32-bit integer vectors to a single 16-bit integer
234            // vector with signed saturation.
235            "packssdw" => {
236                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
237
238                packssdw(this, left, right, dest)?;
239            }
240            // Used to implement the _mm256_packus_epi16 function.
241            // Converts two 16-bit signed integer vectors to a single 8-bit
242            // unsigned integer vector with saturation.
243            "packuswb" => {
244                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
245
246                packuswb(this, left, right, dest)?;
247            }
248            // Used to implement the _mm256_packus_epi32 function.
249            // Concatenates two 32-bit signed integer vectors and converts
250            // the result to a 16-bit unsigned integer vector with saturation.
251            "packusdw" => {
252                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
253
254                packusdw(this, left, right, dest)?;
255            }
256            // Used to implement the _mm256_permutevar8x32_epi32 and
257            // _mm256_permutevar8x32_ps function.
258            // Shuffles `left` using the three low bits of each element of `right`
259            // as indices.
260            "permd" | "permps" => {
261                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
262
263                let (left, left_len) = this.project_to_simd(left)?;
264                let (right, right_len) = this.project_to_simd(right)?;
265                let (dest, dest_len) = this.project_to_simd(dest)?;
266
267                assert_eq!(dest_len, left_len);
268                assert_eq!(dest_len, right_len);
269
270                for i in 0..dest_len {
271                    let dest = this.project_index(&dest, i)?;
272                    let right = this.read_scalar(&this.project_index(&right, i)?)?.to_u32()?;
273                    let left = this.project_index(&left, (right & 0b111).into())?;
274
275                    this.copy_op(&left, &dest)?;
276                }
277            }
278            // Used to implement the _mm256_permute2x128_si256 function.
279            // Shuffles 128-bit blocks of `a` and `b` using `imm` as pattern.
280            "vperm2i128" => {
281                let [left, right, imm] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
282
283                assert_eq!(left.layout.size.bits(), 256);
284                assert_eq!(right.layout.size.bits(), 256);
285                assert_eq!(dest.layout.size.bits(), 256);
286
287                // Transmute to `[i128; 2]`
288
289                let array_layout =
290                    this.layout_of(Ty::new_array(this.tcx.tcx, this.tcx.types.i128, 2))?;
291                let left = left.transmute(array_layout, this)?;
292                let right = right.transmute(array_layout, this)?;
293                let dest = dest.transmute(array_layout, this)?;
294
295                let imm = this.read_scalar(imm)?.to_u8()?;
296
297                for i in 0..2 {
298                    let dest = this.project_index(&dest, i)?;
299                    let src = match (imm >> i.strict_mul(4)) & 0b11 {
300                        0 => this.project_index(&left, 0)?,
301                        1 => this.project_index(&left, 1)?,
302                        2 => this.project_index(&right, 0)?,
303                        3 => this.project_index(&right, 1)?,
304                        _ => unreachable!(),
305                    };
306
307                    this.copy_op(&src, &dest)?;
308                }
309            }
310            // Used to implement the _mm256_sad_epu8 function.
311            // Compute the absolute differences of packed unsigned 8-bit integers
312            // in `left` and `right`, then horizontally sum each consecutive 8
313            // differences to produce four unsigned 16-bit integers, and pack
314            // these unsigned 16-bit integers in the low 16 bits of 64-bit elements
315            // in `dest`.
316            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_sad_epu8
317            "psad.bw" => {
318                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
319
320                let (left, left_len) = this.project_to_simd(left)?;
321                let (right, right_len) = this.project_to_simd(right)?;
322                let (dest, dest_len) = this.project_to_simd(dest)?;
323
324                assert_eq!(left_len, right_len);
325                assert_eq!(left_len, dest_len.strict_mul(8));
326
327                for i in 0..dest_len {
328                    let dest = this.project_index(&dest, i)?;
329
330                    let mut acc: u16 = 0;
331                    for j in 0..8 {
332                        let src_index = i.strict_mul(8).strict_add(j);
333
334                        let left = this.project_index(&left, src_index)?;
335                        let left = this.read_scalar(&left)?.to_u8()?;
336
337                        let right = this.project_index(&right, src_index)?;
338                        let right = this.read_scalar(&right)?.to_u8()?;
339
340                        acc = acc.strict_add(left.abs_diff(right).into());
341                    }
342
343                    this.write_scalar(Scalar::from_u64(acc.into()), &dest)?;
344                }
345            }
346            // Used to implement the _mm256_shuffle_epi8 intrinsic.
347            // Shuffles bytes from `left` using `right` as pattern.
348            // Each 128-bit block is shuffled independently.
349            "pshuf.b" => {
350                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
351
352                let (left, left_len) = this.project_to_simd(left)?;
353                let (right, right_len) = this.project_to_simd(right)?;
354                let (dest, dest_len) = this.project_to_simd(dest)?;
355
356                assert_eq!(dest_len, left_len);
357                assert_eq!(dest_len, right_len);
358
359                for i in 0..dest_len {
360                    let right = this.read_scalar(&this.project_index(&right, i)?)?.to_u8()?;
361                    let dest = this.project_index(&dest, i)?;
362
363                    let res = if right & 0x80 == 0 {
364                        // Shuffle each 128-bit (16-byte) block independently.
365                        let j = u64::from(right % 16).strict_add(i & !15);
366                        this.read_scalar(&this.project_index(&left, j)?)?
367                    } else {
368                        // If the highest bit in `right` is 1, write zero.
369                        Scalar::from_u8(0)
370                    };
371
372                    this.write_scalar(res, &dest)?;
373                }
374            }
375            // Used to implement the _mm256_sign_epi{8,16,32} functions.
376            // Negates elements from `left` when the corresponding element in
377            // `right` is negative. If an element from `right` is zero, zero
378            // is writen to the corresponding output element.
379            // Basically, we multiply `left` with `right.signum()`.
380            "psign.b" | "psign.w" | "psign.d" => {
381                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
382
383                psign(this, left, right, dest)?;
384            }
385            // Used to implement the _mm256_{sll,srl,sra}_epi{16,32,64} functions
386            // (except _mm256_sra_epi64, which is not available in AVX2).
387            // Shifts N-bit packed integers in left by the amount in right.
388            // `right` is as 128-bit vector. but it is interpreted as a single
389            // 64-bit integer (remaining bits are ignored).
390            // For logic shifts, when right is larger than N - 1, zero is produced.
391            // For arithmetic shifts, when right is larger than N - 1, the sign bit
392            // is copied to remaining bits.
393            "psll.w" | "psrl.w" | "psra.w" | "psll.d" | "psrl.d" | "psra.d" | "psll.q"
394            | "psrl.q" => {
395                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
396
397                let which = match unprefixed_name {
398                    "psll.w" | "psll.d" | "psll.q" => ShiftOp::Left,
399                    "psrl.w" | "psrl.d" | "psrl.q" => ShiftOp::RightLogic,
400                    "psra.w" | "psra.d" => ShiftOp::RightArith,
401                    _ => unreachable!(),
402                };
403
404                shift_simd_by_scalar(this, left, right, which, dest)?;
405            }
406            // Used to implement the _mm{,256}_{sllv,srlv,srav}_epi{32,64} functions
407            // (except _mm{,256}_srav_epi64, which are not available in AVX2).
408            "psllv.d" | "psllv.d.256" | "psllv.q" | "psllv.q.256" | "psrlv.d" | "psrlv.d.256"
409            | "psrlv.q" | "psrlv.q.256" | "psrav.d" | "psrav.d.256" => {
410                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
411
412                let which = match unprefixed_name {
413                    "psllv.d" | "psllv.d.256" | "psllv.q" | "psllv.q.256" => ShiftOp::Left,
414                    "psrlv.d" | "psrlv.d.256" | "psrlv.q" | "psrlv.q.256" => ShiftOp::RightLogic,
415                    "psrav.d" | "psrav.d.256" => ShiftOp::RightArith,
416                    _ => unreachable!(),
417                };
418
419                shift_simd_by_simd(this, left, right, which, dest)?;
420            }
421            _ => return interp_ok(EmulateItemResult::NotSupported),
422        }
423        interp_ok(EmulateItemResult::NeedsReturn)
424    }
425}