Cranelift: fix use of pinned reg with SysV calling convention. (#4176)

Previously, the pinned register (enabled by the `enable_pinned_reg`
Cranelift setting and used via the `get_pinned_reg` and `set_pinned_reg`
CLIF ops) was only used when Cranelift was embedded in SpiderMonkey, in
order to support a pinned heap register. SpiderMonkey has its own
calling convention in Cranelift (named after the integration layer,
"Baldrdash").

However, the feature is more general, and should be usable with the
default system calling convention too, e.g. SysV or Windows Fastcall.

This PR fixes the ABI code to properly treat the pinned register as a
globally allocated register -- and hence an implicit input and output to
every function, not saved/restored in the prologue/epilogue -- for SysV
on x86-64 and aarch64, and Fastcall on x86-64.

Fixes #4170.

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

@@ -936,7 +936,8 @@ impl ABIMachineSpec for AArch64MachineDeps {
         _outgoing_args_size: u32,
     ) -> SmallVec<[Inst; 16]> {
         let mut insts = SmallVec::new();
-        let (clobbered_int, clobbered_vec) = get_regs_restored_in_epilogue(call_conv, clobbers);
+        let (clobbered_int, clobbered_vec) =
+            get_regs_restored_in_epilogue(call_conv, flags, clobbers);
 
         // Free the fixed frame if necessary.
         if fixed_frame_storage_size > 0 {
@@ -1189,12 +1190,13 @@ impl ABIMachineSpec for AArch64MachineDeps {
 
     fn get_clobbered_callee_saves(
         call_conv: isa::CallConv,
+        flags: &settings::Flags,
         regs: &[Writable<RealReg>],
     ) -> Vec<Writable<RealReg>> {
         let mut regs: Vec<Writable<RealReg>> = regs
             .iter()
             .cloned()
-            .filter(|r| is_reg_saved_in_prologue(call_conv, r.to_reg()))
+            .filter(|r| is_reg_saved_in_prologue(call_conv, flags.enable_pinned_reg(), r.to_reg()))
             .collect();
 
         // Sort registers for deterministic code output. We can do an unstable
@@ -1229,7 +1231,7 @@ fn legal_type_for_machine(ty: Type) -> bool {
 
 /// 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, r: RealReg) -> bool {
+fn is_reg_saved_in_prologue(call_conv: isa::CallConv, enable_pinned_reg: bool, r: RealReg) -> bool {
     if call_conv.extends_baldrdash() {
         match r.class() {
             RegClass::Int => {
@@ -1246,7 +1248,14 @@ fn is_reg_saved_in_prologue(call_conv: isa::CallConv, r: RealReg) -> bool {
     match r.class() {
         RegClass::Int => {
             // x19 - x28 inclusive are callee-saves.
-            r.hw_enc() >= 19 && r.hw_enc() <= 28
+            // However, x21 is the pinned reg if `enable_pinned_reg`
+            // is set, and is implicitly globally-allocated, hence not
+            // callee-saved in prologues.
+            if enable_pinned_reg && r.hw_enc() == PINNED_REG {
+                false
+            } else {
+                r.hw_enc() >= 19 && r.hw_enc() <= 28
+            }
         }
         RegClass::Float => {
             // v8 - v15 inclusive are callee-saves.
@@ -1260,12 +1269,13 @@ fn is_reg_saved_in_prologue(call_conv: isa::CallConv, r: RealReg) -> bool {
 /// written by the function's body.
 fn get_regs_restored_in_epilogue(
     call_conv: isa::CallConv,
+    flags: &settings::Flags,
     regs: &[Writable<RealReg>],
 ) -> (Vec<Writable<RealReg>>, Vec<Writable<RealReg>>) {
     let mut int_saves = vec![];
     let mut vec_saves = vec![];
     for &reg in regs {
-        if is_reg_saved_in_prologue(call_conv, reg.to_reg()) {
+        if is_reg_saved_in_prologue(call_conv, flags.enable_pinned_reg(), reg.to_reg()) {
             match reg.to_reg().class() {
                 RegClass::Int => int_saves.push(reg),
                 RegClass::Float => vec_saves.push(reg),

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

@@ -719,8 +719,14 @@ impl ABIMachineSpec for S390xMachineDeps {
 
     fn get_clobbered_callee_saves(
         call_conv: isa::CallConv,
+        flags: &settings::Flags,
         regs: &[Writable<RealReg>],
     ) -> Vec<Writable<RealReg>> {
+        assert!(
+            !flags.enable_pinned_reg(),
+            "Pinned register not supported on s390x"
+        );
+
         let mut regs: Vec<Writable<RealReg>> = regs
             .iter()
             .cloned()

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

@@ -397,7 +397,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
 
     fn get_stacklimit_reg() -> Reg {
         debug_assert!(
-            !is_callee_save_systemv(regs::r10().to_real_reg().unwrap())
+            !is_callee_save_systemv(regs::r10().to_real_reg().unwrap(), false)
                 && !is_callee_save_baldrdash(regs::r10().to_real_reg().unwrap())
         );
 
@@ -577,7 +577,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
     ) -> SmallVec<[Self::I; 16]> {
         let mut insts = SmallVec::new();
 
-        let clobbered_callee_saves = Self::get_clobbered_callee_saves(call_conv, clobbers);
+        let clobbered_callee_saves = Self::get_clobbered_callee_saves(call_conv, flags, clobbers);
         let stack_size = fixed_frame_storage_size + compute_clobber_size(&clobbered_callee_saves);
 
         // Restore regs by loading from offsets of RSP. RSP will be
@@ -788,6 +788,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
 
     fn get_clobbered_callee_saves(
         call_conv: CallConv,
+        flags: &settings::Flags,
         regs: &[Writable<RealReg>],
     ) -> Vec<Writable<RealReg>> {
         let mut regs: Vec<Writable<RealReg>> = match call_conv {
@@ -802,12 +803,12 @@ impl ABIMachineSpec for X64ABIMachineSpec {
             CallConv::Fast | CallConv::Cold | CallConv::SystemV | CallConv::WasmtimeSystemV => regs
                 .iter()
                 .cloned()
-                .filter(|r| is_callee_save_systemv(r.to_reg()))
+                .filter(|r| is_callee_save_systemv(r.to_reg(), flags.enable_pinned_reg()))
                 .collect(),
             CallConv::WindowsFastcall | CallConv::WasmtimeFastcall => regs
                 .iter()
                 .cloned()
-                .filter(|r| is_callee_save_fastcall(r.to_reg()))
+                .filter(|r| is_callee_save_fastcall(r.to_reg(), flags.enable_pinned_reg()))
                 .collect(),
             CallConv::Probestack => todo!("probestack?"),
             CallConv::AppleAarch64 | CallConv::WasmtimeAppleAarch64 => unreachable!(),
@@ -969,11 +970,15 @@ fn get_fltreg_for_retval(
     }
 }
 
-fn is_callee_save_systemv(r: RealReg) -> bool {
+fn is_callee_save_systemv(r: RealReg, enable_pinned_reg: bool) -> bool {
     use regs::*;
     match r.class() {
         RegClass::Int => match r.hw_enc() {
-            ENC_RBX | ENC_RBP | ENC_R12 | ENC_R13 | ENC_R14 | ENC_R15 => true,
+            ENC_RBX | ENC_RBP | ENC_R12 | ENC_R13 | ENC_R14 => true,
+            // R15 is the pinned register; if we're using it that way,
+            // it is effectively globally-allocated, and is not
+            // callee-saved.
+            ENC_R15 => !enable_pinned_reg,
             _ => false,
         },
         RegClass::Float => false,
@@ -989,18 +994,20 @@ fn is_callee_save_baldrdash(r: RealReg) -> bool {
                 false
             } else {
                 // Defer to native for the other ones.
-                is_callee_save_systemv(r)
+                is_callee_save_systemv(r, /* enable_pinned_reg = */ true)
             }
         }
         RegClass::Float => false,
     }
 }
 
-fn is_callee_save_fastcall(r: RealReg) -> bool {
+fn is_callee_save_fastcall(r: RealReg, enable_pinned_reg: bool) -> bool {
     use regs::*;
     match r.class() {
         RegClass::Int => match r.hw_enc() {
-            ENC_RBX | ENC_RBP | ENC_RSI | ENC_RDI | ENC_R12 | ENC_R13 | ENC_R14 | ENC_R15 => true,
+            ENC_RBX | ENC_RBP | ENC_RSI | ENC_RDI | ENC_R12 | ENC_R13 | ENC_R14 => true,
+            // See above for SysV: we must treat the pinned reg specially.
+            ENC_R15 => !enable_pinned_reg,
             _ => false,
         },
         RegClass::Float => match r.hw_enc() {

cranelift/codegen/src/machinst/abi_impl.rs

@@ -428,6 +428,7 @@ pub trait ABIMachineSpec {
     /// contains the registers in a sorted order.
     fn get_clobbered_callee_saves(
         call_conv: isa::CallConv,
+        flags: &settings::Flags,
         regs: &[Writable<RealReg>],
     ) -> Vec<Writable<RealReg>>;
 
@@ -1224,7 +1225,8 @@ impl<M: ABIMachineSpec> ABICallee for ABICalleeImpl<M> {
         }
         let mask = M::stack_align(self.call_conv) - 1;
         let total_stacksize = (total_stacksize + mask) & !mask; // 16-align the stack.
-        let clobbered_callee_saves = M::get_clobbered_callee_saves(self.call_conv, &self.clobbered);
+        let clobbered_callee_saves =
+            M::get_clobbered_callee_saves(self.call_conv, &self.flags, &self.clobbered);
         let mut insts = smallvec![];
 
         if !self.call_conv.extends_baldrdash() {

cranelift/filetests/filetests/isa/aarch64/pinned-reg.clif

@@ -0,0 +1,16 @@
+test compile precise-output
+set enable_pinned_reg=true
+target aarch64
+
+function %f0() {
+block0:
+    v1 = get_pinned_reg.i64
+    v2 = iadd_imm v1, 1
+    set_pinned_reg v2
+    return
+}
+
+; block0:
+;   add x21, x21, #1
+;   ret
+

cranelift/filetests/filetests/isa/x64/pinned-reg.clif

@@ -0,0 +1,36 @@
+test compile precise-output
+set enable_pinned_reg=true
+target x86_64
+
+function %f0() {
+block0:
+    v1 = get_pinned_reg.i64
+    v2 = iadd_imm v1, 1
+    set_pinned_reg v2
+    return
+}
+
+;   pushq   %rbp
+;   movq    %rsp, %rbp
+; block0:
+;   addq    %r15, $1, %r15
+;   movq    %rbp, %rsp
+;   popq    %rbp
+;   ret
+
+function %f1() windows_fastcall {
+block0:
+    v1 = get_pinned_reg.i64
+    v2 = iadd_imm v1, 1
+    set_pinned_reg v2
+    return
+}
+
+;   pushq   %rbp
+;   movq    %rsp, %rbp
+; block0:
+;   addq    %r15, $1, %r15
+;   movq    %rbp, %rsp
+;   popq    %rbp
+;   ret
+