ensure byval allocas are sufficiently aligned

erikdesjardins · erikdesjardins · commit 0e76446a9ff8 · 2023-07-10T19:19:38.000-04:00
diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -23,6 +23,8 @@ use rustc_target::abi::call::{ArgAbi, FnAbi, PassMode, Reg};
 use rustc_target::abi::{self, HasDataLayout, WrappingRange};
 use rustc_target::spec::abi::Abi;
 
+use std::cmp;
+
 // Indicates if we are in the middle of merging a BB's successor into it. This
 // can happen when BB jumps directly to its successor and the successor has no
 // other predecessors.
@@ -1360,36 +1362,58 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         // Force by-ref if we have to load through a cast pointer.
         let (mut llval, align, by_ref) = match op.val {
             Immediate(_) | Pair(..) => match arg.mode {
-                PassMode::Indirect { .. } | PassMode::Cast(..) => {
+                PassMode::Indirect { attrs, .. } => {
+                    // Indirect argument may have higher alignment requirements than the type's alignment.
+                    // This can happen, e.g. when passing types with <4 byte alignment on the stack on x86.
+                    let required_align = match attrs.pointee_align {
+                        Some(pointee_align) => cmp::max(pointee_align, arg.layout.align.abi),
+                        None => arg.layout.align.abi,
+                    };
+                    let scratch = PlaceRef::alloca_aligned(bx, arg.layout, required_align);
+                    op.val.store(bx, scratch);
+                    (scratch.llval, scratch.align, true)
+                }
+                PassMode::Cast(..) => {
                     let scratch = PlaceRef::alloca(bx, arg.layout);
                     op.val.store(bx, scratch);
                     (scratch.llval, scratch.align, true)
                 }
                 _ => (op.immediate_or_packed_pair(bx), arg.layout.align.abi, false),
             },
-            Ref(llval, _, align) => {
-                if arg.is_indirect() && align < arg.layout.align.abi {
-                    // `foo(packed.large_field)`. We can't pass the (unaligned) field directly. I
-                    // think that ATM (Rust 1.16) we only pass temporaries, but we shouldn't
-                    // have scary latent bugs around.
-
-                    let scratch = PlaceRef::alloca(bx, arg.layout);
-                    base::memcpy_ty(
-                        bx,
-                        scratch.llval,
-                        scratch.align,
-                        llval,
-                        align,
-                        op.layout,
-                        MemFlags::empty(),
-                    );
-                    (scratch.llval, scratch.align, true)
-                } else {
-                    (llval, align, true)
+            Ref(llval, _, align) => match arg.mode {
+                PassMode::Indirect { attrs, .. } => {
+                    let required_align = match attrs.pointee_align {
+                        Some(pointee_align) => cmp::max(pointee_align, arg.layout.align.abi),
+                        None => arg.layout.align.abi,
+                    };
+                    if align < required_align {
+                        // For `foo(packed.large_field)`, and types with <4 byte alignment on x86,
+                        // alignment requirements may be higher than the type's alignment, so copy
+                        // to a higher-aligned alloca.
+                        let scratch = PlaceRef::alloca_aligned(bx, arg.layout, required_align);
+                        base::memcpy_ty(
+                            bx,
+                            scratch.llval,
+                            scratch.align,
+                            llval,
+                            align,
+                            op.layout,
+                            MemFlags::empty(),
+                        );
+                        (scratch.llval, scratch.align, true)
+                    } else {
+                        (llval, align, true)
+                    }
                 }
-            }
+                _ => (llval, align, true),
+            },
             ZeroSized => match arg.mode {
-                PassMode::Indirect { .. } => {
+                PassMode::Indirect { on_stack, .. } => {
+                    if on_stack {
+                        // It doesn't seem like any target can have `byval` ZSTs, so this assert
+                        // is here to replace a would-be untested codepath.
+                        bug!("ZST {op:?} passed on stack with abi {arg:?}");
+                    }
                     // Though `extern "Rust"` doesn't pass ZSTs, some ABIs pass
                     // a pointer for `repr(C)` structs even when empty, so get
                     // one from an `alloca` (which can be left uninitialized).
diff --git a/compiler/rustc_codegen_ssa/src/mir/place.rs b/compiler/rustc_codegen_ssa/src/mir/place.rs
@@ -47,10 +47,18 @@ impl<'a, 'tcx, V: CodegenObject> PlaceRef<'tcx, V> {
     pub fn alloca<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
         bx: &mut Bx,
         layout: TyAndLayout<'tcx>,
+    ) -> Self {
+        Self::alloca_aligned(bx, layout, layout.align.abi)
+    }
+
+    pub fn alloca_aligned<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
+        bx: &mut Bx,
+        layout: TyAndLayout<'tcx>,
+        align: Align,
     ) -> Self {
         assert!(layout.is_sized(), "tried to statically allocate unsized place");
-        let tmp = bx.alloca(bx.cx().backend_type(layout), layout.align.abi);
-        Self::new_sized(tmp, layout)
+        let tmp = bx.alloca(bx.cx().backend_type(layout), align);
+        Self::new_sized_aligned(tmp, layout, align)
     }
 
     /// Returns a place for an indirect reference to an unsized place.
diff --git a/tests/codegen/align-byval.rs b/tests/codegen/align-byval.rs
@@ -1,3 +1,4 @@
+// ignore-tidy-linelength
 // revisions:m68k wasm x86_64-linux x86_64-windows i686-linux i686-windows
 
 //[m68k] compile-flags: --target m68k-unknown-linux-gnu
@@ -29,6 +30,16 @@
 impl Copy for i32 {}
 impl Copy for i64 {}
 
+// This struct can be represented as a pair, so it exercises the OperandValue::Pair
+// codepath in `codegen_argument`.
+#[repr(C)]
+pub struct NaturalAlign1 {
+    a: i8,
+    b: i8,
+}
+
+// This struct cannot be represented as an immediate, so it exercises the OperandValue::Ref
+// codepath in `codegen_argument`.
 #[repr(C)]
 pub struct NaturalAlign2 {
     a: [i16; 16],
@@ -67,7 +78,93 @@ pub struct ForceAlign16 {
     b: i8
 }
 
+// CHECK-LABEL: @call_na1
+#[no_mangle]
+pub unsafe fn call_na1(x: NaturalAlign1) {
+    // CHECK: start:
+
+    // m68k: [[ALLOCA:%[a-z0-9+]]] = alloca { i8, i8 }, align 1
+    // m68k: call void @natural_align_1({{.*}}byval({ i8, i8 }) align 1{{.*}} [[ALLOCA]])
+
+    // wasm: [[ALLOCA:%[a-z0-9+]]] = alloca { i8, i8 }, align 1
+    // wasm: call void @natural_align_1({{.*}}byval({ i8, i8 }) align 1{{.*}} [[ALLOCA]])
+
+    // x86_64-linux: call void @natural_align_1(i16
+
+    // x86_64-windows: call void @natural_align_1(i16
+
+    // i686-linux: [[ALLOCA:%[a-z0-9+]]] = alloca { i8, i8 }, align 4
+    // i686-linux: call void @natural_align_1({{.*}}byval({ i8, i8 }) align 4{{.*}} [[ALLOCA]])
+
+    // i686-windows: [[ALLOCA:%[a-z0-9+]]] = alloca { i8, i8 }, align 4
+    // i686-windows: call void @natural_align_1({{.*}}byval({ i8, i8 }) align 4{{.*}} [[ALLOCA]])
+    natural_align_1(x);
+}
+
+// CHECK-LABEL: @call_na2
+#[no_mangle]
+pub unsafe fn call_na2(x: NaturalAlign2) {
+    // CHECK: start:
+
+    // m68k-NEXT: call void @natural_align_2
+    // wasm-NEXT: call void @natural_align_2
+    // x86_64-linux-NEXT: call void @natural_align_2
+    // x86_64-windows-NEXT: call void @natural_align_2
+
+    // i686-linux: [[ALLOCA:%[0-9]+]] = alloca %NaturalAlign2, align 4
+    // i686-linux: call void @natural_align_2({{.*}}byval(%NaturalAlign2) align 4{{.*}} [[ALLOCA]])
+
+    // i686-windows: [[ALLOCA:%[0-9]+]] = alloca %NaturalAlign2, align 4
+    // i686-windows: call void @natural_align_2({{.*}}byval(%NaturalAlign2) align 4{{.*}} [[ALLOCA]])
+    natural_align_2(x);
+}
+
+// CHECK-LABEL: @call_fa4
+#[no_mangle]
+pub unsafe fn call_fa4(x: ForceAlign4) {
+    // CHECK: start:
+    // CHECK-NEXT: call void @force_align_4
+    force_align_4(x);
+}
+
+// CHECK-LABEL: @call_na8
+#[no_mangle]
+pub unsafe fn call_na8(x: NaturalAlign8) {
+    // CHECK: start:
+    // CHECK-NEXT: call void @natural_align_8
+    natural_align_8(x);
+}
+
+// CHECK-LABEL: @call_fa8
+#[no_mangle]
+pub unsafe fn call_fa8(x: ForceAlign8) {
+    // CHECK: start:
+    // CHECK-NEXT: call void @force_align_8
+    force_align_8(x);
+}
+
+// CHECK-LABEL: @call_fa16
+#[no_mangle]
+pub unsafe fn call_fa16(x: ForceAlign16) {
+    // CHECK: start:
+    // CHECK-NEXT: call void @force_align_16
+    force_align_16(x);
+}
+
 extern "C" {
+    // m68k: declare void @natural_align_1({{.*}}byval({ i8, i8 }) align 1{{.*}})
+
+    // wasm: declare void @natural_align_1({{.*}}byval({ i8, i8 }) align 1{{.*}})
+
+    // x86_64-linux: declare void @natural_align_1(i16)
+
+    // x86_64-windows: declare void @natural_align_1(i16)
+
+    // i686-linux: declare void @natural_align_1({{.*}}byval({ i8, i8 }) align 4{{.*}})
+
+    // i686-windows: declare void @natural_align_1({{.*}}byval({ i8, i8 }) align 4{{.*}})
+    fn natural_align_1(x: NaturalAlign1);
+
     // m68k: declare void @natural_align_2({{.*}}byval(%NaturalAlign2) align 2{{.*}})
 
     // wasm: declare void @natural_align_2({{.*}}byval(%NaturalAlign2) align 2{{.*}})
@@ -81,7 +178,7 @@ extern "C" {
     // i686-linux: declare void @natural_align_2({{.*}}byval(%NaturalAlign2) align 4{{.*}})
 
     // i686-windows: declare void @natural_align_2({{.*}}byval(%NaturalAlign2) align 4{{.*}})
-    fn natural_align_2(a: NaturalAlign2);
+    fn natural_align_2(x: NaturalAlign2);
 
     // m68k: declare void @force_align_4({{.*}}byval(%ForceAlign4) align 4{{.*}})
 
@@ -96,7 +193,7 @@ extern "C" {
     // i686-linux: declare void @force_align_4({{.*}}byval(%ForceAlign4) align 4{{.*}})
 
     // i686-windows: declare void @force_align_4({{.*}}byval(%ForceAlign4) align 4{{.*}})
-    fn force_align_4(b: ForceAlign4);
+    fn force_align_4(x: ForceAlign4);
 
     // m68k: declare void @natural_align_8({{.*}}byval(%NaturalAlign8) align 4{{.*}})
 
@@ -128,7 +225,7 @@ extern "C" {
     // i686-windows: declare void @force_align_8(
     // i686-windows-NOT: byval
     // i686-windows-SAME: align 8{{.*}})
-    fn force_align_8(y: ForceAlign8);
+    fn force_align_8(x: ForceAlign8);
 
     // m68k: declare void @force_align_16({{.*}}byval(%ForceAlign16) align 16{{.*}})
 
@@ -145,16 +242,5 @@ extern "C" {
     // i686-windows: declare void @force_align_16(
     // i686-windows-NOT: byval
     // i686-windows-SAME: align 16{{.*}})
-    fn force_align_16(z: ForceAlign16);
-}
-
-pub unsafe fn main(
-    a: NaturalAlign2, b: ForceAlign4,
-    x: NaturalAlign8, y: ForceAlign8, z: ForceAlign16
-) {
-    natural_align_2(a);
-    force_align_4(b);
-    natural_align_8(x);
-    force_align_8(y);
-    force_align_16(z);
+    fn force_align_16(x: ForceAlign16);
 }
