Repository: bytecodealliance/wasmtime
Pull request #8475 (2024-04-25 17:42:29)

cranelift: Simplify allocating ABI temp VRegs
This merges several existing passes over the ABI signature that handle incoming arguments:

1. Callee::temps_needed counted how many and what type of temporaries would be needed in some later steps, and Lower::lower allocated them.
2. Callee::init duplicated the same logic to determine where to save those temporaries for later.
3. Various places used the saved temporaries.

There were two kinds of temporaries allocated. One kind was the return-area pointer, which must be allocated before any instructions are lowered so that the return instructions will all use the same VReg that the prologue defines. After this commit, that allocation still happens at the same point, but now it follows the pattern of other ABI glue and uses a VRegAllocator directly.

The other kind was a truly local kind of temporary used in only one uncommon path in gen_copy_arg_to_regs. That is unnecessary because we already allocate temporaries as-needed in all other paths in that method, so this commit just follows the same pattern.

This change avoids heap-allocating two temporary Vecs during init, as well as a third Vec that was live for the entire lifetime of the Callee. The latter almost always contained only a bunch of `None` values, and was only used at the end of lowering.


Patch:
diff --git a/cranelift/codegen/src/machinst/abi.rs b/cranelift/codegen/src/machinst/abi.rs
index a413ae98e..a5de3824d 100644
--- a/cranelift/codegen/src/machinst/abi.rs
+++ b/cranelift/codegen/src/machinst/abi.rs
@@ -1097,9 +1097,7 @@ pub struct Callee<M: ABIMachineSpec> {
     /// Finalized frame layout for this function.
     frame_layout: Option<FrameLayout>,
     /// The register holding the return-area pointer, if needed.
-    ret_area_ptr: Option<Writable<Reg>>,
-    /// Temp registers required for argument setup, if needed.
-    arg_temp_reg: Vec<Option<Writable<Reg>>>,
+    ret_area_ptr: Option<Reg>,
     /// Calling convention this function expects.
     call_conv: isa::CallConv,
     /// The settings controlling this function's compilation.
@@ -1241,7 +1239,6 @@ impl<M: ABIMachineSpec> Callee<M> {
             reg_args: vec![],
             frame_layout: None,
             ret_area_ptr: None,
-            arg_temp_reg: vec![],
             call_conv,
             flags,
             isa_flags: isa_flags.clone(),
@@ -1406,45 +1403,19 @@ impl<M: ABIMachineSpec> Callee<M> {
         &self.ir_sig
     }
 
-    /// Does the ABI-body code need temp registers (and if so, of what type)?
-    /// They will be provided to `init()` as the `temps` arg if so.
-    pub fn temps_needed(&self, sigs: &SigSet) -> Vec<Type> {
-        let mut temp_tys = vec![];
-        for arg in sigs.args(self.sig) {
-            match arg {
-                &ABIArg::ImplicitPtrArg { pointer, .. } => match &pointer {
-                    &ABIArgSlot::Reg { .. } => {}
-                    &ABIArgSlot::Stack { ty, .. } => {
-                        temp_tys.push(ty);
-                    }
-                },
-                _ => {}
-            }
-        }
-        if sigs[self.sig].stack_ret_arg.is_some() {
-            temp_tys.push(M::word_type());
-        }
-        temp_tys
-    }
-
     /// Initialize. This is called after the Callee is constructed because it
-    /// may be provided with a vector of temp vregs, which can only be allocated
-    /// once the lowering context exists.
-    pub fn init(&mut self, sigs: &SigSet, temps: Vec<Writable<Reg>>) {
-        let mut temps_iter = temps.into_iter();
-        for arg in sigs.args(self.sig) {
-            let temp = match arg {
-                &ABIArg::ImplicitPtrArg { pointer, .. } => match &pointer {
-                    &ABIArgSlot::Reg { .. } => None,
-                    &ABIArgSlot::Stack { .. } => Some(temps_iter.next().unwrap()),
-                },
-                _ => None,
-            };
-            self.arg_temp_reg.push(temp);
-        }
+    /// may allocate a temp vreg, which can only be allocated once the lowering
+    /// context exists.
+    pub fn init_retval_area(
+        &mut self,
+        sigs: &SigSet,
+        vregs: &mut VRegAllocator<M::I>,
+    ) -> CodegenResult<()> {
         if sigs[self.sig].stack_ret_arg.is_some() {
-            self.ret_area_ptr = Some(temps_iter.next().unwrap());
+            let ret_area_ptr = vregs.alloc(M::word_type())?;
+            self.ret_area_ptr = Some(ret_area_ptr.only_reg().unwrap());
         }
+        Ok(())
     }
 
     /// Accumulate outgoing arguments.
@@ -1574,9 +1545,9 @@ impl<M: ABIMachineSpec> Callee<M> {
                         tmp
                     }
                     &ABIArgSlot::Stack { offset, ty, .. } => {
-                        // In this case we need a temp register to hold the address.
-                        // This was allocated in the `init` routine.
-                        let addr_reg = self.arg_temp_reg[idx].unwrap();
+                        let addr_reg = writable_value_regs(vregs.alloc_with_deferred_error(ty))
+                            .only_reg()
+                            .unwrap();
                         insts.push(M::gen_load_stack(
                             StackAMode::IncomingArg(offset, sigs[self.sig].sized_stack_arg_space),
                             addr_reg,
@@ -1676,7 +1647,7 @@ impl<M: ABIMachineSpec> Callee<M> {
                                 _ => {}
                             };
                             ret.push(M::gen_store_base_offset(
-                                self.ret_area_ptr.unwrap().to_reg(),
+                                self.ret_area_ptr.unwrap(),
                                 off,
                                 from_reg,
                                 ty,
@@ -1706,17 +1677,14 @@ impl<M: ABIMachineSpec> Callee<M> {
         vregs: &mut VRegAllocator<M::I>,
     ) -> Option<M::I> {
         if let Some(i) = sigs[self.sig].stack_ret_arg {
-            let insts = self.gen_copy_arg_to_regs(
-                sigs,
-                i.into(),
-                ValueRegs::one(self.ret_area_ptr.unwrap()),
-                vregs,
-            );
+            let ret_area_ptr = Writable::from_reg(self.ret_area_ptr.unwrap());
+            let insts =
+                self.gen_copy_arg_to_regs(sigs, i.into(), ValueRegs::one(ret_area_ptr), vregs);
             insts.into_iter().next().map(|inst| {
                 trace!(
                     "gen_retval_area_setup: inst {:?}; ptr reg is {:?}",
                     inst,
-                    self.ret_area_ptr.unwrap().to_reg()
+                    ret_area_ptr.to_reg()
                 );
                 inst
             })
@@ -2390,7 +2358,7 @@ impl<M: ABIMachineSpec> CallSite<M> {
                 "if the tail callee has a return pointer, then the tail caller \
                  must as well",
             );
-            self.gen_arg(ctx, i.into(), ValueRegs::one(ret_area_ptr.to_reg()));
+            self.gen_arg(ctx, i.into(), ValueRegs::one(ret_area_ptr));
         }
     }
 
diff --git a/cranelift/codegen/src/machinst/lower.rs b/cranelift/codegen/src/machinst/lower.rs
index a6b3bc501..cbd0491fc 100644
--- a/cranelift/codegen/src/machinst/lower.rs
+++ b/cranelift/codegen/src/machinst/lower.rs
@@ -1019,15 +1019,7 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
     ) -> CodegenResult<VCode<I>> {
         trace!("about to lower function: {:?}", self.f);
 
-        // Initialize the ABI object, giving it temps if requested.
-        let temps = self
-            .vcode
-            .abi()
-            .temps_needed(self.sigs())
-            .into_iter()
-            .map(|temp_ty| self.alloc_tmp(temp_ty).only_reg().unwrap())
-            .collect::<Vec<_>>();
-        self.vcode.init_abi(temps);
+        self.vcode.init_retval_area(&mut self.vregs)?;
 
         // Get the pinned reg here (we only parameterize this function on `B`,
         // not the whole `Lower` impl).
diff --git a/cranelift/codegen/src/machinst/vcode.rs b/cranelift/codegen/src/machinst/vcode.rs
index 7f4fab821..39010b9da 100644
--- a/cranelift/codegen/src/machinst/vcode.rs
+++ b/cranelift/codegen/src/machinst/vcode.rs
@@ -291,8 +291,8 @@ impl<I: VCodeInst> VCodeBuilder<I> {
         }
     }
 
-    pub fn init_abi(&mut self, temps: Vec<Writable<Reg>>) {
-        self.vcode.abi.init(&self.vcode.sigs, temps);
+    pub fn init_retval_area(&mut self, vregs: &mut VRegAllocator<I>) -> CodegenResult<()> {
+        self.vcode.abi.init_retval_area(&self.vcode.sigs, vregs)
     }
 
     /// Access the ABI object.
@@ -1259,26 +1259,8 @@ impl<I: VCodeInst> VCode<I> {
     }
 
     #[inline]
-    fn assert_no_vreg_aliases<'a>(&self, list: &'a [VReg]) -> &'a [VReg] {
-        for vreg in list {
-            self.assert_not_vreg_alias(*vreg);
-        }
-        list
-    }
-
-    #[inline]
-    fn assert_not_vreg_alias(&self, vreg: VReg) -> VReg {
-        debug_assert!(self.resolve_vreg_alias(vreg) == vreg);
-        vreg
-    }
-
-    #[inline]
-    fn assert_operand_not_vreg_alias(&self, op: Operand) -> Operand {
-        // It should be true by construction that `Operand`s do not contain any
-        // aliased vregs since they're all collected and mapped when the VCode
-        // is itself constructed.
-        self.assert_not_vreg_alias(op.vreg());
-        op
+    fn debug_assert_no_vreg_aliases(&self, mut list: impl Iterator<Item = VReg>) {
+        debug_assert!(list.all(|vreg| !self.vreg_aliases.contains_key(&vreg)));
     }
 
     /// Get the type of a VReg.
@@ -1355,7 +1337,8 @@ impl<I: VCodeInst> RegallocFunction for VCode<I> {
         let ret = &self.block_params[start as usize..end as usize];
         // Currently block params are never aliased to another vreg, but
         // double-check just to be sure.
-        self.assert_no_vreg_aliases(ret)
+        self.debug_assert_no_vreg_aliases(ret.iter().copied());
+        ret
     }
 
     fn branch_blockparams(&self, block: BlockIndex, _insn: InsnIndex, succ_idx: usize) -> &[VReg] {
@@ -1365,7 +1348,8 @@ impl<I: VCodeInst> RegallocFunction for VCode<I> {
         let (branch_block_args_start, branch_block_args_end) = succ_ranges[succ_idx];
         let ret = &self.branch_block_args
             [branch_block_args_start as usize..branch_block_args_end as usize];
-        self.assert_no_vreg_aliases(ret)
+        self.debug_assert_no_vreg_aliases(ret.iter().copied());
+        ret
     }
 
     fn is_ret(&self, insn: InsnIndex) -> bool {
@@ -1391,9 +1375,10 @@ impl<I: VCodeInst> RegallocFunction for VCode<I> {
     fn inst_operands(&self, insn: InsnIndex) -> &[Operand] {
         let (start, end) = self.operand_ranges[insn.index()];
         let ret = &self.operands[start as usize..end as usize];
-        for op in ret {
-            self.assert_operand_not_vreg_alias(*op);
-        }
+        // It should be true by construction that `Operand`s do not contain any
+        // aliased vregs since they're all collected and mapped when the VCode
+        // is itself constructed.
+        self.debug_assert_no_vreg_aliases(ret.iter().map(|op| op.vreg()));
         ret
     }
 
@@ -1406,17 +1391,18 @@ impl<I: VCodeInst> RegallocFunction for VCode<I> {
     }
 
     fn reftype_vregs(&self) -> &[VReg] {
-        self.assert_no_vreg_aliases(&self.reftyped_vregs[..])
+        let ret = &self.reftyped_vregs;
+        self.debug_assert_no_vreg_aliases(ret.iter().copied());
+        ret
     }
 
     fn debug_value_labels(&self) -> &[(VReg, InsnIndex, InsnIndex, u32)] {
         // VRegs here are inserted into `debug_value_labels` after code is
-        // generated and aliases are fully defined, so no double-check that
+        // generated and aliases are fully defined, so double-check that
         // aliases are not lingering.
-        for (vreg, ..) in self.debug_value_labels.iter() {
-            self.assert_not_vreg_alias(*vreg);
-        }
-        &self.debug_value_labels[..]
+        let ret = &self.debug_value_labels;
+        self.debug_assert_no_vreg_aliases(ret.iter().map(|&(vreg, ..)| vreg));
+        ret
     }
 
     fn spillslot_size(&self, regclass: RegClass) -> usize {


Test patch:
diff --git a/cranelift/filetests/filetests/isa/s390x/call.clif b/cranelift/filetests/filetests/isa/s390x/call.clif
index 1d72253b1..41bf58ea7 100644
--- a/cranelift/filetests/filetests/isa/s390x/call.clif
+++ b/cranelift/filetests/filetests/isa/s390x/call.clif
@@ -332,14 +332,14 @@ block0(v0: i128, v1: i128, v2: i128, v3: i128, v4: i128, v5: i128, v6: i128, v7:
 ;   vl %v24, 0(%r5)
 ;   lg %r4, 184(%r15)
 ;   vl %v27, 0(%r4)
-;   vaq %v16, %v1, %v3
-;   vaq %v17, %v5, %v7
-;   vaq %v18, %v18, %v21
-;   vaq %v19, %v24, %v27
-;   vaq %v16, %v16, %v17
-;   vaq %v17, %v18, %v19
-;   vaq %v16, %v16, %v17
-;   vst %v16, 0(%r2)
+;   vaq %v4, %v1, %v3
+;   vaq %v5, %v5, %v7
+;   vaq %v6, %v18, %v21
+;   vaq %v7, %v24, %v27
+;   vaq %v4, %v4, %v5
+;   vaq %v5, %v6, %v7
+;   vaq %v4, %v4, %v5
+;   vst %v4, 0(%r2)
 ;   br %r14
 ;
 ; Disassembled:
@@ -356,14 +356,14 @@ block0(v0: i128, v1: i128, v2: i128, v3: i128, v4: i128, v5: i128, v6: i128, v7:
 ;   vl %v24, 0(%r5)
 ;   lg %r4, 0xb8(%r15)
 ;   vl %v27, 0(%r4)
-;   vaq %v16, %v1, %v3
-;   vaq %v17, %v5, %v7
-;   vaq %v18, %v18, %v21
-;   vaq %v19, %v24, %v27
-;   vaq %v16, %v16, %v17
-;   vaq %v17, %v18, %v19
-;   vaq %v16, %v16, %v17
-;   vst %v16, 0(%r2)
+;   vaq %v4, %v1, %v3
+;   vaq %v5, %v5, %v7
+;   vaq %v6, %v18, %v21
+;   vaq %v7, %v24, %v27
+;   vaq %v4, %v4, %v5
+;   vaq %v5, %v6, %v7
+;   vaq %v4, %v4, %v5
+;   vst %v4, 0(%r2)
 ;   br %r14
 
 function %call_sret() -> i64 {