Cranelift: use SP-offset amodes for stack_addr+load/store. (#11727)

We provide `stack_load`/ `stack_store` / `stack_addr` instructions in
Cranelift to operate on stack slots, and the first two are legalized
to a `stack_addr` plus an ordinary load or store instruction.

We currently have lowerings for `stack_addr` that materialize an
SP-relative address into a register: for example, `leaq 8(%rsp), %rax`
on x86-64 or `add x0, sp, #8` on aarch64.

Taken together, we see sequences like (aarch64 / x86-64)

```
    add x0, sp, #8       /   leaq 8(%rsp), %rax
    str x1, [x0]         /   movq %rdx, (%rax)
```

when using `stack_store`s. In particular, we do *not* use the direct
SP-relative form, which would look like

```
    str x1, [sp, #8]     /   movq %rdx, 8(%rsp)
```

and which we can already generate in other cases, e.g. spillslot
moves (spills/reloads) and clobber saves/restores.

This inefficiency is undesirable whenever the embedder is using
stackslots, but in particular when we expect to have high memory
traffic to stack slots (e.g., I am seeing this now when implementing
debug instrumentation in Wasmtime, and user stack map instrumentation
for GC will also benefit).

This PR adds new lowerings that use the existing synthetic address
mode we already use for spillslots to emit loads/stores to stackslots
directly when possible. The PR does this for x86-64 and aarch64;
others could be updated later.

Author: cfallin

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

@@ -3899,6 +3899,11 @@
         (if-let new_offset (i32_checked_add x offset))
         (amode_no_more_iconst ty y new_offset))
 
+(rule 3
+      (amode ty (stack_addr slot offset1) offset2)
+      (AMode.SlotOffset
+       (abi_stackslot_offset_into_slot_region slot offset1 offset2)))
+
 (decl amode_no_more_iconst (Type Value i32) AMode)
 ;; Base case: move the `offset` into a register and add it to `val` via the
 ;; amode

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

@@ -403,6 +403,15 @@
 (decl amode_to_synthetic_amode (Amode) SyntheticAmode)
 (extern constructor amode_to_synthetic_amode amode_to_synthetic_amode)
 
+(decl synthetic_amode_slot (i32) SyntheticAmode)
+(extern constructor synthetic_amode_slot synthetic_amode_slot)
+
+;; Helper for loads/stores to/from stackslots.
+(decl stackslot_amode (StackSlot Offset32 Offset32) SyntheticAmode)
+(rule (stackslot_amode slot offset1 offset2)
+      (let ((slot_offset i32 (abi_stackslot_offset_into_slot_region slot offset1 offset2)))
+        (synthetic_amode_slot slot_offset)))
+
 ;; An `Amode` represents a possible addressing mode that can be used
 ;; in instructions. These denote a 64-bit value only.
 (type Amode (enum
@@ -494,12 +503,16 @@
       (provide (= result (concat flags (bvadd val (sign_ext 64 offset)))))
       (require
             (= (widthof val) 64)))
-(decl to_amode (MemFlags Value Offset32) Amode)
+(decl to_amode (MemFlags Value Offset32) SyntheticAmode)
 (rule 0 (to_amode flags base offset)
         (amode_imm_reg flags base offset))
 (rule 1 (to_amode flags (iadd x y) offset)
         (to_amode_add flags x y offset))
 
+(rule 2
+      (to_amode flags (stack_addr slot offset1) offset2)
+      (stackslot_amode slot offset1 offset2))
+
 ;; Same as `to_amode`, except that the base address is computed via the addition
 ;; of the two `Value` arguments provided.
 ;;
@@ -572,7 +585,7 @@
 
 ;; Offsetting an Amode. Used when we need to do consecutive
 ;; loads/stores to adjacent addresses.
-(decl amode_offset (Amode i32) Amode)
+(decl amode_offset (SyntheticAmode i32) SyntheticAmode)
 (extern constructor amode_offset amode_offset)
 
 ;; Return a zero offset as an `Offset32`.
@@ -1368,7 +1381,7 @@
 (rule 0 (x64_load (multi_lane _bits _lanes) addr _ext_kind)
       (x64_movdqu_load addr))
 
-(decl x64_mov (Amode) Reg)
+(decl x64_mov (SyntheticAmode) Reg)
 (spec (x64_mov addr)
       (provide (= result (conv_to 64 (load_effect (extract 79 64 addr) 64 (extract 63 0 addr))))))
 (rule (x64_mov addr) (x64_movq_rm addr))
@@ -2701,25 +2714,25 @@
 (decl x64_pextrb (Xmm u8) Gpr)
 (rule (x64_pextrb src lane) (x64_pextrb_a_or_avx src lane))
 
-(decl x64_pextrb_store (Amode Xmm u8) SideEffectNoResult)
+(decl x64_pextrb_store (SyntheticAmode Xmm u8) SideEffectNoResult)
 (rule (x64_pextrb_store addr src lane) (x64_pextrb_a_mem_or_avx addr src lane))
 
 (decl x64_pextrw (Xmm u8) Gpr)
 (rule (x64_pextrw src lane) (x64_pextrw_a_or_avx src lane))
 
-(decl x64_pextrw_store (Amode Xmm u8) SideEffectNoResult)
+(decl x64_pextrw_store (SyntheticAmode Xmm u8) SideEffectNoResult)
 (rule (x64_pextrw_store addr src lane) (x64_pextrw_b_mem_or_avx addr src lane))
 
 (decl x64_pextrd (Xmm u8) Gpr)
 (rule (x64_pextrd src lane) (x64_pextrd_a_or_avx src lane))
 
-(decl x64_pextrd_store (Amode Xmm u8) SideEffectNoResult)
+(decl x64_pextrd_store (SyntheticAmode Xmm u8) SideEffectNoResult)
 (rule (x64_pextrd_store addr src lane) (x64_pextrd_a_mem_or_avx addr src lane))
 
 (decl x64_pextrq (Xmm u8) Gpr)
 (rule (x64_pextrq src lane) (x64_pextrq_a_or_avx src lane))
 
-(decl x64_pextrq_store (Amode Xmm u8) SideEffectNoResult)
+(decl x64_pextrq_store (SyntheticAmode Xmm u8) SideEffectNoResult)
 (rule (x64_pextrq_store addr src lane) (x64_pextrq_a_mem_or_avx addr src lane))
 
 ;; Helper for creating `pmovmskb` instructions.
@@ -3181,7 +3194,7 @@
           )
           (x64_por low_gt_and_high_eq high_halves_gt)))
 
-(decl x64_add_mem (Type Amode Value) SideEffectNoResult)
+(decl x64_add_mem (Type SyntheticAmode Value) SideEffectNoResult)
 (spec (x64_add_mem ty addr val)
       (provide (= result (store_effect
                         (extract 79 64 addr)
@@ -3207,7 +3220,7 @@
 (rule 2 (x64_add_mem $I32 addr (i8_from_iconst val)) (x64_addl_mi_sxb_mem addr val))
 (rule 2 (x64_add_mem $I64 addr (i8_from_iconst val)) (x64_addq_mi_sxb_mem addr val))
 
-(decl x64_sub_mem (Type Amode Value) SideEffectNoResult)
+(decl x64_sub_mem (Type SyntheticAmode Value) SideEffectNoResult)
 
 ;; `sub mem, reg`
 (rule 0 (x64_sub_mem $I8 addr val) (x64_subb_mr_mem addr val))
@@ -3223,7 +3236,7 @@
 (rule 2 (x64_sub_mem $I32 addr (i8_from_iconst val)) (x64_subl_mi_sxb_mem addr val))
 (rule 2 (x64_sub_mem $I64 addr (i8_from_iconst val)) (x64_subq_mi_sxb_mem addr val))
 
-(decl x64_and_mem (Type Amode Value) SideEffectNoResult)
+(decl x64_and_mem (Type SyntheticAmode Value) SideEffectNoResult)
 
 ;; `and mem, imm`
 (rule (x64_and_mem $I8 addr val) (x64_andb_mr_mem addr val))
@@ -3241,7 +3254,7 @@
 (rule 2 (x64_and_mem $I32 addr (i8_from_iconst val)) (x64_andl_mi_sxb_mem addr val))
 (rule 2 (x64_and_mem $I64 addr (i8_from_iconst val)) (x64_andq_mi_sxb_mem addr val))
 
-(decl x64_or_mem (Type Amode Value) SideEffectNoResult)
+(decl x64_or_mem (Type SyntheticAmode Value) SideEffectNoResult)
 
 ;; `or mem, reg`
 (rule 0 (x64_or_mem $I8 addr val) (x64_orb_mr_mem addr val))
@@ -3259,7 +3272,7 @@
 (rule 2 (x64_or_mem $I32 addr (i8_from_iconst val)) (x64_orl_mi_sxb_mem addr val))
 (rule 2 (x64_or_mem $I64 addr (i8_from_iconst val)) (x64_orq_mi_sxb_mem addr val))
 
-(decl x64_xor_mem (Type Amode Value) SideEffectNoResult)
+(decl x64_xor_mem (Type SyntheticAmode Value) SideEffectNoResult)
 
 ;; `xor mem, reg`
 (rule 0 (x64_xor_mem $I8 addr val) (x64_xorb_mr_mem addr val))
@@ -3644,31 +3657,31 @@
 (rule (x64_xchg $I32 addr operand) (x64_xchgl_rm operand addr))
 (rule (x64_xchg $I64 addr operand) (x64_xchgq_rm operand addr))
 
-(decl x64_lock_add (OperandSize Amode Gpr) SideEffectNoResult)
+(decl x64_lock_add (OperandSize SyntheticAmode Gpr) SideEffectNoResult)
 (rule (x64_lock_add (OperandSize.Size8) addr reg)   (x64_lock_addb_mr_mem addr reg))
 (rule (x64_lock_add (OperandSize.Size16) addr reg)  (x64_lock_addw_mr_mem addr reg))
 (rule (x64_lock_add (OperandSize.Size32) addr reg)  (x64_lock_addl_mr_mem addr reg))
 (rule (x64_lock_add (OperandSize.Size64) addr reg)  (x64_lock_addq_mr_mem addr reg))
 
-(decl x64_lock_sub (OperandSize Amode Gpr) SideEffectNoResult)
+(decl x64_lock_sub (OperandSize SyntheticAmode Gpr) SideEffectNoResult)
 (rule (x64_lock_sub (OperandSize.Size8) addr reg)   (x64_lock_subb_mr_mem addr reg))
 (rule (x64_lock_sub (OperandSize.Size16) addr reg)  (x64_lock_subw_mr_mem addr reg))
 (rule (x64_lock_sub (OperandSize.Size32) addr reg)  (x64_lock_subl_mr_mem addr reg))
 (rule (x64_lock_sub (OperandSize.Size64) addr reg)  (x64_lock_subq_mr_mem addr reg))
 
-(decl x64_lock_and (OperandSize Amode Gpr) SideEffectNoResult)
+(decl x64_lock_and (OperandSize SyntheticAmode Gpr) SideEffectNoResult)
 (rule (x64_lock_and (OperandSize.Size8) addr reg)   (x64_lock_andb_mr_mem addr reg))
 (rule (x64_lock_and (OperandSize.Size16) addr reg)  (x64_lock_andw_mr_mem addr reg))
 (rule (x64_lock_and (OperandSize.Size32) addr reg)  (x64_lock_andl_mr_mem addr reg))
 (rule (x64_lock_and (OperandSize.Size64) addr reg)  (x64_lock_andq_mr_mem addr reg))
 
-(decl x64_lock_or (OperandSize Amode Gpr) SideEffectNoResult)
+(decl x64_lock_or (OperandSize SyntheticAmode Gpr) SideEffectNoResult)
 (rule (x64_lock_or (OperandSize.Size8) addr reg)   (x64_lock_orb_mr_mem addr reg))
 (rule (x64_lock_or (OperandSize.Size16) addr reg)  (x64_lock_orw_mr_mem addr reg))
 (rule (x64_lock_or (OperandSize.Size32) addr reg)  (x64_lock_orl_mr_mem addr reg))
 (rule (x64_lock_or (OperandSize.Size64) addr reg)  (x64_lock_orq_mr_mem addr reg))
 
-(decl x64_lock_xor (OperandSize Amode Gpr) SideEffectNoResult)
+(decl x64_lock_xor (OperandSize SyntheticAmode Gpr) SideEffectNoResult)
 (rule (x64_lock_xor (OperandSize.Size8) addr reg)   (x64_lock_xorb_mr_mem addr reg))
 (rule (x64_lock_xor (OperandSize.Size16) addr reg)  (x64_lock_xorw_mr_mem addr reg))
 (rule (x64_lock_xor (OperandSize.Size32) addr reg)  (x64_lock_xorl_mr_mem addr reg))

cranelift/codegen/src/isa/x64/inst/args.rs

@@ -573,6 +573,21 @@ impl SyntheticAmode {
             | SyntheticAmode::ConstantOffset { .. } => true,
         }
     }
+
+    /// Offset the synthetic amode by a fixed offset.
+    pub(crate) fn offset(&self, offset: i32) -> Self {
+        let mut ret = self.clone();
+        match &mut ret {
+            SyntheticAmode::Real(amode) => *amode = amode.offset(offset),
+            SyntheticAmode::SlotOffset { simm32 } => *simm32 += offset,
+            // `amode_offset` is used only in i128.load/store which
+            // takes a synthetic amode from `to_amode`; `to_amode` can
+            // only produce Real or SlotOffset amodes, never
+            // IncomingArg or ConstantOffset.
+            _ => panic!("Cannot offset SyntheticAmode: {self:?}"),
+        }
+        ret
+    }
 }
 
 impl From<Amode> for SyntheticAmode {

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

@@ -3056,8 +3056,8 @@
 ;; We can load an I128 by doing two 64-bit loads.
 (rule -3 (lower (has_type $I128
                        (load (little_or_native_endian flags) address offset)))
-      (let ((addr_lo Amode (to_amode flags address offset))
-            (addr_hi Amode (amode_offset addr_lo 8))
+      (let ((addr_lo SyntheticAmode (to_amode flags address offset))
+            (addr_hi SyntheticAmode (amode_offset addr_lo 8))
             (value_lo Reg (x64_mov addr_lo))
             (value_hi Reg (x64_mov addr_hi)))
         (value_regs value_lo value_hi)))
@@ -3187,8 +3187,8 @@
       (let ((value_reg ValueRegs value)
             (value_lo Gpr (value_regs_get_gpr value_reg 0))
             (value_hi Gpr (value_regs_get_gpr value_reg 1))
-            (addr_lo Amode (to_amode flags address offset))
-            (addr_hi Amode (amode_offset addr_lo 8)))
+            (addr_lo SyntheticAmode (to_amode flags address offset))
+            (addr_hi SyntheticAmode (amode_offset addr_lo 8)))
       (side_effect
        (side_effect_concat
         (x64_movrm $I64 addr_lo value_lo)

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

@@ -409,6 +409,11 @@ impl Context for IsleContext<'_, '_, MInst, X64Backend> {
         amode.clone().into()
     }
 
+    #[inline]
+    fn synthetic_amode_slot(&mut self, offset: i32) -> SyntheticAmode {
+        SyntheticAmode::SlotOffset { simm32: offset }
+    }
+
     #[inline]
     fn const_to_synthetic_amode(&mut self, c: VCodeConstant) -> SyntheticAmode {
         SyntheticAmode::ConstantOffset(c)
@@ -633,7 +638,7 @@ impl Context for IsleContext<'_, '_, MInst, X64Backend> {
     }
 
     #[inline]
-    fn amode_offset(&mut self, addr: &Amode, offset: i32) -> Amode {
+    fn amode_offset(&mut self, addr: &SyntheticAmode, offset: i32) -> SyntheticAmode {
         addr.offset(offset)
     }
 

cranelift/codegen/src/machinst/abi.rs

@@ -2139,6 +2139,11 @@ impl<M: ABIMachineSpec> Callee<M> {
         }
     }
 
+    /// Get the raw offset of a sized stackslot in the slot region.
+    pub fn sized_stackslot_offset(&self, slot: StackSlot) -> u32 {
+        self.sized_stackslots[slot]
+    }
+
     /// Produce an instruction that computes a sized stackslot address.
     pub fn sized_stackslot_addr(
         &self,

cranelift/codegen/src/machinst/isle.rs

@@ -507,6 +507,22 @@ macro_rules! isle_lower_prelude_methods {
                 .into()
         }
 
+        fn abi_stackslot_offset_into_slot_region(
+            &mut self,
+            stack_slot: StackSlot,
+            offset1: Offset32,
+            offset2: Offset32,
+        ) -> i32 {
+            let offset1 = i32::from(offset1);
+            let offset2 = i32::from(offset2);
+            i32::try_from(self.lower_ctx.abi().sized_stackslot_offset(stack_slot))
+                .expect("Stack slot region cannot be larger than 2GiB")
+                .checked_add(offset1)
+                .expect("Stack slot region cannot be larger than 2GiB")
+                .checked_add(offset2)
+                .expect("Stack slot region cannot be larger than 2GiB")
+        }
+
         fn abi_dynamic_stackslot_addr(
             &mut self,
             dst: WritableReg,

cranelift/codegen/src/prelude_lower.isle

@@ -1119,6 +1119,10 @@
 (decl abi_stackslot_addr (WritableReg StackSlot Offset32) MInst)
 (extern constructor abi_stackslot_addr abi_stackslot_addr)
 
+;; StackSlot raw offset into slot region
+(decl abi_stackslot_offset_into_slot_region (StackSlot Offset32 Offset32) i32)
+(extern constructor abi_stackslot_offset_into_slot_region abi_stackslot_offset_into_slot_region)
+
 ;; DynamicStackSlot addr
 (decl abi_dynamic_stackslot_addr (WritableReg DynamicStackSlot) MInst)
 (extern constructor abi_dynamic_stackslot_addr abi_dynamic_stackslot_addr)

cranelift/filetests/filetests/isa/aarch64/dynamic-slot.clif

@@ -16,9 +16,8 @@ block0:
 ;   mov fp, sp
 ;   sub sp, sp, #16
 ; block0:
-;   movz x1, #1
-;   mov x2, sp
-;   str x1, [x2]
+;   movz x0, #1
+;   str x0, [sp]
 ;   add sp, sp, #16
 ;   ldp fp, lr, [sp], #16
 ;   ret
@@ -29,9 +28,8 @@ block0:
 ;   mov x29, sp
 ;   sub sp, sp, #0x10
 ; block1: ; offset 0xc
-;   mov x1, #1
-;   mov x2, sp
-;   str x1, [x2]
+;   mov x0, #1
+;   stur x0, [sp]
 ;   add sp, sp, #0x10
 ;   ldp x29, x30, [sp], #0x10
 ;   ret
@@ -51,9 +49,8 @@ block0:
 ;   mov fp, sp
 ;   sub sp, sp, #16
 ; block0:
-;   movz x1, #1
-;   mov x2, sp
-;   str x1, [x2]
+;   movz x0, #1
+;   str x0, [sp]
 ;   add sp, sp, #16
 ;   ldp fp, lr, [sp], #16
 ;   ret
@@ -64,9 +61,8 @@ block0:
 ;   mov x29, sp
 ;   sub sp, sp, #0x10
 ; block1: ; offset 0xc
-;   mov x1, #1
-;   mov x2, sp
-;   str x1, [x2]
+;   mov x0, #1
+;   stur x0, [sp]
 ;   add sp, sp, #0x10
 ;   ldp x29, x30, [sp], #0x10
 ;   ret