Add support to array-based SIMD (#2633)

Originally, repr(simd) supported only multi-field form. An array based version was later added and it's likely to become the only supported way (rust-lang/compiler-team#621).

The array-based version is currently used in the standard library, and it is used to implement `portable-simd`. This change adds support to instantiating and using the array-based version.

Author: celinval

kani-compiler/src/codegen_cprover_gotoc/codegen/place.rs

@@ -23,7 +23,7 @@ use tracing::{debug, trace, warn};
 
 /// A projection in Kani can either be to a type (the normal case),
 /// or a variant in the case of a downcast.
-#[derive(Debug)]
+#[derive(Copy, Clone, Debug)]
 pub enum TypeOrVariant<'tcx> {
     Type(Ty<'tcx>),
     Variant(&'tcx VariantDef),
@@ -235,15 +235,21 @@ impl<'tcx> TypeOrVariant<'tcx> {
 }
 
 impl<'tcx> GotocCtx<'tcx> {
+    /// Codegen field access for types that allow direct field projection.
+    ///
+    /// I.e.: Algebraic data types, closures, and generators.
+    ///
+    /// Other composite types such as array only support index projection.
     fn codegen_field(
         &mut self,
-        res: Expr,
-        t: TypeOrVariant<'tcx>,
-        f: &FieldIdx,
+        parent_expr: Expr,
+        parent_ty_or_var: TypeOrVariant<'tcx>,
+        field: &FieldIdx,
+        field_ty_or_var: TypeOrVariant<'tcx>,
     ) -> Result<Expr, UnimplementedData> {
-        match t {
-            TypeOrVariant::Type(t) => {
-                match t.kind() {
+        match parent_ty_or_var {
+            TypeOrVariant::Type(parent_ty) => {
+                match parent_ty.kind() {
                     ty::Alias(..)
                     | ty::Bool
                     | ty::Char
@@ -254,56 +260,98 @@ impl<'tcx> GotocCtx<'tcx> {
                     | ty::Never
                     | ty::FnDef(..)
                     | ty::GeneratorWitness(..)
+                    | ty::GeneratorWitnessMIR(..)
                     | ty::Foreign(..)
                     | ty::Dynamic(..)
                     | ty::Bound(..)
                     | ty::Placeholder(..)
                     | ty::Param(_)
                     | ty::Infer(_)
-                    | ty::Error(_) => unreachable!("type {:?} does not have a field", t),
-                    ty::Tuple(_) => {
-                        Ok(res.member(&Self::tuple_fld_name(f.index()), &self.symbol_table))
-                    }
-                    ty::Adt(def, _) if def.repr().simd() => {
-                        // this is a SIMD vector - the index represents one
-                        // of the elements, so we want to index as an array
-                        // Example:
-                        // pub struct i64x2(i64, i64);
-                        // fn main() {
-                        //   let v = i64x2(1, 2);
-                        //   assert!(v.0 == 1); // refers to the first i64
-                        //   assert!(v.1 == 2);
-                        // }
-                        let size_index = Expr::int_constant(f.index(), Type::size_t());
-                        Ok(res.index_array(size_index))
-                    }
+                    | ty::Error(_) => unreachable!("type {parent_ty:?} does not have a field"),
+                    ty::Tuple(_) => Ok(parent_expr
+                        .member(&Self::tuple_fld_name(field.index()), &self.symbol_table)),
+                    ty::Adt(def, _) if def.repr().simd() => Ok(self.codegen_simd_field(
+                        parent_expr,
+                        *field,
+                        field_ty_or_var.expect_type(),
+                    )),
                     // if we fall here, then we are handling either a struct or a union
                     ty::Adt(def, _) => {
-                        let field = &def.variants().raw[0].fields[*f];
-                        Ok(res.member(&field.name.to_string(), &self.symbol_table))
+                        let field = &def.variants().raw[0].fields[*field];
+                        Ok(parent_expr.member(&field.name.to_string(), &self.symbol_table))
+                    }
+                    ty::Closure(..) => {
+                        Ok(parent_expr.member(&field.index().to_string(), &self.symbol_table))
                     }
-                    ty::Closure(..) => Ok(res.member(&f.index().to_string(), &self.symbol_table)),
                     ty::Generator(..) => {
-                        let field_name = self.generator_field_name(f.as_usize());
-                        Ok(res
+                        let field_name = self.generator_field_name(field.as_usize());
+                        Ok(parent_expr
                             .member("direct_fields", &self.symbol_table)
                             .member(field_name, &self.symbol_table))
                     }
-                    _ => unimplemented!(),
+                    ty::Str | ty::Array(_, _) | ty::Slice(_) | ty::RawPtr(_) | ty::Ref(_, _, _) => {
+                        unreachable!(
+                            "element of {parent_ty:?} is not accessed via field projection"
+                        )
+                    }
                 }
             }
             // if we fall here, then we are handling an enum
-            TypeOrVariant::Variant(v) => {
-                let field = &v.fields[*f];
-                Ok(res.member(&field.name.to_string(), &self.symbol_table))
+            TypeOrVariant::Variant(parent_var) => {
+                let field = &parent_var.fields[*field];
+                Ok(parent_expr.member(&field.name.to_string(), &self.symbol_table))
             }
             TypeOrVariant::GeneratorVariant(_var_idx) => {
-                let field_name = self.generator_field_name(f.index());
-                Ok(res.member(field_name, &self.symbol_table))
+                let field_name = self.generator_field_name(field.index());
+                Ok(parent_expr.member(field_name, &self.symbol_table))
             }
         }
     }
 
+    /// This is a SIMD vector, which has 2 possible internal representations:
+    /// 1- Multi-field representation (original and currently deprecated)
+    ///    In this case, a field is one lane (i.e.: one element)
+    ///    Example:
+    /// ```ignore
+    ///    pub struct i64x2(i64, i64);
+    ///    fn main() {
+    ///      let v = i64x2(1, 2);
+    ///      assert!(v.0 == 1); // refers to the first i64
+    ///      assert!(v.1 == 2);
+    ///    }
+    /// ```
+    /// 2- Array-based representation
+    ///    In this case, the projection refers to the entire array.
+    /// ```ignore
+    ///    pub struct i64x2([i64; 2]);
+    ///    fn main() {
+    ///      let v = i64x2([1, 2]);
+    ///      assert!(v.0 == [1, 2]); // refers to the entire array
+    ///    }
+    /// ```
+    /// * Note that projection inside SIMD structs may eventually become illegal.
+    /// See <https://github.com/rust-lang/stdarch/pull/1422#discussion_r1176415609> thread.
+    ///
+    /// Since the goto representation for both is the same, we use the expected type to decide
+    /// what to return.
+    fn codegen_simd_field(
+        &mut self,
+        parent_expr: Expr,
+        field: FieldIdx,
+        field_ty: Ty<'tcx>,
+    ) -> Expr {
+        if matches!(field_ty.kind(), ty::Array { .. }) {
+            // Array based
+            assert_eq!(field.index(), 0);
+            let field_typ = self.codegen_ty(field_ty);
+            parent_expr.reinterpret_cast(field_typ)
+        } else {
+            // Return the given field.
+            let index_expr = Expr::int_constant(field.index(), Type::size_t());
+            parent_expr.index_array(index_expr)
+        }
+    }
+
     /// If a local is a function definition, ignore the local variable name and
     /// generate a function call based on the def id.
     ///
@@ -424,7 +472,8 @@ impl<'tcx> GotocCtx<'tcx> {
             }
             ProjectionElem::Field(f, t) => {
                 let typ = TypeOrVariant::Type(t);
-                let expr = self.codegen_field(before.goto_expr, before.mir_typ_or_variant, &f)?;
+                let expr =
+                    self.codegen_field(before.goto_expr, before.mir_typ_or_variant, &f, typ)?;
                 ProjectedPlace::try_new(
                     expr,
                     typ,

kani-compiler/src/codegen_cprover_gotoc/codegen/rvalue.rs

@@ -582,24 +582,28 @@ impl<'tcx> GotocCtx<'tcx> {
             AggregateKind::Adt(_, _, _, _, _) if res_ty.is_simd() => {
                 let typ = self.codegen_ty(res_ty);
                 let layout = self.layout_of(res_ty);
-                let vector_element_type = typ.base_type().unwrap().clone();
-                Expr::vector_expr(
-                    typ,
-                    layout
-                        .fields
-                        .index_by_increasing_offset()
-                        .map(|idx| {
-                            let cgo = self.codegen_operand(&operands[idx.into()]);
-                            // The input operand might actually be a one-element array, as seen
-                            // when running assess on firecracker.
-                            if *cgo.typ() == vector_element_type {
-                                cgo
-                            } else {
-                                cgo.transmute_to(vector_element_type.clone(), &self.symbol_table)
-                            }
-                        })
-                        .collect(),
-                )
+                trace!(shape=?layout.fields, "codegen_rvalue_aggregate");
+                assert!(operands.len() > 0, "SIMD vector cannot be empty");
+                if operands.len() == 1 {
+                    let data = self.codegen_operand(&operands[0u32.into()]);
+                    if data.typ().is_array() {
+                        // Array-based SIMD representation.
+                        data.transmute_to(typ, &self.symbol_table)
+                    } else {
+                        // Multi field-based representation with one field.
+                        Expr::vector_expr(typ, vec![data])
+                    }
+                } else {
+                    // Multi field SIMD representation.
+                    Expr::vector_expr(
+                        typ,
+                        layout
+                            .fields
+                            .index_by_increasing_offset()
+                            .map(|idx| self.codegen_operand(&operands[idx.into()]))
+                            .collect(),
+                    )
+                }
             }
             AggregateKind::Adt(_, variant_index, ..) if res_ty.is_enum() => {
                 self.codegen_rvalue_enum_aggregate(variant_index, operands, res_ty, loc)

tests/kani/SIMD/array_simd_repr.rs

@@ -0,0 +1,36 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//! Verify that Kani can properly handle SIMD declaration and field access using array syntax.
+
+#![allow(non_camel_case_types)]
+#![feature(repr_simd)]
+
+#[repr(simd)]
+#[derive(Clone, PartialEq, Eq, PartialOrd, kani::Arbitrary)]
+pub struct i64x2([i64; 2]);
+
+#[kani::proof]
+fn check_diff() {
+    let x = i64x2([1, 2]);
+    let y = i64x2([3, 4]);
+    assert!(x != y);
+}
+
+#[kani::proof]
+fn check_ge() {
+    let x: i64x2 = kani::any();
+    kani::assume(x.0[0] > 0);
+    kani::assume(x.0[1] > 0);
+    assert!(x > i64x2([0, 0]));
+}
+
+#[derive(Clone, Debug)]
+#[repr(simd)]
+struct CustomSimd<T, const LANES: usize>([T; LANES]);
+
+#[kani::proof]
+fn simd_vec() {
+    let simd = CustomSimd([0u8; 10]);
+    let idx: usize = kani::any_where(|x: &usize| *x < 10);
+    assert_eq!(simd.0[idx], 0);
+}

tests/kani/SIMD/multi_field_simd.rs

@@ -0,0 +1,28 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//! Verify that Kani can properly handle SIMD declaration and field access using multi-field syntax.
+//! Note: Multi-field SIMD is actually being deprecated, but until it's removed, we might
+//! as well keep supporting it.
+//! See <https://github.com/rust-lang/compiler-team/issues/621> for more details.
+
+#![allow(non_camel_case_types)]
+#![feature(repr_simd)]
+
+#[repr(simd)]
+#[derive(PartialEq, Eq, PartialOrd, kani::Arbitrary)]
+pub struct i64x2(i64, i64);
+
+#[kani::proof]
+fn check_diff() {
+    let x = i64x2(1, 2);
+    let y = i64x2(3, 4);
+    assert!(x != y);
+}
+
+#[kani::proof]
+fn check_ge() {
+    let x: i64x2 = kani::any();
+    kani::assume(x.0 > 0);
+    kani::assume(x.1 > 0);
+    assert!(x > i64x2(0, 0));
+}

tests/kani/SIMD/portable_simd.rs

@@ -0,0 +1,15 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+//! Ensure we have basic support of portable SIMD.
+#![feature(portable_simd)]
+
+use std::simd::u64x16;
+
+#[kani::proof]
+fn check_sum_any() {
+    let a = u64x16::splat(0);
+    let b = u64x16::from_array(kani::any());
+    // Cannot compare them directly: https://github.com/model-checking/kani/issues/2632
+    assert_eq!((a + b).as_array(), b.as_array());
+}

tests/kani/SIMD/simd_float_ops_fixme.rs

@@ -0,0 +1,39 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+//! Ensure we can handle SIMD defined in the standard library
+//! FIXME: <https://github.com/model-checking/kani/issues/2631>
+#![allow(non_camel_case_types)]
+#![feature(repr_simd, platform_intrinsics, portable_simd)]
+use std::simd::f32x4;
+
+extern "platform-intrinsic" {
+    fn simd_add<T>(x: T, y: T) -> T;
+    fn simd_eq<T, U>(x: T, y: T) -> U;
+}
+
+#[repr(simd)]
+#[derive(Clone, PartialEq, kani::Arbitrary)]
+pub struct f32x2(f32, f32);
+
+impl f32x2 {
+    fn as_array(&self) -> &[f32; 2] {
+        unsafe { &*(self as *const f32x2 as *const [f32; 2]) }
+    }
+}
+
+#[kani::proof]
+fn check_sum() {
+    let a = f32x2(0.0, 0.0);
+    let b = kani::any::<f32x2>();
+    let sum = unsafe { simd_add(a.clone(), b) };
+    assert_eq!(sum.as_array(), a.as_array());
+}
+
+#[kani::proof]
+fn check_sum_portable() {
+    let a = f32x4::splat(0.0);
+    let b = f32x4::from_array(kani::any());
+    // Cannot compare them directly: https://github.com/model-checking/kani/issues/2632
+    assert_eq!((a + b).as_array(), b.as_array());
+}