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/aarch64/inst.isle

@@ -165,6 +165,12 @@
         (rd WritableReg)
         (rm Reg))
 
+       ;; Like `Move` but with a particular `PReg` source (for implementing CLIF
+       ;; instructions like `get_stack_pointer`).
+       (MovPReg
+        (rd WritableReg)
+        (rm PReg))
+
        ;; A MOV[Z,N,K] with a 16-bit immediate.
        (MovWide
         (op MoveWideOp)
@@ -2421,3 +2427,31 @@
         ;; And finally, copy the preordained AtomicCASLoop output reg to its destination.
         ;; Also, x24 and x28 are trashed.
         (mov64_from_real 27)))
+
+;; Helper for emitting `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_sp () PReg)
+(extern constructor preg_sp preg_sp)
+
+(decl preg_fp () PReg)
+(extern constructor preg_fp preg_fp)
+
+(decl preg_link () PReg)
+(extern constructor preg_link preg_link)
+
+(decl aarch64_sp () Reg)
+(rule (aarch64_sp)
+      (mov_preg (preg_sp)))
+
+(decl aarch64_fp () Reg)
+(rule (aarch64_fp)
+      (mov_preg (preg_fp)))
+
+(decl aarch64_link () Reg)
+(rule (aarch64_link)
+      (mov_preg (preg_link)))

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

@@ -1334,6 +1334,15 @@ impl MachInstEmit for Inst {
                     }
                 }
             }
+            &Inst::MovPReg { rd, rm } => {
+                let rd = allocs.next_writable(rd);
+                let rm: Reg = rm.into();
+                debug_assert!([regs::fp_reg(), regs::stack_reg(), regs::link_reg()].contains(&rm));
+                assert!(rm.class() == RegClass::Int);
+                assert!(rd.to_reg().class() == rm.class());
+                let size = OperandSize::Size64;
+                Inst::Mov { size, rd, rm }.emit(&[], sink, emit_info, state);
+            }
             &Inst::MovWide { op, rd, imm, size } => {
                 let rd = allocs.next_writable(rd);
                 sink.put4(enc_move_wide(op, rd, imm, size));

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

@@ -649,6 +649,13 @@ fn aarch64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
             collector.reg_def(rd);
             collector.reg_use(rm);
         }
+        &Inst::MovPReg { rd, rm } => {
+            debug_assert!(
+                [regs::fp_reg(), regs::stack_reg(), regs::link_reg()].contains(&rm.into())
+            );
+            debug_assert!(rd.to_reg().is_virtual());
+            collector.reg_def(rd);
+        }
         &Inst::MovWide { op, rd, .. } => match op {
             MoveWideOp::MovK => collector.reg_mod(rd),
             _ => collector.reg_def(rd),
@@ -1474,6 +1481,11 @@ impl Inst {
                 let rm = pretty_print_ireg(rm, size, allocs);
                 format!("mov {}, {}", rd, rm)
             }
+            &Inst::MovPReg { rd, rm } => {
+                let rd = pretty_print_ireg(rd.to_reg(), OperandSize::Size64, allocs);
+                let rm = show_ireg_sized(rm.into(), OperandSize::Size64);
+                format!("mov {}, {}", rd, rm)
+            }
             &Inst::MovWide {
                 op,
                 rd,

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

@@ -1697,6 +1697,7 @@
 
 
 ;;;; Rules for `uunarrow` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
 (rule (lower (has_type (ty_vec128_int ty) (uunarrow x y)))
       (if (zero_value y))
       (uqxtn x (lane_size ty)))
@@ -1733,3 +1734,14 @@
 
 (rule (lower (debugtrap))
       (side_effect (brk)))
+
+;;; Rules for `get_{frame,stack}_pointer` and `get_return_address` ;;;;;;;;;;;;;
+
+(rule (lower (get_frame_pointer))
+      (aarch64_fp))
+
+(rule (lower (get_stack_pointer))
+      (aarch64_sp))
+
+(rule (lower (get_return_address))
+      (aarch64_link))

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

@@ -26,6 +26,7 @@ use crate::{
     isa::unwind::UnwindInst,
     machinst::{ty_bits, InsnOutput, LowerCtx, VCodeConstant, VCodeConstantData},
 };
+use regalloc2::PReg;
 use std::boxed::Box;
 use std::convert::TryFrom;
 use std::vec::Vec;
@@ -466,4 +467,16 @@ where
 
         rd.to_reg()
     }
+
+    fn preg_sp(&mut self) -> PReg {
+        super::regs::stack_reg().to_real_reg().unwrap().into()
+    }
+
+    fn preg_fp(&mut self) -> PReg {
+        super::regs::fp_reg().to_real_reg().unwrap().into()
+    }
+
+    fn preg_link(&mut self) -> PReg {
+        super::regs::link_reg().to_real_reg().unwrap().into()
+    }
 }

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

@@ -53,6 +53,10 @@ pub(crate) fn lower_insn_to_regs<C: LowerCtx<I = Inst>>(
             point, as constants are rematerialized at use-sites"
         ),
 
+        Opcode::GetFramePointer | Opcode::GetStackPointer | Opcode::GetReturnAddress => {
+            implemented_in_isle(ctx)
+        }
+
         Opcode::Iadd => implemented_in_isle(ctx),
         Opcode::Isub => implemented_in_isle(ctx),
         Opcode::UaddSat | Opcode::SaddSat | Opcode::UsubSat | Opcode::SsubSat => {