miri/intrinsics/

aarch64.rs

use rustc_middle::mir::BinOp;
use rustc_span::Symbol;

use crate::intrinsics::math::{compute_crc32, sha256};
use crate::*;

impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
    fn emulate_aarch64_intrinsic(
        &mut self,
        link_name: Symbol,
        args: &[OpTy<'tcx>],
        dest: &MPlaceTy<'tcx>,
    ) -> InterpResult<'tcx, EmulateItemResult> {
        let ecx = self.eval_context_mut();
        // Prefix should have already been checked.
        let unprefixed_name = link_name.as_str().strip_prefix("llvm.aarch64.").unwrap();
        match unprefixed_name {
            // Used to implement the vpmaxq_u8 function.
            // Computes the maximum of adjacent pairs; the first half of the output is produced from the
            // `left` input, the second half of the output from the `right` input.
            // https://developer.arm.com/architectures/instruction-sets/intrinsics/vpmaxq_u8
            "neon.umaxp.v16i8" => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let (left, left_len) = ecx.project_to_simd(left)?;
                let (right, right_len) = ecx.project_to_simd(right)?;
                let (dest, lane_count) = ecx.project_to_simd(dest)?;
                assert_eq!(left_len, right_len);
                assert_eq!(lane_count, left_len);

                for lane_idx in 0..lane_count {
                    let src = if lane_idx < (lane_count / 2) { &left } else { &right };
                    let src_idx = lane_idx.strict_rem(lane_count / 2);

                    let lhs_lane =
                        ecx.read_immediate(&ecx.project_index(src, src_idx.strict_mul(2))?)?;
                    let rhs_lane = ecx.read_immediate(
                        &ecx.project_index(src, src_idx.strict_mul(2).strict_add(1))?,
                    )?;

                    // Compute `if lhs > rhs { lhs } else { rhs }`, i.e., `max`.
                    let res_lane =
                        if ecx.binary_op(BinOp::Gt, &lhs_lane, &rhs_lane)?.to_scalar().to_bool()? {
                            lhs_lane
                        } else {
                            rhs_lane
                        };

                    let dest = ecx.project_index(&dest, lane_idx)?;
                    ecx.write_immediate(*res_lane, &dest)?;
                }
            }

            // Wrapping pairwise addition.
            //
            // Concatenates the two input vectors and adds adjacent elements. For input vectors `v`
            // and `w` this computes `[v0 + v1, v2 + v3, ..., w0 + w1, w2 + w3, ...]`, using
            // wrapping addition for `+`.
            //
            // Used by `vpadd_{s8, u8, s16, u16, s32, u32}`.
            name if name.starts_with("neon.addp.") => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let (left, left_len) = ecx.project_to_simd(left)?;
                let (right, right_len) = ecx.project_to_simd(right)?;
                let (dest, dest_len) = ecx.project_to_simd(dest)?;

                assert_eq!(left_len, right_len);
                assert_eq!(left_len, dest_len);

                assert_eq!(left.layout, right.layout);
                assert_eq!(left.layout, dest.layout);

                assert!(dest_len.is_multiple_of(2));
                let half_len = dest_len.strict_div(2);

                for lane_idx in 0..dest_len {
                    // The left and right vectors are concatenated.
                    let (src, src_pair_idx) = if lane_idx < half_len {
                        (&left, lane_idx)
                    } else {
                        (&right, lane_idx.strict_sub(half_len))
                    };
                    // Convert "pair index" into "index of first element of the pair".
                    let i = src_pair_idx.strict_mul(2);

                    let lhs = ecx.read_immediate(&ecx.project_index(src, i)?)?;
                    let rhs = ecx.read_immediate(&ecx.project_index(src, i.strict_add(1))?)?;

                    // Wrapping addition on the element type.
                    let sum = ecx.binary_op(BinOp::Add, &lhs, &rhs)?;

                    let dst_lane = ecx.project_index(&dest, lane_idx)?;
                    ecx.write_immediate(*sum, &dst_lane)?;
                }
            }

            // Widening pairwise addition.
            //
            // Takes a single input vector, and an output vector with half as many lanes and double
            // the element width. Takes adjacent pairs of elements, widens both, and then adds them
            // together.
            //
            // Used by `vpaddl_{u8, u16, u32}` and `vpaddlq_{u8, u16, u32}`.
            name if name.starts_with("neon.uaddlp.") => {
                let [src] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let (src, src_len) = ecx.project_to_simd(src)?;
                let (dest, dest_len) = ecx.project_to_simd(dest)?;

                // Operates pairwise, so src has twice as many lanes.
                assert_eq!(src_len, dest_len.strict_mul(2));

                let src_elem_size = src.layout.field(ecx, 0).size;
                let dest_elem_size = dest.layout.field(ecx, 0).size;

                // Widens, so dest elements must be exactly twice as wide.
                assert_eq!(dest_elem_size.bytes(), src_elem_size.bytes().strict_mul(2));

                for dest_idx in 0..dest_len {
                    let src_idx = dest_idx.strict_mul(2);

                    let a_scalar = ecx.read_scalar(&ecx.project_index(&src, src_idx)?)?;
                    let b_scalar =
                        ecx.read_scalar(&ecx.project_index(&src, src_idx.strict_add(1))?)?;

                    let a_val = a_scalar.to_uint(src_elem_size)?;
                    let b_val = b_scalar.to_uint(src_elem_size)?;

                    // Use addition on u128 to simulate widening addition for the destination type.
                    // This cannot wrap since the element type is at most u64.
                    let sum = a_val.strict_add(b_val);

                    let dst_lane = ecx.project_index(&dest, dest_idx)?;
                    ecx.write_scalar(Scalar::from_uint(sum, dest_elem_size), &dst_lane)?;
                }
            }

            // Signed saturating doubling multiply returning the high half.
            //
            // Used by the `vqdmulh*` functions.
            //
            // This LLVM intrinsic multiplies the values of corresponding elements of the two source
            // vector registers (which are signed integers), doubles the results, places the most significant half of the
            // final results (using a saturating cast to fit the element type) into a vector, and writes the vector to the destination register.
            //
            // https://developer.arm.com/architectures/instruction-sets/intrinsics#f:@navigationhierarchiessimdisa=[Neon]&q=vqdmulh
            name if name.starts_with("neon.sqdmulh.") => {
                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let (left, left_len) = ecx.project_to_simd(left)?;
                let (right, right_len) = ecx.project_to_simd(right)?;
                let (dest, dest_len) = ecx.project_to_simd(dest)?;
                assert_eq!(left_len, right_len);
                assert_eq!(left_len, dest_len);

                let elem_size = dest.layout.field(ecx, 0).size;
                let bits = elem_size.bits();
                let min = elem_size.signed_int_min();
                let max = elem_size.signed_int_max();

                for i in 0..dest_len {
                    let a = ecx.read_scalar(&ecx.project_index(&left, i)?)?.to_int(elem_size)?;
                    let b = ecx.read_scalar(&ecx.project_index(&right, i)?)?.to_int(elem_size)?;

                    // Uses i128 arithmetic, which cannot overflow because the intrinsic takes at most i32.
                    let doubled = a.strict_mul(b).strict_mul(2);
                    let res = (doubled >> bits).clamp(min, max);

                    ecx.write_scalar(
                        Scalar::from_int(res, elem_size),
                        &ecx.project_index(&dest, i)?,
                    )?;
                }
            }

            // Vector table lookup: each index selects a byte from the table,
            // out-of-range -> 0.
            //
            // Used to implement the vtblN, vqtblN and vqtblNq set of functions, e.g.:
            //
            // - https://developer.arm.com/architectures/instruction-sets/intrinsics/vtbl1_u8
            // - https://developer.arm.com/architectures/instruction-sets/intrinsics/vqtbl1_u8
            // - https://developer.arm.com/architectures/instruction-sets/intrinsics/vqtbl1q_s8
            // - https://developer.arm.com/architectures/instruction-sets/intrinsics/vtbl2_u8
            // - https://developer.arm.com/architectures/instruction-sets/intrinsics/vqtbl4q_s8
            //
            // LLVM does not have a portable shuffle that takes non-const indices
            // so we need to implement this ourselves.
            _ if unprefixed_name.starts_with("neon.tbl") => {
                // The table segments always have 16 elements, the index vector and output vector
                // have either 8 or 16 elements.
                let (table_segments, indices) = match unprefixed_name {
                    "neon.tbl1.v8i8" | "neon.tbl1.v16i8" => {
                        let [table, indices] =
                            ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                        let (table, len) = ecx.project_to_simd(table)?;
                        assert_eq!(len, 16);
                        (vec![table], indices)
                    }
                    "neon.tbl2.v8i8" | "neon.tbl2.v16i8" => {
                        let [table0, table1, indices] =
                            ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                        let (table0, len0) = ecx.project_to_simd(table0)?;
                        let (table1, len1) = ecx.project_to_simd(table1)?;
                        assert_eq!([len0, len1], [16; 2]);
                        (vec![table0, table1], indices)
                    }
                    "neon.tbl3.v8i8" | "neon.tbl3.v16i8" => {
                        let [table0, table1, table2, indices] =
                            ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                        let (table0, len0) = ecx.project_to_simd(table0)?;
                        let (table1, len1) = ecx.project_to_simd(table1)?;
                        let (table2, len2) = ecx.project_to_simd(table2)?;
                        assert_eq!([len0, len1, len2], [16; 3]);
                        (vec![table0, table1, table2], indices)
                    }
                    "neon.tbl4.v8i8" | "neon.tbl4.v16i8" => {
                        let [table0, table1, table2, table3, indices] =
                            ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                        let (table0, len0) = ecx.project_to_simd(table0)?;
                        let (table1, len1) = ecx.project_to_simd(table1)?;
                        let (table2, len2) = ecx.project_to_simd(table2)?;
                        let (table3, len3) = ecx.project_to_simd(table3)?;
                        assert_eq!([len0, len1, len2, len3], [16; 4]);
                        (vec![table0, table1, table2, table3], indices)
                    }
                    _ => unreachable!(),
                };

                let (indices, idx_len) = ecx.project_to_simd(indices)?;
                let (dest, dest_len) = ecx.project_to_simd(dest)?;
                assert_eq!(idx_len, dest_len);

                for i in 0..dest_len {
                    let idx = ecx.read_immediate(&ecx.project_index(&indices, i)?)?;
                    let idx = idx.to_scalar().to_u8()?;

                    // The LLVM intrinsic table segments always have 16 elements.
                    let val = if usize::from(idx) < table_segments.len().strict_mul(16) {
                        let table = &table_segments[usize::from(idx.strict_div(16))];
                        let lane = u64::from(idx.strict_rem(16));
                        let t = ecx.read_immediate(&ecx.project_index(table, lane)?)?;
                        t.to_scalar()
                    } else {
                        Scalar::from_u8(0)
                    };

                    ecx.write_scalar(val, &ecx.project_index(&dest, i)?)?;
                }
            }
            // Used to implement the __crc32{b,h,w,x} and __crc32c{b,h,w,x} functions.
            // Polynomial 0x04C11DB7 (standard CRC-32):
            // https://developer.arm.com/documentation/ddi0602/latest/Base-Instructions/CRC32B--CRC32H--CRC32W--CRC32X--CRC32-checksum-
            // Polynomial 0x1EDC6F41 (CRC-32C / Castagnoli):
            // https://developer.arm.com/documentation/ddi0602/latest/Base-Instructions/CRC32CB--CRC32CH--CRC32CW--CRC32CX--CRC32C-checksum-
            "crc32b" | "crc32h" | "crc32w" | "crc32x" | "crc32cb" | "crc32ch" | "crc32cw"
            | "crc32cx" => {
                ecx.expect_target_feature_for_intrinsic(link_name, "crc")?;
                // The polynomial constants below include the leading 1 bit
                // (e.g. 0x104C11DB7 instead of 0x04C11DB7) which the ARM docs
                // omit but the polynomial division algorithm requires.
                let (bit_size, polynomial): (u32, u128) = match unprefixed_name {
                    "crc32b" => (8, 0x104C11DB7),
                    "crc32h" => (16, 0x104C11DB7),
                    "crc32w" => (32, 0x104C11DB7),
                    "crc32x" => (64, 0x104C11DB7),
                    "crc32cb" => (8, 0x11EDC6F41),
                    "crc32ch" => (16, 0x11EDC6F41),
                    "crc32cw" => (32, 0x11EDC6F41),
                    "crc32cx" => (64, 0x11EDC6F41),
                    _ => unreachable!(),
                };

                let [crc, data] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                let crc = ecx.read_scalar(crc)?;
                let data = ecx.read_scalar(data)?;

                // The CRC accumulator is always u32. The data argument is u32 for
                // b/h/w variants and u64 for the x variant, per the LLVM intrinsic
                // definitions (all b/h/w take i32, only x takes i64).
                // https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/IntrinsicsAArch64.td
                // If the higher bits are non-zero, `compute_crc32` will panic. We should probably
                // raise a proper error instead, but outside stdarch nobody can trigger this anyway.
                let crc = crc.to_u32()?;
                let data = if bit_size == 64 { data.to_u64()? } else { u64::from(data.to_u32()?) };

                let result = compute_crc32(crc, data, bit_size, polynomial);
                ecx.write_scalar(Scalar::from_u32(result), dest)?;
            }
            // Polynomial multiply long (64-bit x 64-bit -> 128-bit).
            //
            // This is the same as "carryless" multiplication, see
            // <https://en.wikipedia.org/wiki/Carry-less_product#Multiplication_of_polynomials>.
            //
            // Used to implement the vmull_p64 and vmull_high_p64 functions.
            // https://developer.arm.com/architectures/instruction-sets/intrinsics/vmull_p64
            "neon.pmull64" => {
                // LLVM and GCC group pmull with the AES intrinsics.
                // Also see <https://gcc.gnu.org/pipermail/gcc-patches/2023-February/612088.html>.
                ecx.expect_target_feature_for_intrinsic(link_name, "aes")?;

                let [left, right] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;
                let left = ecx.read_scalar(left)?.to_u64()?;
                let right = ecx.read_scalar(right)?.to_u64()?;

                let result = left.widening_carryless_mul(right);

                // dest is int8x16_t, transmute to u128 for the write.
                let dest = dest.transmute(ecx.machine.layouts.u128, ecx)?;
                ecx.write_scalar(Scalar::from_u128(result), &dest)?;
            }

            // Used to implement the vsha256hq_u32 function.
            // https://developer.arm.com/architectures/instruction-sets/intrinsics/vsha256hq_u32
            "crypto.sha256h" => {
                ecx.expect_target_feature_for_intrinsic(link_name, "sha2")?;

                let [abcd, efgh, wk] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let (abcd, abcd_len) = ecx.project_to_simd(abcd)?;
                let (efgh, efgh_len) = ecx.project_to_simd(efgh)?;
                let (wk, wk_len) = ecx.project_to_simd(wk)?;
                let (dest, dest_len) = ecx.project_to_simd(dest)?;

                assert_eq!(abcd_len, 4);
                assert_eq!(efgh_len, 4);
                assert_eq!(wk_len, 4);
                assert_eq!(dest_len, 4);

                let abcd: [u32; 4] = read_u32x4(ecx, &abcd)?;
                let efgh: [u32; 4] = read_u32x4(ecx, &efgh)?;
                let wk: [u32; 4] = read_u32x4(ecx, &wk)?;

                let result = sha256h(abcd, efgh, wk);

                write_u32x4(ecx, &dest, result)?;
            }
            // Used to implement the vsha256h2q_u32 function.
            // https://developer.arm.com/architectures/instruction-sets/intrinsics/vsha256h2q_u32
            "crypto.sha256h2" => {
                ecx.expect_target_feature_for_intrinsic(link_name, "sha2")?;

                let [efgh, abcd, wk] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let (efgh, efgh_len) = ecx.project_to_simd(efgh)?;
                let (abcd, abcd_len) = ecx.project_to_simd(abcd)?;
                let (wk, wk_len) = ecx.project_to_simd(wk)?;
                let (dest, dest_len) = ecx.project_to_simd(dest)?;

                assert_eq!(efgh_len, 4);
                assert_eq!(abcd_len, 4);
                assert_eq!(wk_len, 4);
                assert_eq!(dest_len, 4);

                let efgh: [u32; 4] = read_u32x4(ecx, &efgh)?;
                let abcd: [u32; 4] = read_u32x4(ecx, &abcd)?;
                let wk: [u32; 4] = read_u32x4(ecx, &wk)?;

                let result = sha256h2(efgh, abcd, wk);

                write_u32x4(ecx, &dest, result)?;
            }
            // Used to implement the vsha256su0q_u32 function.
            // https://developer.arm.com/architectures/instruction-sets/intrinsics/vsha256su0q_u32
            "crypto.sha256su0" => {
                ecx.expect_target_feature_for_intrinsic(link_name, "sha2")?;

                let [a, b] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let (a, a_len) = ecx.project_to_simd(a)?;
                let (b, b_len) = ecx.project_to_simd(b)?;
                let (dest, dest_len) = ecx.project_to_simd(dest)?;

                assert_eq!(a_len, 4);
                assert_eq!(b_len, 4);
                assert_eq!(dest_len, 4);

                let a: [u32; 4] = read_u32x4(ecx, &a)?;
                let b: [u32; 4] = read_u32x4(ecx, &b)?;

                let result = sha256su0(a, b);

                write_u32x4(ecx, &dest, result)?;
            }
            // Used to implement the vsha256su1q_u32 function.
            // https://developer.arm.com/architectures/instruction-sets/intrinsics/vsha256su1q_u32
            "crypto.sha256su1" => {
                ecx.expect_target_feature_for_intrinsic(link_name, "sha2")?;

                let [a, b, c] = ecx.check_shim_sig_llvm_intrinsic(link_name, args)?;

                let (a, a_len) = ecx.project_to_simd(a)?;
                let (b, b_len) = ecx.project_to_simd(b)?;
                let (c, c_len) = ecx.project_to_simd(c)?;
                let (dest, dest_len) = ecx.project_to_simd(dest)?;

                assert_eq!(a_len, 4);
                assert_eq!(b_len, 4);
                assert_eq!(c_len, 4);
                assert_eq!(dest_len, 4);

                let a: [u32; 4] = read_u32x4(ecx, &a)?;
                let b: [u32; 4] = read_u32x4(ecx, &b)?;
                let c: [u32; 4] = read_u32x4(ecx, &c)?;

                let result = sha256su1(a, b, c);

                write_u32x4(ecx, &dest, result)?;
            }
            _ => return interp_ok(EmulateItemResult::NotSupported),
        }
        interp_ok(EmulateItemResult::NeedsReturn)
    }
}

/// Reads a `[u32; 4]` array.
fn read_u32x4<'c>(ecx: &mut MiriInterpCx<'c>, vec: &OpTy<'c>) -> InterpResult<'c, [u32; 4]> {
    let mut res = [0; 4];
    for (i, dst) in res.iter_mut().enumerate() {
        let projected = &ecx.project_index(vec, i.try_into().unwrap())?;
        *dst = ecx.read_scalar(projected)?.to_u32()?;
    }
    interp_ok(res)
}

fn write_u32x4<'c>(
    ecx: &mut MiriInterpCx<'c>,
    dest: &MPlaceTy<'c>,
    val: [u32; 4],
) -> InterpResult<'c, ()> {
    for (i, part) in val.into_iter().enumerate() {
        let projected = &ecx.project_index(dest, i.to_u64())?;
        ecx.write_scalar(Scalar::from_u32(part), projected)?;
    }
    interp_ok(())
}

fn sha256su0(v0: [u32; 4], v1: [u32; 4]) -> [u32; 4] {
    [
        v0[0].wrapping_add(sha256::sigma0(v0[1])),
        v0[1].wrapping_add(sha256::sigma0(v0[2])),
        v0[2].wrapping_add(sha256::sigma0(v0[3])),
        v0[3].wrapping_add(sha256::sigma0(v1[0])),
    ]
}

fn sha256su1(v0: [u32; 4], v1: [u32; 4], v2: [u32; 4]) -> [u32; 4] {
    let r0 = v0[0].wrapping_add(v1[1]).wrapping_add(sha256::sigma1(v2[2]));
    let r1 = v0[1].wrapping_add(v1[2]).wrapping_add(sha256::sigma1(v2[3]));
    let r2 = v0[2].wrapping_add(v1[3]).wrapping_add(sha256::sigma1(r0));
    let r3 = v0[3].wrapping_add(v2[0]).wrapping_add(sha256::sigma1(r1));
    [r0, r1, r2, r3]
}

// SHA256H/SHA256H2 do four compression rounds on the abcd/efgh layout.
// https://developer.arm.com/architectures/instruction-sets/intrinsics/#f:@navigationhierarchiesinstructiongroup=[Cryptography,SHA256]
fn sha256hash(abcd: [u32; 4], efgh: [u32; 4], wk: [u32; 4]) -> ([u32; 4], [u32; 4]) {
    let mut state = [abcd[0], abcd[1], abcd[2], abcd[3], efgh[0], efgh[1], efgh[2], efgh[3]];
    for &wk_i in &wk {
        state = sha256::round(state, wk_i);
    }
    ([state[0], state[1], state[2], state[3]], [state[4], state[5], state[6], state[7]])
}

fn sha256h(abcd: [u32; 4], efgh: [u32; 4], wk: [u32; 4]) -> [u32; 4] {
    sha256hash(abcd, efgh, wk).0
}

// sha256h2 takes efgh as the first argument. abcd and efgh are swapped when calling sha256hash.
fn sha256h2(efgh: [u32; 4], abcd: [u32; 4], wk: [u32; 4]) -> [u32; 4] {
    sha256hash(abcd, efgh, wk).1
}