Skip to content

Commit

Permalink
Align stack handling with other architectures
Browse files Browse the repository at this point in the history 
In the past the stack in ZKASM grew from lower to upper addresses, but
there is no real reason for this, as we can always start from the
largest offset and go down.

This aligns the behaviour with other architectures which makes some
common logic like spill stack allocation work correctly.

We also now track all offsets in bytes and only do conversion to slots on
accesses to make sure that beforementioned spill slock logic works
correctly.
  • Loading branch information
@aborg-dev
aborg-dev committed Jan 9, 2024
1 parent cefafb6 commit b47542d
Show file tree
Hide file tree
Showing 734 changed files with 12,785 additions and 12,055 deletions.
36 changes: 18 additions & 18 deletions cranelift/codegen/src/isa/zkasm/abi.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 {
Expand Down Expand Up @@ -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() {
Expand All @@ -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
}
Expand All @@ -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() {
Expand All @@ -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
}
Expand Down
22 changes: 12 additions & 10 deletions cranelift/codegen/src/isa/zkasm/inst/emit.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
);
Expand Down Expand Up @@ -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,
);
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion cranelift/codegen/src/isa/zkasm/inst/emit_tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
10 changes: 6 additions & 4 deletions cranelift/codegen/src/isa/zkasm/inst/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 } => {
Expand Down
3 changes: 3 additions & 0 deletions cranelift/filetests/src/test_zkasm.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
13 changes: 7 additions & 6 deletions cranelift/zkasm_data/benchmarks/fibonacci/generated/from_rust.zkasm
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
19 changes: 10 additions & 9 deletions cranelift/zkasm_data/benchmarks/fibonacci/generated/handwritten_wat.zkasm
View file
Original file line number Diff line number Diff line change
@@ -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)
Expand All @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions cranelift/zkasm_data/benchmarks/fibonacci/state.csv
View file
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit b47542d

Please sign in to comment.