Delay argument location use in CallSite::gen_arg (#8151)

elliottt · jameysharp · web-flow · commit e3a08d4c40f6 · 2024-03-15T20:36:36.000Z
* Delay argument location use in `CallSite::gen_arg`

Co-authored-by: Jamey Sharp &lt;jsharp@fastly.com&gt;

* Directly emit the instructions produced by gen_arg

---------

Co-authored-by: Jamey Sharp &lt;jsharp@fastly.com&gt;
diff --git a/cranelift/codegen/src/isa/x64/lower.rs b/cranelift/codegen/src/isa/x64/lower.rs
@@ -179,9 +179,8 @@ fn emit_vm_call(
     assert_eq!(inputs.len(), abi.num_args(ctx.sigs()));
 
     for (i, input) in inputs.iter().enumerate() {
-        for inst in abi.gen_arg(ctx, i, ValueRegs::one(*input)) {
-            ctx.emit(inst);
-        }
+        let moves = abi.gen_arg(ctx, i, ValueRegs::one(*input));
+        abi.emit_arg_moves(ctx, moves);
     }
 
     let mut retval_insts: SmallInstVec<_> = smallvec![];
diff --git a/cranelift/codegen/src/machinst/abi.rs b/cranelift/codegen/src/machinst/abi.rs
@@ -301,6 +301,14 @@ impl StackAMode {
             StackAMode::SPOffset(off, ty) => StackAMode::SPOffset(off + addend, ty),
         }
     }
+
+    pub fn get_type(&self) -> ir::Type {
+        match self {
+            &StackAMode::FPOffset(_, ty) => ty,
+            &StackAMode::NominalSPOffset(_, ty) => ty,
+            &StackAMode::SPOffset(_, ty) => ty,
+        }
+    }
 }
 
 /// Trait implemented by machine-specific backend to represent ISA flags.
@@ -2041,6 +2049,13 @@ impl<M: ABIMachineSpec> Callee<M> {
     }
 }
 
+/// The register or stack slot location of an argument.
+#[derive(Clone, Debug)]
+pub enum ArgLoc {
+    Reg(PReg),
+    Stack(StackAMode),
+}
+
 /// An input argument to a call instruction: the vreg that is used,
 /// and the preg it is constrained to (per the ABI).
 #[derive(Clone, Debug)]
@@ -2289,6 +2304,20 @@ impl<M: ABIMachineSpec> CallSite<M> {
         }
     }
 
+    /// Emit moves or uses for the moves list generated by [`Self::gen_arg`].
+    pub fn emit_arg_moves(&mut self, ctx: &mut Lower<M::I>, moves: SmallVec<[(VReg, ArgLoc); 2]>) {
+        for (vreg, loc) in moves {
+            let vreg = vreg.into();
+            match loc {
+                ArgLoc::Reg(preg) => self.uses.push(CallArgPair {
+                    vreg,
+                    preg: preg.into(),
+                }),
+                ArgLoc::Stack(amode) => ctx.emit(M::gen_store_stack(amode, vreg, amode.get_type())),
+            }
+        }
+    }
+
     /// Add a constraint for an argument value from a source register.
     /// For large arguments with associated stack buffer, this may
     /// load the address of the buffer into the argument register, if
@@ -2298,8 +2327,9 @@ impl<M: ABIMachineSpec> CallSite<M> {
         ctx: &mut Lower<M::I>,
         idx: usize,
         from_regs: ValueRegs<Reg>,
-    ) -> SmallInstVec<M::I> {
-        let mut insts = smallvec![];
+    ) -> SmallVec<[(VReg, ArgLoc); 2]> {
+        let mut insts = SmallInstVec::new();
+        let mut locs = smallvec![];
         let word_rc = M::word_reg_class();
         let word_bits = M::word_bits() as usize;
 
@@ -2357,10 +2387,7 @@ impl<M: ABIMachineSpec> CallSite<M> {
                                     ty_bits(ty) as u8,
                                     word_bits as u8,
                                 ));
-                                self.uses.push(CallArgPair {
-                                    vreg: extend_result.to_reg(),
-                                    preg: reg.into(),
-                                });
+                                locs.push((extend_result.to_reg().into(), ArgLoc::Reg(reg.into())));
                             } else if ty.is_ref() {
                                 // Reference-typed args need to be
                                 // passed as a copy; the original vreg
@@ -2369,15 +2396,10 @@ impl<M: ABIMachineSpec> CallSite<M> {
                                 let ref_copy =
                                     temps.pop().expect("Must have allocated enough temps");
                                 insts.push(M::gen_move(ref_copy, *from_reg, M::word_type()));
-                                self.uses.push(CallArgPair {
-                                    vreg: ref_copy.to_reg(),
-                                    preg: reg.into(),
-                                });
+
+                                locs.push((ref_copy.to_reg().into(), ArgLoc::Reg(reg.into())));
                             } else {
-                                self.uses.push(CallArgPair {
-                                    vreg: *from_reg,
-                                    preg: reg.into(),
-                                });
+                                locs.push((from_reg.into(), ArgLoc::Reg(reg.into())));
                             }
                         }
                         &ABIArgSlot::Stack {
@@ -2409,10 +2431,9 @@ impl<M: ABIMachineSpec> CallSite<M> {
                                 } else {
                                     (*from_reg, ty)
                                 };
-                            insts.push(M::gen_store_stack(
-                                StackAMode::SPOffset(offset, ty),
-                                data,
-                                ty,
+                            locs.push((
+                                data.into(),
+                                ArgLoc::Stack(StackAMode::SPOffset(offset, ty)),
                             ));
                         }
                     }
@@ -2434,25 +2455,22 @@ impl<M: ABIMachineSpec> CallSite<M> {
                 insts.push(M::gen_get_stack_addr(amode, tmp, ty));
                 let tmp = tmp.to_reg();
                 insts.push(M::gen_store_base_offset(tmp, 0, vreg, ty));
-                match pointer {
-                    ABIArgSlot::Reg { reg, .. } => {
-                        self.uses.push(CallArgPair {
-                            vreg: tmp,
-                            preg: reg.into(),
-                        });
-                    }
+                let loc = match pointer {
+                    ABIArgSlot::Reg { reg, .. } => ArgLoc::Reg(reg.into()),
                     ABIArgSlot::Stack { offset, .. } => {
                         let ty = M::word_type();
-                        insts.push(M::gen_store_stack(
-                            StackAMode::SPOffset(offset, ty),
-                            tmp,
-                            ty,
-                        ));
+                        ArgLoc::Stack(StackAMode::SPOffset(offset, ty))
                     }
                 };
+                locs.push((tmp.into(), loc));
             }
         }
-        insts
+
+        for inst in insts {
+            ctx.emit(inst);
+        }
+
+        locs
     }
 
     /// Call `gen_arg` for each non-hidden argument and emit all instructions
@@ -2470,9 +2488,8 @@ impl<M: ABIMachineSpec> CallSite<M> {
             self.emit_copy_regs_to_buffer(ctx, i, *arg_regs);
         }
         for (i, value_regs) in arg_value_regs.iter().enumerate() {
-            for inst in self.gen_arg(ctx, i, *value_regs) {
-                ctx.emit(inst);
-            }
+            let moves = self.gen_arg(ctx, i, *value_regs);
+            self.emit_arg_moves(ctx, moves);
         }
     }
 
@@ -2484,9 +2501,8 @@ impl<M: ABIMachineSpec> CallSite<M> {
                 "if the tail callee has a return pointer, then the tail caller \
                  must as well",
             );
-            for inst in self.gen_arg(ctx, i.into(), ValueRegs::one(ret_area_ptr.to_reg())) {
-                ctx.emit(inst);
-            }
+            let moves = self.gen_arg(ctx, i.into(), ValueRegs::one(ret_area_ptr.to_reg()));
+            self.emit_arg_moves(ctx, moves);
         }
     }
 
@@ -2592,9 +2608,8 @@ impl<M: ABIMachineSpec> CallSite<M> {
                 rd,
                 I8,
             ));
-            for inst in self.gen_arg(ctx, i.into(), ValueRegs::one(rd.to_reg())) {
-                ctx.emit(inst);
-            }
+            let moves = self.gen_arg(ctx, i.into(), ValueRegs::one(rd.to_reg()));
+            self.emit_arg_moves(ctx, moves);
         }
 
         let (uses, defs) = (
diff --git a/cranelift/codegen/src/machinst/isle.rs b/cranelift/codegen/src/machinst/isle.rs
@@ -825,9 +825,8 @@ macro_rules! isle_prelude_method_helpers {
                 call_site.emit_copy_regs_to_buffer(self.lower_ctx, i, *arg_regs);
             }
             for (i, arg_regs) in arg_regs.iter().enumerate() {
-                for inst in call_site.gen_arg(self.lower_ctx, i, *arg_regs) {
-                    self.lower_ctx.emit(inst);
-                }
+                let moves = call_site.gen_arg(self.lower_ctx, i, *arg_regs);
+                call_site.emit_arg_moves(self.lower_ctx, moves);
             }
         }
 

Original file line number	Diff line number	Diff line change
`@@ -825,9 +825,8 @@ macro_rules! isle_prelude_method_helpers {`
`825`	`825`	`call_site.emit_copy_regs_to_buffer(self.lower_ctx, i, *arg_regs);`
`826`	`826`	`}`
`827`	`827`	`for (i, arg_regs) in arg_regs.iter().enumerate() {`
`828`		`- for inst in call_site.gen_arg(self.lower_ctx, i, *arg_regs) {`
`829`		`- self.lower_ctx.emit(inst);`
`830`		`- }`
	`828`	`+ let moves = call_site.gen_arg(self.lower_ctx, i, *arg_regs);`
	`829`	`+ call_site.emit_arg_moves(self.lower_ctx, moves);`
`831`	`830`	`}`
`832`	`831`	`}`
`833`	`832`