Merge pull request #4678 from folkertdev/ternary-logic

add shim for avx512 ternarylogic functions

Author: RalfJung

src/tools/miri/src/shims/x86/avx512.rs

@@ -0,0 +1,85 @@
+use rustc_abi::CanonAbi;
+use rustc_middle::ty::Ty;
+use rustc_span::Symbol;
+use rustc_target::callconv::FnAbi;
+
+use crate::*;
+
+impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
+pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
+    fn emulate_x86_avx512_intrinsic(
+        &mut self,
+        link_name: Symbol,
+        abi: &FnAbi<'tcx, Ty<'tcx>>,
+        args: &[OpTy<'tcx>],
+        dest: &MPlaceTy<'tcx>,
+    ) -> InterpResult<'tcx, EmulateItemResult> {
+        let this = self.eval_context_mut();
+        // Prefix should have already been checked.
+        let unprefixed_name = link_name.as_str().strip_prefix("llvm.x86.avx512.").unwrap();
+
+        match unprefixed_name {
+            // Used by the ternarylogic functions.
+            "pternlog.d.128" | "pternlog.d.256" | "pternlog.d.512" => {
+                this.expect_target_feature_for_intrinsic(link_name, "avx512f")?;
+                if matches!(unprefixed_name, "pternlog.d.128" | "pternlog.d.256") {
+                    this.expect_target_feature_for_intrinsic(link_name, "avx512vl")?;
+                }
+
+                let [a, b, c, imm8] =
+                    this.check_shim_sig_lenient(abi, CanonAbi::C, link_name, args)?;
+
+                assert_eq!(dest.layout, a.layout);
+                assert_eq!(dest.layout, b.layout);
+                assert_eq!(dest.layout, c.layout);
+
+                // The signatures of these operations are:
+                //
+                // ```
+                // fn vpternlogd(a: i32x16, b: i32x16, c: i32x16, imm8: i32) -> i32x16;
+                // fn vpternlogd256(a: i32x8, b: i32x8, c: i32x8, imm8: i32) -> i32x8;
+                // fn vpternlogd128(a: i32x4, b: i32x4, c: i32x4, imm8: i32) -> i32x4;
+                // ```
+                //
+                // The element type is always a 32-bit integer, the width varies.
+
+                let (a, _a_len) = this.project_to_simd(a)?;
+                let (b, _b_len) = this.project_to_simd(b)?;
+                let (c, _c_len) = this.project_to_simd(c)?;
+                let (dest, dest_len) = this.project_to_simd(dest)?;
+
+                // Compute one lane with ternary table.
+                let tern = |xa: u32, xb: u32, xc: u32, imm: u32| -> u32 {
+                    let mut out = 0u32;
+                    // At each bit position, select bit from imm8 at index = (a << 2) | (b << 1) | c
+                    for bit in 0..32 {
+                        let ia = (xa >> bit) & 1;
+                        let ib = (xb >> bit) & 1;
+                        let ic = (xc >> bit) & 1;
+                        let idx = (ia << 2) | (ib << 1) | ic;
+                        let v = (imm >> idx) & 1;
+                        out |= v << bit;
+                    }
+                    out
+                };
+
+                let imm8 = this.read_scalar(imm8)?.to_u32()? & 0xFF;
+                for i in 0..dest_len {
+                    let a_lane = this.project_index(&a, i)?;
+                    let b_lane = this.project_index(&b, i)?;
+                    let c_lane = this.project_index(&c, i)?;
+                    let d_lane = this.project_index(&dest, i)?;
+
+                    let va = this.read_scalar(&a_lane)?.to_u32()?;
+                    let vb = this.read_scalar(&b_lane)?.to_u32()?;
+                    let vc = this.read_scalar(&c_lane)?.to_u32()?;
+
+                    let r = tern(va, vb, vc, imm8);
+                    this.write_scalar(Scalar::from_u32(r), &d_lane)?;
+                }
+            }
+            _ => return interp_ok(EmulateItemResult::NotSupported),
+        }
+        interp_ok(EmulateItemResult::NeedsReturn)
+    }
+}

src/tools/miri/src/shims/x86/mod.rs

@@ -13,6 +13,7 @@ use crate::*;
 mod aesni;
 mod avx;
 mod avx2;
+mod avx512;
 mod bmi;
 mod gfni;
 mod sha;
@@ -152,6 +153,11 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     this, link_name, abi, args, dest,
                 );
             }
+            name if name.starts_with("avx512.") => {
+                return avx512::EvalContextExt::emulate_x86_avx512_intrinsic(
+                    this, link_name, abi, args, dest,
+                );
+            }
 
             _ => return interp_ok(EmulateItemResult::NotSupported),
         }

src/tools/miri/tests/pass/shims/x86/intrinsics-x86-avx512.rs

@@ -17,6 +17,7 @@ fn main() {
     unsafe {
         test_avx512bitalg();
         test_avx512vpopcntdq();
+        test_avx512ternarylogic();
     }
 }
 
@@ -191,6 +192,77 @@ unsafe fn test_avx512vpopcntdq() {
     test_mm_popcnt_epi64();
 }
 
+#[target_feature(enable = "avx512f,avx512vl")]
+unsafe fn test_avx512ternarylogic() {
+    #[target_feature(enable = "avx512f")]
+    unsafe fn test_mm512_ternarylogic_epi32() {
+        let a = _mm512_set4_epi32(0b100, 0b110, 0b001, 0b101);
+        let b = _mm512_set4_epi32(0b010, 0b011, 0b001, 0b110);
+        let c = _mm512_set4_epi32(0b001, 0b000, 0b001, 0b111);
+
+        // Identity of A.
+        let r = _mm512_ternarylogic_epi32::<0b1111_0000>(a, b, c);
+        assert_eq_m512i(r, a);
+
+        // Bitwise xor.
+        let r = _mm512_ternarylogic_epi32::<0b10010110>(a, b, c);
+        let e = _mm512_set4_epi32(0b111, 0b101, 0b001, 0b100);
+        assert_eq_m512i(r, e);
+
+        // Majority (2 or more bits set).
+        let r = _mm512_ternarylogic_epi32::<0b1110_1000>(a, b, c);
+        let e = _mm512_set4_epi32(0b000, 0b010, 0b001, 0b111);
+        assert_eq_m512i(r, e);
+    }
+    test_mm512_ternarylogic_epi32();
+
+    #[target_feature(enable = "avx512f,avx512vl")]
+    unsafe fn test_mm256_ternarylogic_epi32() {
+        let _mm256_set4_epi32 = |a, b, c, d| _mm256_setr_epi32(a, b, c, d, a, b, c, d);
+
+        let a = _mm256_set4_epi32(0b100, 0b110, 0b001, 0b101);
+        let b = _mm256_set4_epi32(0b010, 0b011, 0b001, 0b110);
+        let c = _mm256_set4_epi32(0b001, 0b000, 0b001, 0b111);
+
+        // Identity of A.
+        let r = _mm256_ternarylogic_epi32::<0b1111_0000>(a, b, c);
+        assert_eq_m256i(r, a);
+
+        // Bitwise xor.
+        let r = _mm256_ternarylogic_epi32::<0b10010110>(a, b, c);
+        let e = _mm256_set4_epi32(0b111, 0b101, 0b001, 0b100);
+        assert_eq_m256i(r, e);
+
+        // Majority (2 or more bits set).
+        let r = _mm256_ternarylogic_epi32::<0b1110_1000>(a, b, c);
+        let e = _mm256_set4_epi32(0b000, 0b010, 0b001, 0b111);
+        assert_eq_m256i(r, e);
+    }
+    test_mm256_ternarylogic_epi32();
+
+    #[target_feature(enable = "avx512f,avx512vl")]
+    unsafe fn test_mm_ternarylogic_epi32() {
+        let a = _mm_setr_epi32(0b100, 0b110, 0b001, 0b101);
+        let b = _mm_setr_epi32(0b010, 0b011, 0b001, 0b110);
+        let c = _mm_setr_epi32(0b001, 0b000, 0b001, 0b111);
+
+        // Identity of A.
+        let r = _mm_ternarylogic_epi32::<0b1111_0000>(a, b, c);
+        assert_eq_m128i(r, a);
+
+        // Bitwise xor.
+        let r = _mm_ternarylogic_epi32::<0b10010110>(a, b, c);
+        let e = _mm_setr_epi32(0b111, 0b101, 0b001, 0b100);
+        assert_eq_m128i(r, e);
+
+        // Majority (2 or more bits set).
+        let r = _mm_ternarylogic_epi32::<0b1110_1000>(a, b, c);
+        let e = _mm_setr_epi32(0b000, 0b010, 0b001, 0b111);
+        assert_eq_m128i(r, e);
+    }
+    test_mm_ternarylogic_epi32();
+}
+
 #[track_caller]
 unsafe fn assert_eq_m512i(a: __m512i, b: __m512i) {
     assert_eq!(transmute::<_, [i32; 16]>(a), transmute::<_, [i32; 16]>(b))