Skip to content

Commit

Permalink
Merge pull request bytecodealliance#186 from near/stack_fixes
Browse files Browse the repository at this point in the history 
Align stack handling with other architectures
  • Loading branch information
@akashin
akashin committed Jan 9, 2024
2 parents cefafb6 + b47542d commit f1231ff
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
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
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
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
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
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
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
@@ -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
@@ -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

0 comments on commit f1231ff

Please sign in to comment.