aarch64-linux: properly handle 128bit aligned aggregates

Author: erikdesjardins

compiler/rustc_abi/src/layout.rs

@@ -41,6 +41,7 @@ pub trait LayoutCalculator {
             align,
             size,
             repr_align: None,
+            unadjusted_abi_align: align.abi,
         }
     }
 
@@ -124,6 +125,7 @@ pub trait LayoutCalculator {
             align: dl.i8_align,
             size: Size::ZERO,
             repr_align: None,
+            unadjusted_abi_align: dl.i8_align.abi,
         }
     }
 
@@ -291,13 +293,16 @@ pub trait LayoutCalculator {
             }
 
             let mut align = dl.aggregate_align;
+            let mut unadjusted_abi_align = align.abi;
+
             let mut variant_layouts = variants
                 .iter_enumerated()
                 .map(|(j, v)| {
                     let mut st = self.univariant(dl, v, repr, StructKind::AlwaysSized)?;
                     st.variants = Variants::Single { index: j };
 
                     align = align.max(st.align);
+                    unadjusted_abi_align = unadjusted_abi_align.max(st.unadjusted_abi_align);
 
                     Some(st)
                 })
@@ -425,6 +430,7 @@ pub trait LayoutCalculator {
                 size,
                 align,
                 repr_align: repr.align,
+                unadjusted_abi_align,
             };
 
             Some(TmpLayout { layout, variants: variant_layouts })
@@ -459,6 +465,8 @@ pub trait LayoutCalculator {
         let (min_ity, signed) = discr_range_of_repr(min, max); //Integer::repr_discr(tcx, ty, &repr, min, max);
 
         let mut align = dl.aggregate_align;
+        let mut unadjusted_abi_align = align.abi;
+
         let mut size = Size::ZERO;
 
         // We're interested in the smallest alignment, so start large.
@@ -501,6 +509,7 @@ pub trait LayoutCalculator {
                 }
                 size = cmp::max(size, st.size);
                 align = align.max(st.align);
+                unadjusted_abi_align = unadjusted_abi_align.max(st.unadjusted_abi_align);
                 Some(st)
             })
             .collect::<Option<IndexVec<VariantIdx, _>>>()?;
@@ -695,6 +704,7 @@ pub trait LayoutCalculator {
             align,
             size,
             repr_align: repr.align,
+            unadjusted_abi_align,
         };
 
         let tagged_layout = TmpLayout { layout: tagged_layout, variants: layout_variants };
@@ -735,10 +745,6 @@ pub trait LayoutCalculator {
         let dl = dl.borrow();
         let mut align = if repr.pack.is_some() { dl.i8_align } else { dl.aggregate_align };
 
-        if let Some(repr_align) = repr.align {
-            align = align.max(AbiAndPrefAlign::new(repr_align));
-        }
-
         // If all the non-ZST fields have the same ABI and union ABI optimizations aren't
         // disabled, we can use that common ABI for the union as a whole.
         struct AbiMismatch;
@@ -791,6 +797,14 @@ pub trait LayoutCalculator {
         if let Some(pack) = repr.pack {
             align = align.min(AbiAndPrefAlign::new(pack));
         }
+        // The unadjusted ABI alignment does not include repr(align), but does include repr(pack).
+        // See documentation on `LayoutS::unadjusted_abi_align`.
+        let unadjusted_abi_align = align.abi;
+        if let Some(repr_align) = repr.align {
+            align = align.max(AbiAndPrefAlign::new(repr_align));
+        }
+        // `align` must not be modified after this, or `unadjusted_abi_align` could be inaccurate.
+        let align = align;
 
         // If all non-ZST fields have the same ABI, we may forward that ABI
         // for the union as a whole, unless otherwise inhibited.
@@ -814,6 +828,7 @@ pub trait LayoutCalculator {
             align,
             size: size.align_to(align.abi),
             repr_align: repr.align,
+            unadjusted_abi_align,
         })
     }
 }
@@ -1023,9 +1038,16 @@ fn univariant(
 
         offset = offset.checked_add(field.size(), dl)?;
     }
+
+    // The unadjusted ABI alignment does not include repr(align), but does include repr(pack).
+    // See documentation on `LayoutS::unadjusted_abi_align`.
+    let unadjusted_abi_align = align.abi;
     if let Some(repr_align) = repr.align {
         align = align.max(AbiAndPrefAlign::new(repr_align));
     }
+    // `align` must not be modified after this point, or `unadjusted_abi_align` could be inaccurate.
+    let align = align;
+
     debug!("univariant min_size: {:?}", offset);
     let min_size = offset;
     // As stated above, inverse_memory_index holds field indices by increasing offset.
@@ -1111,9 +1133,29 @@ fn univariant(
         abi = Abi::Uninhabited;
     }
 
-    let repr_align = repr.align.or_else(|| {
-        if repr.transparent() { layout_of_single_non_zst_field?.repr_align() } else { None }
-    });
+    let (repr_align, unadjusted_abi_align) = if repr.transparent() {
+        match layout_of_single_non_zst_field {
+            Some(l) => (l.repr_align(), l.unadjusted_abi_align()),
+            None => {
+                // `repr(transparent)` with all ZST fields.
+                //
+                // Using `None` for `repr_align` here is technically incorrect, since one of
+                // the ZSTs could have `repr(align(1))`. It's an interesting question, if you have
+                // `#{repr(transparent)] struct Foo((), ZstWithReprAlign1)`, which of those ZSTs'
+                // ABIs is forwarded by `repr(transparent)`? The answer to that question determines
+                // whether we should use `None` or `Some(align 1)` here. Thanksfully, two things
+                // together mean this doesn't matter:
+                // - You're not allowed to have a `repr(transparent)` struct that contains
+                //   `repr(align)` > 1 ZSTs. See error E0691.
+                // - MSVC never treats requested align 1 differently from natural align 1.
+                //   (And the `repr_align` field is only used on i686-windows, see `LayoutS` docs.)
+                // So just use `None` for now.
+                (None, align.abi)
+            }
+        }
+    } else {
+        (repr.align, unadjusted_abi_align)
+    };
 
     Some(LayoutS {
         variants: Variants::Single { index: FIRST_VARIANT },
@@ -1123,6 +1165,7 @@ fn univariant(
         align,
         size,
         repr_align,
+        unadjusted_abi_align,
     })
 }
 

compiler/rustc_abi/src/lib.rs

@@ -1536,6 +1536,11 @@ pub struct LayoutS {
     /// Only used on i686-windows, where the argument passing ABI is different when alignment is
     /// requested, even if the requested alignment is equal to the natural alignment.
     pub repr_align: Option<Align>,
+
+    /// The alignment the type would have, ignoring any `repr(align)` but including `repr(packed)`.
+    /// Only used on aarch64-linux, where the argument passing ABI ignores the requested alignment
+    /// in some cases.
+    pub unadjusted_abi_align: Align,
 }
 
 impl LayoutS {
@@ -1551,6 +1556,7 @@ impl LayoutS {
             size,
             align,
             repr_align: None,
+            unadjusted_abi_align: align.abi,
         }
     }
 }
@@ -1560,7 +1566,16 @@ impl fmt::Debug for LayoutS {
         // This is how `Layout` used to print before it become
         // `Interned<LayoutS>`. We print it like this to avoid having to update
         // expected output in a lot of tests.
-        let LayoutS { size, align, abi, fields, largest_niche, variants, repr_align } = self;
+        let LayoutS {
+            size,
+            align,
+            abi,
+            fields,
+            largest_niche,
+            variants,
+            repr_align,
+            unadjusted_abi_align,
+        } = self;
         f.debug_struct("Layout")
             .field("size", size)
             .field("align", align)
@@ -1569,6 +1584,7 @@ impl fmt::Debug for LayoutS {
             .field("largest_niche", largest_niche)
             .field("variants", variants)
             .field("repr_align", repr_align)
+            .field("unadjusted_abi_align", unadjusted_abi_align)
             .finish()
     }
 }
@@ -1613,6 +1629,10 @@ impl<'a> Layout<'a> {
         self.0.0.repr_align
     }
 
+    pub fn unadjusted_abi_align(self) -> Align {
+        self.0.0.unadjusted_abi_align
+    }
+
     /// Whether the layout is from a type that implements [`std::marker::PointerLike`].
     ///
     /// Currently, that means that the type is pointer-sized, pointer-aligned,

compiler/rustc_middle/src/ty/layout.rs

@@ -756,6 +756,7 @@ where
                     align: tcx.data_layout.i8_align,
                     size: Size::ZERO,
                     repr_align: None,
+                    unadjusted_abi_align: tcx.data_layout.i8_align.abi,
                 })
             }
 

compiler/rustc_target/src/abi/call/aarch64.rs

@@ -1,25 +1,15 @@
 use crate::abi::call::{ArgAbi, FnAbi, Reg, RegKind, Uniform};
 use crate::abi::{HasDataLayout, TyAbiInterface};
 
-/// Given integer-types M and register width N (e.g. M=u16 and N=32 bits), the
-/// `ParamExtension` policy specifies how a uM value should be treated when
-/// passed via register or stack-slot of width N. See also rust-lang/rust#97463.
+/// Indicates the variant of the AArch64 ABI we are compiling for.
+/// Used to accommodate Apple and Microsoft's deviations from the usual AAPCS ABI.
+///
+/// Corresponds to Clang's `AArch64ABIInfo::ABIKind`.
 #[derive(Copy, Clone, PartialEq)]
-pub enum ParamExtension {
-    /// Indicates that when passing an i8/i16, either as a function argument or
-    /// as a return value, it must be sign-extended to 32 bits, and likewise a
-    /// u8/u16 must be zero-extended to 32-bits. (This variant is here to
-    /// accommodate Apple's deviation from the usual AArch64 ABI as defined by
-    /// ARM.)
-    ///
-    /// See also: <https://developer.apple.com/documentation/xcode/writing-arm64-code-for-apple-platforms#Pass-Arguments-to-Functions-Correctly>
-    ExtendTo32Bits,
-
-    /// Indicates that no sign- nor zero-extension is performed: if a value of
-    /// type with bitwidth M is passed as function argument or return value,
-    /// then M bits are copied into the least significant M bits, and the
-    /// remaining bits of the register (or word of memory) are untouched.
-    NoExtension,
+pub enum AbiKind {
+    AAPCS,
+    DarwinPCS,
+    Win64,
 }
 
 fn is_homogeneous_aggregate<'a, Ty, C>(cx: &C, arg: &mut ArgAbi<'a, Ty>) -> Option<Uniform>
@@ -45,15 +35,17 @@ where
     })
 }
 
-fn classify_ret<'a, Ty, C>(cx: &C, ret: &mut ArgAbi<'a, Ty>, param_policy: ParamExtension)
+fn classify_ret<'a, Ty, C>(cx: &C, ret: &mut ArgAbi<'a, Ty>, kind: AbiKind)
 where
     Ty: TyAbiInterface<'a, C> + Copy,
     C: HasDataLayout,
 {
     if !ret.layout.is_aggregate() {
-        match param_policy {
-            ParamExtension::ExtendTo32Bits => ret.extend_integer_width_to(32),
-            ParamExtension::NoExtension => {}
+        if kind == AbiKind::DarwinPCS {
+            // On Darwin, when returning an i8/i16, it must be sign-extended to 32 bits,
+            // and likewise a u8/u16 must be zero-extended to 32-bits.
+            // See also: <https://developer.apple.com/documentation/xcode/writing-arm64-code-for-apple-platforms#Pass-Arguments-to-Functions-Correctly>
+            ret.extend_integer_width_to(32)
         }
         return;
     }
@@ -70,15 +62,17 @@ where
     ret.make_indirect();
 }
 
-fn classify_arg<'a, Ty, C>(cx: &C, arg: &mut ArgAbi<'a, Ty>, param_policy: ParamExtension)
+fn classify_arg<'a, Ty, C>(cx: &C, arg: &mut ArgAbi<'a, Ty>, kind: AbiKind)
 where
     Ty: TyAbiInterface<'a, C> + Copy,
     C: HasDataLayout,
 {
     if !arg.layout.is_aggregate() {
-        match param_policy {
-            ParamExtension::ExtendTo32Bits => arg.extend_integer_width_to(32),
-            ParamExtension::NoExtension => {}
+        if kind == AbiKind::DarwinPCS {
+            // On Darwin, when passing an i8/i16, it must be sign-extended to 32 bits,
+            // and likewise a u8/u16 must be zero-extended to 32-bits.
+            // See also: <https://developer.apple.com/documentation/xcode/writing-arm64-code-for-apple-platforms#Pass-Arguments-to-Functions-Correctly>
+            arg.extend_integer_width_to(32);
         }
         return;
     }
@@ -87,27 +81,39 @@ where
         return;
     }
     let size = arg.layout.size;
-    let bits = size.bits();
-    if bits <= 128 {
-        arg.cast_to(Uniform { unit: Reg::i64(), total: size });
+    let align = if kind == AbiKind::AAPCS {
+        // When passing small aggregates by value, the AAPCS ABI mandates using the unadjusted
+        // alignment of the type (not including `repr(align)`).
+        // This matches behavior of `AArch64ABIInfo::classifyArgumentType` in Clang.
+        // See: <https://github.com/llvm/llvm-project/blob/5e691a1c9b0ad22689d4a434ddf4fed940e58dec/clang/lib/CodeGen/TargetInfo.cpp#L5816-L5823>
+        arg.layout.unadjusted_abi_align
+    } else {
+        arg.layout.align.abi
+    };
+    if size.bits() <= 128 {
+        if align.bits() == 128 {
+            arg.cast_to(Uniform { unit: Reg::i128(), total: size });
+        } else {
+            arg.cast_to(Uniform { unit: Reg::i64(), total: size });
+        }
         return;
     }
     arg.make_indirect();
 }
 
-pub fn compute_abi_info<'a, Ty, C>(cx: &C, fn_abi: &mut FnAbi<'a, Ty>, param_policy: ParamExtension)
+pub fn compute_abi_info<'a, Ty, C>(cx: &C, fn_abi: &mut FnAbi<'a, Ty>, kind: AbiKind)
 where
     Ty: TyAbiInterface<'a, C> + Copy,
     C: HasDataLayout,
 {
     if !fn_abi.ret.is_ignore() {
-        classify_ret(cx, &mut fn_abi.ret, param_policy);
+        classify_ret(cx, &mut fn_abi.ret, kind);
     }
 
     for arg in fn_abi.args.iter_mut() {
         if arg.is_ignore() {
             continue;
         }
-        classify_arg(cx, arg, param_policy);
+        classify_arg(cx, arg, kind);
     }
 }

compiler/rustc_target/src/abi/call/mod.rs

@@ -679,12 +679,14 @@ impl<'a, Ty> FnAbi<'a, Ty> {
                 }
             },
             "aarch64" => {
-                let param_policy = if cx.target_spec().is_like_osx {
-                    aarch64::ParamExtension::ExtendTo32Bits
+                let kind = if cx.target_spec().is_like_osx {
+                    aarch64::AbiKind::DarwinPCS
+                } else if cx.target_spec().is_like_windows {
+                    aarch64::AbiKind::Win64
                 } else {
-                    aarch64::ParamExtension::NoExtension
+                    aarch64::AbiKind::AAPCS
                 };
-                aarch64::compute_abi_info(cx, self, param_policy)
+                aarch64::compute_abi_info(cx, self, kind)
             }
             "amdgpu" => amdgpu::compute_abi_info(cx, self),
             "arm" => arm::compute_abi_info(cx, self),