x86: Correctly pass large integers and structs in registers with -Zregparm (v2)

compiler/rustc_target/src/callconv/x86.rs

@@ -2,7 +2,7 @@ use rustc_abi::{
     AddressSpace, Align, BackendRepr, HasDataLayout, Primitive, Reg, RegKind, TyAndLayout,
 };
 
-use crate::callconv::{ArgAttribute, FnAbi, PassMode, TyAbiInterface};
+use crate::callconv::{ArgAttribute, FnAbi, PassMode, TyAbiInterface, Uniform};
 use crate::spec::{HasTargetSpec, RustcAbi};
 
 #[derive(PartialEq)]
@@ -58,6 +58,13 @@ where
         }
     }
 
+    // Some calling conversions pass arguments in registers.
+    // This is set with -Zregparm, or the fastcall/vectorcall (always 2 registers)
+    let mut free_regs = opts.regparm.unwrap_or(0).into();
+    if opts.flavor == Flavor::FastcallOrVectorcall {
+        free_regs = 2;
+    }
+
     for arg in fn_abi.args.iter_mut() {
         if arg.is_ignore() || !arg.layout.is_sized() {
             continue;
@@ -72,7 +79,21 @@ where
         let align_4 = Align::from_bytes(4).unwrap();
         let align_16 = Align::from_bytes(16).unwrap();
 
+        let size_in_regs = arg.layout.size.bits().div_ceil(32);
+        let mut pass_in_reg = free_regs >= size_in_regs;
+
+        // In fastcall/vectorcall only 32-bit integers can be passed in registers.
+        if opts.flavor == Flavor::FastcallOrVectorcall && size_in_regs > 1 {
+            pass_in_reg = false;
+            // Once we've had an argument which doesn't fit, don't try to fit any more.
+            free_regs = 0;
+        }
+
         if arg.layout.is_aggregate() {
+            if opts.flavor == Flavor::FastcallOrVectorcall && size_in_regs > 1 {
+                pass_in_reg = false;
+            }
+
             // We need to compute the alignment of the `byval` argument. The rules can be found in
             // `X86_32ABIInfo::getTypeStackAlignInBytes` in Clang's `TargetInfo.cpp`. Summarized
             // here, they are:
@@ -119,83 +140,40 @@ where
                 align_4
             };
 
-            arg.pass_by_stack_offset(Some(byval_align));
+            if pass_in_reg {
+                arg.cast_to(Uniform::new(Reg::i32(), arg.layout.size));
+            } else {
+                arg.pass_by_stack_offset(Some(byval_align))
+            }
         } else {
-            arg.extend_integer_width_to(32);
-        }
-    }
-
-    fill_inregs(cx, fn_abi, opts, false);
-}
-
-pub(crate) fn fill_inregs<'a, Ty, C>(
-    cx: &C,
-    fn_abi: &mut FnAbi<'a, Ty>,
-    opts: X86Options,
-    rust_abi: bool,
-) where
-    Ty: TyAbiInterface<'a, C> + Copy,
-{
-    if opts.flavor != Flavor::FastcallOrVectorcall && opts.regparm.is_none_or(|x| x == 0) {
-        return;
-    }
-    // Mark arguments as InReg like clang does it,
-    // so our fastcall/vectorcall is compatible with C/C++ fastcall/vectorcall.
-
-    // Clang reference: lib/CodeGen/TargetInfo.cpp
-    // See X86_32ABIInfo::shouldPrimitiveUseInReg(), X86_32ABIInfo::updateFreeRegs()
-
-    // IsSoftFloatABI is only set to true on ARM platforms,
-    // which in turn can't be x86?
-
-    // 2 for fastcall/vectorcall, regparm limited by 3 otherwise
-    let mut free_regs = opts.regparm.unwrap_or(2).into();
+            let unit = arg.layout.homogeneous_aggregate(cx).unwrap().unit().unwrap();
 
-    // For types generating PassMode::Cast, InRegs will not be set.
-    // Maybe, this is a FIXME
-    let has_casts = fn_abi.args.iter().any(|arg| matches!(arg.mode, PassMode::Cast { .. }));
-    if has_casts && rust_abi {
-        return;
-    }
-
-    for arg in fn_abi.args.iter_mut() {
-        let attrs = match arg.mode {
-            PassMode::Ignore | PassMode::Indirect { attrs: _, meta_attrs: None, on_stack: _ } => {
-                continue;
+            // Fastcall/regparm won't pass floats in registers
+            // (though regparm _will_ pass f]oat-containing structs in registers)
+            if unit.kind != RegKind::Integer {
+                pass_in_reg = false;
             }
-            PassMode::Direct(ref mut attrs) => attrs,
-            PassMode::Pair(..)
-            | PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ }
-            | PassMode::Cast { .. } => {
-                unreachable!("x86 shouldn't be passing arguments by {:?}", arg.mode)
-            }
-        };
 
-        // At this point we know this must be a primitive of sorts.
-        let unit = arg.layout.homogeneous_aggregate(cx).unwrap().unit().unwrap();
-        assert_eq!(unit.size, arg.layout.size);
-        if matches!(unit.kind, RegKind::Float | RegKind::Vector { .. }) {
-            continue;
-        }
-
-        let size_in_regs = arg.layout.size.bits().div_ceil(32);
-
-        if size_in_regs == 0 {
-            continue;
-        }
-
-        if size_in_regs > free_regs {
-            break;
-        }
-
-        free_regs -= size_in_regs;
-
-        if arg.layout.size.bits() <= 32 && unit.kind == RegKind::Integer {
-            attrs.set(ArgAttribute::InReg);
+            arg.extend_integer_width_to(32);
         }
 
-        if free_regs == 0 {
-            break;
+        // Set the InReg annotation if we're passing by register
+        if pass_in_reg {
+            match arg.mode {
+                PassMode::Ignore
+                | PassMode::Indirect { attrs: _, meta_attrs: None, on_stack: _ } => {}
+                PassMode::Cast { pad_i32: _, ref mut cast } => {
+                    cast.attrs.set(ArgAttribute::InReg);
+                }
+                PassMode::Direct(ref mut attrs) => {
+                    attrs.set(ArgAttribute::InReg);
+                }
+                PassMode::Pair(..)
+                | PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+                    unreachable!("x86 shouldn't be passing arguments by {:?}", arg.mode)
+                }
+            };
+            free_regs -= size_in_regs;
         }
     }
 }

compiler/rustc_target/src/callconv/x86_win32.rs

@@ -1,6 +1,7 @@
-use rustc_abi::{Align, HasDataLayout, Reg, TyAbiInterface};
+use rustc_abi::{Align, HasDataLayout, Reg, RegKind};
 
-use crate::callconv::FnAbi;
+use crate::callconv::x86::Flavor;
+use crate::callconv::{ArgAttribute, FnAbi, PassMode, TyAbiInterface, Uniform};
 use crate::spec::HasTargetSpec;
 
 pub(crate) fn compute_abi_info<'a, Ty, C>(
@@ -41,6 +42,13 @@ pub(crate) fn compute_abi_info<'a, Ty, C>(
         }
     }
 
+    // Some calling conversions pass arguments in registers.
+    // This is set with -Zregparm, or the fastcall/vectorcall (always 2 registers)
+    let mut free_regs = opts.regparm.unwrap_or(0).into();
+    if opts.flavor == Flavor::FastcallOrVectorcall {
+        free_regs = 2;
+    }
+
     for arg in fn_abi.args.iter_mut() {
         if arg.is_ignore() || !arg.layout.is_sized() {
             continue;
@@ -51,6 +59,16 @@ pub(crate) fn compute_abi_info<'a, Ty, C>(
             continue;
         }
 
+        let size_in_regs = arg.layout.size.bits().div_ceil(32);
+        let mut pass_in_reg = free_regs >= size_in_regs;
+
+        // In fastcall/vectorcall only 32-bit integers can be passed in registers.
+        if opts.flavor == Flavor::FastcallOrVectorcall && size_in_regs > 1 {
+            pass_in_reg = false;
+            // Once we've had an argument which doesn't fit, don't try to fit any more.
+            free_regs = 0;
+        }
+
         // FIXME: MSVC 2015+ will pass the first 3 vector arguments in [XYZ]MM0-2
         // See https://reviews.llvm.org/D72114 for Clang behavior
 
@@ -73,14 +91,45 @@ pub(crate) fn compute_abi_info<'a, Ty, C>(
             );
             arg.make_indirect();
         } else if arg.layout.is_aggregate() {
+            if opts.flavor == Flavor::FastcallOrVectorcall && size_in_regs > 1 {
+                pass_in_reg = false;
+            }
             // Alignment of the `byval` argument.
             // The rules can be found in `X86_32ABIInfo::getTypeStackAlignInBytes` in Clang's `TargetInfo.cpp`.
             let byval_align = align_4;
-            arg.pass_by_stack_offset(Some(byval_align));
+            if pass_in_reg {
+                arg.cast_to(Uniform::new(Reg::i32(), arg.layout.size));
+            } else {
+                arg.pass_by_stack_offset(Some(byval_align))
+            }
         } else {
+            let unit = arg.layout.homogeneous_aggregate(cx).unwrap().unit().unwrap();
+
+            // Fastcall/regparm won't pass floats in registers
+            // (though regparm _will_ pass f]oat-containing structs in registers)
+            if unit.kind != RegKind::Integer {
+                pass_in_reg = false;
+            }
+
             arg.extend_integer_width_to(32);
         }
+        // Set the InReg annotation if we're passing by register
+        if pass_in_reg {
+            match arg.mode {
+                PassMode::Ignore
+                | PassMode::Indirect { attrs: _, meta_attrs: None, on_stack: _ } => {}
+                PassMode::Cast { pad_i32: _, ref mut cast } => {
+                    cast.attrs.set(ArgAttribute::InReg);
+                }
+                PassMode::Direct(ref mut attrs) => {
+                    attrs.set(ArgAttribute::InReg);
+                }
+                PassMode::Pair(..)
+                | PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+                    unreachable!("x86 shouldn't be passing arguments by {:?}", arg.mode)
+                }
+            };
+            free_regs -= size_in_regs;
+        }
     }
-
-    super::x86::fill_inregs(cx, fn_abi, opts, false);
 }

tests/assembly-llvm/regparm-module-flag.rs

@@ -16,9 +16,58 @@
 extern crate minicore;
 use minicore::*;
 
+#[repr(C)]
+struct ThreeRegStruct {
+    a: i64,
+    b: i32,
+}
+
 unsafe extern "C" {
     fn memset(p: *mut c_void, val: i32, len: usize) -> *mut c_void;
     fn non_builtin_memset(p: *mut c_void, val: i32, len: usize) -> *mut c_void;
+    fn test_i64_arg(s: i64) -> i64;
+    fn test_struct_arg(s: ThreeRegStruct);
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn test_i64() -> i64 {
+    // REGPARM1-LABEL: test_i64
+    // REGPARM1: pushl
+    // REGPARM1: pushl
+    // REGPARM1: calll test_i64_arg
+
+    // REGPARM2-LABEL: test_i64
+    // REGPARM2: movl $42, %eax
+    // REGPARM2: xorl %edx, %edx
+    // REGPARM2: jmp test_i64_arg
+
+    // REGPARM3-LABEL: test_i64
+    // REGPARM3: movl $42, %eax
+    // REGPARM3: xorl %edx, %edx
+    // REGPARM3: jmp test_i64_arg
+    unsafe { test_i64_arg(42) }
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn test_struct() {
+    // REGPARM1-LABEL: test_struct
+    // REGPARM1: movl $0, {{.*}}(%esp)
+    // REGPARM1: movl $42, {{.*}}(%esp)
+    // REGPARM1: movl $1, {{.*}}(%esp)
+    // REGPARM1: calll test_struct_arg
+
+    // REGPARM2-LABEL: test_struct
+    // REGPARM2: movl $0, {{.*}}(%esp)
+    // REGPARM2: movl $42, {{.*}}(%esp)
+    // REGPARM2: movl $1, {{.*}}(%esp)
+    // REGPARM2: calll test_struct_arg
+
+    // REGPARM3-LABEL: test_struct
+    // REGPARM3: movl $42, %eax
+    // REGPARM3: xorl %edx, %edx
+    // REGPARM3: movl $1, %ecx
+    // REGPARM3: jmp test_struct_arg
+    unsafe { test_struct_arg(ThreeRegStruct { a: 42, b: 1 }) }
 }
 
 #[unsafe(no_mangle)]

tests/codegen-llvm/regparm-inreg.rs

@@ -121,4 +121,31 @@ pub mod tests {
     // regparm3-SAME: i32 inreg noundef %_4)
     #[no_mangle]
     pub extern "C" fn f12(_: i32, _: __m256, _: i32, _: i32) {}
+
+    // regparm0: @f13(i64 noundef %_1)
+    // regparm1: @f13(i64 noundef %_1)
+    // regparm2: @f13(i64 inreg noundef %_1)
+    // regparm3: @f13(i64 inreg noundef %_1)
+    #[no_mangle]
+    pub extern "C" fn f13(_: i64) {}
+
+    // regparm0: @f14(ptr noalias nofree noundef byval([8 x i8]) align 4 captures(address) dereferenceable(8) %_1)
+    // regparm1: @f14(ptr noalias nofree noundef byval([8 x i8]) align 4 captures(address) dereferenceable(8) %_1)
+    // regparm2: @f14([2 x i32] inreg %0)
+    // regparm3: @f14([2 x i32] inreg %0)
+    #[no_mangle]
+    pub extern "C" fn f14(_: S2) {}
+
+    #[repr(C)]
+    struct S3 {
+        x1: i32,
+        x2: i32,
+        x3: i32,
+    }
+    // regparm0: @f15(ptr noalias nofree noundef byval([12 x i8]) align 4 captures(address) dereferenceable(12) %_1)
+    // regparm1: @f15(ptr noalias nofree noundef byval([12 x i8]) align 4 captures(address) dereferenceable(12) %_1)
+    // regparm2: @f15(ptr noalias nofree noundef byval([12 x i8]) align 4 captures(address) dereferenceable(12) %_1)
+    // regparm3: @f15([3 x i32] inreg %0)
+    #[no_mangle]
+    pub extern "C" fn f15(_: S3) {}
 }