cranelift: Always consider sret arguments used (#8664) (#8665)

* cranelift: Always consider sret arguments used

In #8438 we stopped emitting register bindings for unused arguments,
based on the use-counts from `compute_use_states`. However, that doesn't
count the use of the struct-return argument that's automatically added
after lowering when the `rets` instruction is generated in the epilogue.
As a result, using a struct-return argument caused register allocation
to panic due to the VReg not being defined anywhere.

This commit adds a use to the struct-return argument so that it's always
available in the epilogue.

Fixes #8659

* Review comments

Co-authored-by: Jamey Sharp <[email protected]>

cranelift/codegen/src/machinst/lower.rs

@@ -398,7 +398,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
         }
 
         // Find the sret register, if it's used.
-        let mut sret_reg = None;
+        let mut sret_param = None;
         for ret in vcode.abi().signature().returns.iter() {
             if ret.purpose == ArgumentPurpose::StructReturn {
                 let entry_bb = f.stencil.layout.entry_block().unwrap();
@@ -409,18 +409,21 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
                     .zip(vcode.abi().signature().params.iter())
                 {
                     if sig_param.purpose == ArgumentPurpose::StructReturn {
-                        let regs = value_regs[param];
-                        assert!(regs.len() == 1);
-
-                        assert!(sret_reg.is_none());
-                        sret_reg = Some(regs);
+                        assert!(sret_param.is_none());
+                        sret_param = Some(param);
                     }
                 }
 
-                assert!(sret_reg.is_some());
+                assert!(sret_param.is_some());
             }
         }
 
+        let sret_reg = sret_param.map(|param| {
+            let regs = value_regs[param];
+            assert!(regs.len() == 1);
+            regs
+        });
+
         // Compute instruction colors, find constant instructions, and find instructions with
         // side-effects, in one combined pass.
         let mut cur_color = 0;
@@ -449,7 +452,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
             block_end_colors[bb] = InstColor::new(cur_color);
         }
 
-        let value_ir_uses = Self::compute_use_states(f);
+        let value_ir_uses = Self::compute_use_states(f, sret_param);
 
         Ok(Lower {
             f,
@@ -482,7 +485,10 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
     /// Pre-analysis: compute `value_ir_uses`. See comment on
     /// `ValueUseState` for a description of what this analysis
     /// computes.
-    fn compute_use_states<'a>(f: &'a Function) -> SecondaryMap<Value, ValueUseState> {
+    fn compute_use_states<'a>(
+        f: &'a Function,
+        sret_param: Option<Value>,
+    ) -> SecondaryMap<Value, ValueUseState> {
         // We perform the analysis without recursion, so we don't
         // overflow the stack on long chains of ops in the input.
         //
@@ -502,6 +508,12 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
 
         let mut value_ir_uses = SecondaryMap::with_default(ValueUseState::Unused);
 
+        if let Some(sret_param) = sret_param {
+            // There's an implicit use of the struct-return parameter in each
+            // copy of the function epilogue, which we count here.
+            value_ir_uses[sret_param] = ValueUseState::Multiple;
+        }
+
         // Stack of iterators over Values as we do DFS to mark
         // Multiple-state subtrees. The iterator type is whatever is
         // returned by `uses` below.

cranelift/filetests/filetests/isa/aarch64/issue-8659.clif

@@ -0,0 +1,16 @@
+test compile precise-output
+target aarch64
+
+function u0:11(i64 sret) system_v {
+block0(v0: i64):
+    return
+}
+
+; VCode:
+; block0:
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   ret
+

cranelift/filetests/filetests/isa/x64/issue-8659.clif

@@ -0,0 +1,27 @@
+test compile precise-output
+target x86_64
+
+function u0:11(i64 sret) system_v {
+block0(v0: i64):
+    return
+}
+
+; VCode:
+;   pushq   %rbp
+;   movq    %rsp, %rbp
+; block0:
+;   movq    %rdi, %rax
+;   movq    %rbp, %rsp
+;   popq    %rbp
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   pushq %rbp
+;   movq %rsp, %rbp
+; block1: ; offset 0x4
+;   movq %rdi, %rax
+;   movq %rbp, %rsp
+;   popq %rbp
+;   retq
+