Introduce simd_shuffle intrinsic (#417)

* introduce simd_shuffle intrinsic

* add links to cbmc issues in simd_shuffle comments

* add assert for arguments length to simd instrinsic functions

* typo fix

Co-authored-by: Adrian Palacios <[email protected]>

* separate simd_shuffle test from other tests

Co-authored-by: Adrian Palacios <[email protected]>
Co-authored-by: Daniel Schwartz-Narbonne <[email protected]>

Author: 3 people

compiler/rustc_codegen_llvm/src/gotoc/cbmc/goto_program/expr.rs

@@ -146,6 +146,10 @@ pub enum ExprValue {
         op: UnaryOperand,
         e: Expr,
     },
+    /// `vec_typ x = >>> {elems0, elems1 ...} <<<`
+    Vector {
+        elems: Vec<Expr>,
+    },
 }
 
 /// Binary operators. The names are the same as in the Irep representation.
@@ -373,6 +377,21 @@ impl Expr {
         expr!(Array { elems }, typ)
     }
 
+    pub fn vector_expr(typ: Type, elems: Vec<Expr>) -> Self {
+        if let Type::Vector { size, typ: value_typ } = typ.clone() {
+            assert_eq!(size as usize, elems.len());
+            assert!(
+                elems.iter().all(|x| x.typ == *value_typ),
+                "Vector type and value types don't match: \n{:?}\n{:?}",
+                typ,
+                elems
+            );
+        } else {
+            unreachable!("Can't make a vector_val with non-vector target type {:?}", typ);
+        }
+        expr!(Vector { elems }, typ)
+    }
+
     /// `(__CPROVER_bool) >>> true/false <<<`. True/False as a single bit boolean.
     pub fn bool_constant(c: bool) -> Self {
         expr!(BoolConstant(c), Type::bool())

compiler/rustc_codegen_llvm/src/gotoc/cbmc/goto_program/symtab_transformer/transformer.rs

@@ -266,6 +266,7 @@ pub trait Transformer: Sized {
             ExprValue::Typecast(child) => self.transform_expr_typecast(typ, child),
             ExprValue::Union { value, field } => self.transform_expr_union(typ, value, field),
             ExprValue::UnOp { op, e } => self.transform_expr_un_op(typ, op, e),
+            ExprValue::Vector { elems } => self.transform_expr_vector(typ, elems),
         }
         .with_location(e.location().clone())
     }
@@ -282,6 +283,13 @@ pub trait Transformer: Sized {
         Expr::array_expr(transformed_typ, transformed_elems)
     }
 
+    /// Transform a vector expr (`vec_typ x[] = >>> {elems0, elems1 ...} <<<`)
+    fn transform_expr_vector(&self, typ: &Type, elems: &[Expr]) -> Expr {
+        let transformed_typ = self.transform_type(typ);
+        let transformed_elems = elems.iter().map(|elem| self.transform_expr(elem)).collect();
+        Expr::vector_expr(transformed_typ, transformed_elems)
+    }
+
     /// Transforms an array of expr (`typ x[width] = >>> {elem} <<<`)
     fn transform_expr_array_of(&self, typ: &Type, elem: &Expr) -> Expr {
         let transformed_typ = self.transform_type(typ);

compiler/rustc_codegen_llvm/src/gotoc/cbmc/irep/to_irep.rs

@@ -175,7 +175,6 @@ impl ToIrep for ExprValue {
                 sub: vec![e.to_irep(mm), Expr::int_constant(*offset, Type::ssize_t()).to_irep(mm)],
                 named_sub: btree_map![],
             },
-
             ExprValue::CBoolConstant(i) => Irep {
                 id: IrepId::Constant,
                 sub: vec![],
@@ -292,6 +291,11 @@ impl ToIrep for ExprValue {
             ExprValue::UnOp { op, e } => {
                 Irep { id: op.to_irep_id(), sub: vec![e.to_irep(mm)], named_sub: btree_map![] }
             }
+            ExprValue::Vector { elems } => Irep {
+                id: IrepId::Vector,
+                sub: elems.iter().map(|x| x.to_irep(mm)).collect(),
+                named_sub: btree_map![],
+            },
         }
     }
 }

compiler/rustc_codegen_llvm/src/gotoc/mir_to_goto/codegen/intrinsic.rs

@@ -217,6 +217,11 @@ impl<'tcx> GotocCtx<'tcx> {
             }};
         }
 
+        if let Some(stripped) = intrinsic.strip_prefix("simd_shuffle") {
+            let n: u64 = stripped.parse().unwrap();
+            return self.codegen_intrinsic_simd_shuffle(fargs, p, cbmc_ret_ty, n);
+        }
+
         match intrinsic {
             "abort" => Stmt::assert_false("abort intrinsic", loc),
             "add_with_overflow" => codegen_op_with_overflow!(add_overflow),
@@ -401,6 +406,7 @@ impl<'tcx> GotocCtx<'tcx> {
                     codegen_intrinsic_binop!(lshr)
                 }
             }
+            // "simd_shuffle#" => handled in an `if` preceding this match
             "simd_sub" => codegen_intrinsic_binop!(sub),
             "simd_xor" => codegen_intrinsic_binop!(bitxor),
             "size_of" => codegen_intrinsic_const!(),
@@ -814,6 +820,7 @@ impl<'tcx> GotocCtx<'tcx> {
         cbmc_ret_ty: Type,
         loc: Location,
     ) -> Stmt {
+        assert!(fargs.len() == 3, "simd_insert had unexpected arguments {:?}", fargs);
         let vec = fargs.remove(0);
         let index = fargs.remove(0);
         let newval = fargs.remove(0);
@@ -829,4 +836,47 @@ impl<'tcx> GotocCtx<'tcx> {
             loc,
         )
     }
+
+    /// simd_shuffle constructs a new vector from the elements of two input vectors,
+    /// choosing values according to an input array of indexes.
+    ///
+    /// This code mimics CBMC's `shuffle_vector_exprt::lower()` here:
+    /// https://github.com/diffblue/cbmc/blob/develop/src/ansi-c/c_expr.cpp
+    ///
+    /// We can't use shuffle_vector_exprt because it's not understood by the CBMC backend,
+    /// it's immediately lowered by the C frontend.
+    /// Issue: https://github.com/diffblue/cbmc/issues/6297
+    pub fn codegen_intrinsic_simd_shuffle(
+        &mut self,
+        mut fargs: Vec<Expr>,
+        p: &Place<'tcx>,
+        cbmc_ret_ty: Type,
+        n: u64,
+    ) -> Stmt {
+        assert!(fargs.len() == 3, "simd_shuffle had unexpected arguments {:?}", fargs);
+        // vector, size n: translated as vector types which cbmc treats as arrays
+        let vec1 = fargs.remove(0);
+        let vec2 = fargs.remove(0);
+        // [u32; n]: translated wrapped in a struct
+        let indexes = fargs.remove(0);
+
+        // An unsigned type here causes an invariant violation in CBMC.
+        // Issue: https://github.com/diffblue/cbmc/issues/6298
+        let st_rep = Type::ssize_t();
+        let n_rep = Expr::int_constant(n, st_rep.clone());
+
+        // P = indexes.expanded_map(v -> if v < N then vec1[v] else vec2[v-N])
+        let elems = (0..n)
+            .map(|i| {
+                let i = Expr::int_constant(i, st_rep.clone());
+                // Must not use `indexes.index(i)` directly, because codegen wraps arrays in struct
+                let v = self.codegen_idx_array(indexes.clone(), i).cast_to(st_rep.clone());
+                let cond = v.clone().lt(n_rep.clone());
+                let t = vec1.clone().index(v.clone());
+                let e = vec2.clone().index(v.sub(n_rep.clone()));
+                cond.ternary(t, e)
+            })
+            .collect();
+        self.codegen_expr_to_place(p, Expr::vector_expr(cbmc_ret_ty, elems))
+    }
 }

src/test/cbmc/SIMD/Shuffle/main.rs

@@ -0,0 +1,36 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+#![feature(repr_simd, platform_intrinsics)]
+
+#[repr(simd)]
+#[allow(non_camel_case_types)]
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub struct i64x2(i64, i64);
+
+#[repr(simd)]
+#[allow(non_camel_case_types)]
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub struct i64x4(i64, i64, i64, i64);
+
+extern "platform-intrinsic" {
+    fn simd_shuffle2<T, U>(x: T, y: T, idx: [u32; 2]) -> U;
+    fn simd_shuffle4<T, U>(x: T, y: T, idx: [u32; 4]) -> U;
+}
+
+fn main() {
+    {
+        let y = i64x2(0, 1);
+        let z = i64x2(1, 2);
+        const I: [u32; 2] = [1, 2];
+        let x: i64x2 = unsafe { simd_shuffle2(y, z, I) };
+        assert!(x.0 == 1);
+        assert!(x.1 == 1);
+    }
+    {
+        let a = i64x4(1, 2, 3, 4);
+        let b = i64x4(5, 6, 7, 8);
+        const I: [u32; 4] = [1, 3, 5, 7];
+        let c: i64x4 = unsafe { simd_shuffle4(a, b, I) };
+        assert!(c == i64x4(2, 4, 6, 8));
+    }
+}