From c38f75c21f016bffbf841ed5abce338f94201bde Mon Sep 17 00:00:00 2001 From: Scott McMurray Date: Sun, 14 Apr 2024 00:51:49 -0700 Subject: [PATCH] Make SSA aggregates without needing an alloca --- Cargo.lock | 2 + compiler/rustc_codegen_ssa/Cargo.toml | 4 +- compiler/rustc_codegen_ssa/src/mir/operand.rs | 30 +++++ compiler/rustc_codegen_ssa/src/mir/rvalue.rs | 81 +++++++++++-- tests/codegen/mir-aggregate-no-alloca.rs | 114 +++++++++++++++--- 5 files changed, 203 insertions(+), 28 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f1304735f8f..f2c65c7cadb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3746,8 +3746,10 @@ name = "rustc_codegen_ssa" version = "0.0.0" dependencies = [ "ar_archive_writer", + "arrayvec", "bitflags 2.5.0", "cc", + "either", "itertools 0.12.1", "jobserver", "libc", diff --git a/compiler/rustc_codegen_ssa/Cargo.toml b/compiler/rustc_codegen_ssa/Cargo.toml index 1562a09f5c1..3771fc6b0a2 100644 --- a/compiler/rustc_codegen_ssa/Cargo.toml +++ b/compiler/rustc_codegen_ssa/Cargo.toml @@ -6,8 +6,10 @@ edition = "2021" [dependencies] # tidy-alphabetical-start ar_archive_writer = "0.2.0" +arrayvec = { version = "0.7", default-features = false } bitflags = "2.4.1" -cc = "1.0.97" +cc = "1.0.90" +either = "1.5.0" itertools = "0.12" jobserver = "0.1.28" pathdiff = "0.2.0" diff --git a/compiler/rustc_codegen_ssa/src/mir/operand.rs b/compiler/rustc_codegen_ssa/src/mir/operand.rs index 38f77f2e646..4cff6ebd35d 100644 --- a/compiler/rustc_codegen_ssa/src/mir/operand.rs +++ b/compiler/rustc_codegen_ssa/src/mir/operand.rs @@ -14,6 +14,9 @@ use rustc_target::abi::{self, Abi, Align, Size}; use std::fmt; +use arrayvec::ArrayVec; +use either::Either; + /// The representation of a Rust value. The enum variant is in fact /// uniquely determined by the value's type, but is kept as a /// safety check. @@ -58,6 +61,33 @@ pub enum OperandValue { ZeroSized, } +impl OperandValue { + /// If this is ZeroSized/Immediate/Pair, return an array of the 0/1/2 values. + /// If this is Ref, return the place. + #[inline] + pub fn immediates_or_place(self) -> Either, PlaceValue> { + match self { + OperandValue::ZeroSized => Either::Left(ArrayVec::new()), + OperandValue::Immediate(a) => Either::Left(ArrayVec::from_iter([a])), + OperandValue::Pair(a, b) => Either::Left([a, b].into()), + OperandValue::Ref(p) => Either::Right(p), + } + } + + /// Given an array of 0/1/2 immediate values, return ZeroSized/Immediate/Pair. + #[inline] + pub fn from_immediates(immediates: ArrayVec) -> Self { + let mut it = immediates.into_iter(); + let Some(a) = it.next() else { + return OperandValue::ZeroSized; + }; + let Some(b) = it.next() else { + return OperandValue::Immediate(a); + }; + OperandValue::Pair(a, b) + } +} + /// An `OperandRef` is an "SSA" reference to a Rust value, along with /// its type. /// diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs index 6be473ffa1f..06663a0f61c 100644 --- a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs +++ b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs @@ -8,14 +8,16 @@ use crate::traits::*; use crate::MemFlags; use rustc_hir as hir; -use rustc_middle::mir::{self, AggregateKind, Operand}; +use rustc_middle::mir; use rustc_middle::ty::cast::{CastTy, IntTy}; use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, TyAndLayout}; use rustc_middle::ty::{self, adjustment::PointerCoercion, Instance, Ty, TyCtxt}; use rustc_middle::{bug, span_bug}; use rustc_session::config::OptLevel; use rustc_span::{Span, DUMMY_SP}; -use rustc_target::abi::{self, FIRST_VARIANT}; +use rustc_target::abi::{self, FieldIdx, FIRST_VARIANT}; + +use arrayvec::ArrayVec; impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { #[instrument(level = "trace", skip(self, bx))] @@ -581,7 +583,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { self.codegen_place_to_pointer(bx, place, mk_ref) } - mir::Rvalue::CopyForDeref(place) => self.codegen_operand(bx, &Operand::Copy(place)), + mir::Rvalue::CopyForDeref(place) => { + self.codegen_operand(bx, &mir::Operand::Copy(place)) + } mir::Rvalue::AddressOf(mutability, place) => { let mk_ptr = move |tcx: TyCtxt<'tcx>, ty: Ty<'tcx>| Ty::new_ptr(tcx, ty, mutability); @@ -739,11 +743,40 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { } } mir::Rvalue::Repeat(..) => bug!("{rvalue:?} in codegen_rvalue_operand"), - mir::Rvalue::Aggregate(..) => { - // According to `rvalue_creates_operand`, only ZST - // aggregate rvalues are allowed to be operands. + mir::Rvalue::Aggregate(_, ref fields) => { let ty = rvalue.ty(self.mir, self.cx.tcx()); - OperandRef::zero_sized(self.cx.layout_of(self.monomorphize(ty))) + let ty = self.monomorphize(ty); + let layout = self.cx.layout_of(ty); + + // `rvalue_creates_operand` has arranged that we only get here if + // we can build the aggregate immediate from the field immediates. + let mut inputs = ArrayVec::::new(); + let mut input_scalars = ArrayVec::::new(); + for field_idx in layout.fields.index_by_increasing_offset() { + let field_idx = FieldIdx::from_usize(field_idx); + let op = self.codegen_operand(bx, &fields[field_idx]); + let values = op.val.immediates_or_place().left_or_else(|p| { + bug!("Field {field_idx:?} is {p:?} making {layout:?}"); + }); + inputs.extend(values); + let scalars = self.value_kind(op.layout).scalars().unwrap(); + input_scalars.extend(scalars); + } + + let output_scalars = self.value_kind(layout).scalars().unwrap(); + itertools::izip!(&mut inputs, input_scalars, output_scalars).for_each( + |(v, in_s, out_s)| { + if in_s != out_s { + // We have to be really careful about bool here, because + // `(bool,)` stays i1 but `Cell` becomes i8. + *v = bx.from_immediate(*v); + *v = bx.to_immediate_scalar(*v, out_s); + } + }, + ); + + let val = OperandValue::from_immediates(inputs); + OperandRef { val, layout } } mir::Rvalue::ShallowInitBox(ref operand, content_ty) => { let operand = self.codegen_operand(bx, operand); @@ -1051,16 +1084,29 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { mir::Rvalue::ThreadLocalRef(_) | mir::Rvalue::Use(..) => // (*) true, - // This always produces a `ty::RawPtr`, so will be Immediate or Pair - mir::Rvalue::Aggregate(box AggregateKind::RawPtr(..), ..) => 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(..) => { + mir::Rvalue::Aggregate(ref kind, _) => { + let allowed_kind = match **kind { + // This always produces a `ty::RawPtr`, so will be Immediate or Pair + mir::AggregateKind::RawPtr(..) => true, + mir::AggregateKind::Array(..) => false, + mir::AggregateKind::Tuple => true, + mir::AggregateKind::Adt(def_id, ..) => { + let adt_def = self.cx.tcx().adt_def(def_id); + adt_def.is_struct() && !adt_def.repr().simd() + } + mir::AggregateKind::Closure(..) => true, + // FIXME: Can we do this for simple coroutines too? + mir::AggregateKind::Coroutine(..) | mir::AggregateKind::CoroutineClosure(..) => false, + }; + allowed_kind && { let ty = rvalue.ty(self.mir, self.cx.tcx()); let ty = self.monomorphize(ty); - // For ZST this can be `OperandValueKind::ZeroSized`. - self.cx.spanned_layout_of(ty, span).is_zst() + let layout = self.cx.spanned_layout_of(ty, span); + !self.cx.is_backend_ref(layout) + } } } @@ -1102,3 +1148,14 @@ enum OperandValueKind { Pair(abi::Scalar, abi::Scalar), ZeroSized, } + +impl OperandValueKind { + fn scalars(self) -> Option> { + Some(match self { + OperandValueKind::ZeroSized => ArrayVec::new(), + OperandValueKind::Immediate(a) => ArrayVec::from_iter([a]), + OperandValueKind::Pair(a, b) => [a, b].into(), + OperandValueKind::Ref => return None, + }) + } +} diff --git a/tests/codegen/mir-aggregate-no-alloca.rs b/tests/codegen/mir-aggregate-no-alloca.rs index 8714fc56980..a7752d714fe 100644 --- a/tests/codegen/mir-aggregate-no-alloca.rs +++ b/tests/codegen/mir-aggregate-no-alloca.rs @@ -1,17 +1,15 @@ -//@ compile-flags: -O -C no-prepopulate-passes +//@ compile-flags: -O -C no-prepopulate-passes -Z randomize-layout=no #![crate_type = "lib"] #[repr(transparent)] -struct Transparent32(u32); +pub struct Transparent32(u32); // CHECK: i32 @make_transparent(i32 noundef %x) #[no_mangle] pub fn make_transparent(x: u32) -> Transparent32 { - // CHECK: %a = alloca i32 - // CHECK: store i32 %x, ptr %a - // CHECK: %[[TEMP:.+]] = load i32, ptr %a - // CHECK: ret i32 %[[TEMP]] + // CHECK-NOT: alloca + // CHECK: ret i32 %x let a = Transparent32(x); a } @@ -19,21 +17,107 @@ pub fn make_transparent(x: u32) -> Transparent32 { // CHECK: i32 @make_closure(i32 noundef %x) #[no_mangle] pub fn make_closure(x: i32) -> impl Fn(i32) -> i32 { - // CHECK: %[[ALLOCA:.+]] = alloca i32 - // CHECK: store i32 %x, ptr %[[ALLOCA]] - // CHECK: %[[TEMP:.+]] = load i32, ptr %[[ALLOCA]] - // CHECK: ret i32 %[[TEMP]] + // CHECK-NOT: alloca + // CHECK: ret i32 %x move |y| x + y } +#[repr(transparent)] +pub struct TransparentPair((), (u16, u16), ()); + +// CHECK: { i16, i16 } @make_transparent_pair(i16 noundef %x.0, i16 noundef %x.1) +#[no_mangle] +pub fn make_transparent_pair(x: (u16, u16)) -> TransparentPair { + // CHECK-NOT: alloca + // CHECK: %[[TEMP0:.+]] = insertvalue { i16, i16 } poison, i16 %x.0, 0 + // CHECK: %[[TEMP1:.+]] = insertvalue { i16, i16 } %[[TEMP0]], i16 %x.1, 1 + // CHECK: ret { i16, i16 } %[[TEMP1]] + let a = TransparentPair((), x, ()); + a +} + // CHECK-LABEL: { i32, i32 } @make_2_tuple(i32 noundef %x) #[no_mangle] pub fn make_2_tuple(x: u32) -> (u32, u32) { - // CHECK: %pair = alloca { i32, i32 } - // CHECK: store i32 - // CHECK: store i32 - // CHECK: load i32 - // CHECK: load i32 + // CHECK-NOT: alloca + // CHECK: %[[TEMP0:.+]] = insertvalue { i32, i32 } poison, i32 %x, 0 + // CHECK: %[[TEMP1:.+]] = insertvalue { i32, i32 } %[[TEMP0]], i32 %x, 1 + // CHECK: ret { i32, i32 } %[[TEMP1]] let pair = (x, x); pair } + +// CHECK-LABEL: i8 @make_cell_of_bool(i1 noundef zeroext %b) +#[no_mangle] +pub fn make_cell_of_bool(b: bool) -> std::cell::Cell { + // CHECK: %[[BYTE:.+]] = zext i1 %b to i8 + // CHECK: ret i8 %[[BYTE]] + std::cell::Cell::new(b) +} + +// CHECK-LABLE: { i8, i16 } @make_cell_of_bool_and_short(i1 noundef zeroext %b, i16 noundef %s) +#[no_mangle] +pub fn make_cell_of_bool_and_short(b: bool, s: u16) -> std::cell::Cell<(bool, u16)> { + // CHECK-NOT: alloca + // CHECK: %[[BYTE:.+]] = zext i1 %b to i8 + // CHECK: %[[TEMP0:.+]] = insertvalue { i8, i16 } poison, i8 %[[BYTE]], 0 + // CHECK: %[[TEMP1:.+]] = insertvalue { i8, i16 } %[[TEMP0]], i16 %s, 1 + // CHECK: ret { i8, i16 } %[[TEMP1]] + std::cell::Cell::new((b, s)) +} + +// CHECK-LABEL: { i1, i1 } @make_tuple_of_bools(i1 noundef zeroext %a, i1 noundef zeroext %b) +#[no_mangle] +pub fn make_tuple_of_bools(a: bool, b: bool) -> (bool, bool) { + // CHECK-NOT: alloca + // CHECK: %[[TEMP0:.+]] = insertvalue { i1, i1 } poison, i1 %a, 0 + // CHECK: %[[TEMP1:.+]] = insertvalue { i1, i1 } %[[TEMP0]], i1 %b, 1 + // CHECK: ret { i1, i1 } %[[TEMP1]] + (a, b) +} + +pub struct Struct0(); + +// CHECK-LABEL: void @make_struct_0() +#[no_mangle] +pub fn make_struct_0() -> Struct0 { + // CHECK: ret void + let s = Struct0(); + s +} + +pub struct Struct1(i32); + +// CHECK-LABEL: i32 @make_struct_1(i32 noundef %a) +#[no_mangle] +pub fn make_struct_1(a: i32) -> Struct1 { + // CHECK: ret i32 %a + let s = Struct1(a); + s +} + +pub struct Struct2Asc(i16, i64); + +// CHECK-LABEL: { i64, i16 } @make_struct_2_asc(i16 noundef %a, i64 noundef %b) +#[no_mangle] +pub fn make_struct_2_asc(a: i16, b: i64) -> Struct2Asc { + // CHECK-NOT: alloca + // CHECK: %[[TEMP0:.+]] = insertvalue { i64, i16 } poison, i64 %b, 0 + // CHECK: %[[TEMP1:.+]] = insertvalue { i64, i16 } %[[TEMP0]], i16 %a, 1 + // CHECK: ret { i64, i16 } %[[TEMP1]] + let s = Struct2Asc(a, b); + s +} + +pub struct Struct2Desc(i64, i16); + +// CHECK-LABEL: { i64, i16 } @make_struct_2_desc(i64 noundef %a, i16 noundef %b) +#[no_mangle] +pub fn make_struct_2_desc(a: i64, b: i16) -> Struct2Desc { + // CHECK-NOT: alloca + // CHECK: %[[TEMP0:.+]] = insertvalue { i64, i16 } poison, i64 %a, 0 + // CHECK: %[[TEMP1:.+]] = insertvalue { i64, i16 } %[[TEMP0]], i16 %b, 1 + // CHECK: ret { i64, i16 } %[[TEMP1]] + let s = Struct2Desc(a, b); + s +}