miri/intrinsics/x86/

avx2.rs

use rustc_span::Symbol;

use super::{
    ShiftOp, mpsadbw, packssdw, packsswb, packusdw, packuswb, pavg, permute, pmaddbw, pmaddwd,
    pmulh, pmulhrsw, psadbw, pshufb, psign, shift_simd_by_scalar, shift_simd_by_simd,
};
use crate::*;

impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
    fn emulate_x86_avx2_intrinsic(
        &mut self,
        link_name: Symbol,
        args: &[OpTy<'tcx>],
        dest: &MPlaceTy<'tcx>,
    ) -> InterpResult<'tcx, EmulateItemResult> {
        let ecx = self.eval_context_mut();
        ecx.expect_target_feature_for_intrinsic(link_name, "avx2")?;
        // Prefix should have already been checked.
        let unprefixed_name = link_name.as_str().strip_prefix("llvm.x86.avx2.").unwrap();

        match unprefixed_name {
            // Used to implement the _mm256_avg_epu8 and _mm256_avg_epu16 functions.
            // Averages packed unsigned 8/16-bit integers in `left` and `right`.
            "pavg.b" | "pavg.w" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                pavg(ecx, left, right, dest)?;
            }
            // Used to implement the _mm256_mulhi_epi16 and _mm256_mulhi_epu16 functions.
            "pmulh.w" | "pmulhu.w" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                pmulh(ecx, left, right, dest)?;
            }
            // Used to implement `_mm{,_mask}_{i32,i64}gather_{epi32,epi64,pd,ps}` functions
            // Gathers elements from `slice` using `offsets * scale` as indices.
            // When the highest bit of the corresponding element of `mask` is 0,
            // the value is copied from `src` instead.
            "gather.d.d" | "gather.d.d.256" | "gather.d.q" | "gather.d.q.256" | "gather.q.d"
            | "gather.q.d.256" | "gather.q.q" | "gather.q.q.256" | "gather.d.pd"
            | "gather.d.pd.256" | "gather.q.pd" | "gather.q.pd.256" | "gather.d.ps"
            | "gather.d.ps.256" | "gather.q.ps" | "gather.q.ps.256" => {
                let [src, slice, offsets, mask, scale] =
                    ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                assert_eq!(dest.layout, src.layout);

                let (src, _) = ecx.project_to_simd(src)?;
                let (offsets, offsets_len) = ecx.project_to_simd(offsets)?;
                let (mask, mask_len) = ecx.project_to_simd(mask)?;
                let (dest, dest_len) = ecx.project_to_simd(dest)?;

                // There are cases like dest: i32x4, offsets: i64x2
                // If dest has more elements than offset, extra dest elements are filled with zero.
                // If offsets has more elements than dest, extra offsets are ignored.
                let actual_len = dest_len.min(offsets_len);

                assert_eq!(dest_len, mask_len);

                let mask_item_size = mask.layout.field(ecx, 0).size;
                let high_bit_offset = mask_item_size.bits().strict_sub(1);

                let scale = ecx.read_scalar(scale)?.to_i8()?;
                if !matches!(scale, 1 | 2 | 4 | 8) {
                    panic!("invalid gather scale {scale}");
                }
                let scale = i64::from(scale);

                let slice = ecx.read_pointer(slice)?;
                for i in 0..actual_len {
                    let mask = ecx.project_index(&mask, i)?;
                    let dest = ecx.project_index(&dest, i)?;

                    if ecx.read_scalar(&mask)?.to_uint(mask_item_size)? >> high_bit_offset != 0 {
                        let offset = ecx.project_index(&offsets, i)?;
                        let offset =
                            i64::try_from(ecx.read_scalar(&offset)?.to_int(offset.layout.size)?)
                                .unwrap();
                        let ptr = slice.wrapping_signed_offset(offset.strict_mul(scale), &ecx.tcx);
                        // Unaligned copy, which is what we want.
                        ecx.mem_copy(
                            ptr,
                            dest.ptr(),
                            dest.layout.size,
                            /*nonoverlapping*/ true,
                        )?;
                    } else {
                        ecx.copy_op(&ecx.project_index(&src, i)?, &dest)?;
                    }
                }
                for i in actual_len..dest_len {
                    let dest = ecx.project_index(&dest, i)?;
                    ecx.write_scalar(Scalar::from_int(0, dest.layout.size), &dest)?;
                }
            }
            // Used to implement the _mm256_maddubs_epi16 function.
            "pmadd.ub.sw" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                pmaddbw(ecx, left, right, dest)?;
            }
            // Used to implement the _mm256_mpsadbw_epu8 function.
            // Compute the sum of absolute differences of quadruplets of unsigned
            // 8-bit integers in `left` and `right`, and store the 16-bit results
            // in `right`. Quadruplets are selected from `left` and `right` with
            // offsets specified in `imm`.
            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_mpsadbw_epu8
            "mpsadbw" => {
                let [left, right, imm] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                mpsadbw(ecx, left, right, imm, dest)?;
            }
            // Used to implement the _mm256_mulhrs_epi16 function.
            // Multiplies packed 16-bit signed integer values, truncates the 32-bit
            // product to the 18 most significant bits by right-shifting, and then
            // divides the 18-bit value by 2 (rounding to nearest) by first adding
            // 1 and then taking the bits `1..=16`.
            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_mulhrs_epi16
            "pmul.hr.sw" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                pmulhrsw(ecx, left, right, dest)?;
            }
            // Used to implement the _mm256_packs_epi16 function.
            // Converts two 16-bit integer vectors to a single 8-bit integer
            // vector with signed saturation.
            "packsswb" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                packsswb(ecx, left, right, dest)?;
            }
            // Used to implement the _mm256_packs_epi32 function.
            // Converts two 32-bit integer vectors to a single 16-bit integer
            // vector with signed saturation.
            "packssdw" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                packssdw(ecx, left, right, dest)?;
            }
            // Used to implement the _mm256_packus_epi16 function.
            // Converts two 16-bit signed integer vectors to a single 8-bit
            // unsigned integer vector with saturation.
            "packuswb" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                packuswb(ecx, left, right, dest)?;
            }
            // Used to implement the _mm256_packus_epi32 function.
            // Concatenates two 32-bit signed integer vectors and converts
            // the result to a 16-bit unsigned integer vector with saturation.
            "packusdw" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                packusdw(ecx, left, right, dest)?;
            }
            // Used to implement _mm256_permutevar8x32_epi32 and _mm256_permutevar8x32_ps.
            "permd" | "permps" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                permute(ecx, left, right, dest)?;
            }
            // Used to implement the _mm256_sad_epu8 function.
            "psad.bw" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                psadbw(ecx, left, right, dest)?
            }
            // Used to implement the _mm256_shuffle_epi8 intrinsic.
            // Shuffles bytes from `left` using `right` as pattern.
            // Each 128-bit block is shuffled independently.
            "pshuf.b" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                pshufb(ecx, left, right, dest)?;
            }
            // Used to implement the _mm256_sign_epi{8,16,32} functions.
            // Negates elements from `left` when the corresponding element in
            // `right` is negative. If an element from `right` is zero, zero
            // is written to the corresponding output element.
            // Basically, we multiply `left` with `right.signum()`.
            "psign.b" | "psign.w" | "psign.d" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                psign(ecx, left, right, dest)?;
            }
            // Used to implement the _mm256_{sll,srl,sra}_epi{16,32,64} functions
            // (except _mm256_sra_epi64, which is not available in AVX2).
            // Shifts N-bit packed integers in left by the amount in right.
            // `right` is as 128-bit vector. but it is interpreted as a single
            // 64-bit integer (remaining bits are ignored).
            // For logic shifts, when right is larger than N - 1, zero is produced.
            // For arithmetic shifts, when right is larger than N - 1, the sign bit
            // is copied to remaining bits.
            "psll.w" | "psrl.w" | "psra.w" | "psll.d" | "psrl.d" | "psra.d" | "psll.q"
            | "psrl.q" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let which = match unprefixed_name {
                    "psll.w" | "psll.d" | "psll.q" => ShiftOp::Left,
                    "psrl.w" | "psrl.d" | "psrl.q" => ShiftOp::RightLogic,
                    "psra.w" | "psra.d" => ShiftOp::RightArith,
                    _ => unreachable!(),
                };

                shift_simd_by_scalar(ecx, left, right, which, dest)?;
            }
            // Used to implement the _mm256_madd_epi16 function.
            // Multiplies packed signed 16-bit integers in `left` and `right`, producing
            // intermediate signed 32-bit integers. Horizontally add adjacent pairs of
            // intermediate 32-bit integers, and pack the results in `dest`.
            "pmadd.wd" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                pmaddwd(ecx, left, right, dest)?;
            }
            // Used to implement the _mm{,256}_{sllv,srlv,srav}_epi{32,64} functions
            // (except _mm{,256}_srav_epi64, which are not available in AVX2).
            "psllv.d" | "psllv.d.256" | "psllv.q" | "psllv.q.256" | "psrlv.d" | "psrlv.d.256"
            | "psrlv.q" | "psrlv.q.256" | "psrav.d" | "psrav.d.256" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let which = match unprefixed_name {
                    "psllv.d" | "psllv.d.256" | "psllv.q" | "psllv.q.256" => ShiftOp::Left,
                    "psrlv.d" | "psrlv.d.256" | "psrlv.q" | "psrlv.q.256" => ShiftOp::RightLogic,
                    "psrav.d" | "psrav.d.256" => ShiftOp::RightArith,
                    _ => unreachable!(),
                };

                shift_simd_by_simd(ecx, left, right, which, dest)?;
            }

            _ => return interp_ok(EmulateItemResult::NotSupported),
        }
        interp_ok(EmulateItemResult::NeedsReturn)
    }
}