ABI: implement register arguments with constraints. (#4858)

* ABI: implement register arguments with constraints.

Currently, Cranelift's ABI code emits a sequence of moves from physical
registers into vregs at the top of the function body, one for every
register-carried argument.

For a number of reasons, we want to move to operand constraints instead,
and remove the use of explicitly-named "pinned vregs"; this allows for
better regalloc in theory, as it removes the need to "reverse-engineer"
the sequence of moves.

This PR alters the ABI code so that it generates a single "args"
pseudo-instruction as the first instruction in the function body. This
pseudo-inst defs all register arguments, and constrains them to the
appropriate registers at the def-point. Subsequently the regalloc can
move them wherever it needs to.

Some care was taken not to have this pseudo-inst show up in
post-regalloc disassemblies, but the change did cause a general regalloc
"shift" in many tests, so the precise-output updates are a bit noisy.
Sorry about that!

A subsequent PR will handle the other half of the ABI code, namely, the
callsite case, with a similar preg-to-constraint conversion.

* Update based on review feedback.

* Review feedback.

cranelift/codegen/src/machinst/abi.rs

@@ -124,6 +124,17 @@ use std::mem;
 /// a small fixed sequence implementing one operation.
 pub type SmallInstVec<I> = SmallVec<[I; 4]>;
 
+/// A type used by backends to track argument-binding info in the "args"
+/// pseudoinst. The pseudoinst holds a vec of `ArgPair` structs.
+#[derive(Clone, Debug)]
+pub struct ArgPair {
+    /// The vreg that is defined by this args pseudoinst.
+    pub vreg: Writable<Reg>,
+    /// The preg that the arg arrives in; this constrains the vreg's
+    /// placement at the pseudoinst.
+    pub preg: Reg,
+}
+
 /// A location for (part of) an argument or return value. These "storage slots"
 /// are specified for each register-sized part of an argument.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -364,6 +375,10 @@ pub trait ABIMachineSpec {
         to_bits: u8,
     ) -> Self::I;
 
+    /// Generate an "args" pseudo-instruction to capture input args in
+    /// registers.
+    fn gen_args(isa_flags: &Self::F, args: Vec<ArgPair>) -> Self::I;
+
     /// Generate a return instruction.
     fn gen_ret(setup_frame: bool, isa_flags: &Self::F, rets: Vec<Reg>) -> Self::I;
 
@@ -856,6 +871,9 @@ pub struct Callee<M: ABIMachineSpec> {
     stackslots_size: u32,
     /// Stack size to be reserved for outgoing arguments.
     outgoing_args_size: u32,
+    /// Register-argument defs, to be provided to the `args`
+    /// pseudo-inst, and pregs to constrain them to.
+    reg_args: Vec<ArgPair>,
     /// Clobbered registers, from regalloc.
     clobbered: Vec<Writable<RealReg>>,
     /// Total number of spillslots, including for 'dynamic' types, from regalloc.
@@ -1012,6 +1030,7 @@ impl<M: ABIMachineSpec> Callee<M> {
             sized_stackslots,
             stackslots_size,
             outgoing_args_size: 0,
+            reg_args: vec![],
             clobbered: vec![],
             spillslots: None,
             fixed_frame_storage_size: 0,
@@ -1283,7 +1302,7 @@ impl<M: ABIMachineSpec> Callee<M> {
     /// Generate an instruction which copies an argument to a destination
     /// register.
     pub fn gen_copy_arg_to_regs(
-        &self,
+        &mut self,
         sigs: &SigSet,
         idx: usize,
         into_regs: ValueRegs<Writable<Reg>>,
@@ -1291,10 +1310,16 @@ impl<M: ABIMachineSpec> Callee<M> {
         let mut insts = smallvec![];
         let mut copy_arg_slot_to_reg = |slot: &ABIArgSlot, into_reg: &Writable<Reg>| {
             match slot {
-                &ABIArgSlot::Reg { reg, ty, .. } => {
-                    // Extension mode doesn't matter (we're copying out, not in; we
-                    // ignore high bits by convention).
-                    insts.push(M::gen_move(*into_reg, reg.into(), ty));
+                &ABIArgSlot::Reg { reg, .. } => {
+                    // Add a preg -> def pair to the eventual `args`
+                    // instruction.  Extension mode doesn't matter
+                    // (we're copying out, not in; we ignore high bits
+                    // by convention).
+                    let arg = ArgPair {
+                        vreg: *into_reg,
+                        preg: reg.into(),
+                    };
+                    self.reg_args.push(arg);
                 }
                 &ABIArgSlot::Stack {
                     offset,
@@ -1481,17 +1506,18 @@ impl<M: ABIMachineSpec> Callee<M> {
     /// values or an otherwise large return value that must be passed on the
     /// stack; typically the ABI specifies an extra hidden argument that is a
     /// pointer to that memory.
-    pub fn gen_retval_area_setup(&self, sigs: &SigSet) -> Option<M::I> {
+    pub fn gen_retval_area_setup(&mut self, sigs: &SigSet) -> Option<M::I> {
         if let Some(i) = sigs[self.sig].stack_ret_arg {
             let insts =
                 self.gen_copy_arg_to_regs(sigs, i, ValueRegs::one(self.ret_area_ptr.unwrap()));
-            let inst = insts.into_iter().next().unwrap();
-            trace!(
-                "gen_retval_area_setup: inst {:?}; ptr reg is {:?}",
-                inst,
-                self.ret_area_ptr.unwrap().to_reg()
-            );
-            Some(inst)
+            insts.into_iter().next().map(|inst| {
+                trace!(
+                    "gen_retval_area_setup: inst {:?}; ptr reg is {:?}",
+                    inst,
+                    self.ret_area_ptr.unwrap().to_reg()
+                );
+                inst
+            })
         } else {
             trace!("gen_retval_area_setup: not needed");
             None
@@ -1573,6 +1599,23 @@ impl<M: ABIMachineSpec> Callee<M> {
 
         gen_store_stack_multi::<M>(StackAMode::NominalSPOffset(sp_off, ty), from_regs, ty)
     }
+
+    /// Get an `args` pseudo-inst, if any, that should appear at the
+    /// very top of the function body prior to regalloc.
+    pub fn take_args(&mut self) -> Option<M::I> {
+        if self.reg_args.len() > 0 {
+            // Very first instruction is an `args` pseudo-inst that
+            // establishes live-ranges for in-register arguments and
+            // constrains them at the start of the function to the
+            // locations defined by the ABI.
+            Some(M::gen_args(
+                &self.isa_flags,
+                std::mem::take(&mut self.reg_args),
+            ))
+        } else {
+            None
+        }
+    }
 }
 
 /// ### Post-Regalloc Functions

cranelift/codegen/src/machinst/lower.rs

@@ -586,8 +586,9 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
                 let regs = writable_value_regs(self.value_regs[*param]);
                 for insn in self
                     .vcode
-                    .abi()
-                    .gen_copy_arg_to_regs(self.sigs(), i, regs)
+                    .vcode
+                    .abi
+                    .gen_copy_arg_to_regs(&self.vcode.vcode.sigs, i, regs)
                     .into_iter()
                 {
                     self.emit(insn);
@@ -611,7 +612,22 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
                     ));
                 }
             }
-            if let Some(insn) = self.vcode.abi().gen_retval_area_setup(self.sigs()) {
+            if let Some(insn) = self
+                .vcode
+                .vcode
+                .abi
+                .gen_retval_area_setup(&self.vcode.vcode.sigs)
+            {
+                self.emit(insn);
+            }
+
+            // The `args` instruction below must come first. Finish
+            // the current "IR inst" (with a default source location,
+            // as for other special instructions inserted during
+            // lowering) and continue the scan backward.
+            self.finish_ir_inst(Default::default());
+
+            if let Some(insn) = self.vcode.vcode.abi.take_args() {
                 self.emit(insn);
             }
         }

cranelift/codegen/src/machinst/mod.rs

@@ -100,6 +100,9 @@ pub trait MachInst: Clone + Debug {
     /// (ret/uncond/cond) and target if applicable.
     fn is_term(&self) -> MachTerminator;
 
+    /// Is this an "args" pseudoinst?
+    fn is_args(&self) -> bool;
+
     /// Should this instruction be included in the clobber-set?
     fn is_included_in_clobbers(&self) -> bool {
         true

cranelift/codegen/src/machinst/vcode.rs

@@ -158,7 +158,7 @@ pub struct VCode<I: VCodeInst> {
     block_order: BlockLoweringOrder,
 
     /// ABI object.
-    abi: Callee<I::ABIMachineSpec>,
+    pub(crate) abi: Callee<I::ABIMachineSpec>,
 
     /// Constant information used during code emission. This should be
     /// immutable across function compilations within the same module.
@@ -179,7 +179,7 @@ pub struct VCode<I: VCodeInst> {
     /// Value labels for debuginfo attached to vregs.
     debug_value_labels: Vec<(VReg, InsnIndex, InsnIndex, u32)>,
 
-    sigs: SigSet,
+    pub(crate) sigs: SigSet,
 }
 
 /// The result of `VCode::emit`. Contains all information computed
@@ -244,7 +244,7 @@ pub struct EmitResult<I: VCodeInst> {
 /// terminator instructions with successor blocks.)
 pub struct VCodeBuilder<I: VCodeInst> {
     /// In-progress VCode.
-    vcode: VCode<I>,
+    pub(crate) vcode: VCode<I>,
 
     /// In what direction is the build occuring?
     direction: VCodeBuildDirection,
@@ -862,7 +862,7 @@ impl<I: VCodeInst> VCode<I> {
                            disasm: &mut String,
                            buffer: &mut MachBuffer<I>,
                            state: &mut I::State| {
-                if want_disasm {
+                if want_disasm && !inst.is_args() {
                     let mut s = state.clone();
                     writeln!(disasm, "  {}", inst.pretty_print_inst(allocs, &mut s)).unwrap();
                 }