Merge pull request #186 from near/stack_fixes

Align stack handling with other architectures

cranelift/codegen/src/isa/zkasm/abi.rs

@@ -347,14 +347,14 @@ impl ABIMachineSpec for ZkAsmMachineDeps {
         let mut insts = SmallVec::new();
 
         if frame_layout.setup_area_size > 0 {
-            insts.push(Inst::ReserveSp {
-                amount: frame_layout.setup_area_size.into(),
-            });
             insts.push(Self::gen_store_stack(
-                StackAMode::SPOffset(-1, I64),
+                StackAMode::SPOffset(0, I64),
                 link_reg(),
                 I64,
             ));
+            insts.push(Inst::ReserveSp {
+                amount: frame_layout.setup_area_size.into(),
+            });
         }
 
         insts
@@ -369,14 +369,14 @@ impl ABIMachineSpec for ZkAsmMachineDeps {
         let mut insts = SmallVec::new();
 
         if frame_layout.setup_area_size > 0 {
+            insts.push(Inst::ReleaseSp {
+                amount: frame_layout.setup_area_size.into(),
+            });
             insts.push(Self::gen_load_stack(
-                StackAMode::SPOffset(-1, I64),
+                StackAMode::SPOffset(0, I64),
                 writable_link_reg(),
                 I64,
             ));
-            insts.push(Inst::ReleaseSp {
-                amount: frame_layout.setup_area_size.into(),
-            });
         }
 
         if call_conv == isa::CallConv::Tail {
@@ -415,10 +415,7 @@ impl ABIMachineSpec for ZkAsmMachineDeps {
         // Store each clobbered register in order at offsets from SP,
         // placing them above the fixed frame slots.
         if stack_size > 0 {
-            insts.push(Inst::ReserveSp {
-                amount: stack_size.into(),
-            });
-            let mut cur_offset = -1;
+            let mut cur_offset = -8;
             for reg in &frame_layout.clobbered_callee_saves {
                 let r_reg = reg.to_reg();
                 let ty = match r_reg.class() {
@@ -431,8 +428,11 @@ impl ABIMachineSpec for ZkAsmMachineDeps {
                     real_reg_to_reg(reg.to_reg()),
                     ty,
                 ));
-                cur_offset -= 1
+                cur_offset -= 8
             }
+            insts.push(Inst::ReserveSp {
+                amount: stack_size.into(),
+            });
         }
         insts
     }
@@ -445,7 +445,10 @@ impl ABIMachineSpec for ZkAsmMachineDeps {
         let mut insts = SmallVec::new();
         let stack_size = frame_layout.fixed_frame_storage_size + frame_layout.clobber_size;
         if stack_size > 0 {
-            let mut cur_offset = -1;
+            insts.push(Inst::ReleaseSp {
+                amount: stack_size.into(),
+            });
+            let mut cur_offset = -8;
             for reg in &frame_layout.clobbered_callee_saves {
                 let rreg = reg.to_reg();
                 let ty = match rreg.class() {
@@ -458,11 +461,8 @@ impl ABIMachineSpec for ZkAsmMachineDeps {
                     Writable::from_reg(real_reg_to_reg(reg.to_reg())),
                     ty,
                 ));
-                cur_offset -= 1
+                cur_offset -= 8
             }
-            insts.push(Inst::ReleaseSp {
-                amount: stack_size.into(),
-            });
         }
         insts
     }

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

@@ -601,11 +601,12 @@ impl MachInstEmit for Inst {
                         );
                     }
                     AMode::SPOffset(off, _) | AMode::NominalSPOffset(off, _) => {
+                        assert_eq!(off % 8, 0);
                         put_string(
                             &format!(
                                 "$ => {} :MLOAD({})\n",
                                 reg_name(rd.to_reg()),
-                                access_reg_with_offset(stack_reg(), off),
+                                access_reg_with_offset(stack_reg(), off / 8),
                             ),
                             sink,
                         );
@@ -647,11 +648,12 @@ impl MachInstEmit for Inst {
                         );
                     }
                     AMode::SPOffset(off, _) | AMode::NominalSPOffset(off, _) => {
+                        assert_eq!(off % 8, 0);
                         put_string(
                             &format!(
                                 "{} :MSTORE({})\n",
                                 reg_name(src),
-                                access_reg_with_offset(stack_reg(), off),
+                                access_reg_with_offset(stack_reg(), off / 8),
                             ),
                             sink,
                         );
@@ -720,20 +722,20 @@ impl MachInstEmit for Inst {
                 }
             }
             &Inst::ReleaseSp { amount } => {
-                // Stack is growing "up" in zkASM contrary to traditional architectures.
-                // Furthermore, addressing is done in slots rather than bytes.
+                // Stack addressing is done in slots rather than bytes.
                 //
                 // FIXME: add helper functions to implement these conversions.
-                let amount = amount.checked_div(8).unwrap();
-                put_string(&format!("SP - {amount} => SP\n"), sink);
+                let (q, r) = (amount / 8, amount % 8);
+                assert_eq!(r, 0);
+                put_string(&format!("SP + {q} => SP\n"), sink);
             }
             &Inst::ReserveSp { amount } => {
-                // Stack is growing "up" in zkASM contrary to traditional architectures.
-                // Furthermore, addressing is done in slots rather than bytes.
+                // Stack addressing is done in slots rather than bytes.
                 //
                 // FIXME: add helper functions to implement these conversions.
-                let amount = amount.checked_div(8).unwrap();
-                put_string(&format!("SP + {amount} => SP\n"), sink);
+                let (q, r) = (amount / 8, amount % 8);
+                assert_eq!(r, 0);
+                put_string(&format!("SP - {q} => SP\n"), sink);
             }
             &Inst::Call { ref info } => {
                 // call

cranelift/codegen/src/isa/zkasm/inst/emit_tests.rs

@@ -32,7 +32,7 @@ fn test_zkasm_binemit() {
             rets: vec![],
             stack_bytes_to_pop: 16,
         },
-        "SP - 2 => SP\n  :JMP(RR)",
+        "SP + 2 => SP\n  :JMP(RR)",
     ));
     // TODO: a test with `rets`.
     insns.push(TestUnit::new(

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

@@ -1122,13 +1122,15 @@ impl Inst {
             }
             &Inst::ReserveSp { amount } => {
                 // FIXME: conversion methods
-                let amount = amount.checked_div(8).unwrap();
-                format!("  SP + {} => SP", amount)
+                let (q, r) = (amount / 8, amount % 8);
+                assert_eq!(r, 0);
+                format!("  SP - {q} => SP")
             }
             &Inst::ReleaseSp { amount } => {
                 // FIXME: conversion methods
-                let amount = amount.checked_div(8).unwrap();
-                format!("  SP - {} => SP", amount)
+                let (q, r) = (amount / 8, amount % 8);
+                assert_eq!(r, 0);
+                format!("  SP + {q} => SP")
             }
             &Inst::Call { ref info } => format!("call {}", info.dest.display(None)),
             &Inst::CallInd { ref info } => {

cranelift/filetests/src/test_zkasm.rs

@@ -73,6 +73,9 @@ mod tests {
             }
         }
 
+        // The total amount of stack available on ZKASM processor is 2^16 of 8-byte words.
+        // Stack memory is a separate region that is independent from the heap.
+        program.push("  0xffff => SP".to_string());
         program.push("  zkPC + 2 => RR".to_string());
         program.push(format!("  :JMP(function_{})", start_func_index));
         program.push("  :JMP(finalizeExecution)".to_string());

cranelift/zkasm_data/benchmarks/fibonacci/generated/from_rust.zkasm

@@ -5,17 +5,18 @@ start:
   1048576 :MSTORE(global_0)  ;; Global32(1048576)
   1048576 :MSTORE(global_1)  ;; Global32(1048576)
   1048576 :MSTORE(global_2)  ;; Global32(1048576)
+  0xffff => SP
   zkPC + 2 => RR
   :JMP(function_1)
   :JMP(finalizeExecution)
 function_1:
-  SP + 1 => SP
-  RR :MSTORE(SP - 1)
-  SP + 4 => SP
+  RR :MSTORE(SP)
+  SP - 1 => SP
   C :MSTORE(SP - 1)
   D :MSTORE(SP - 2)
   E :MSTORE(SP - 3)
   B :MSTORE(SP - 4)
+  SP - 4 => SP
   1n => A  ;; LoadConst64
   0n => B  ;; LoadConst64
   10000n => C  ;; LoadConst32
@@ -75,13 +76,13 @@ label_1_3:
   15574651946073070043n => B  ;; LoadConst64
   D => A
   B :ASSERT
+  SP + 4 => SP
   $ => C :MLOAD(SP - 1)
   $ => D :MLOAD(SP - 2)
   $ => E :MLOAD(SP - 3)
   $ => B :MLOAD(SP - 4)
-  SP - 4 => SP
-  $ => RR :MLOAD(SP - 1)
-  SP - 1 => SP
+  SP + 1 => SP
+  $ => RR :MLOAD(SP)
   :JMP(RR)
 finalizeExecution:
   ${beforeLast()}  :JMPN(finalizeExecution)

cranelift/zkasm_data/benchmarks/fibonacci/generated/handwritten_wat.zkasm

@@ -1,29 +1,30 @@
 start:
+  0xffff => SP
   zkPC + 2 => RR
   :JMP(function_1)
   :JMP(finalizeExecution)
 function_1:
-  SP + 1 => SP
-  RR :MSTORE(SP - 1)
-  SP + 6 => SP
+  RR :MSTORE(SP)
+  SP - 1 => SP
   C :MSTORE(SP - 1)
   D :MSTORE(SP - 2)
   E :MSTORE(SP - 3)
   B :MSTORE(SP - 4)
+  SP - 6 => SP
   0n => A  ;; LoadConst32
-  A :MSTORE(SP + 8)
+  A :MSTORE(SP + 1)
   0n => A  ;; LoadConst64
   1n => B  ;; LoadConst64
   :JMP(label_1_1)
 label_1_1:
   $ => E :ADD
   B :MSTORE(SP)
   1n => B  ;; LoadConst32
-  $ => A :MLOAD(SP + 8)
+  $ => A :MLOAD(SP + 1)
   $ => A :ADD
   4294967295n => B  ;; LoadConst64
   $ => A :AND
-  A :MSTORE(SP + 8)
+  A :MSTORE(SP + 1)
   10000n => B  ;; LoadConst32
   $ => A :EQ
   A :JMPNZ(label_1_3)
@@ -34,13 +35,13 @@ label_1_3:
   15574651946073070043n => B  ;; LoadConst64
   $ => A :MLOAD(SP)
   B :ASSERT
+  SP + 6 => SP
   $ => C :MLOAD(SP - 1)
   $ => D :MLOAD(SP - 2)
   $ => E :MLOAD(SP - 3)
   $ => B :MLOAD(SP - 4)
-  SP - 6 => SP
-  $ => RR :MLOAD(SP - 1)
-  SP - 1 => SP
+  SP + 1 => SP
+  $ => RR :MLOAD(SP)
   :JMP(RR)
 finalizeExecution:
   ${beforeLast()}  :JMPN(finalizeExecution)

cranelift/zkasm_data/benchmarks/fibonacci/state.csv

@@ -1,4 +1,4 @@
 Test,Status,Cycles
-from_rust,pass,48025
+from_rust,pass,48026
 handwritten,pass,50008
-handwritten_wat,pass,140023
+handwritten_wat,pass,140024