ty_utils: lower tuples to ScalableVector repr

davidtwco · davidtwco · commit a2f7f3c1ebb7 · 2026-04-03T10:27:30.000Z
Instead of just using regular struct lowering for these types, which
results in an incorrect ABI (e.g. returning indirectly), use
`BackendRepr::ScalableVector` which will lower to the correct type and
be passed in registers.

This also enables some simplifications for generating alloca of scalable
vectors and greater re-use of `scalable_vector_parts`.

A LLVM codegen test demonstrating the changed IR this generates is
included in the next commit alongside some intrinsics that make these
tuples usable.
diff --git a/compiler/rustc_abi/src/layout.rs b/compiler/rustc_abi/src/layout.rs
@@ -10,8 +10,8 @@ use tracing::{debug, trace};
 
 use crate::{
     AbiAlign, Align, BackendRepr, FieldsShape, HasDataLayout, IndexSlice, IndexVec, Integer,
-    LayoutData, Niche, NonZeroUsize, Primitive, ReprOptions, Scalar, Size, StructKind, TagEncoding,
-    TargetDataLayout, Variants, WrappingRange,
+    LayoutData, Niche, NonZeroUsize, NumScalableVectors, Primitive, ReprOptions, Scalar, Size,
+    StructKind, TagEncoding, TargetDataLayout, Variants, WrappingRange,
 };
 
 mod coroutine;
@@ -204,13 +204,19 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         &self,
         element: F,
         count: u64,
+        number_of_vectors: NumScalableVectors,
     ) -> LayoutCalculatorResult<FieldIdx, VariantIdx, F>
     where
         FieldIdx: Idx,
         VariantIdx: Idx,
         F: AsRef<LayoutData<FieldIdx, VariantIdx>> + fmt::Debug,
     {
-        vector_type_layout(SimdVectorKind::Scalable, self.cx.data_layout(), element, count)
+        vector_type_layout(
+            SimdVectorKind::Scalable(number_of_vectors),
+            self.cx.data_layout(),
+            element,
+            count,
+        )
     }
 
     pub fn simd_type<FieldIdx, VariantIdx, F>(
@@ -1526,7 +1532,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
 
 enum SimdVectorKind {
     /// `#[rustc_scalable_vector]`
-    Scalable,
+    Scalable(NumScalableVectors),
     /// `#[repr(simd, packed)]`
     PackedFixed,
     /// `#[repr(simd)]`
@@ -1559,9 +1565,10 @@ where
     let size =
         elt.size.checked_mul(count, dl).ok_or_else(|| LayoutCalculatorError::SizeOverflow)?;
     let (repr, align) = match kind {
-        SimdVectorKind::Scalable => {
-            (BackendRepr::SimdScalableVector { element, count }, dl.llvmlike_vector_align(size))
-        }
+        SimdVectorKind::Scalable(number_of_vectors) => (
+            BackendRepr::SimdScalableVector { element, count, number_of_vectors },
+            dl.llvmlike_vector_align(size),
+        ),
         // Non-power-of-two vectors have padding up to the next power-of-two.
         // If we're a packed repr, remove the padding while keeping the alignment as close
         // to a vector as possible.
diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs
@@ -1696,6 +1696,28 @@ impl AddressSpace {
     pub const ZERO: Self = AddressSpace(0);
 }
 
+/// How many scalable vectors are in a `BackendRepr::ScalableVector`?
+#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
+#[cfg_attr(feature = "nightly", derive(HashStable_Generic))]
+pub struct NumScalableVectors(pub u8);
+
+impl NumScalableVectors {
+    /// Returns a `NumScalableVector` for a non-tuple scalable vector (e.g. a single vector).
+    pub fn for_non_tuple() -> Self {
+        NumScalableVectors(1)
+    }
+
+    // Returns `NumScalableVectors` for values of two through eight, which are a valid number of
+    // fields for a tuple of scalable vectors to have. `1` is a valid value of `NumScalableVectors`
+    // but not for a tuple which would have a field count.
+    pub fn from_field_count(count: usize) -> Option<Self> {
+        match count {
+            2..8 => Some(NumScalableVectors(count as u8)),
+            _ => None,
+        }
+    }
+}
+
 /// The way we represent values to the backend
 ///
 /// Previously this was conflated with the "ABI" a type is given, as in the platform-specific ABI.
@@ -1714,6 +1736,7 @@ pub enum BackendRepr {
     SimdScalableVector {
         element: Scalar,
         count: u64,
+        number_of_vectors: NumScalableVectors,
     },
     SimdVector {
         element: Scalar,
@@ -1820,8 +1843,12 @@ impl BackendRepr {
                 BackendRepr::SimdVector { element: element.to_union(), count }
             }
             BackendRepr::Memory { .. } => BackendRepr::Memory { sized: true },
-            BackendRepr::SimdScalableVector { element, count } => {
-                BackendRepr::SimdScalableVector { element: element.to_union(), count }
+            BackendRepr::SimdScalableVector { element, count, number_of_vectors } => {
+                BackendRepr::SimdScalableVector {
+                    element: element.to_union(),
+                    count,
+                    number_of_vectors,
+                }
             }
         }
     }
@@ -2161,7 +2188,7 @@ impl<FieldIdx: Idx, VariantIdx: Idx> LayoutData<FieldIdx, VariantIdx> {
     }
 
     /// Returns `true` if the size of the type is only known at runtime.
-    pub fn is_runtime_sized(&self) -> bool {
+    pub fn is_scalable_vector(&self) -> bool {
         matches!(self.backend_repr, BackendRepr::SimdScalableVector { .. })
     }
 
diff --git a/compiler/rustc_codegen_gcc/src/builder.rs b/compiler/rustc_codegen_gcc/src/builder.rs
@@ -24,7 +24,8 @@ use rustc_data_structures::fx::FxHashSet;
 use rustc_middle::bug;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
 use rustc_middle::ty::layout::{
-    FnAbiError, FnAbiOfHelpers, FnAbiRequest, HasTyCtxt, HasTypingEnv, LayoutError, LayoutOfHelpers,
+    FnAbiError, FnAbiOfHelpers, FnAbiRequest, HasTyCtxt, HasTypingEnv, LayoutError,
+    LayoutOfHelpers, TyAndLayout,
 };
 use rustc_middle::ty::{self, AtomicOrdering, Instance, Ty, TyCtxt};
 use rustc_span::Span;
@@ -943,8 +944,8 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
             .get_address(self.location)
     }
 
-    fn scalable_alloca(&mut self, _elt: u64, _align: Align, _element_ty: Ty<'_>) -> RValue<'gcc> {
-        todo!()
+    fn alloca_with_ty(&mut self, ty: TyAndLayout<'tcx>) -> RValue<'gcc> {
+        self.alloca(ty.layout.size, ty.layout.align.abi)
     }
 
     fn load(&mut self, pointee_ty: Type<'gcc>, ptr: RValue<'gcc>, align: Align) -> RValue<'gcc> {
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -7,8 +7,7 @@ pub(crate) mod autodiff;
 pub(crate) mod gpu_offload;
 
 use libc::{c_char, c_uint};
-use rustc_abi as abi;
-use rustc_abi::{Align, Size, WrappingRange};
+use rustc_abi::{self as abi, Align, Size, WrappingRange};
 use rustc_codegen_ssa::MemFlags;
 use rustc_codegen_ssa::common::{IntPredicate, RealPredicate, SynchronizationScope, TypeKind};
 use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
@@ -616,21 +615,14 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
     }
 
-    fn scalable_alloca(&mut self, elt: u64, align: Align, element_ty: Ty<'_>) -> Self::Value {
+    fn alloca_with_ty(&mut self, layout: TyAndLayout<'tcx>) -> Self::Value {
         let mut bx = Builder::with_cx(self.cx);
         bx.position_at_start(unsafe { llvm::LLVMGetFirstBasicBlock(self.llfn()) });
-        let llvm_ty = match element_ty.kind() {
-            ty::Bool => bx.type_i1(),
-            ty::Int(int_ty) => self.cx.type_int_from_ty(*int_ty),
-            ty::Uint(uint_ty) => self.cx.type_uint_from_ty(*uint_ty),
-            ty::Float(float_ty) => self.cx.type_float_from_ty(*float_ty),
-            _ => unreachable!("scalable vectors can only contain a bool, int, uint or float"),
-        };
+        let scalable_vector_ty = layout.llvm_type(self.cx);
 
         unsafe {
-            let ty = llvm::LLVMScalableVectorType(llvm_ty, elt.try_into().unwrap());
-            let alloca = llvm::LLVMBuildAlloca(&bx.llbuilder, ty, UNNAMED);
-            llvm::LLVMSetAlignment(alloca, align.bytes() as c_uint);
+            let alloca = llvm::LLVMBuildAlloca(&bx.llbuilder, scalable_vector_ty, UNNAMED);
+            llvm::LLVMSetAlignment(alloca, layout.align.abi.bytes() as c_uint);
             alloca
         }
     }
diff --git a/compiler/rustc_codegen_llvm/src/type_of.rs b/compiler/rustc_codegen_llvm/src/type_of.rs
@@ -24,14 +24,54 @@ fn uncached_llvm_type<'a, 'tcx>(
             let element = layout.scalar_llvm_type_at(cx, element);
             return cx.type_vector(element, count);
         }
-        BackendRepr::SimdScalableVector { ref element, count } => {
+        BackendRepr::SimdScalableVector { ref element, count, number_of_vectors } => {
             let element = if element.is_bool() {
                 cx.type_i1()
             } else {
                 layout.scalar_llvm_type_at(cx, *element)
             };
 
-            return cx.type_scalable_vector(element, count);
+            let vector_type = cx.type_scalable_vector(element, count);
+            return match number_of_vectors.0 {
+                1 => vector_type,
+                2 => cx.type_struct(&[vector_type, vector_type], false),
+                3 => cx.type_struct(&[vector_type, vector_type, vector_type], false),
+                4 => cx.type_struct(&[vector_type, vector_type, vector_type, vector_type], false),
+                5 => cx.type_struct(
+                    &[vector_type, vector_type, vector_type, vector_type, vector_type],
+                    false,
+                ),
+                6 => cx.type_struct(
+                    &[vector_type, vector_type, vector_type, vector_type, vector_type, vector_type],
+                    false,
+                ),
+                7 => cx.type_struct(
+                    &[
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                    ],
+                    false,
+                ),
+                8 => cx.type_struct(
+                    &[
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                        vector_type,
+                    ],
+                    false,
+                ),
+                _ => bug!("`#[rustc_scalable_vector]` tuple struct with too many fields"),
+            };
         }
         BackendRepr::Memory { .. } | BackendRepr::ScalarPair(..) => {}
     }
diff --git a/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs b/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
@@ -438,8 +438,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 if operand.layout.ty.is_scalable_vector()
                     && bx.sess().target.arch == rustc_target::spec::Arch::AArch64
                 {
-                    let (count, element_ty) =
-                        operand.layout.ty.scalable_vector_element_count_and_type(bx.tcx());
+                    let (count, element_ty, _) =
+                        operand.layout.ty.scalable_vector_parts(bx.tcx()).unwrap();
                     // i.e. `<vscale x N x i1>` when `N != 16`
                     if element_ty.is_bool() && count != 16 {
                         return;
diff --git a/compiler/rustc_codegen_ssa/src/mir/place.rs b/compiler/rustc_codegen_ssa/src/mir/place.rs
@@ -1,3 +1,5 @@
+use std::ops::Deref as _;
+
 use rustc_abi::{
     Align, BackendRepr, FieldIdx, FieldsShape, Size, TagEncoding, VariantIdx, Variants,
 };
@@ -109,8 +111,8 @@ impl<'a, 'tcx, V: CodegenObject> PlaceRef<'tcx, V> {
         bx: &mut Bx,
         layout: TyAndLayout<'tcx>,
     ) -> Self {
-        if layout.is_runtime_sized() {
-            Self::alloca_runtime_sized(bx, layout)
+        if layout.deref().is_scalable_vector() {
+            Self::alloca_scalable(bx, layout)
         } else {
             Self::alloca_size(bx, layout.size, layout)
         }
@@ -151,16 +153,11 @@ impl<'a, 'tcx, V: CodegenObject> PlaceRef<'tcx, V> {
         }
     }
 
-    fn alloca_runtime_sized<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
+    fn alloca_scalable<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
         bx: &mut Bx,
         layout: TyAndLayout<'tcx>,
     ) -> Self {
-        let (element_count, ty) = layout.ty.scalable_vector_element_count_and_type(bx.tcx());
-        PlaceValue::new_sized(
-            bx.scalable_alloca(element_count as u64, layout.align.abi, ty),
-            layout.align.abi,
-        )
-        .with_type(layout)
+        PlaceValue::new_sized(bx.alloca_with_ty(layout), layout.align.abi).with_type(layout)
     }
 }
 
diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
@@ -235,7 +235,7 @@ pub trait BuilderMethods<'a, 'tcx>:
     fn to_immediate_scalar(&mut self, val: Self::Value, scalar: Scalar) -> Self::Value;
 
     fn alloca(&mut self, size: Size, align: Align) -> Self::Value;
-    fn scalable_alloca(&mut self, elt: u64, align: Align, element_ty: Ty<'_>) -> Self::Value;
+    fn alloca_with_ty(&mut self, layout: TyAndLayout<'tcx>) -> Self::Value;
 
     fn load(&mut self, ty: Self::Type, ptr: Self::Value, align: Align) -> Self::Value;
     fn volatile_load(&mut self, ty: Self::Type, ptr: Self::Value) -> Self::Value;
diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs
@@ -7,7 +7,7 @@ use std::debug_assert_matches;
 use std::ops::{ControlFlow, Range};
 
 use hir::def::{CtorKind, DefKind};
-use rustc_abi::{FIRST_VARIANT, FieldIdx, ScalableElt, VariantIdx};
+use rustc_abi::{FIRST_VARIANT, FieldIdx, NumScalableVectors, ScalableElt, VariantIdx};
 use rustc_errors::{ErrorGuaranteed, MultiSpan};
 use rustc_hir as hir;
 use rustc_hir::LangItem;
@@ -1261,17 +1261,27 @@ impl<'tcx> Ty<'tcx> {
         }
     }
 
-    pub fn scalable_vector_element_count_and_type(self, tcx: TyCtxt<'tcx>) -> (u16, Ty<'tcx>) {
+    pub fn scalable_vector_parts(
+        self,
+        tcx: TyCtxt<'tcx>,
+    ) -> Option<(u16, Ty<'tcx>, NumScalableVectors)> {
         let Adt(def, args) = self.kind() else {
-            bug!("`scalable_vector_size_and_type` called on invalid type")
+            return None;
         };
-        let Some(ScalableElt::ElementCount(element_count)) = def.repr().scalable else {
-            bug!("`scalable_vector_size_and_type` called on non-scalable vector type");
+        let (num_vectors, vec_def) = match def.repr().scalable? {
+            ScalableElt::ElementCount(_) => (NumScalableVectors::for_non_tuple(), *def),
+            ScalableElt::Container => (
+                NumScalableVectors::from_field_count(def.non_enum_variant().fields.len())?,
+                def.non_enum_variant().fields[FieldIdx::ZERO].ty(tcx, args).ty_adt_def()?,
+            ),
         };
-        let variant = def.non_enum_variant();
+        let Some(ScalableElt::ElementCount(element_count)) = vec_def.repr().scalable else {
+            return None;
+        };
+        let variant = vec_def.non_enum_variant();
         assert_eq!(variant.fields.len(), 1);
         let field_ty = variant.fields[FieldIdx::ZERO].ty(tcx, args);
-        (element_count, field_ty)
+        Some((element_count, field_ty, num_vectors))
     }
 
     pub fn simd_size_and_type(self, tcx: TyCtxt<'tcx>) -> (u64, Ty<'tcx>) {
diff --git a/compiler/rustc_public/src/abi.rs b/compiler/rustc_public/src/abi.rs
@@ -232,6 +232,10 @@ pub enum TagEncoding {
     },
 }
 
+/// How many scalable vectors are in a `ValueAbi::ScalableVector`?
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)]
+pub struct NumScalableVectors(pub(crate) u8);
+
 /// Describes how values of the type are passed by target ABIs,
 /// in terms of categories of C types there are ABI rules for.
 #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)]
@@ -245,6 +249,7 @@ pub enum ValueAbi {
     ScalableVector {
         element: Scalar,
         count: u64,
+        number_of_vectors: NumScalableVectors,
     },
     Aggregate {
         /// If true, the size is exact, otherwise it's only a lower bound.
diff --git a/compiler/rustc_public/src/unstable/convert/stable/abi.rs b/compiler/rustc_public/src/unstable/convert/stable/abi.rs
@@ -10,8 +10,9 @@ use rustc_target::callconv;
 
 use crate::abi::{
     AddressSpace, ArgAbi, CallConvention, FieldsShape, FloatLength, FnAbi, IntegerLength,
-    IntegerType, Layout, LayoutShape, PassMode, Primitive, ReprFlags, ReprOptions, Scalar,
-    TagEncoding, TyAndLayout, ValueAbi, VariantFields, VariantsShape, WrappingRange,
+    IntegerType, Layout, LayoutShape, NumScalableVectors, PassMode, Primitive, ReprFlags,
+    ReprOptions, Scalar, TagEncoding, TyAndLayout, ValueAbi, VariantFields, VariantsShape,
+    WrappingRange,
 };
 use crate::compiler_interface::BridgeTys;
 use crate::target::MachineSize as Size;
@@ -249,6 +250,18 @@ impl<'tcx> Stable<'tcx> for rustc_abi::TagEncoding<rustc_abi::VariantIdx> {
     }
 }
 
+impl<'tcx> Stable<'tcx> for rustc_abi::NumScalableVectors {
+    type T = NumScalableVectors;
+
+    fn stable<'cx>(
+        &self,
+        _tables: &mut Tables<'cx, BridgeTys>,
+        _cx: &CompilerCtxt<'cx, BridgeTys>,
+    ) -> Self::T {
+        NumScalableVectors(self.0)
+    }
+}
+
 impl<'tcx> Stable<'tcx> for rustc_abi::BackendRepr {
     type T = ValueAbi;
 
@@ -265,8 +278,12 @@ impl<'tcx> Stable<'tcx> for rustc_abi::BackendRepr {
             rustc_abi::BackendRepr::SimdVector { element, count } => {
                 ValueAbi::Vector { element: element.stable(tables, cx), count }
             }
-            rustc_abi::BackendRepr::SimdScalableVector { element, count } => {
-                ValueAbi::ScalableVector { element: element.stable(tables, cx), count }
+            rustc_abi::BackendRepr::SimdScalableVector { element, count, number_of_vectors } => {
+                ValueAbi::ScalableVector {
+                    element: element.stable(tables, cx),
+                    count,
+                    number_of_vectors: number_of_vectors.stable(tables, cx),
+                }
             }
             rustc_abi::BackendRepr::Memory { sized } => ValueAbi::Aggregate { sized },
         }
diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
