fix is_single_fp_element for s390x and x86

compiler/rustc_abi/src/callconv/reg.rs

@@ -38,6 +38,7 @@ impl Reg {
     reg_ctor!(i64, Integer, 64);
     reg_ctor!(i128, Integer, 128);
 
+    reg_ctor!(f16, Float, 16);
     reg_ctor!(f32, Float, 32);
     reg_ctor!(f64, Float, 64);
     reg_ctor!(f128, Float, 128);

compiler/rustc_abi/src/layout/ty.rs

@@ -155,26 +155,6 @@ impl<'a, Ty> TyAndLayout<'a, Ty> {
         Ty::ty_and_layout_pointee_info_at(self, cx, offset)
     }
 
-    pub fn is_single_fp_element<C>(self, cx: &C) -> bool
-    where
-        Ty: TyAbiInterface<'a, C>,
-        C: HasDataLayout,
-    {
-        match self.backend_repr {
-            BackendRepr::Scalar(scalar) => {
-                matches!(scalar.primitive(), Primitive::Float(Float::F32 | Float::F64))
-            }
-            BackendRepr::Memory { .. } => {
-                if self.fields.count() == 1 && self.fields.offset(0).bytes() == 0 {
-                    self.field(cx, 0).is_single_fp_element(cx)
-                } else {
-                    false
-                }
-            }
-            _ => false,
-        }
-    }
-
     pub fn is_single_vector_element<C>(self, cx: &C, expected_size: Size) -> bool
     where
         Ty: TyAbiInterface<'a, C>,
@@ -301,6 +281,29 @@ impl<'a, Ty> TyAndLayout<'a, Ty> {
         found
     }
 
+    /// Finds the one field that is not a ZST.
+    /// Returns `None` if there are multiple non-ZST fields or only ZST-fields.
+    ///
+    /// Note that this function checks for ZSTs, not just 1-ZSTs.
+    pub fn non_zst_field_ignore_alignment<C>(&self, cx: &C) -> Option<(FieldIdx, Self)>
+    where
+        Ty: TyAbiInterface<'a, C> + Copy,
+    {
+        let mut found = None;
+        for field_idx in 0..self.fields.count() {
+            let field = self.field(cx, field_idx);
+            if field.is_zst() {
+                continue;
+            }
+            if found.is_some() {
+                // More than one non-ZST field.
+                return None;
+            }
+            found = Some((FieldIdx::from_usize(field_idx), field));
+        }
+        found
+    }
+
     /// If this type should match the ABI of the C `_Complex` type, returns the primitive that is
     /// used for its components.
     ///

compiler/rustc_codegen_llvm/src/va_arg.rs

@@ -474,7 +474,19 @@ fn emit_s390x_va_arg<'ll, 'tcx>(
     let padded_size = 8;
     let padding = padded_size - unpadded_size;
 
-    let gpr_type = indirect || !layout.is_single_fp_element(bx.cx);
+    // NOTE: if we ever allow aggregate types, this should handle structs with a single fp element.
+    let is_single_fp_element = |layout: TyAndLayout<'_>| -> bool {
+        match layout.layout.backend_repr() {
+            BackendRepr::Scalar(scalar) => match scalar.primitive() {
+                Primitive::Float(Float::F16 | Float::F32 | Float::F64) => true,
+                Primitive::Float(Float::F128) => false,
+                Primitive::Int(_, _) | Primitive::Pointer(_) => false,
+            },
+            _ => false,
+        }
+    };
+
+    let gpr_type = indirect || !is_single_fp_element(layout);
     let (max_regs, reg_count, reg_save_index, reg_padding) =
         if gpr_type { (5, gpr, 2, padding) } else { (4, fpr, 16, 0) };
 

compiler/rustc_target/src/callconv/s390x.rs

@@ -1,11 +1,44 @@
 // Reference: ELF Application Binary Interface s390x Supplement
 // https://github.com/IBM/s390x-abi
 
-use rustc_abi::{BackendRepr, HasDataLayout, TyAbiInterface};
+use rustc_abi::{BackendRepr, FieldsShape, HasDataLayout, Primitive, TyAbiInterface, TyAndLayout};
 
 use crate::callconv::{ArgAbi, FnAbi, Reg};
 use crate::spec::{Env, HasTargetSpec, Os};
 
+/// Is this a struct with a single float field?
+fn is_single_fp_element<'a, Ty, C>(mut layout: TyAndLayout<'a, Ty>, cx: &C) -> bool
+where
+    Ty: TyAbiInterface<'a, C> + Copy,
+    C: HasDataLayout,
+{
+    // Contrary to X86, trailing padding is allowed on s390x.
+
+    loop {
+        layout = layout.peel_transparent_wrappers(cx);
+
+        return match layout.backend_repr {
+            BackendRepr::Scalar(scalar) => match scalar.primitive() {
+                Primitive::Float(_) => true,
+                Primitive::Int(_, _) | Primitive::Pointer(_) => false,
+            },
+            BackendRepr::Memory { .. } => {
+                // A single-element array or union does not qualify.
+                if let FieldsShape::Arbitrary { .. } = layout.fields
+                    && layout.fields.count() == 1
+                    && layout.fields.offset(0).bytes() == 0
+                {
+                    layout = layout.field(cx, 0);
+                    continue;
+                } else {
+                    false
+                }
+            }
+            _ => false,
+        };
+    }
+}
+
 fn classify_ret<Ty>(ret: &mut ArgAbi<'_, Ty>) {
     let size = ret.layout.size;
     if size.bits() <= 128 && matches!(ret.layout.backend_repr, BackendRepr::SimdVector { .. }) {
@@ -65,8 +98,19 @@ where
         return;
     }
 
-    if arg.layout.is_single_fp_element(cx) {
+    if is_single_fp_element(arg.layout, cx) {
+        // Match GCC and Clang by explicitly passing padding, even though their behavior violates
+        // (our reading of) the specification, which says that:
+        //
+        // > Structures equivalent to a floating point type are passed in floating point registers.
+        // > A structure is equivalent to a floating point type if and only if it has exactly one
+        // > member, which is either of floating point type of itself a structure equivalent to a
+        // > floating point type.
+        //
+        // When the alignment is higher than 8, we pass the argument indirectly, which violates
+        // the specification but is consistent with GCC and Clang.
         match size.bytes() {
+            2 => arg.cast_to(Reg::f16()),
             4 => arg.cast_to(Reg::f32()),
             8 => arg.cast_to(Reg::f64()),
             _ => arg.make_indirect(),

compiler/rustc_target/src/callconv/x86.rs

@@ -5,6 +5,38 @@ use rustc_abi::{
 use crate::callconv::{ArgAttribute, FnAbi, PassMode, TyAbiInterface};
 use crate::spec::{HasTargetSpec, RustcAbi};
 
+/// Is this a struct with a single float field?
+fn is_single_fp_element<'a, Ty, C>(mut layout: TyAndLayout<'a, Ty>, cx: &C) -> bool
+where
+    Ty: TyAbiInterface<'a, C> + Copy,
+    C: HasDataLayout,
+{
+    // On X86 over-aligned structs are disqualified.
+    let outer_size = layout.layout.size();
+
+    loop {
+        layout = layout.peel_transparent_wrappers(cx);
+
+        return match layout.backend_repr {
+            BackendRepr::Scalar(scalar) => match scalar.primitive() {
+                Primitive::Float(float) => float.size() == outer_size,
+                Primitive::Int(_, _) | Primitive::Pointer(_) => false,
+            },
+            BackendRepr::Memory { .. } => {
+                // Structs, unions and arrays all qualify.
+                if let Some((_idx, field)) = layout.non_zst_field_ignore_alignment(cx) {
+                    // NOTE: alignment is not relevant here, checking for 1-ZST is incorrect.
+                    layout = field;
+                    continue;
+                } else {
+                    false
+                }
+            }
+            _ => false,
+        };
+    }
+}
+
 #[derive(PartialEq)]
 pub(crate) enum Flavor {
     General,
@@ -42,8 +74,9 @@ where
             {
                 // According to Clang, everyone but MSVC returns single-element
                 // float aggregates directly in a floating-point register.
-                if fn_abi.ret.layout.is_single_fp_element(cx) {
+                if is_single_fp_element(fn_abi.ret.layout, cx) {
                     match fn_abi.ret.layout.size.bytes() {
+                        2 => fn_abi.ret.cast_to(Reg::f16()),
                         4 => fn_abi.ret.cast_to(Reg::f32()),
                         8 => fn_abi.ret.cast_to(Reg::f64()),
                         _ => fn_abi.ret.make_indirect(),

tests/codegen-llvm/s390x-abi/single-fp-element.rs

@@ -0,0 +1,129 @@
+//@ add-minicore
+//@ needs-llvm-components: systemz
+//@ compile-flags: --target=s390x-unknown-linux-gnu -Copt-level=3 -Zmerge-functions=disabled
+#![crate_type = "lib"]
+#![feature(no_core, f16, f128)]
+#![no_core]
+
+extern crate minicore;
+use minicore::hint::black_box;
+use minicore::*;
+
+#[repr(C)]
+struct Wrapper<T>(T);
+
+// CHECK-LABEL: define void @plain_f16(half noundef %x)
+#[unsafe(no_mangle)]
+extern "C" fn plain_f16(x: f16) {
+    black_box(x);
+}
+
+// CHECK-LABEL: define void @wrapped_f16(half %0)
+#[unsafe(no_mangle)]
+extern "C" fn wrapped_f16(x: Wrapper<f16>) {
+    black_box(x);
+}
+
+// CHECK-LABEL: define void @plain_f32(float noundef %x)
+#[unsafe(no_mangle)]
+extern "C" fn plain_f32(x: f32) {
+    black_box(x);
+}
+
+// CHECK-LABEL: define void @wrapped_f32(float %0)
+#[unsafe(no_mangle)]
+extern "C" fn wrapped_f32(x: Wrapper<f32>) {
+    black_box(x);
+}
+
+// CHECK-LABEL: define void @plain_f64(double noundef %x)
+#[unsafe(no_mangle)]
+extern "C" fn plain_f64(x: f64) {
+    black_box(x);
+}
+
+// CHECK-LABEL: define void @wrapped_f64(double %0)
+#[unsafe(no_mangle)]
+extern "C" fn wrapped_f64(x: Wrapper<f64>) {
+    black_box(x);
+}
+
+// CHECK-LABEL: define void @plain_f128(ptr {{.*}}dereferenceable(16) %x)
+#[unsafe(no_mangle)]
+extern "C" fn plain_f128(x: f128) {
+    black_box(x);
+}
+
+// CHECK-LABEL: define void @wrapped_f128(ptr {{.*}}dereferenceable(16) %x)
+#[unsafe(no_mangle)]
+extern "C" fn wrapped_f128(x: Wrapper<f128>) {
+    black_box(x);
+}
+
+#[repr(transparent)]
+struct Transparent<T>(T);
+
+// CHECK-LABEL: define void @transparent_wrapped_f32(float %0)
+#[unsafe(no_mangle)]
+extern "C" fn transparent_wrapped_f32(x: Transparent<Wrapper<f32>>) {
+    black_box(x);
+}
+
+// CHECK-LABEL: define void @transparent_transparent_wrapped_f32(float %0)
+#[unsafe(no_mangle)]
+extern "C" fn transparent_transparent_wrapped_f32(x: Transparent<Transparent<Wrapper<f32>>>) {
+    black_box(x);
+}
+
+#[repr(C, align(8))]
+struct Aligned8Wrapper<T>(T);
+
+// CHECK-LABEL: define void @aligned_8_wrapped_f16(double %0)
+#[unsafe(no_mangle)]
+extern "C" fn aligned_8_wrapped_f16(x: Aligned8Wrapper<f16>) {
+    black_box(x);
+}
+
+// CHECK-LABEL: define void @aligned_8_wrapped_f32(double %0)
+#[unsafe(no_mangle)]
+extern "C" fn aligned_8_wrapped_f32(x: Aligned8Wrapper<f32>) {
+    black_box(x);
+}
+
+#[repr(C, align(16))]
+struct Aligned16Wrapper<T>(T);
+
+// CHECK-LABEL: define void @aligned_16_wrapped_f32(ptr {{.*}}dereferenceable(16)
+#[unsafe(no_mangle)]
+extern "C" fn aligned_16_wrapped_f32(x: Aligned16Wrapper<f32>) {
+    black_box(x);
+}
+
+#[repr(C)]
+union UnionWrapper<T: Copy> {
+    a: T,
+}
+
+// A repr(C) union does not count.
+//
+// CHECK-LABEL: define void @union_wrapped_f32(i32 %0)
+#[unsafe(no_mangle)]
+extern "C" fn union_wrapped_f32(x: UnionWrapper<f32>) {
+    black_box(x);
+}
+
+// But a repr(transparent) union does.
+//
+// CHECK-LABEL: define void @maybe_uninit_f32(float %x)
+#[unsafe(no_mangle)]
+extern "C" fn maybe_uninit_f32(x: MaybeUninit<f32>) {
+    black_box(x);
+}
+
+// A single-element array also does not count.
+//
+// CHECK-LABEL: define void @array_f32(i32 %0)
+#[unsafe(no_mangle)]
+extern "C" fn array_f32(x: [f32; 1]) {
+    black_box(x);
+}