rust_for_linux: -Zreg-struct-return commandline flag for X86 (#116973)

Author: azhogin

compiler/rustc_codegen_gcc/src/context.rs

@@ -542,7 +542,10 @@ impl<'gcc, 'tcx> HasWasmCAbiOpt for CodegenCx<'gcc, 'tcx> {
 
 impl<'gcc, 'tcx> HasX86AbiOpt for CodegenCx<'gcc, 'tcx> {
     fn x86_abi_opt(&self) -> X86Abi {
-        X86Abi { regparm: self.tcx.sess.opts.unstable_opts.regparm }
+        X86Abi {
+            regparm: self.tcx.sess.opts.unstable_opts.regparm,
+            reg_struct_return: self.tcx.sess.opts.unstable_opts.reg_struct_return,
+        }
     }
 }
 

compiler/rustc_interface/src/tests.rs

@@ -836,6 +836,7 @@ fn test_unstable_options_tracking_hash() {
     tracked!(profile_emit, Some(PathBuf::from("abc")));
     tracked!(profile_sample_use, Some(PathBuf::from("abc")));
     tracked!(profiler_runtime, "abc".to_string());
+    tracked!(reg_struct_return, true);
     tracked!(regparm, Some(3));
     tracked!(relax_elf_relocations, Some(true));
     tracked!(remap_cwd_prefix, Some(PathBuf::from("abc")));

compiler/rustc_middle/src/ty/layout.rs

@@ -548,7 +548,10 @@ impl<'tcx> HasWasmCAbiOpt for TyCtxt<'tcx> {
 
 impl<'tcx> HasX86AbiOpt for TyCtxt<'tcx> {
     fn x86_abi_opt(&self) -> X86Abi {
-        X86Abi { regparm: self.sess.opts.unstable_opts.regparm }
+        X86Abi {
+            regparm: self.sess.opts.unstable_opts.regparm,
+            reg_struct_return: self.sess.opts.unstable_opts.reg_struct_return,
+        }
     }
 }
 

compiler/rustc_session/src/options.rs

@@ -2000,6 +2000,9 @@ options! {
         "enable queries of the dependency graph for regression testing (default: no)"),
     randomize_layout: bool = (false, parse_bool, [TRACKED],
         "randomize the layout of types (default: no)"),
+    reg_struct_return: bool = (false, parse_bool, [TRACKED],
+        "On x86-32 targets, it overrides the default ABI to return small structs in registers.
+        It is UNSOUND to link together crates that use different values for this flag!"),
     regparm: Option<u32> = (None, parse_opt_number, [TRACKED],
         "On x86-32 targets, setting this to N causes the compiler to pass N arguments \
         in registers EAX, EDX, and ECX instead of on the stack.\

compiler/rustc_target/src/callconv/mod.rs

@@ -654,7 +654,9 @@ impl<'a, Ty> FnAbi<'a, Ty> {
                     }
                     _ => (x86::Flavor::General, None),
                 };
-                x86::compute_abi_info(cx, self, x86::X86Options { flavor, regparm });
+                let reg_struct_return = cx.x86_abi_opt().reg_struct_return;
+                let opts = x86::X86Options { flavor, regparm, reg_struct_return };
+                x86::compute_abi_info(cx, self, opts);
             }
             "x86_64" => match abi {
                 spec::abi::Abi::SysV64 { .. } => x86_64::compute_abi_info(cx, self),
@@ -735,6 +737,7 @@ impl<'a, Ty> FnAbi<'a, Ty> {
                     x86::X86Options {
                         flavor: x86::Flavor::General,
                         regparm: cx.x86_abi_opt().regparm,
+                        reg_struct_return: cx.x86_abi_opt().reg_struct_return,
                     },
                     true,
                 );

compiler/rustc_target/src/callconv/x86.rs

@@ -11,6 +11,7 @@ pub(crate) enum Flavor {
 pub(crate) struct X86Options {
     pub flavor: Flavor,
     pub regparm: Option<u32>,
+    pub reg_struct_return: bool,
 }
 
 pub(crate) fn compute_abi_info<'a, Ty, C>(cx: &C, fn_abi: &mut FnAbi<'a, Ty>, opts: X86Options)
@@ -28,7 +29,7 @@ where
             // https://www.angelcode.com/dev/callconv/callconv.html
             // Clang's ABI handling is in lib/CodeGen/TargetInfo.cpp
             let t = cx.target_spec();
-            if t.abi_return_struct_as_int {
+            if t.abi_return_struct_as_int || opts.reg_struct_return {
                 // According to Clang, everyone but MSVC returns single-element
                 // float aggregates directly in a floating-point register.
                 if !t.is_like_msvc && fn_abi.ret.layout.is_single_fp_element(cx) {

compiler/rustc_target/src/spec/mod.rs

@@ -2102,6 +2102,8 @@ pub struct X86Abi {
     /// On x86-32 targets, the regparm N causes the compiler to pass arguments
     /// in registers EAX, EDX, and ECX instead of on the stack.
     pub regparm: Option<u32>,
+    /// Override the default ABI to return small structs in registers
+    pub reg_struct_return: bool,
 }
 
 pub trait HasX86AbiOpt {

compiler/rustc_ty_utils/src/abi.rs

@@ -730,6 +730,14 @@ fn fn_abi_adjust_for_abi<'tcx>(
 
             match arg.layout.abi {
                 Abi::Aggregate { .. } => {}
+                Abi::ScalarPair(..) => {
+                    // FIXME: It is strange that small struct optimizations are enabled
+                    // for 3-fields and more, but disabled for 2-fields (represented by ScalarPair)
+                    // Here 2-fields optimizations are enabled only for return value
+                    if !tcx.sess.opts.unstable_opts.reg_struct_return || arg_idx.is_some() {
+                        return;
+                    }
+                }
 
                 // This is a fun case! The gist of what this is doing is
                 // that we want callers and callees to always agree on the
@@ -761,12 +769,22 @@ fn fn_abi_adjust_for_abi<'tcx>(
             }
             // Compute `Aggregate` ABI.
 
-            let is_indirect_not_on_stack =
-                matches!(arg.mode, PassMode::Indirect { on_stack: false, .. });
-            assert!(is_indirect_not_on_stack, "{:?}", arg);
+            let is_indirect_not_on_stack = !matches!(arg.layout.abi, Abi::Aggregate { .. })
+                || matches!(arg.mode, PassMode::Indirect { on_stack: false, .. });
+            assert!(is_indirect_not_on_stack, "is_indirect_not_on_stack: {:?}", arg);
 
             let size = arg.layout.size;
-            if !arg.layout.is_unsized() && size <= Pointer(AddressSpace::DATA).size(cx) {
+            let ptr_size = Pointer(AddressSpace::DATA).size(cx);
+
+            // In x86 we may return 2x pointer sized struct as i64
+            let reg_struct_return_case = tcx.sess.target.arch == "x86"
+                && arg_idx.is_none()
+                && tcx.sess.opts.unstable_opts.reg_struct_return
+                && abi != SpecAbi::RustIntrinsic;
+
+            if !arg.layout.is_unsized()
+                && (size <= ptr_size || (reg_struct_return_case && (size <= 2 * ptr_size)))
+            {
                 // We want to pass small aggregates as immediates, but using
                 // an LLVM aggregate type for this leads to bad optimizations,
                 // so we pick an appropriately sized integer type instead.

tests/codegen/reg-struct-return.rs

@@ -0,0 +1,63 @@
+// Checks how `reg-struct-return` flag works with different calling conventions:
+// Return struct with 8/16/32/64 bit size will be converted into i8/i16/i32/i64
+// (like abi_return_struct_as_int target spec).
+// x86 only.
+
+//@ compile-flags: --target i686-unknown-linux-gnu -Zreg-struct-return -O -C no-prepopulate-passes
+//@ needs-llvm-components: x86
+
+#![crate_type = "lib"]
+#![no_std]
+#![no_core]
+#![feature(no_core, lang_items)]
+
+#[lang = "sized"]
+trait Sized {}
+#[lang = "copy"]
+trait Copy {}
+
+#[repr(C)]
+pub struct Foo {
+    x: u32,
+    y: u32,
+}
+
+pub mod tests {
+    use Foo;
+
+    // CHECK: i64 @f1()
+    #[no_mangle]
+    pub extern "fastcall" fn f1() -> Foo {
+        Foo { x: 1, y: 2 }
+    }
+
+    // CHECK: i64 @f2()
+    #[no_mangle]
+    pub extern "Rust" fn f2() -> Foo {
+        Foo { x: 1, y: 2 }
+    }
+
+    // CHECK: i64 @f3()
+    #[no_mangle]
+    pub extern "C" fn f3() -> Foo {
+        Foo { x: 1, y: 2 }
+    }
+
+    // CHECK: i64 @f4()
+    #[no_mangle]
+    pub extern "cdecl" fn f4() -> Foo {
+        Foo { x: 1, y: 2 }
+    }
+
+    // CHECK: i64 @f5()
+    #[no_mangle]
+    pub extern "stdcall" fn f5() -> Foo {
+        Foo { x: 1, y: 2 }
+    }
+
+    // CHECK: i64 @f6()
+    #[no_mangle]
+    pub extern "thiscall" fn f6() -> Foo {
+        Foo { x: 1, y: 2 }
+    }
+}