Cranelift: Add instructions for getting the current stack/frame/return pointers (#4573)

* Cranelift: Add instructions for getting the current stack/frame pointers and return address

This is the initial part of https://github.com/bytecodealliance/wasmtime/issues/4535

* x64: Remove `Amode::RbpOffset` and use `Amode::ImmReg` instead

We just special case getting operands from `Amode`s now.

* Fix s390x `get_return_address`; require `preserve_frame_pointers=true`

* Assert that `Amode::ImmRegRegShift` doesn't use rbp/rsp

* Handle non-allocatable registers in Amode::with_allocs

* Use "stack" instead of "r15" on s390x

* r14 is an allocatable register on s390x, so it shouldn't be used with `MovPReg`

cranelift/codegen/src/isa/s390x/inst.isle

@@ -373,6 +373,11 @@
       (rd WritableReg)
       (rm Reg))
 
+    ;; Like `Mov64` but with a particular physical register source.
+    (MovPReg
+      (rd WritableReg)
+      (rm PReg))
+
     ;; A 32-bit move instruction with a full 32-bit immediate.
     (Mov32Imm
       (rd WritableReg)
@@ -1556,6 +1561,10 @@
 (decl memarg_stack_off (i64 i64) MemArg)
 (extern constructor memarg_stack_off memarg_stack_off)
 
+;; Create a `MemArg` referring to an offset from the initial SP.
+(decl memarg_initial_sp_offset (i64) MemArg)
+(extern constructor memarg_initial_sp_offset memarg_initial_sp_offset)
+
 ;; Form the sum of two offset values, and check that the result is
 ;; a valid `MemArg::Symbol` offset (i.e. is even and fits into i32).
 (decl pure memarg_symbol_offset_sum (i64 i64) i32)
@@ -2469,6 +2478,20 @@
 (rule (emit_load $I64 dst addr)
       (emit (MInst.Load64 dst addr)))
 
+;; Helper for creating `MInst.MovPReg` instructions.
+(decl mov_preg (PReg) Reg)
+(rule (mov_preg src)
+      (let ((dst WritableReg (temp_writable_reg $I64))
+            (_ Unit (emit (MInst.MovPReg dst src))))
+        dst))
+
+(decl preg_stack () PReg)
+(extern constructor preg_stack preg_stack)
+
+;; Copy the physical stack register into a virtual register.
+(decl sp () Reg)
+(rule (sp)
+      (mov_preg (preg_stack)))
 
 ;; Helpers for accessing argument / return value slots ;;;;;;;;;;;;;;;;;;;;;;;;;
 

cranelift/codegen/src/isa/s390x/inst/emit.rs

@@ -2081,6 +2081,12 @@ impl MachInstEmit for Inst {
                 let opcode = 0xb904; // LGR
                 put(sink, &enc_rre(opcode, rd.to_reg(), rm));
             }
+            &Inst::MovPReg { rd, rm } => {
+                let rm: Reg = rm.into();
+                debug_assert!([regs::gpr(15)].contains(&rm));
+                let rd = allocs.next_writable(rd);
+                Inst::Mov64 { rd, rm }.emit(&[], sink, emit_info, state);
+            }
             &Inst::Mov32 { rd, rm } => {
                 let rd = allocs.next_writable(rd);
                 let rm = allocs.next(rm);

cranelift/codegen/src/isa/s390x/inst/mod.rs

@@ -144,6 +144,7 @@ impl Inst {
             | Inst::StoreMultiple64 { .. }
             | Inst::Mov32 { .. }
             | Inst::Mov64 { .. }
+            | Inst::MovPReg { .. }
             | Inst::Mov32Imm { .. }
             | Inst::Mov32SImm16 { .. }
             | Inst::Mov64SImm16 { .. }
@@ -623,6 +624,11 @@ fn s390x_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut OperandC
             collector.reg_def(rd);
             collector.reg_use(rm);
         }
+        &Inst::MovPReg { rd, rm } => {
+            debug_assert!([regs::gpr(14), regs::gpr(15)].contains(&rm.into()));
+            debug_assert!(rd.to_reg().is_virtual());
+            collector.reg_def(rd);
+        }
         &Inst::Mov32 { rd, rm } => {
             collector.reg_def(rd);
             collector.reg_use(rm);
@@ -1778,6 +1784,11 @@ impl Inst {
                 let rm = pretty_print_reg(rm, allocs);
                 format!("lgr {}, {}", rd, rm)
             }
+            &Inst::MovPReg { rd, rm } => {
+                let rd = pretty_print_reg(rd.to_reg(), allocs);
+                let rm = show_reg(rm.into());
+                format!("lgr {}, {}", rd, rm)
+            }
             &Inst::Mov32 { rd, rm } => {
                 let rd = pretty_print_reg(rd.to_reg(), allocs);
                 let rm = pretty_print_reg(rm, allocs);

cranelift/codegen/src/isa/s390x/lower.isle

@@ -3616,3 +3616,15 @@
             (_ Unit (output_builder_push builder ret)))
         (lower_call_rets abi tail builder)))
 
+;;;; Rules for `get_{frame,stack}_pointer` and `get_return_address` ;;;;;;;;;;;;
+
+(rule (lower (get_stack_pointer))
+      (sp))
+
+(rule (lower (get_frame_pointer))
+      (load64 (memarg_stack_off 0 0)))
+
+(rule (lower (get_return_address))
+      ;; The return address is 14 pointer-sized slots above the initial SP. So
+      ;; our offset is `14 * 8 = 112`.
+      (load64 (memarg_initial_sp_offset 112)))

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

@@ -180,7 +180,10 @@ impl LowerBackend for S390xBackend {
             | Opcode::Return
             | Opcode::StackAddr
             | Opcode::FuncAddr
-            | Opcode::SymbolValue => {
+            | Opcode::SymbolValue
+            | Opcode::GetFramePointer
+            | Opcode::GetStackPointer
+            | Opcode::GetReturnAddress => {
                 unreachable!(
                     "implemented in ISLE: inst = `{}`, type = `{:?}`",
                     ctx.dfg().display_inst(ir_inst),

cranelift/codegen/src/isa/s390x/lower/isle.rs

@@ -21,6 +21,7 @@ use crate::{
     isa::unwind::UnwindInst,
     machinst::{InsnOutput, LowerCtx, VCodeConstant, VCodeConstantData},
 };
+use regalloc2::PReg;
 use std::boxed::Box;
 use std::cell::Cell;
 use std::convert::TryFrom;
@@ -603,6 +604,11 @@ where
         MemArg::reg_plus_off(stack_reg(), base + off, MemFlags::trusted())
     }
 
+    #[inline]
+    fn memarg_initial_sp_offset(&mut self, off: i64) -> MemArg {
+        MemArg::InitialSPOffset { off }
+    }
+
     #[inline]
     fn memarg_symbol(&mut self, name: ExternalName, offset: i32, flags: MemFlags) -> MemArg {
         MemArg::Symbol {
@@ -670,6 +676,11 @@ where
     fn emit(&mut self, inst: &MInst) -> Unit {
         self.lower_ctx.emit(inst.clone());
     }
+
+    #[inline]
+    fn preg_stack(&mut self) -> PReg {
+        stack_reg().to_real_reg().unwrap().into()
+    }
 }
 
 /// Zero-extend the low `from_bits` bits of `value` to a full u64.