Cranelift: add a "patchable call" ABI. (#12061)

This ABI is intended for use in scenarios where we want a very
lightweight callsite that can be turned on and off by patching in one
instruction. (The actual patchable call instruction is not in this PR;
that will be a separate PR.)

The idea is that we define a call to clobber *no* registers -- not
even the arguments! And we restrict signatures such that on all of our
supported architectures, all arguments go into registers only. Those
two requirements together mean that all callsites for this ABI should
have only a raw call instruction, with no loads/stores to stackslots;
and have the minimum possible impact on regalloc, by only imposing
constraints on args to ensure they are in certain registers but not
altering those registers.

Given this, we could implement, e.g., breakpoints with patchable
callsites (off by default) at every sequence point in compiled
code. In a typical use-case with Wasmtime-compiled Wasm, that would
put a bunch of uses of vmctx constrained to the first argument
register in every code path, but vmctx likely already sits there most
of the time anyway (for any call to other Wasm functions or for
libcalls). Thus, the impact is just the one instruction and nothing
else.

This PR adds the calling convention itself and tests that show
that *two* consecutive callsites can be compiled with no register
setup re-occurring from one call to the next (thus demonstrating no
clobbers).

Author: cfallin

cranelift/codegen/src/isa/aarch64/abi.rs

@@ -1102,6 +1102,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
             (isa::CallConv::Winch, true) => ALL_CLOBBERS,
             (isa::CallConv::Winch, false) => WINCH_CLOBBERS,
             (isa::CallConv::SystemV, _) => DEFAULT_AAPCS_CLOBBERS,
+            (isa::CallConv::Patchable, _) => NO_CLOBBERS,
             (_, false) => DEFAULT_AAPCS_CLOBBERS,
             (_, true) => panic!("unimplemented clobbers for exn abi of {call_conv:?}"),
         }
@@ -1265,11 +1266,15 @@ impl AArch64MachineDeps {
 /// Is the given register saved in the prologue if clobbered, i.e., is it a
 /// callee-save?
 fn is_reg_saved_in_prologue(
-    _call_conv: isa::CallConv,
+    call_conv: isa::CallConv,
     enable_pinned_reg: bool,
     sig: &Signature,
     r: RealReg,
 ) -> bool {
+    if call_conv == isa::CallConv::Patchable {
+        return true;
+    }
+
     // FIXME: We need to inspect whether a function is returning Z or P regs too.
     let save_z_regs = sig
         .params
@@ -1518,6 +1523,7 @@ const fn all_clobbers() -> PRegSet {
 const DEFAULT_AAPCS_CLOBBERS: PRegSet = default_aapcs_clobbers();
 const WINCH_CLOBBERS: PRegSet = winch_clobbers();
 const ALL_CLOBBERS: PRegSet = all_clobbers();
+const NO_CLOBBERS: PRegSet = PRegSet::empty();
 
 fn create_reg_env(enable_pinned_reg: bool) -> MachineEnv {
     fn preg(r: Reg) -> PReg {

cranelift/codegen/src/isa/call_conv.rs

@@ -45,6 +45,26 @@ pub enum CallConv {
     /// defines no callee-save registers, and restricts the number of return
     /// registers to one integer, and one floating point.
     Winch,
+    /// Calling convention for patchable-call instructions.
+    ///
+    /// This is designed for a very specific need: we want a *single*
+    /// call instruction at our callsite, with no other setup, and we
+    /// don't want any registers clobbered. This allows patchable
+    /// callsites to be as unobtrusive as possible.
+    ///
+    /// The ABI is based on the native register-argument ABI on each
+    /// respective platform, but puts severe restrictions on allowable
+    /// signatures: only up to four arguments of integer type, and no
+    /// return values. It does not support tail-calls, and disallows
+    /// any extension modes on arguments.
+    ///
+    /// The ABI specifies that *no* registers, not even argument
+    /// registers, are clobbered. This is pretty unique: it means that
+    /// the call instruction will constrain regalloc to have any args
+    /// in the right registers, but those registers will be preserved,
+    /// so multiple patchable callsites can reuse those values. This
+    /// further reduces the cost of the callsites.
+    Patchable,
 }
 
 impl CallConv {
@@ -123,6 +143,7 @@ impl fmt::Display for CallConv {
             Self::AppleAarch64 => "apple_aarch64",
             Self::Probestack => "probestack",
             Self::Winch => "winch",
+            Self::Patchable => "patchable",
         })
     }
 }
@@ -139,6 +160,7 @@ impl str::FromStr for CallConv {
             "apple_aarch64" => Ok(Self::AppleAarch64),
             "probestack" => Ok(Self::Probestack),
             "winch" => Ok(Self::Winch),
+            "patchable" => Ok(Self::Patchable),
             _ => Err(()),
         }
     }

cranelift/codegen/src/isa/pulley_shared/abi.rs

@@ -480,18 +480,20 @@ where
     }
 
     fn get_regs_clobbered_by_call(
-        _call_conv_of_callee: isa::CallConv,
+        call_conv_of_callee: isa::CallConv,
         is_exception: bool,
     ) -> PRegSet {
-        if is_exception {
+        if call_conv_of_callee == isa::CallConv::Patchable {
+            NO_CLOBBERS
+        } else if is_exception {
             ALL_CLOBBERS
         } else {
             DEFAULT_CLOBBERS
         }
     }
 
     fn compute_frame_layout(
-        _call_conv: isa::CallConv,
+        call_conv: isa::CallConv,
         flags: &settings::Flags,
         _sig: &Signature,
         regs: &[Writable<RealReg>],
@@ -502,11 +504,12 @@ where
         fixed_frame_storage_size: u32,
         outgoing_args_size: u32,
     ) -> FrameLayout {
-        let mut regs: Vec<Writable<RealReg>> = regs
-            .iter()
-            .cloned()
-            .filter(|r| DEFAULT_CALLEE_SAVES.contains(r.to_reg().into()))
-            .collect();
+        let is_callee_save = |reg: &Writable<RealReg>| match call_conv {
+            isa::CallConv::Patchable => true,
+            _ => DEFAULT_CALLEE_SAVES.contains(reg.to_reg().into()),
+        };
+        let mut regs: Vec<Writable<RealReg>> =
+            regs.iter().cloned().filter(is_callee_save).collect();
 
         regs.sort_unstable();
 
@@ -720,7 +723,7 @@ impl FrameLayout {
                     I64
                 }
                 RegClass::Float => F64,
-                RegClass::Vector => unreachable!("no vector registers are callee-save"),
+                RegClass::Vector => I8X16,
             };
             let offset = i32::try_from(offset).unwrap();
             Some((offset, ty, Reg::from(reg.to_reg())))
@@ -759,7 +762,11 @@ fn compute_clobber_size(clobbers: &[Writable<RealReg>]) -> u32 {
             RegClass::Float => {
                 clobbered_size += 8;
             }
-            RegClass::Vector => unimplemented!("Vector Size Clobbered"),
+            RegClass::Vector => {
+                // No alignment concerns: the Pulley virtual CPU
+                // supports unaligned vector load/stores.
+                clobbered_size += 16;
+            }
         }
     }
     align_to(clobbered_size, 16)
@@ -948,6 +955,8 @@ const ALL_CLOBBERS: PRegSet = PRegSet::empty()
     .with(pv_reg(30))
     .with(pv_reg(31));
 
+const NO_CLOBBERS: PRegSet = PRegSet::empty();
+
 fn create_reg_environment() -> MachineEnv {
     // Prefer caller-saved registers over callee-saved registers, because that
     // way we don't need to emit code to save and restore them if we don't

cranelift/codegen/src/isa/riscv64/abi.rs

@@ -489,16 +489,17 @@ impl ABIMachineSpec for Riscv64MachineDeps {
         if stack_size > 0 {
             insts.extend(Self::gen_sp_reg_adjust(-(stack_size as i32)));
 
-            let mut cur_offset = 8;
+            let mut cur_offset = 0;
             for reg in &frame_layout.clobbered_callee_saves {
                 let r_reg = reg.to_reg();
                 let ty = match r_reg.class() {
                     RegClass::Int => I64,
                     RegClass::Float => F64,
-                    RegClass::Vector => unimplemented!("Vector Clobber Saves"),
+                    RegClass::Vector => I8X16,
                 };
+                cur_offset = align_to(cur_offset, ty.bytes());
                 insts.push(Inst::gen_store(
-                    AMode::SPOffset((stack_size - cur_offset) as i64),
+                    AMode::SPOffset(i64::from(stack_size - cur_offset - ty.bytes())),
                     Reg::from(reg.to_reg()),
                     ty,
                     MemFlags::trusted(),
@@ -507,13 +508,14 @@ impl ABIMachineSpec for Riscv64MachineDeps {
                 if flags.unwind_info() {
                     insts.push(Inst::Unwind {
                         inst: UnwindInst::SaveReg {
-                            clobber_offset: frame_layout.clobber_size - cur_offset,
+                            clobber_offset: frame_layout.clobber_size - cur_offset - ty.bytes(),
                             reg: r_reg,
                         },
                     });
                 }
 
-                cur_offset += 8
+                cur_offset += ty.bytes();
+                assert!(cur_offset <= stack_size);
             }
         }
         insts
@@ -529,22 +531,23 @@ impl ABIMachineSpec for Riscv64MachineDeps {
         let stack_size = frame_layout.clobber_size
             + frame_layout.fixed_frame_storage_size
             + frame_layout.outgoing_args_size;
+        let mut cur_offset = 0;
 
-        let mut cur_offset = 8;
         for reg in &frame_layout.clobbered_callee_saves {
             let rreg = reg.to_reg();
             let ty = match rreg.class() {
                 RegClass::Int => I64,
                 RegClass::Float => F64,
-                RegClass::Vector => unimplemented!("Vector Clobber Restores"),
+                RegClass::Vector => I8X16,
             };
+            cur_offset = align_to(cur_offset, ty.bytes());
             insts.push(Inst::gen_load(
                 reg.map(Reg::from),
-                AMode::SPOffset(i64::from(stack_size - cur_offset)),
+                AMode::SPOffset(i64::from(stack_size - cur_offset - ty.bytes())),
                 ty,
                 MemFlags::trusted(),
             ));
-            cur_offset += 8
+            cur_offset += ty.bytes();
         }
 
         if stack_size > 0 {
@@ -619,12 +622,13 @@ impl ABIMachineSpec for Riscv64MachineDeps {
     ) -> PRegSet {
         match call_conv_of_callee {
             isa::CallConv::Tail if is_exception => ALL_CLOBBERS,
+            isa::CallConv::Patchable => NO_CLOBBERS,
             _ => DEFAULT_CLOBBERS,
         }
     }
 
     fn compute_frame_layout(
-        _call_conv: isa::CallConv,
+        call_conv: isa::CallConv,
         flags: &settings::Flags,
         _sig: &Signature,
         regs: &[Writable<RealReg>],
@@ -635,11 +639,12 @@ impl ABIMachineSpec for Riscv64MachineDeps {
         fixed_frame_storage_size: u32,
         outgoing_args_size: u32,
     ) -> FrameLayout {
-        let mut regs: Vec<Writable<RealReg>> = regs
-            .iter()
-            .cloned()
-            .filter(|r| DEFAULT_CALLEE_SAVES.contains(r.to_reg().into()))
-            .collect();
+        let is_callee_saved = |reg: &Writable<RealReg>| match call_conv {
+            isa::CallConv::Patchable => true,
+            _ => DEFAULT_CALLEE_SAVES.contains(reg.to_reg().into()),
+        };
+        let mut regs: Vec<Writable<RealReg>> =
+            regs.iter().cloned().filter(is_callee_saved).collect();
 
         regs.sort_unstable();
 
@@ -760,7 +765,10 @@ fn compute_clobber_size(clobbers: &[Writable<RealReg>]) -> u32 {
             RegClass::Float => {
                 clobbered_size += 8;
             }
-            RegClass::Vector => unimplemented!("Vector Size Clobbered"),
+            RegClass::Vector => {
+                clobbered_size = align_to(clobbered_size, 16);
+                clobbered_size += 16;
+            }
         }
     }
     align_to(clobbered_size, 16)
@@ -937,6 +945,8 @@ const ALL_CLOBBERS: PRegSet = PRegSet::empty()
     .with(pv_reg(30))
     .with(pv_reg(31));
 
+const NO_CLOBBERS: PRegSet = PRegSet::empty();
+
 fn create_reg_environment() -> MachineEnv {
     // Some C Extension instructions can only use a subset of the registers.
     // x8 - x15, f8 - f15, v8 - v15 so we should prefer to use those since

cranelift/codegen/src/isa/s390x/abi.rs

@@ -880,6 +880,7 @@ impl ABIMachineSpec for S390xMachineDeps {
         match call_conv_of_callee {
             isa::CallConv::Tail if is_exception => ALL_CLOBBERS,
             isa::CallConv::Tail => TAIL_CLOBBERS,
+            isa::CallConv::Patchable => NO_CLOBBERS,
             _ => SYSV_CLOBBERS,
         }
     }
@@ -1103,6 +1104,7 @@ fn is_reg_saved_in_prologue(call_conv: isa::CallConv, r: RealReg) -> bool {
             // r8 - r14 inclusive are callee-saves.
             r.hw_enc() >= 8 && r.hw_enc() <= 14
         }
+        (isa::CallConv::Patchable, _) => true,
         (_, RegClass::Int) => {
             // r6 - r15 inclusive are callee-saves.
             r.hw_enc() >= 6 && r.hw_enc() <= 15
@@ -1380,8 +1382,11 @@ const fn all_clobbers() -> PRegSet {
         .with(vr_preg(30))
         .with(vr_preg(31))
 }
+
 const ALL_CLOBBERS: PRegSet = all_clobbers();
 
+const NO_CLOBBERS: PRegSet = PRegSet::empty();
+
 fn sysv_create_machine_env() -> MachineEnv {
     MachineEnv {
         preferred_regs_by_class: [

cranelift/codegen/src/isa/x64/abi.rs

@@ -891,6 +891,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
             (isa::CallConv::Winch, _) => ALL_CLOBBERS,
             (isa::CallConv::SystemV, _) => SYSV_CLOBBERS,
             (isa::CallConv::WindowsFastcall, false) => WINDOWS_CLOBBERS,
+            (isa::CallConv::Patchable, _) => NO_CLOBBERS,
             (_, false) => SYSV_CLOBBERS,
             (call_conv, true) => panic!("unimplemented clobbers for exn abi of {call_conv:?}"),
         }
@@ -931,6 +932,8 @@ impl ABIMachineSpec for X64ABIMachineSpec {
                 .cloned()
                 .filter(|r| is_callee_save_fastcall(r.to_reg(), flags.enable_pinned_reg()))
                 .collect(),
+            // The `patchable` calling convention makes every reg a callee-save reg.
+            CallConv::Patchable => regs.iter().cloned().collect(),
             CallConv::Probestack => todo!("probestack?"),
             CallConv::AppleAarch64 => unreachable!(),
         };
@@ -1085,6 +1088,8 @@ fn get_intreg_for_retval(
         },
 
         CallConv::Winch => is_last.then(|| regs::rax()),
+        // The patchable ABI does not support any return values.
+        CallConv::Patchable => None,
         CallConv::Probestack => todo!(),
         CallConv::AppleAarch64 => unreachable!(),
     }
@@ -1113,6 +1118,8 @@ fn get_fltreg_for_retval(call_conv: CallConv, fltreg_idx: usize, is_last: bool)
             _ => None,
         },
         CallConv::Winch => is_last.then(|| regs::xmm0()),
+        // The patchable ABI does not support any return values.
+        CallConv::Patchable => None,
         CallConv::Probestack => todo!(),
         CallConv::AppleAarch64 => unreachable!(),
     }
@@ -1174,6 +1181,7 @@ fn compute_clobber_size(clobbers: &[Writable<RealReg>]) -> u32 {
 const WINDOWS_CLOBBERS: PRegSet = windows_clobbers();
 const SYSV_CLOBBERS: PRegSet = sysv_clobbers();
 pub(crate) const ALL_CLOBBERS: PRegSet = all_clobbers();
+const NO_CLOBBERS: PRegSet = PRegSet::empty();
 
 const fn windows_clobbers() -> PRegSet {
     use asm::gpr::enc::*;

cranelift/codegen/src/machinst/abi.rs

@@ -1222,7 +1222,8 @@ impl<M: ABIMachineSpec> Callee<M> {
                 || call_conv == isa::CallConv::Cold
                 || call_conv == isa::CallConv::WindowsFastcall
                 || call_conv == isa::CallConv::AppleAarch64
-                || call_conv == isa::CallConv::Winch,
+                || call_conv == isa::CallConv::Winch
+                || call_conv == isa::CallConv::Patchable,
             "Unsupported calling convention: {call_conv:?}"
         );
 
@@ -2490,10 +2491,14 @@ impl<T> CallInfo<T> {
         }
 
         let temp = M::retval_temp_reg(self.callee_conv);
-        // The temporary must be noted as clobbered.
+        // The temporary must be noted as clobbered unless there are
+        // no returns (hence it isn't needed). The latter can only be
+        // the case statically for an ABI when the ABI doesn't allow
+        // any returns at all (e.g., patchable-call ABI).
         debug_assert!(
-            M::get_regs_clobbered_by_call(self.callee_conv, self.try_call_info.is_some())
-                .contains(PReg::from(temp.to_reg().to_real_reg().unwrap()))
+            self.defs.is_empty()
+                || M::get_regs_clobbered_by_call(self.callee_conv, self.try_call_info.is_some())
+                    .contains(PReg::from(temp.to_reg().to_real_reg().unwrap()))
         );
 
         for CallRetPair { vreg, location } in &self.defs {