Add correct ISLE rules for memory operations

This is an intermediate step that makes ISLE rules support generating memory loads and stores.
The test file `memory.wat` still fails because we actually need to implement ZK ASM instruction generation correctly.

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

@@ -1267,32 +1267,75 @@
       (_ Unit (emit (MInst.Popcnt sum step tmp rs ty))))
     (writable_reg_to_reg sum)))
 
-(decl gen_amode (Reg Offset32 Type) AMode)
-(extern constructor gen_amode gen_amode)
+;; Generates a AMode that points to a register plus an offset.
+(decl gen_reg_offset_amode (Reg i64 Type) AMode)
+(extern constructor gen_reg_offset_amode gen_reg_offset_amode)
+
+;; Generates a AMode that an offset from the stack pointer.
+(decl gen_sp_offset_amode (i64 Type) AMode)
+(extern constructor gen_sp_offset_amode gen_sp_offset_amode)
+
+;; Generates a AMode that an offset from the frame pointer.
+(decl gen_fp_offset_amode (i64 Type) AMode)
+(extern constructor gen_fp_offset_amode gen_fp_offset_amode)
+
+;; Generates an AMode that points to a stack slot + offset.
+(decl gen_stack_slot_amode (StackSlot i64 Type) AMode)
+(extern constructor gen_stack_slot_amode gen_stack_slot_amode)
 
 ;; Generates a AMode that points to a constant in the constant pool.
 (decl gen_const_amode (VCodeConstant) AMode)
 (extern constructor gen_const_amode gen_const_amode)
 
-(decl offset32_imm (i32) Offset32)
-(extern constructor offset32_imm offset32_imm)
+;; Tries to match a Value + Offset into an AMode
+(decl amode (Value i32 Type) AMode)
+(rule 0 (amode addr offset ty) (amode_inner addr offset ty))
 
-;; helper function to load from memory.
-(decl gen_load (Reg Offset32 LoadOP MemFlags Type) Reg)
-(rule
-  (gen_load p offset op flags ty)
-  (let
-    ((tmp WritableReg (temp_writable_reg ty))
-      (_ Unit (emit (MInst.Load tmp op flags (gen_amode p offset $I64)))))
-    tmp))
+;; If we are adding a constant offset with an iadd we can instead make that
+;; offset part of the amode offset.
+;;
+;; We can't recurse into `amode` again since that could cause stack overflows.
+;; See: https://github.com/bytecodealliance/wasmtime/pull/6968
+(rule 1 (amode (iadd addr (iconst (simm32 y))) offset ty)
+  (if-let new_offset (s32_add_fallible y offset))
+  (amode_inner addr new_offset ty))
+(rule 2 (amode (iadd (iconst (simm32 x)) addr) offset ty)
+  (if-let new_offset (s32_add_fallible x offset))
+  (amode_inner addr new_offset ty))
+
+
+;; These are the normal rules for generating an AMode.
+(decl amode_inner (Value i32 Type) AMode)
+
+;; In the simplest case we just lower into a Reg+Offset
+(rule 0 (amode_inner r @ (value_type (ty_addr64 _)) offset ty)
+  (gen_reg_offset_amode r offset ty))
+
+;; If the value is a `get_frame_pointer`, we can just use the offset from that.
+(rule 1 (amode_inner (get_frame_pointer) offset ty)
+  (gen_fp_offset_amode offset ty))
+
+;; If the value is a `get_stack_pointer`, we can just use the offset from that.
+(rule 1 (amode_inner (get_stack_pointer) offset ty)
+  (gen_sp_offset_amode offset ty))
+
+;; Similarly if the value is a `stack_addr` we can also turn that into an sp offset.
+(rule 1 (amode_inner (stack_addr ss ss_offset) amode_offset ty)
+  (if-let combined_offset (s32_add_fallible ss_offset amode_offset))
+  (gen_stack_slot_amode ss combined_offset ty))
+
+;; Returns a canonical type for a LoadOP. We only return I64 or F64.
+(decl load_op_reg_type (LoadOP) Type)
+(rule 1 (load_op_reg_type (LoadOP.Fld)) $F64)
+(rule 1 (load_op_reg_type (LoadOP.Flw)) $F64)
+(rule 0 (load_op_reg_type _) $I64)
 
-(decl gen_load_128 (Reg Offset32 MemFlags) ValueRegs)
-(rule
-  (gen_load_128 p offset flags)
-  (let
-    ((low Reg (gen_load p offset (LoadOP.Ld) flags $I64))
-      (high Reg (gen_load p (offset32_add offset 8) (LoadOP.Ld) flags $I64)))
-    (value_regs low high)))
+;; helper function to load from memory.
+(decl gen_load (AMode LoadOP MemFlags) Reg)
+(rule (gen_load amode op flags)
+  (let ((dst WritableReg (temp_writable_reg (load_op_reg_type op)))
+      (_ Unit (emit (MInst.Load dst op flags amode))))
+    dst))
 
 (decl default_memflags () MemFlags)
 (extern constructor default_memflags default_memflags)
@@ -1301,19 +1344,9 @@
 (extern constructor offset32_add offset32_add)
 
 ;; helper function to store to memory.
-(decl gen_store (Reg Offset32 StoreOP MemFlags Reg) InstOutput)
-(rule
-  (gen_store base offset op flags src)
-  (side_effect (SideEffectNoResult.Inst (MInst.Store (gen_amode base offset $I64) op flags src)))
-)
-
-(decl gen_store_128 (Reg Offset32 MemFlags ValueRegs) InstOutput)
-(rule
-  (gen_store_128 p offset flags src)
-  (side_effect
-    (SideEffectNoResult.Inst2
-      (MInst.Store (gen_amode p offset $I64) (StoreOP.Sd) flags (value_regs_get src 0))
-      (MInst.Store (gen_amode p (offset32_add offset 8) $I64) (StoreOP.Sd) flags (value_regs_get src 1)))))
+(decl gen_store (AMode StoreOP MemFlags Reg) InstOutput)
+(rule (gen_store amode op flags src)
+  (side_effect (SideEffectNoResult.Inst (MInst.Store amode op flags src))))
 
 (decl valid_atomic_transaction (Type) Type)
 (extern extractor valid_atomic_transaction valid_atomic_transaction)
@@ -1372,33 +1405,6 @@
 (decl store_op (Type) StoreOP)
 (extern constructor store_op store_op)
 
-;; bool is "is_signed"
-(decl int_load_op (bool u8) LoadOP)
-(rule
-  (int_load_op $false 8)
-  (LoadOP.Lbu))
-
-(rule
-  (int_load_op $true 8)
-  (LoadOP.Lb))
-
-(rule
-  (int_load_op $false 16)
-  (LoadOP.Lhu))
-(rule
-  (int_load_op $true 16)
-  (LoadOP.Lh))
-(rule
-  (int_load_op $false 32)
-  (LoadOP.Lwu))
-(rule
-  (int_load_op $true 32)
-  (LoadOP.Lw))
-
-(rule
-  (int_load_op _ 64)
-  (LoadOP.Ld))
-
 ;;;; load extern name
 (decl load_ext_name (ExternalName i64) Reg)
 (extern constructor load_ext_name load_ext_name)

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

@@ -305,56 +305,48 @@
   (udf code))
 
 ;;;;;  Rules for `uload8`;;;;;;;;;
-(rule
-  (lower (uload8 flags p @ (value_type (ty_addr64 _)) offset))
-  (gen_load p offset (int_load_op $false 8) flags $I64))
+(rule (lower (uload8 flags addr offset))
+  (gen_load (amode addr offset $I8) (LoadOP.Lbu) flags))
+
 ;;;;;  Rules for `sload8`;;;;;;;;;
-(rule
-  (lower (sload8 flags p @ (value_type (ty_addr64 _)) offset))
-  (gen_load p offset (int_load_op $true 8) flags $I64))
+(rule (lower (sload8 flags addr offset))
+  (gen_load (amode addr offset $I8) (LoadOP.Lb) flags))
+
 ;;;;;  Rules for `uload16`;;;;;;;;;
-(rule
-  (lower (uload16 flags p @ (value_type (ty_addr64 _)) offset))
-  (gen_load p offset (int_load_op $false 16) flags $I64))
+(rule (lower (uload16 flags addr offset))
+  (gen_load (amode addr offset $I16) (LoadOP.Lhu) flags))
 
 ;;;;;  Rules for `iload16`;;;;;;;;;
-(rule
-  (lower (sload16 flags p @ (value_type (ty_addr64 _)) offset))
-  (gen_load p offset (int_load_op $true 16) flags $I64))
+(rule (lower (sload16 flags addr offset))
+  (gen_load (amode addr offset $I16) (LoadOP.Lh) flags))
 
 ;;;;;  Rules for `uload32`;;;;;;;;;
-(rule
-  (lower (uload32 flags p @ (value_type (ty_addr64 _)) offset))
-  (gen_load p offset (int_load_op $false 32) flags $I64))
+(rule (lower (uload32 flags addr offset))
+  (gen_load (amode addr offset $I32) (LoadOP.Lwu) flags))
 
-;;;;;  Rules for `iload32`;;;;;;;;;
-(rule
-  (lower (sload32 flags p @ (value_type (ty_addr64 _)) offset))
-  (gen_load p offset (int_load_op $true 32) flags $I64))
+;;;;;  Rules for `sload32`;;;;;;;;;
+(rule (lower (sload32 flags addr offset))
+  (gen_load (amode addr offset $I32) (LoadOP.Lw) flags))
 
-(rule
-  (lower (has_type ty (load flags p @ (value_type (ty_addr32 _)) offset)))
-  (gen_load p offset (load_op ty) flags ty)
-)
+;;;;;  Rules for `load`;;;;;;;;;
+(rule (lower (has_type ty (load flags addr offset)))
+  (gen_load (amode addr offset ty) (load_op ty) flags))
 
 ;;;;;  Rules for `istore8`;;;;;;;;;
-(rule
-  (lower (istore8 flags x p @ (value_type (ty_addr64 _)) offset))
-  (gen_store p offset (StoreOP.Sb) flags x))
+(rule (lower (istore8 flags src addr offset))
+  (gen_store (amode addr offset $I8) (StoreOP.Sb) flags src))
+
 ;;;;;  Rules for `istore16`;;;;;;;;;
-(rule
-  (lower (istore16 flags x p @ (value_type (ty_addr64 _)) offset))
-  (gen_store p offset (StoreOP.Sh) flags x))
+(rule (lower (istore16 flags src addr offset))
+  (gen_store (amode addr offset $I16) (StoreOP.Sh) flags src))
 
 ;;;;;  Rules for `istore32`;;;;;;;;;
-(rule
-  (lower (istore32 flags x p @ (value_type (ty_addr64 _)) offset))
-  (gen_store p offset (StoreOP.Sw) flags x))
+(rule (lower (istore32 flags src addr offset))
+  (gen_store (amode addr offset $I32) (StoreOP.Sw) flags src))
 
 ;;;;;  Rules for `store`;;;;;;;;;
-(rule
-  (lower (store flags x @ (value_type ty) p @ (value_type (ty_addr32 _)) offset))
-  (gen_store p offset (store_op ty) flags x))
+(rule (lower (store flags src @ (value_type ty) addr offset))
+  (gen_store (amode addr offset ty) (store_op ty) flags src))
 
 (decl gen_icmp (IntCC ValueRegs ValueRegs Type) XReg)
 (rule

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

@@ -323,15 +323,28 @@ impl generated_code::Context for ZkAsmIsleContext<'_, '_, MInst, ZkAsmBackend> {
         todo!()
     }
 
-    fn offset32_imm(&mut self, offset: i32) -> Offset32 {
-        Offset32::new(offset)
-    }
     fn default_memflags(&mut self) -> MemFlags {
         MemFlags::new()
     }
 
-    fn gen_amode(&mut self, base: Reg, offset: Offset32, ty: Type) -> AMode {
-        AMode::RegOffset(base, i64::from(offset), ty)
+    fn gen_reg_offset_amode(&mut self, base: Reg, offset: i64, ty: Type) -> AMode {
+        AMode::RegOffset(base, offset, ty)
+    }
+
+    fn gen_sp_offset_amode(&mut self, offset: i64, ty: Type) -> AMode {
+        AMode::SPOffset(offset, ty)
+    }
+
+    fn gen_fp_offset_amode(&mut self, offset: i64, ty: Type) -> AMode {
+        AMode::FPOffset(offset, ty)
+    }
+
+    fn gen_stack_slot_amode(&mut self, ss: StackSlot, offset: i64, ty: Type) -> AMode {
+        // Offset from beginning of stackslot area, which is at nominal SP (see
+        // [MemArg::NominalSPOffset] for more details on nominal SP tracking).
+        let stack_off = self.lower_ctx.abi().sized_stackslot_offsets()[ss] as i64;
+        let sp_off: i64 = stack_off + offset;
+        AMode::NominalSPOffset(sp_off, ty)
     }
 
     fn gen_const_amode(&mut self, c: VCodeConstant) -> AMode {

cranelift/filetests/src/test_zkasm.rs

@@ -321,6 +321,7 @@ mod tests {
         add,
         locals,
         locals_simple,
+        memory,
         counter,
         add_func,
         mul,

cranelift/zkasm_data/generated/memory.zkasm

@@ -0,0 +1,45 @@
+start:
+  zkPC + 2 => RR
+  :JMP(function_1)
+  :JMP(finalizeExecution)
+function_1:
+  SP + 1 => SP
+  RR :MSTORE(SP - 1)
+  SP + 4 => SP
+  D :MSTORE(SP - 1)
+  E :MSTORE(SP - 2)
+  B :MSTORE(SP - 3)
+  0n => B
+  8589934592n => D
+  $ => A :MLOAD(CTX)
+  $ => E :ADD
+  D :MSTORE(E)
+  34359738368n => B
+  12884901888n => E
+  $ => A :MLOAD(CTX)
+  $ => A :ADD
+  E :MSTORE(A)
+  0n => B
+  $ => A :MLOAD(CTX)
+  $ => A :ADD
+  $ => A :MLOAD(A)
+  A => E
+  34359738368n => B
+  $ => A :MLOAD(CTX)
+  $ => A :ADD
+  $ => B :MLOAD(A)
+  E => A
+  $ => A :ADD
+  21474836480n => B
+  B :ASSERT
+  $ => D :MLOAD(SP - 1)
+  $ => E :MLOAD(SP - 2)
+  $ => B :MLOAD(SP - 3)
+  SP - 4 => SP
+  $ => RR :MLOAD(SP - 1)
+  SP - 1 => SP
+  :JMP(RR)
+finalizeExecution:
+  ${beforeLast()}  :JMPN(finalizeExecution)
+                   :JMP(start)
+INCLUDE "helpers/2-exp.zkasm"

cranelift/zkasm_data/memory.wat

@@ -0,0 +1,12 @@
+(module
+ (import "env" "assert_eq" (func $assert_eq (param i32) (param i32)))
+ (memory (export "memory") 100)
+ (func $main
+	(i32.store (i32.const 0) (i32.const 2))
+	(i32.store (i32.const 8) (i32.const 3))
+	(i32.load (i32.const 0))
+	(i32.load (i32.const 8))
+	i32.add
+	i32.const 5
+	call $assert_eq)
+ (start $main))

cranelift/zkasm_data/state.csv

@@ -19,6 +19,7 @@ locals,pass,17
 locals_simple,pass,15
 lt_s,pass,34
 lt_u,pass,34
+memory,runtime error,
 mul,pass,31
 ne,pass,36
 nop,pass,15

docs/zkasm/test_summary.csv

@@ -1,5 +1,5 @@
 Suite path,Passing count,Total count,Total cycles
-cranelift/zkasm_data,25,26,1021
+cranelift/zkasm_data,25,27,1021
 cranelift/zkasm_data/benchmarks/fibonacci,3,3,282056
 cranelift/zkasm_data/benchmarks/sha256,0,1,0
 cranelift/zkasm_data/spectest/conversions,3,24,45