Auto merge of #143502 - scottmcm:aggregate-simd, r=<try>

Let `rvalue_creates_operand` return true for *all* `Rvalue::Aggregate`s

Draft for now because it's built on Ralf's #143291

Inspired by #138759 (comment) where I noticed that we were nearly at this point, plus the comments I was writing in #143410 that reminded me a type-dependent `true` is fine.

This PR splits the `OperandRef::builder` logic out to a separate type, with the updates needed to handle SIMD as well.  In doing so, that makes the existing `Aggregate` path in `codegen_rvalue_operand` capable of handing SIMD values just fine.

As a result, we no longer need to do layout calculations for aggregate result types when running the analysis to determine which things can be SSA in codegen.

Author: bors

compiler/rustc_codegen_ssa/src/mir/analyze.rs

@@ -171,8 +171,7 @@ impl<'a, 'b, 'tcx, Bx: BuilderMethods<'b, 'tcx>> Visitor<'tcx> for LocalAnalyzer
         if let Some(local) = place.as_local() {
             self.define(local, DefLocation::Assignment(location));
             if self.locals[local] != LocalKind::Memory {
-                let decl_span = self.fx.mir.local_decls[local].source_info.span;
-                if !self.fx.rvalue_creates_operand(rvalue, decl_span) {
+                if !self.fx.rvalue_creates_operand(rvalue) {
                     self.locals[local] = LocalKind::Memory;
                 }
             }

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -1,5 +1,6 @@
 use std::fmt;
 
+use itertools::Either;
 use rustc_abi as abi;
 use rustc_abi::{
     Align, BackendRepr, FIRST_VARIANT, FieldIdx, Primitive, Size, TagEncoding, VariantIdx, Variants,
@@ -564,86 +565,127 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
             }
         }
     }
+}
 
-    /// Creates an incomplete operand containing the [`abi::Scalar`]s expected based
-    /// on the `layout` passed. This is for use with [`OperandRef::insert_field`]
-    /// later to set the necessary immediate(s).
-    ///
-    /// Returns `None` for `layout`s which cannot be built this way.
-    pub(crate) fn builder(
-        layout: TyAndLayout<'tcx>,
-    ) -> Option<OperandRef<'tcx, Result<V, abi::Scalar>>> {
-        // Uninhabited types are weird, because for example `Result<!, !>`
-        // shows up as `FieldsShape::Primitive` and we need to be able to write
-        // a field into `(u32, !)`. We'll do that in an `alloca` instead.
-        if layout.uninhabited {
-            return None;
-        }
+/// Each of these variants starts out as `Either::Right` when it's uninitialized,
+/// then setting the field changes that to `Either::Left` with the backend value.
+#[derive(Debug, Copy, Clone)]
+enum OperandValueBuilder<V> {
+    ZeroSized,
+    Immediate(Either<V, abi::Scalar>),
+    Pair(Either<V, abi::Scalar>, Either<V, abi::Scalar>),
+    /// SIMD vectors need special handling because they're the only case where
+    /// a type with a (non-ZST) `Memory`-ABI field can be `Scalar`-ABI.
+    /// (Also you can't represent a vector type with just an `abi::Scalar`.)
+    Vector(Either<V, ()>),
+}
+
+/// Allows building up an `OperandRef` by setting fields one at a time.
+#[derive(Debug, Copy, Clone)]
+pub(super) struct OperandRefBuilder<'tcx, V> {
+    val: OperandValueBuilder<V>,
+    layout: TyAndLayout<'tcx>,
+}
 
+impl<'a, 'tcx, V: CodegenObject> OperandRefBuilder<'tcx, V> {
+    /// Creates an uninitialized builder for an instance of the `layout`.
+    ///
+    /// ICEs for [`BackendRepr::Memory`] types (other than ZSTs), which should
+    /// be built up inside a [`PlaceRef`] instead as they need an allocated place
+    /// into which to write the values of the fields.
+    pub(super) fn new(layout: TyAndLayout<'tcx>) -> Self {
         let val = match layout.backend_repr {
-            BackendRepr::Memory { .. } if layout.is_zst() => OperandValue::ZeroSized,
-            BackendRepr::Scalar(s) => OperandValue::Immediate(Err(s)),
-            BackendRepr::ScalarPair(a, b) => OperandValue::Pair(Err(a), Err(b)),
-            BackendRepr::Memory { .. } | BackendRepr::SimdVector { .. } => return None,
+            BackendRepr::Memory { .. } if layout.is_zst() => OperandValueBuilder::ZeroSized,
+            BackendRepr::Scalar(s) => OperandValueBuilder::Immediate(Either::Right(s)),
+            BackendRepr::ScalarPair(a, b) => {
+                OperandValueBuilder::Pair(Either::Right(a), Either::Right(b))
+            }
+            BackendRepr::SimdVector { .. } => OperandValueBuilder::Vector(Either::Right(())),
+            BackendRepr::Memory { .. } => {
+                bug!("Cannot use non-ZST Memory-ABI type in operand builder: {layout:?}");
+            }
         };
-        Some(OperandRef { val, layout })
+        OperandRefBuilder { val, layout }
     }
-}
 
-impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, Result<V, abi::Scalar>> {
-    pub(crate) fn insert_field<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
+    pub(super) fn insert_field<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
         &mut self,
         bx: &mut Bx,
-        v: VariantIdx,
-        f: FieldIdx,
+        variant: VariantIdx,
+        field: FieldIdx,
         operand: OperandRef<'tcx, V>,
     ) {
-        let (expect_zst, is_zero_offset) = if let abi::FieldsShape::Primitive = self.layout.fields {
+        if let OperandValue::ZeroSized = operand.val {
+            // A ZST never adds any state, so just ignore it.
+            // This special-casing is worth it because of things like
+            // `Result<!, !>` where `Ok(never)` is legal to write,
+            // but the type shows as FieldShape::Primitive so we can't
+            // actually look at the layout for the field being set.
+            return;
+        }
+
+        let is_zero_offset = if let abi::FieldsShape::Primitive = self.layout.fields {
             // The other branch looking at field layouts ICEs for primitives,
             // so we need to handle them separately.
-            // Multiple fields is possible for cases such as aggregating
-            // a thin pointer, where the second field is the unit.
+            // Because we handled ZSTs above (like the metadata in a thin pointer),
+            // the only possibility is that we're setting the one-and-only field.
             assert!(!self.layout.is_zst());
-            assert_eq!(v, FIRST_VARIANT);
-            let first_field = f == FieldIdx::ZERO;
-            (!first_field, first_field)
+            assert_eq!(variant, FIRST_VARIANT);
+            assert_eq!(field, FieldIdx::ZERO);
+            true
         } else {
-            let variant_layout = self.layout.for_variant(bx.cx(), v);
-            let field_layout = variant_layout.field(bx.cx(), f.as_usize());
-            let field_offset = variant_layout.fields.offset(f.as_usize());
-            (field_layout.is_zst(), field_offset == Size::ZERO)
+            let variant_layout = self.layout.for_variant(bx.cx(), variant);
+            let field_offset = variant_layout.fields.offset(field.as_usize());
+            field_offset == Size::ZERO
         };
 
-        let mut update = |tgt: &mut Result<V, abi::Scalar>, src, from_scalar| {
-            let to_scalar = tgt.unwrap_err();
+        let mut update = |tgt: &mut Either<V, abi::Scalar>, src, from_scalar| {
+            let to_scalar = tgt.unwrap_right();
+            // We transmute here (rather than just `from_immediate`) because in
+            // `Result<usize, *const ()>` the field of the `Ok` is an integer,
+            // but the corresponding scalar in the enum is a pointer.
             let imm = transmute_scalar(bx, src, from_scalar, to_scalar);
-            *tgt = Ok(imm);
+            *tgt = Either::Left(imm);
         };
 
         match (operand.val, operand.layout.backend_repr) {
-            (OperandValue::ZeroSized, _) if expect_zst => {}
+            (OperandValue::ZeroSized, _) => unreachable!("Handled above"),
             (OperandValue::Immediate(v), BackendRepr::Scalar(from_scalar)) => match &mut self.val {
-                OperandValue::Immediate(val @ Err(_)) if is_zero_offset => {
+                OperandValueBuilder::Immediate(val @ Either::Right(_)) if is_zero_offset => {
                     update(val, v, from_scalar);
                 }
-                OperandValue::Pair(fst @ Err(_), _) if is_zero_offset => {
+                OperandValueBuilder::Pair(fst @ Either::Right(_), _) if is_zero_offset => {
                     update(fst, v, from_scalar);
                 }
-                OperandValue::Pair(_, snd @ Err(_)) if !is_zero_offset => {
+                OperandValueBuilder::Pair(_, snd @ Either::Right(_)) if !is_zero_offset => {
                     update(snd, v, from_scalar);
                 }
-                _ => bug!("Tried to insert {operand:?} into {v:?}.{f:?} of {self:?}"),
+                _ => bug!("Tried to insert {operand:?} into {variant:?}.{field:?} of {self:?}"),
+            },
+            (OperandValue::Immediate(v), BackendRepr::SimdVector { .. }) => match &mut self.val {
+                OperandValueBuilder::Vector(val @ Either::Right(())) if is_zero_offset => {
+                    *val = Either::Left(v);
+                }
+                _ => bug!("Tried to insert {operand:?} into {variant:?}.{field:?} of {self:?}"),
             },
             (OperandValue::Pair(a, b), BackendRepr::ScalarPair(from_sa, from_sb)) => {
                 match &mut self.val {
-                    OperandValue::Pair(fst @ Err(_), snd @ Err(_)) => {
+                    OperandValueBuilder::Pair(fst @ Either::Right(_), snd @ Either::Right(_)) => {
                         update(fst, a, from_sa);
                         update(snd, b, from_sb);
                     }
-                    _ => bug!("Tried to insert {operand:?} into {v:?}.{f:?} of {self:?}"),
+                    _ => bug!("Tried to insert {operand:?} into {variant:?}.{field:?} of {self:?}"),
                 }
             }
-            _ => bug!("Unsupported operand {operand:?} inserting into {v:?}.{f:?} of {self:?}"),
+            (OperandValue::Ref(place), BackendRepr::Memory { .. }) => match &mut self.val {
+                OperandValueBuilder::Vector(val @ Either::Right(())) => {
+                    let ibty = bx.cx().immediate_backend_type(self.layout);
+                    let simd = bx.load_from_place(ibty, place);
+                    *val = Either::Left(simd);
+                }
+                _ => bug!("Tried to insert {operand:?} into {variant:?}.{field:?} of {self:?}"),
+            },
+            _ => bug!("Operand cannot be used with `insert_field`: {operand:?}"),
         }
     }
 
@@ -656,45 +698,50 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, Result<V, abi::Scalar>> {
         let field_offset = self.layout.fields.offset(f.as_usize());
         let is_zero_offset = field_offset == Size::ZERO;
         match &mut self.val {
-            OperandValue::Immediate(val @ Err(_)) if is_zero_offset => {
-                *val = Ok(imm);
+            OperandValueBuilder::Immediate(val @ Either::Right(_)) if is_zero_offset => {
+                *val = Either::Left(imm);
             }
-            OperandValue::Pair(fst @ Err(_), _) if is_zero_offset => {
-                *fst = Ok(imm);
+            OperandValueBuilder::Pair(fst @ Either::Right(_), _) if is_zero_offset => {
+                *fst = Either::Left(imm);
             }
-            OperandValue::Pair(_, snd @ Err(_)) if !is_zero_offset => {
-                *snd = Ok(imm);
+            OperandValueBuilder::Pair(_, snd @ Either::Right(_)) if !is_zero_offset => {
+                *snd = Either::Left(imm);
             }
             _ => bug!("Tried to insert {imm:?} into field {f:?} of {self:?}"),
         }
     }
 
     /// After having set all necessary fields, this converts the
-    /// `OperandValue<Result<V, _>>` (as obtained from [`OperandRef::builder`])
+    /// `OperandValue<Either<V, _>>` (as obtained from [`OperandRef::builder`])
     /// to the normal `OperandValue<V>`.
     ///
     /// ICEs if any required fields were not set.
-    pub fn build(&self, cx: &impl CodegenMethods<'tcx, Value = V>) -> OperandRef<'tcx, V> {
-        let OperandRef { val, layout } = *self;
+    pub(super) fn build(&self, cx: &impl CodegenMethods<'tcx, Value = V>) -> OperandRef<'tcx, V> {
+        let OperandRefBuilder { val, layout } = *self;
 
         // For something like `Option::<u32>::None`, it's expected that the
         // payload scalar will not actually have been set, so this converts
         // unset scalars to corresponding `undef` values so long as the scalar
         // from the layout allows uninit.
-        let unwrap = |r: Result<V, abi::Scalar>| match r {
-            Ok(v) => v,
-            Err(s) if s.is_uninit_valid() => {
+        let unwrap = |r: Either<V, abi::Scalar>| match r {
+            Either::Left(v) => v,
+            Either::Right(s) if s.is_uninit_valid() => {
                 let bty = cx.type_from_scalar(s);
                 cx.const_undef(bty)
             }
-            Err(_) => bug!("OperandRef::build called while fields are missing {self:?}"),
+            Either::Right(_) => bug!("OperandRef::build called while fields are missing {self:?}"),
         };
 
         let val = match val {
-            OperandValue::ZeroSized => OperandValue::ZeroSized,
-            OperandValue::Immediate(v) => OperandValue::Immediate(unwrap(v)),
-            OperandValue::Pair(a, b) => OperandValue::Pair(unwrap(a), unwrap(b)),
-            OperandValue::Ref(_) => bug!(),
+            OperandValueBuilder::ZeroSized => OperandValue::ZeroSized,
+            OperandValueBuilder::Immediate(v) => OperandValue::Immediate(unwrap(v)),
+            OperandValueBuilder::Pair(a, b) => OperandValue::Pair(unwrap(a), unwrap(b)),
+            OperandValueBuilder::Vector(v) => match v {
+                Either::Left(v) => OperandValue::Immediate(v),
+                Either::Right(()) => {
+                    bug!("OperandRef::build called while fields are missing {self:?}")
+                }
+            },
         };
         OperandRef { val, layout }
     }

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -4,10 +4,9 @@ use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf, TyAndLayout};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_middle::{bug, mir};
 use rustc_session::config::OptLevel;
-use rustc_span::{DUMMY_SP, Span};
 use tracing::{debug, instrument};
 
-use super::operand::{OperandRef, OperandValue};
+use super::operand::{OperandRef, OperandRefBuilder, OperandValue};
 use super::place::{PlaceRef, codegen_tag_value};
 use super::{FunctionCx, LocalRef};
 use crate::common::{IntPredicate, TypeKind};
@@ -181,7 +180,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             }
 
             _ => {
-                assert!(self.rvalue_creates_operand(rvalue, DUMMY_SP));
+                assert!(self.rvalue_creates_operand(rvalue));
                 let temp = self.codegen_rvalue_operand(bx, rvalue);
                 temp.val.store(bx, dest);
             }
@@ -354,10 +353,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         bx: &mut Bx,
         rvalue: &mir::Rvalue<'tcx>,
     ) -> OperandRef<'tcx, Bx::Value> {
-        assert!(
-            self.rvalue_creates_operand(rvalue, DUMMY_SP),
-            "cannot codegen {rvalue:?} to operand",
-        );
+        assert!(self.rvalue_creates_operand(rvalue), "cannot codegen {rvalue:?} to operand",);
 
         match *rvalue {
             mir::Rvalue::Cast(ref kind, ref source, mir_cast_ty) => {
@@ -668,9 +664,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
                 // `rvalue_creates_operand` has arranged that we only get here if
                 // we can build the aggregate immediate from the field immediates.
-                let Some(mut builder) = OperandRef::builder(layout) else {
-                    bug!("Cannot use type in operand builder: {layout:?}")
-                };
+                let mut builder = OperandRefBuilder::new(layout);
                 for (field_idx, field) in fields.iter_enumerated() {
                     let op = self.codegen_operand(bx, field);
                     let fi = active_field_index.unwrap_or(field_idx);
@@ -980,7 +974,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     /// will not actually take the operand path because the result type is such
     /// that it always gets an `alloca`, but where it's not worth re-checking the
     /// layout in this code when the right thing will happen anyway.
-    pub(crate) fn rvalue_creates_operand(&self, rvalue: &mir::Rvalue<'tcx>, span: Span) -> bool {
+    pub(crate) fn rvalue_creates_operand(&self, rvalue: &mir::Rvalue<'tcx>) -> bool {
         match *rvalue {
             mir::Rvalue::Cast(mir::CastKind::Transmute, ref operand, cast_ty) => {
                 let operand_ty = operand.ty(self.mir, self.cx.tcx());
@@ -1025,18 +1019,13 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             mir::Rvalue::NullaryOp(..) |
             mir::Rvalue::ThreadLocalRef(_) |
             mir::Rvalue::Use(..) |
+            mir::Rvalue::Aggregate(..) | // (*)
             mir::Rvalue::WrapUnsafeBinder(..) => // (*)
                 true,
             // Arrays are always aggregates, so it's not worth checking anything here.
             // (If it's really `[(); N]` or `[T; 0]` and we use the place path, fine.)
             mir::Rvalue::Repeat(..) => false,
-            mir::Rvalue::Aggregate(..) => {
-                    let ty = rvalue.ty(self.mir, self.cx.tcx());
-                    let ty = self.monomorphize(ty);
-                    let layout = self.cx.spanned_layout_of(ty, span);
-                    OperandRef::<Bx::Value>::builder(layout).is_some()
-                }
-            }
+        }
 
         // (*) this is only true if the type is suitable
     }

tests/codegen/enum/enum-aggregate.rs

@@ -112,17 +112,14 @@ fn make_uninhabited_err_indirectly(n: Never) -> Result<u32, Never> {
 
 #[no_mangle]
 fn make_fully_uninhabited_result(v: u32, n: Never) -> Result<(u32, Never), (Never, u32)> {
-    // We don't try to do this in SSA form since the whole type is uninhabited.
+    // Actually reaching this would be UB, so we don't actually build a result.
 
     // CHECK-LABEL: { i32, i32 } @make_fully_uninhabited_result(i32 %v)
-    // CHECK: %[[ALLOC_V:.+]] = alloca [4 x i8]
-    // CHECK: %[[RET:.+]] = alloca [8 x i8]
-    // CHECK: store i32 %v, ptr %[[ALLOC_V]]
-    // CHECK: %[[TEMP_V:.+]] = load i32, ptr %[[ALLOC_V]]
-    // CHECK: %[[INNER:.+]] = getelementptr inbounds i8, ptr %[[RET]]
-    // CHECK: store i32 %[[TEMP_V]], ptr %[[INNER]]
-    // CHECK: call void @llvm.trap()
-    // CHECK: unreachable
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: call void @llvm.trap()
+    // CHECK-NEXT: call void @llvm.trap()
+    // CHECK-NEXT: call void @llvm.trap()
+    // CHECK-NEXT: unreachable
     Ok((v, n))
 }