Allow constants using an Abi::Vector layout to be passed to the backend

GeorgeWort · GeorgeWort · commit 7eb36f36b144 · 2023-11-22T12:05:56.000Z
This allows constant vectors using a repr(simd) type to be propagated
through to the backend. This fixes a few issues with LLVM simd intrinsics
where a constant vector is expected but cannot be produced by rust at
opt-level=0.
diff --git a/compiler/rustc_codegen_gcc/src/type_.rs b/compiler/rustc_codegen_gcc/src/type_.rs
@@ -10,22 +10,6 @@ use crate::context::CodegenCx;
 use crate::type_of::LayoutGccExt;
 
 impl<'gcc, 'tcx> CodegenCx<'gcc, 'tcx> {
-    pub fn type_ix(&self, num_bits: u64) -> Type<'gcc> {
-        // gcc only supports 1, 2, 4 or 8-byte integers.
-        // FIXME(antoyo): this is misleading to use the next power of two as rustc_codegen_ssa
-        // sometimes use 96-bit numbers and the following code will give an integer of a different
-        // size.
-        let bytes = (num_bits / 8).next_power_of_two() as i32;
-        match bytes {
-            1 => self.i8_type,
-            2 => self.i16_type,
-            4 => self.i32_type,
-            8 => self.i64_type,
-            16 => self.i128_type,
-            _ => panic!("unexpected num_bits: {}", num_bits),
-        }
-    }
-
     pub fn type_void(&self) -> Type<'gcc> {
         self.context.new_type::<()>()
     }
@@ -90,6 +74,22 @@ impl<'gcc, 'tcx> CodegenCx<'gcc, 'tcx> {
 }
 
 impl<'gcc, 'tcx> BaseTypeMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
+    fn type_ix(&self, num_bits: u64) -> Type<'gcc> {
+        // gcc only supports 1, 2, 4 or 8-byte integers.
+        // FIXME(antoyo): this is misleading to use the next power of two as rustc_codegen_ssa
+        // sometimes use 96-bit numbers and the following code will give an integer of a different
+        // size.
+        let bytes = (num_bits / 8).next_power_of_two() as i32;
+        match bytes {
+            1 => self.i8_type,
+            2 => self.i16_type,
+            4 => self.i32_type,
+            8 => self.i64_type,
+            16 => self.i128_type,
+            _ => panic!("unexpected num_bits: {}", num_bits),
+        }
+    }
+
     fn type_i1(&self) -> Type<'gcc> {
         self.bool_type
     }
diff --git a/compiler/rustc_codegen_llvm/src/type_.rs b/compiler/rustc_codegen_llvm/src/type_.rs
@@ -60,11 +60,6 @@ impl<'ll> CodegenCx<'ll, '_> {
         unsafe { llvm::LLVMMetadataTypeInContext(self.llcx) }
     }
 
-    ///x Creates an integer type with the given number of bits, e.g., i24
-    pub(crate) fn type_ix(&self, num_bits: u64) -> &'ll Type {
-        unsafe { llvm::LLVMIntTypeInContext(self.llcx, num_bits as c_uint) }
-    }
-
     pub(crate) fn type_vector(&self, ty: &'ll Type, len: u64) -> &'ll Type {
         unsafe { llvm::LLVMVectorType(ty, len as c_uint) }
     }
@@ -128,6 +123,11 @@ impl<'ll> CodegenCx<'ll, '_> {
 }
 
 impl<'ll, 'tcx> BaseTypeMethods<'tcx> for CodegenCx<'ll, 'tcx> {
+    /// Creates an integer type with the given number of bits, e.g., i24
+    fn type_ix(&self, num_bits: u64) -> &'ll Type {
+        unsafe { llvm::LLVMIntTypeInContext(self.llcx, num_bits as c_uint) }
+    }
+
     fn type_i1(&self) -> &'ll Type {
         unsafe { llvm::LLVMInt1TypeInContext(self.llcx) }
     }
diff --git a/compiler/rustc_codegen_ssa/src/mir/operand.rs b/compiler/rustc_codegen_ssa/src/mir/operand.rs
@@ -92,13 +92,29 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
         let layout = bx.layout_of(ty);
 
         let val = match val {
-            ConstValue::Scalar(x) => {
-                let Abi::Scalar(scalar) = layout.abi else {
-                    bug!("from_const: invalid ByVal layout: {:#?}", layout);
-                };
-                let llval = bx.scalar_to_backend(x, scalar, bx.immediate_backend_type(layout));
-                OperandValue::Immediate(llval)
-            }
+            ConstValue::Scalar(x) => match layout.abi {
+                Abi::Scalar(scalar) => {
+                    let llval = bx.scalar_to_backend(x, scalar, bx.immediate_backend_type(layout));
+                    OperandValue::Immediate(llval)
+                }
+                Abi::Vector {
+                    element: abi::Scalar::Initialized { value: abi::Primitive::Int(..), .. },
+                    ..
+                } if layout.size.bytes() <= 16 => {
+                    if let Scalar::Int(_) = x {
+                        let llval = bx.const_uint_big(
+                            bx.type_ix(layout.size.bits()),
+                            x.to_bits(layout.size).unwrap(),
+                        );
+                        OperandValue::Immediate(
+                            bx.const_bitcast(llval, bx.immediate_backend_type(layout)),
+                        )
+                    } else {
+                        bug!("Only Scalar::Int constant vectors are supported")
+                    }
+                }
+                _ => bug!("from_const: invalid ByVal layout: {:#?}", layout),
+            },
             ConstValue::ZeroSized => return OperandRef::zero_sized(layout),
             ConstValue::Slice { data, meta } => {
                 let Abi::ScalarPair(a_scalar, _) = layout.abi else {
diff --git a/compiler/rustc_codegen_ssa/src/traits/type_.rs b/compiler/rustc_codegen_ssa/src/traits/type_.rs
@@ -11,6 +11,7 @@ use rustc_target::abi::{AddressSpace, Integer};
 // This depends on `Backend` and not `BackendTypes`, because consumers will probably want to use
 // `LayoutOf` or `HasTyCtxt`. This way, they don't have to add a constraint on it themselves.
 pub trait BaseTypeMethods<'tcx>: Backend<'tcx> {
+    fn type_ix(&self, num_bits: u64) -> Self::Type;
     fn type_i1(&self) -> Self::Type;
     fn type_i8(&self) -> Self::Type;
     fn type_i16(&self) -> Self::Type;
diff --git a/compiler/rustc_const_eval/src/const_eval/eval_queries.rs b/compiler/rustc_const_eval/src/const_eval/eval_queries.rs
@@ -136,6 +136,10 @@ pub(super) fn op_to_const<'tcx>(
         // they can be stored as `ConstValue::Indirect`), but that's not relevant since we barely
         // ever have to do this. (`try_get_slice_bytes_for_diagnostics` exists to provide this
         // functionality.)
+        Abi::Vector {
+            element: abi::Scalar::Initialized { value: abi::Primitive::Int(..), .. },
+            ..
+        } => op.layout.size.bytes() <= 16,
         _ => false,
     };
     let immediate = if force_as_immediate {
diff --git a/compiler/rustc_const_eval/src/interpret/operand.rs b/compiler/rustc_const_eval/src/interpret/operand.rs
@@ -10,12 +10,12 @@ use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::ty::print::{FmtPrinter, PrettyPrinter};
 use rustc_middle::ty::{ConstInt, Ty, TyCtxt};
 use rustc_middle::{mir, ty};
-use rustc_target::abi::{self, Abi, HasDataLayout, Size};
+use rustc_target::abi::{self, Abi, HasDataLayout, Primitive, Size};
 
 use super::{
     alloc_range, from_known_layout, mir_assign_valid_types, AllocId, Frame, InterpCx, InterpResult,
-    MPlaceTy, Machine, MemPlace, MemPlaceMeta, OffsetMode, PlaceTy, Pointer, Projectable,
-    Provenance, Scalar,
+    MPlaceTy, Machine, MemPlace, MemPlaceMeta, MemoryKind, OffsetMode, PlaceTy, Pointer,
+    Projectable, Provenance, Scalar,
 };
 
 /// An `Immediate` represents a single immediate self-contained Rust value.
@@ -261,12 +261,29 @@ impl<'tcx, Prov: Provenance> ImmTy<'tcx, Prov> {
             {
                 Immediate::Uninit
             }
+            (
+                Immediate::Scalar(scalar),
+                Abi::Vector {
+                    element: abi::Scalar::Initialized { value: s @ Primitive::Int(..), .. },
+                    ..
+                },
+            ) => {
+                let size = s.size(cx);
+                assert_eq!(size, layout.size);
+                assert!(self.layout.size.bytes() <= 16);
+                let vector_bits: u128 = scalar.to_bits(self.layout.size).unwrap();
+                Immediate::Scalar(Scalar::Int(
+                    ty::ScalarInt::try_from_uint(size.truncate(vector_bits >> offset.bits()), size)
+                        .unwrap(),
+                ))
+            }
             // the field covers the entire type
             _ if layout.size == self.layout.size => {
                 assert_eq!(offset.bytes(), 0);
                 assert!(
                     match (self.layout.abi, layout.abi) {
                         (Abi::Scalar(..), Abi::Scalar(..)) => true,
+                        (Abi::Aggregate { sized: true }, Abi::Scalar(..)) => true,
                         (Abi::ScalarPair(..), Abi::ScalarPair(..)) => true,
                         _ => false,
                     },
@@ -500,6 +517,21 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 )?;
                 Some(ImmTy::from_immediate(Immediate::ScalarPair(a_val, b_val), mplace.layout))
             }
+            Abi::Vector {
+                element: abi::Scalar::Initialized { value: s @ abi::Primitive::Int(..), .. },
+                ..
+            } => {
+                let size = s.size(self);
+                let stride = size.align_to(s.align(self).abi);
+                if mplace.layout.size.bytes() <= 16 {
+                    assert!(stride.bytes() > 0);
+                    let vector_scalar =
+                        alloc.read_scalar(alloc_range(Size::ZERO, mplace.layout.size), false)?;
+                    Some(ImmTy { imm: Immediate::Scalar(vector_scalar), layout: mplace.layout })
+                } else {
+                    None
+                }
+            }
             _ => {
                 // Neither a scalar nor scalar pair.
                 None
@@ -539,12 +571,23 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         &self,
         op: &impl Readable<'tcx, M::Provenance>,
     ) -> InterpResult<'tcx, ImmTy<'tcx, M::Provenance>> {
-        if !matches!(
-            op.layout().abi,
+        let can_primitive_read = match op.layout().abi {
             Abi::Scalar(abi::Scalar::Initialized { .. })
-                | Abi::ScalarPair(abi::Scalar::Initialized { .. }, abi::Scalar::Initialized { .. })
-        ) {
-            span_bug!(self.cur_span(), "primitive read not possible for type: {}", op.layout().ty);
+            | Abi::ScalarPair(abi::Scalar::Initialized { .. }, abi::Scalar::Initialized { .. }) => {
+                true
+            }
+            Abi::Vector {
+                element: abi::Scalar::Initialized { value: abi::Primitive::Int(..), .. },
+                ..
+            } => op.layout().size.bytes() <= 16,
+            _ => false,
+        };
+        if !can_primitive_read {
+            span_bug!(
+                self.cur_span(),
+                "primitive read not possible for type: {:?}",
+                op.layout().ty
+            );
         }
         let imm = self.read_immediate_raw(op)?.right().unwrap();
         if matches!(*imm, Immediate::Uninit) {
@@ -599,7 +642,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     ///
     /// Can (but does not always) trigger UB if `op` is uninitialized.
     pub fn operand_to_simd(
-        &self,
+        &mut self,
         op: &OpTy<'tcx, M::Provenance>,
     ) -> InterpResult<'tcx, (MPlaceTy<'tcx, M::Provenance>, u64)> {
         // Basically we just transmute this place into an array following simd_size_and_type.
@@ -612,7 +655,15 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                     throw_ub!(InvalidUninitBytes(None))
                 }
                 Immediate::Scalar(..) | Immediate::ScalarPair(..) => {
-                    bug!("arrays/slices can never have Scalar/ScalarPair layout")
+                    if let Abi::Vector {
+                                        element: abi::Scalar::Initialized { value: abi::Primitive::Int(..), .. }, ..
+                                    } = op.layout.abi && op.layout.size.bytes() <= 16 {
+                        let mplace = self.allocate(op.layout, MemoryKind::Stack)?;
+                        self.write_immediate(imm.imm, &mplace).unwrap();
+                        self.mplace_to_simd(&mplace)
+                    } else {
+                        bug!("arrays/slices can never have Scalar/ScalarPair layout")
+                    }
                 }
             },
         }
diff --git a/compiler/rustc_const_eval/src/interpret/place.rs b/compiler/rustc_const_eval/src/interpret/place.rs
@@ -12,7 +12,8 @@ use rustc_middle::mir;
 use rustc_middle::ty;
 use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::ty::Ty;
-use rustc_target::abi::{Abi, Align, FieldIdx, HasDataLayout, Size, FIRST_VARIANT};
+use rustc_target::abi;
+use rustc_target::abi::{Abi, Align, FieldIdx, HasDataLayout, Primitive, Size, FIRST_VARIANT};
 
 use super::{
     alloc_range, mir_assign_valid_types, AllocId, AllocRef, AllocRefMut, CheckAlignMsg, ImmTy,
@@ -696,17 +697,32 @@ where
         };
 
         match value {
-            Immediate::Scalar(scalar) => {
-                let Abi::Scalar(s) = layout.abi else {
-                    span_bug!(
-                        self.cur_span(),
-                        "write_immediate_to_mplace: invalid Scalar layout: {layout:#?}",
-                    )
-                };
-                let size = s.size(&tcx);
-                assert_eq!(size, layout.size, "abi::Scalar size does not match layout size");
-                alloc.write_scalar(alloc_range(Size::ZERO, size), scalar)
-            }
+            Immediate::Scalar(scalar) => match layout.abi {
+                Abi::Scalar(s) => {
+                    let size = s.size(&tcx);
+                    assert_eq!(size, layout.size, "abi::Scalar size does not match layout size");
+                    alloc.write_scalar(alloc_range(Size::ZERO, size), scalar)
+                }
+                Abi::Vector {
+                    element: abi::Scalar::Initialized { value: s @ Primitive::Int(..), .. },
+                    ..
+                } => {
+                    let size = s.size(&tcx);
+                    let stride = size.align_to(s.align(&tcx).abi);
+                    assert!(stride.bytes() > 0);
+                    if layout.size.bytes() > 16 {
+                        span_bug!(
+                            self.cur_span(),
+                            "write_immediate_to_mplace: invalid Vector layout: {layout:#?}",
+                        )
+                    }
+                    alloc.write_scalar(alloc_range(Size::ZERO, layout.size), scalar)
+                }
+                _ => span_bug!(
+                    self.cur_span(),
+                    "write_immediate_to_mplace: invalid Scalar layout: {layout:#?}",
+                ),
+            },
             Immediate::ScalarPair(a_val, b_val) => {
                 let Abi::ScalarPair(a, b) = layout.abi else {
                     span_bug!(

Original file line number	Diff line number	Diff line change
`@@ -60,11 +60,6 @@ impl<'ll> CodegenCx<'ll, '_> {`
`60`	`60`	`unsafe { llvm::LLVMMetadataTypeInContext(self.llcx) }`
`61`	`61`	`}`
`62`	`62`
`63`		`- ///x Creates an integer type with the given number of bits, e.g., i24`
`64`		`- pub(crate) fn type_ix(&self, num_bits: u64) -> &'ll Type {`
`65`		`- unsafe { llvm::LLVMIntTypeInContext(self.llcx, num_bits as c_uint) }`
`66`		`- }`
`67`		`-`
`68`	`63`	`pub(crate) fn type_vector(&self, ty: &'ll Type, len: u64) -> &'ll Type {`
`69`	`64`	`unsafe { llvm::LLVMVectorType(ty, len as c_uint) }`
`70`	`65`	`}`
`@@ -128,6 +123,11 @@ impl<'ll> CodegenCx<'ll, '_> {`
`128`	`123`	`}`
`129`	`124`
`130`	`125`	`impl<'ll, 'tcx> BaseTypeMethods<'tcx> for CodegenCx<'ll, 'tcx> {`
	`126`	`+ /// Creates an integer type with the given number of bits, e.g., i24`
	`127`	`+ fn type_ix(&self, num_bits: u64) -> &'ll Type {`
	`128`	`+ unsafe { llvm::LLVMIntTypeInContext(self.llcx, num_bits as c_uint) }`
	`129`	`+ }`
	`130`	`+`
`131`	`131`	`fn type_i1(&self) -> &'ll Type {`
`132`	`132`	`unsafe { llvm::LLVMInt1TypeInContext(self.llcx) }`
`133`	`133`	`}`