Merge pull request #4912 from folkertdev/aarch64-pairwise-add

RalfJung · web-flow · commit 76621481ee89 · 2026-03-30T12:08:11.000Z
aarch64: add shims for pairwise widening/wrapping addition
diff --git a/src/shims/aarch64.rs b/src/shims/aarch64.rs
@@ -59,6 +59,93 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     this.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] =
+                    this.check_shim_sig_lenient(abi, CanonAbi::C, link_name, args)?;
+
+                let (left, left_len) = this.project_to_simd(left)?;
+                let (right, right_len) = this.project_to_simd(right)?;
+                let (dest, dest_len) = this.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 = this.read_immediate(&this.project_index(src, i)?)?;
+                    let rhs = this.read_immediate(&this.project_index(src, i.strict_add(1))?)?;
+
+                    // Wrapping addition on the element type.
+                    let sum = this.binary_op(BinOp::Add, &lhs, &rhs)?;
+
+                    let dst_lane = this.project_index(&dest, lane_idx)?;
+                    this.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] = this.check_shim_sig_lenient(abi, CanonAbi::C, link_name, args)?;
+
+                let (src, src_len) = this.project_to_simd(src)?;
+                let (dest, dest_len) = this.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(this, 0).size;
+                let dest_elem_size = dest.layout.field(this, 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 = this.read_scalar(&this.project_index(&src, src_idx)?)?;
+                    let b_scalar =
+                        this.read_scalar(&this.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 = this.project_index(&dest, dest_idx)?;
+                    this.write_scalar(Scalar::from_uint(sum, dest_elem_size), &dst_lane)?;
+                }
+            }
+
             // Vector table lookup: each index selects a byte from the 16-byte table, out-of-range -> 0.
             // Used to implement vtbl1_u8 function.
             // LLVM does not have a portable shuffle that takes non-const indices
diff --git a/tests/pass/shims/aarch64/intrinsics-aarch64-neon.rs b/tests/pass/shims/aarch64/intrinsics-aarch64-neon.rs
@@ -12,6 +12,8 @@ fn main() {
     unsafe {
         test_vpmaxq_u8();
         test_tbl1_v16i8_basic();
+        test_vpadd();
+        test_vpaddl();
     }
 }
 
@@ -65,3 +67,93 @@ fn test_tbl1_v16i8_basic() {
         assert_eq!(&got2_arr[3..16], &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12][..]);
     }
 }
+#[target_feature(enable = "neon")]
+unsafe fn test_vpadd() {
+    let a = vld1_s8([1, 2, 3, 4, 5, 6, 7, 8].as_ptr());
+    let b = vld1_s8([9, 10, -1, 2, i8::MIN, i8::MIN, i8::MAX, i8::MAX].as_ptr());
+    let e =
+        [3i8, 7, 11, 15, 19, -1 + 2, i8::MIN.wrapping_add(i8::MIN), i8::MAX.wrapping_add(i8::MAX)];
+    let mut r = [0i8; 8];
+    vst1_s8(r.as_mut_ptr(), vpadd_s8(a, b));
+    assert_eq!(r, e);
+
+    let a = vld1_s16([1, 2, 3, 4].as_ptr());
+    let b = vld1_s16([-1, 2, i16::MAX, i16::MAX].as_ptr());
+    let e = [3i16, 7, -1 + 2, i16::MAX.wrapping_add(i16::MAX)];
+    let mut r = [0i16; 4];
+    vst1_s16(r.as_mut_ptr(), vpadd_s16(a, b));
+    assert_eq!(r, e);
+
+    let a = vld1_s32([1, 2].as_ptr());
+    let b = vld1_s32([i32::MAX, i32::MAX].as_ptr());
+    let e = [3i32, i32::MAX.wrapping_add(i32::MAX)];
+    let mut r = [0i32; 2];
+    vst1_s32(r.as_mut_ptr(), vpadd_s32(a, b));
+    assert_eq!(r, e);
+
+    let a = vld1_u8([1, 2, 3, 4, 5, 6, 7, 8].as_ptr());
+    let b = vld1_u8([9, 10, 11, 12, 13, 14, u8::MAX, u8::MAX].as_ptr());
+    let e = [3u8, 7, 11, 15, 19, 23, 27, 254];
+    let mut r = [0u8; 8];
+    vst1_u8(r.as_mut_ptr(), vpadd_u8(a, b));
+    assert_eq!(r, e);
+
+    let a = vld1_u16([1, 2, 3, 4].as_ptr());
+    let b = vld1_u16([5, 6, u16::MAX, u16::MAX].as_ptr());
+    let e = [3u16, 7, 11, 65534];
+    let mut r = [0u16; 4];
+    vst1_u16(r.as_mut_ptr(), vpadd_u16(a, b));
+    assert_eq!(r, e);
+
+    let a = vld1_u32([1, 2].as_ptr());
+    let b = vld1_u32([u32::MAX, u32::MAX].as_ptr());
+    let e = [3u32, u32::MAX.wrapping_add(u32::MAX)];
+    let mut r = [0u32; 2];
+    vst1_u32(r.as_mut_ptr(), vpadd_u32(a, b));
+    assert_eq!(r, e);
+}
+
+#[target_feature(enable = "neon")]
+unsafe fn test_vpaddl() {
+    let a = vld1_u8([1, 2, 3, 4, 5, 6, u8::MAX, u8::MAX].as_ptr());
+    let e = [3u16, 7, 11, 510];
+    let mut r = [0u16; 4];
+    vst1_u16(r.as_mut_ptr(), vpaddl_u8(a));
+    assert_eq!(r, e);
+
+    let a = vld1q_u8([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, u8::MAX, u8::MAX].as_ptr());
+    let e = [3u16, 7, 11, 15, 19, 23, 27, 510];
+    let mut r = [0u16; 8];
+    vst1q_u16(r.as_mut_ptr(), vpaddlq_u8(a));
+    assert_eq!(r, e);
+
+    let a = vld1_u16([1, 2, u16::MAX, u16::MAX].as_ptr());
+    let e = [3u32, 131070];
+    let mut r = [0u32; 2];
+    vst1_u32(r.as_mut_ptr(), vpaddl_u16(a));
+    assert_eq!(r, e);
+
+    let a = vld1q_u16([1, 2, 3, 4, 5, 6, u16::MAX, u16::MAX].as_ptr());
+    let e = [3u32, 7, 11, 131070];
+    let mut r = [0u32; 4];
+    vst1q_u32(r.as_mut_ptr(), vpaddlq_u16(a));
+    assert_eq!(r, e);
+
+    let a = vld1_u32([1, 2].as_ptr());
+    let e = [3u64];
+    let mut r = [0u64; 1];
+    vst1_u64(r.as_mut_ptr(), vpaddl_u32(a));
+    assert_eq!(r, e);
+
+    let a = vld1_u32([u32::MAX, u32::MAX].as_ptr());
+    let e = [8589934590];
+    let mut r = [0u64; 1];
+    vst1_u64(r.as_mut_ptr(), vpaddl_u32(a));
+    assert_eq!(r, e);
+
+    let a = vld1q_u32([1, 2, u32::MAX, u32::MAX].as_ptr());
+    let e = [3u64, 8589934590];
+    let mut r = [0u64; 2];
+    vst1q_u64(r.as_mut_ptr(), vpaddlq_u32(a));
+    assert_eq!(r, e);
+}
