Skip to content

Commit

Permalink
zkasm: rename some types and structures away from RV64 (bytecodeallia…
Browse files Browse the repository at this point in the history
…nce#15)

* zkasm: rename some types and structures away from RV64

* zkASM: remove the unwinding code

* Revert "zkasm: split AdjustSp to separate variants for Reserve and Release stack space (bytecodealliance#14)"

This reverts commit 0798ff7.
  • Loading branch information
nagisa authored Oct 10, 2023
1 parent 0798ff7 commit bb378cd
Show file tree
Hide file tree
Showing 20 changed files with 260 additions and 534 deletions.
3 changes: 3 additions & 0 deletions .github/workflows/main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,9 @@ jobs:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v3
- uses: actions/setup-node@v3
with:
node-version: '20.8.0'
- run: ./ci/test-zk.sh --install-zkwasm --all

rustfmt:
Expand Down
121 changes: 48 additions & 73 deletions cranelift/codegen/src/isa/zkasm/abi.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
//! Implementation of a standard Riscv64 ABI.
//! Implementation of a standard zkASM ABI.
use std::sync::OnceLock;

Expand All @@ -16,7 +16,7 @@ use crate::machinst::*;
use crate::ir::types::I8;
use crate::ir::LibCall;
use crate::ir::Signature;
use crate::isa::zkasm::settings::Flags as RiscvFlags;
use crate::isa::zkasm::settings::Flags;
use crate::settings;
use crate::CodegenError;
use crate::CodegenResult;
Expand All @@ -27,60 +27,26 @@ use regs::{create_reg_environment, x_reg};

use smallvec::{smallvec, SmallVec};

/// Support for the Riscv64 ABI from the callee side (within a function body).
pub(crate) type Riscv64Callee = Callee<Riscv64MachineDeps>;
/// Support for the zkASM ABI from the callee side (within a function body).
pub(crate) type ZkAsmCallee = Callee<ZkAsmMachineDeps>;

/// Support for the Riscv64 ABI from the caller side (at a callsite).
pub(crate) type Riscv64ABICallSite = CallSite<Riscv64MachineDeps>;
/// Support for the zkASM ABI from the caller side (at a callsite).
pub(crate) type ZkAsmABICallSite = CallSite<ZkAsmMachineDeps>;

/// This is the limit for the size of argument and return-value areas on the
/// stack. We place a reasonable limit here to avoid integer overflow issues
/// with 32-bit arithmetic: for now, 128 MB.
static STACK_ARG_RET_SIZE_LIMIT: u32 = 128 * 1024 * 1024;

/// Riscv64-specific ABI behavior. This struct just serves as an implementation
/// zkASM-specific ABI behavior. This struct just serves as an implementation
/// point for the trait; it is never actually instantiated.
pub struct Riscv64MachineDeps;

impl IsaFlags for RiscvFlags {}

impl RiscvFlags {
pub(crate) fn min_vec_reg_size(&self) -> u64 {
let entries = [
(self.has_zvl65536b(), 65536),
(self.has_zvl32768b(), 32768),
(self.has_zvl16384b(), 16384),
(self.has_zvl8192b(), 8192),
(self.has_zvl4096b(), 4096),
(self.has_zvl2048b(), 2048),
(self.has_zvl1024b(), 1024),
(self.has_zvl512b(), 512),
(self.has_zvl256b(), 256),
// In order to claim the Application Profile V extension, a minimum
// register size of 128 is required. i.e. V implies Zvl128b.
(self.has_v(), 128),
(self.has_zvl128b(), 128),
(self.has_zvl64b(), 64),
(self.has_zvl32b(), 32),
];

for (has_flag, size) in entries.into_iter() {
if !has_flag {
continue;
}
pub struct ZkAsmMachineDeps;

// Due to a limitation in regalloc2, we can't support types
// larger than 1024 bytes. So limit that here.
return std::cmp::min(size, 1024);
}
impl IsaFlags for Flags {}

return 0;
}
}

impl ABIMachineSpec for Riscv64MachineDeps {
impl ABIMachineSpec for ZkAsmMachineDeps {
type I = Inst;
type F = RiscvFlags;
type F = Flags;

fn word_bits() -> u32 {
64
Expand Down Expand Up @@ -328,15 +294,8 @@ impl ABIMachineSpec for Riscv64MachineDeps {
if amount == 0 {
return insts;
}
// FIXME: is this right/sufficient for stack growing up
insts.push(if amount < 0 {
Inst::ReserveSp {
amount: -amount as u64,
}
} else {
Inst::ReleaseSp {
amount: amount as u64,
}
insts.push(Inst::AdjustSp {
amount: amount as i64,
});
insts
}
Expand Down Expand Up @@ -403,9 +362,11 @@ impl ABIMachineSpec for Riscv64MachineDeps {
let mut insts = SmallVec::new();

if frame_layout.setup_area_size > 0 {
insts.push(Inst::ReserveSp {
amount: frame_layout.setup_area_size.into(),
});
// add sp,sp,-16 ;; alloc stack space for fp.
// sd ra,8(sp) ;; save ra.
// sd fp,0(sp) ;; store old fp.
// mv fp,sp ;; set fp to sp.
insts.push(Inst::AdjustSp { amount: -1 });
insts.push(Self::gen_store_stack(
StackAMode::SPOffset(0, I64),
link_reg(),
Expand Down Expand Up @@ -453,9 +414,7 @@ impl ABIMachineSpec for Riscv64MachineDeps {
// writable_fp_reg(),
// I64,
// ));
insts.push(Inst::ReleaseSp {
amount: frame_layout.setup_area_size.into(),
});
insts.push(Inst::AdjustSp { amount: 1 });
}

if call_conv == isa::CallConv::Tail && frame_layout.stack_args_size > 0 {
Expand Down Expand Up @@ -514,16 +473,25 @@ impl ABIMachineSpec for Riscv64MachineDeps {

fn gen_clobber_save(
_call_conv: isa::CallConv,
_flags: &settings::Flags,
flags: &settings::Flags,
frame_layout: &FrameLayout,
) -> SmallVec<[Self::I; 16]> {
let mut insts = SmallVec::new();
// Adjust the stack pointer upward for clobbers and the function fixed
// Adjust the stack pointer downward for clobbers and the function fixed
// frame (spillslots and storage slots).
let stack_size = frame_layout.fixed_frame_storage_size + frame_layout.clobber_size;
// The stack (and memory in general) in zkASM is addressed in slots, rather than bytes.
// Adjust accordingly.
// Each stack slot is 64 bit and can fit a u8 value.
let stack_size = stack_size / 8;
if flags.unwind_info() && frame_layout.setup_area_size > 0 {
// The *unwind* frame (but not the actual frame) starts at the
// clobbers, just below the saved FP/LR pair.
// insts.push(Inst::Unwind {
// inst: UnwindInst::DefineNewFrame {
// offset_downward_to_clobbers: frame_layout.clobber_size,
// offset_upward_to_caller_sp: frame_layout.setup_area_size,
// },
// });
}
// Store each clobbered register in order at offsets from SP,
// placing them above the fixed frame slots.
if stack_size > 0 {
Expand All @@ -536,15 +504,23 @@ impl ABIMachineSpec for Riscv64MachineDeps {
RegClass::Float => F64,
RegClass::Vector => unimplemented!("Vector Clobber Saves"),
};
if flags.unwind_info() {
// insts.push(Inst::Unwind {
// inst: UnwindInst::SaveReg {
// clobber_offset: frame_layout.clobber_size - cur_offset,
// reg: r_reg,
// },
// });
}
insts.push(Self::gen_store_stack(
StackAMode::SPOffset(-(cur_offset as i64), ty),
real_reg_to_reg(reg.to_reg()),
ty,
));
cur_offset += 1
}
insts.push(Inst::ReserveSp {
amount: stack_size.into(),
insts.push(Inst::AdjustSp {
amount: -(stack_size as i64),
});
}
insts
Expand All @@ -558,11 +534,10 @@ impl ABIMachineSpec for Riscv64MachineDeps {
let mut insts = SmallVec::new();
let stack_size = frame_layout.fixed_frame_storage_size + frame_layout.clobber_size;
// Each stack slot is 64 bit and can fit a u8 value.
// FIXME(nagisa): WHAT DOES THIS MEAA~~~AAAN????
let stack_size = stack_size / 8;
if stack_size > 0 {
insts.push(Inst::ReleaseSp {
amount: stack_size.into(),
insts.push(Inst::AdjustSp {
amount: stack_size as i64,
});
}
let mut cur_offset = 1;
Expand Down Expand Up @@ -668,13 +643,13 @@ impl ABIMachineSpec for Riscv64MachineDeps {
fn get_number_of_spillslots_for_value(
rc: RegClass,
_target_vector_bytes: u32,
isa_flags: &RiscvFlags,
_isa_flags: &Flags,
) -> u32 {
// We allocate in terms of 8-byte slots.
match rc {
RegClass::Int => 1,
RegClass::Float => 1,
RegClass::Vector => (isa_flags.min_vec_reg_size() / 8) as u32,
RegClass::Vector => todo!(),
}
}

Expand Down Expand Up @@ -712,7 +687,7 @@ impl ABIMachineSpec for Riscv64MachineDeps {
}
}

impl Riscv64ABICallSite {
impl ZkAsmABICallSite {
pub fn emit_return_call(mut self, ctx: &mut Lower<Inst>, args: isle::ValueSlice) {
let (new_stack_arg_size, old_stack_arg_size) =
self.emit_temporary_tail_call_frame(ctx, args);
Expand Down Expand Up @@ -980,7 +955,7 @@ const fn tail_clobbers() -> PRegSet {

const TAIL_CLOBBERS: PRegSet = tail_clobbers();

impl Riscv64MachineDeps {
impl ZkAsmMachineDeps {
fn gen_probestack_unroll(insts: &mut SmallInstVec<Inst>, guard_size: u32, probe_count: u32) {
insts.reserve(probe_count as usize);
for i in 0..probe_count {
Expand Down
9 changes: 3 additions & 6 deletions cranelift/codegen/src/isa/zkasm/inst.isle
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@
(args VecArgPair))

(Ret (rets VecRetPair)
(stack_bytes_to_pop u64))
(stack_bytes_to_pop u32))

(Extend
(rd WritableReg)
Expand All @@ -61,11 +61,8 @@
(from_bits u8)
(to_bits u8))

;; Adjust the stack pointer as necessary.
(ReserveSp
(amount u64))
(ReleaseSp
(amount u64))
(AdjustSp
(amount i64))

(Call
(info BoxCallInfo))
Expand Down
2 changes: 1 addition & 1 deletion cranelift/codegen/src/isa/zkasm/inst/args.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
//! Riscv64 ISA definitions: instruction arguments.
//! zkASM ISA definitions: instruction arguments.
// Some variants are never constructed, but we still want them as options in the future.
#![allow(dead_code)]
Expand Down
33 changes: 14 additions & 19 deletions cranelift/codegen/src/isa/zkasm/inst/emit.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
//! Riscv64 ISA: binary code emission.
//! zkASM ISA: binary code emission.
use crate::binemit::StackMap;
use crate::ir::{self, RelSourceLoc, TrapCode};
Expand Down Expand Up @@ -139,7 +139,7 @@ impl EmitState {

impl MachInstEmitState<Inst> for EmitState {
fn new(
abi: &Callee<crate::isa::zkasm::abi::Riscv64MachineDeps>,
abi: &Callee<crate::isa::zkasm::abi::ZkAsmMachineDeps>,
ctrl_plane: ControlPlane,
) -> Self {
EmitState {
Expand Down Expand Up @@ -569,8 +569,8 @@ impl MachInstEmit for Inst {
stack_bytes_to_pop, ..
} => {
if stack_bytes_to_pop != 0 {
Inst::ReleaseSp {
amount: stack_bytes_to_pop,
Inst::AdjustSp {
amount: i64::from(stack_bytes_to_pop),
}
.emit(&[], sink, emit_info, state);
}
Expand Down Expand Up @@ -621,21 +621,16 @@ impl MachInstEmit for Inst {
.into_iter()
.for_each(|i| i.emit(&[], sink, emit_info, state)); */
}
&Inst::ReleaseSp { amount } => {
// Stack is growing "up" in zkASM contrary to traditional architectures.
// Furthermore, 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);
}
&Inst::ReserveSp { amount } => {
// Stack is growing "up" in zkASM contrary to traditional architectures.
// Furthermore, 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);
&Inst::AdjustSp { amount } => {
// Stack is growing "up" in zkASM contrary to traditional architectures, so
// we flip it here.
let amount = -amount;
let amount = if amount > 0 {
format!("+ {}", amount)
} else {
format!("- {}", -amount)
};
put_string(&format!("SP {amount} => 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
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 - 16 => SP\n :JMP(RR)",
));
// TODO: a test with `rets`.
insns.push(TestUnit::new(
Expand Down
6 changes: 5 additions & 1 deletion cranelift/codegen/src/isa/zkasm/inst/imms.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
//! Riscv64 ISA definitions: immediate constants.
//! zkASM ISA definitions: immediate constants.
// Some variants are never constructed, but we still want them as options in the future.
use super::Inst;
#[allow(dead_code)]
use std::fmt::{Debug, Display, Formatter, Result};

// TODO: remove
#[derive(Copy, Clone, Debug, Default)]
pub struct Imm12 {
pub bits: i16,
Expand Down Expand Up @@ -98,6 +99,7 @@ impl std::ops::Neg for Imm12 {
}

// singed
// TODO: remove
#[derive(Clone, Copy, Default)]
pub struct Imm20 {
/// The immediate bits.
Expand Down Expand Up @@ -130,6 +132,7 @@ impl Display for Imm20 {
}

/// An unsigned 5-bit immediate.
// TODO: remove
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct UImm5 {
value: u8,
Expand Down Expand Up @@ -158,6 +161,7 @@ impl Display for UImm5 {
}

/// A Signed 5-bit immediate.
// TODO: remove
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Imm5 {
value: i8,
Expand Down
Loading

0 comments on commit bb378cd

Please sign in to comment.