From f9cb07350b0182cfe070ab1b4a5b8ac0c255c828 Mon Sep 17 00:00:00 2001 From: Andrei Kashin Date: Fri, 1 Dec 2023 17:22:59 +0000 Subject: [PATCH] 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 | 124 +++++++++--------- cranelift/codegen/src/isa/zkasm/lower.isle | 62 ++++----- cranelift/codegen/src/isa/zkasm/lower/isle.rs | 23 +++- cranelift/filetests/src/test_zkasm.rs | 1 + cranelift/zkasm_data/generated/memory.zkasm | 45 +++++++ cranelift/zkasm_data/memory.wat | 12 ++ cranelift/zkasm_data/state.csv | 1 + docs/zkasm/test_summary.csv | 2 +- 8 files changed, 170 insertions(+), 100 deletions(-) create mode 100644 cranelift/zkasm_data/generated/memory.zkasm create mode 100644 cranelift/zkasm_data/memory.wat diff --git a/cranelift/codegen/src/isa/zkasm/inst.isle b/cranelift/codegen/src/isa/zkasm/inst.isle index acf5a6b4b614..25e4b5d996d5 100644 --- a/cranelift/codegen/src/isa/zkasm/inst.isle +++ b/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) diff --git a/cranelift/codegen/src/isa/zkasm/lower.isle b/cranelift/codegen/src/isa/zkasm/lower.isle index a6acc64cb41b..760aa4adcbae 100644 --- a/cranelift/codegen/src/isa/zkasm/lower.isle +++ b/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 diff --git a/cranelift/codegen/src/isa/zkasm/lower/isle.rs b/cranelift/codegen/src/isa/zkasm/lower/isle.rs index 56570d74d041..8b29bfb86523 100644 --- a/cranelift/codegen/src/isa/zkasm/lower/isle.rs +++ b/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 { diff --git a/cranelift/filetests/src/test_zkasm.rs b/cranelift/filetests/src/test_zkasm.rs index 408d5de1c9e5..001472eb1b6a 100644 --- a/cranelift/filetests/src/test_zkasm.rs +++ b/cranelift/filetests/src/test_zkasm.rs @@ -321,6 +321,7 @@ mod tests { add, locals, locals_simple, + memory, counter, add_func, mul, diff --git a/cranelift/zkasm_data/generated/memory.zkasm b/cranelift/zkasm_data/generated/memory.zkasm new file mode 100644 index 000000000000..2a7146d230a5 --- /dev/null +++ b/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" \ No newline at end of file diff --git a/cranelift/zkasm_data/memory.wat b/cranelift/zkasm_data/memory.wat new file mode 100644 index 000000000000..70579b8c980f --- /dev/null +++ b/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)) diff --git a/cranelift/zkasm_data/state.csv b/cranelift/zkasm_data/state.csv index 8a228971f0f0..a108032d5cab 100644 --- a/cranelift/zkasm_data/state.csv +++ b/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 diff --git a/docs/zkasm/test_summary.csv b/docs/zkasm/test_summary.csv index a0f5ff4e4cee..985f69552ac0 100644 --- a/docs/zkasm/test_summary.csv +++ b/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