Skip to content

Commit

Permalink
Remove vmctx in ZkasmEnvironment
Browse files Browse the repository at this point in the history 
This commit fully removes "vmctx" pointer in the ZkasmEnvironment, as we don't really have a need for it.
As a result, we can finally get rid of CTX 32-bit register that was causing problems in ROTL tests.
  • Loading branch information
@aborg-dev
aborg-dev committed Dec 11, 2023
1 parent 08dd14e commit 6c9a9b2
Show file tree
Hide file tree
Showing 32 changed files with 6,062 additions and 5,305 deletions.
15 changes: 0 additions & 15 deletions cranelift/codegen/src/isa/zkasm/abi.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,21 +94,6 @@ impl ABIMachineSpec for ZkAsmMachineDeps {
continue;
}

// For now we pin VMContext register to `CTX` register of ZK ASM.
if let ir::ArgumentPurpose::VMContext = param.purpose {
let mut slots = ABIArgSlotVec::new();
slots.push(ABIArgSlot::Reg {
reg: context_reg().to_real_reg().unwrap(),
ty: I32,
extension: param.extension,
});
args.push(ABIArg::Slots {
slots,
purpose: param.purpose,
});
continue;
}

// Find regclass(es) of the register(s) used to store a value of this type.
let (rcs, reg_tys) = Inst::rc_for_type(param.value_type)?;
let mut slots = ABIArgSlotVec::new();
Expand Down
7 changes: 1 addition & 6 deletions cranelift/codegen/src/isa/zkasm/inst/regs.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ pub fn create_reg_environment() -> MachineEnv {
let preferred_regs_by_class: [Vec<PReg>; 3] = {
// Registers are A, B, C, D, E.
let x_registers: Vec<PReg> = (5..=7)
.chain(10..=12)
.chain(10..12)
.map(|i| PReg::new(i, RegClass::Int))
.collect();

Expand All @@ -139,11 +139,6 @@ pub fn create_reg_environment() -> MachineEnv {

let non_preferred_regs_by_class: [Vec<PReg>; 3] = {
let x_registers: Vec<PReg> = Vec::new();
// (9..=9)
// .chain(18..=27)
// .map(|i| PReg::new(i, RegClass::Int))
// .collect();

let f_registers: Vec<PReg> = Vec::new();
let v_registers = vec![];
[x_registers, f_registers, v_registers]
Expand Down
43 changes: 8 additions & 35 deletions cranelift/wasm/src/environ/zkasm.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -227,15 +227,9 @@ impl<'zkasm_environment> ZkasmFuncEnvironment<'zkasm_environment> {
}
}

/// Create a signature for `sigidx` amended with a `vmctx` argument after
/// the standard wasm arguments.
pub fn vmctx_sig(&self, sigidx: TypeIndex) -> ir::Signature {
let mut sig = self.mod_info.signatures[sigidx].clone();
sig.params.push(ir::AbiParam::special(
self.pointer_type(),
ir::ArgumentPurpose::VMContext,
));
sig
/// Create a signature for `sigidx`.
pub fn func_sig(&self, sigidx: TypeIndex) -> ir::Signature {
self.mod_info.signatures[sigidx].clone()
}

fn reference_type(&self) -> ir::Type {
Expand Down Expand Up @@ -318,7 +312,6 @@ impl<'zkasm_environment> FuncEnvironment for ZkasmFuncEnvironment<'zkasm_environ
func: &mut ir::Function,
index: GlobalIndex,
) -> WasmResult<GlobalVariable> {
// Just create a zkasm `vmctx` global.
let offset = i32::try_from((index.index() * 8) + 8).unwrap().into();
Ok(GlobalVariable::Memory {
gv: ZkasmFuncEnvironment::globals_base(func),
Expand Down Expand Up @@ -375,9 +368,7 @@ impl<'zkasm_environment> FuncEnvironment for ZkasmFuncEnvironment<'zkasm_environ
func: &mut ir::Function,
index: TypeIndex,
) -> WasmResult<ir::SigRef> {
// A real implementation would probably change the calling convention and add `vmctx` and
// signature index arguments.
Ok(func.import_signature(self.vmctx_sig(index)))
Ok(func.import_signature(self.func_sig(index)))
}

fn make_direct_func(
Expand All @@ -386,9 +377,7 @@ impl<'zkasm_environment> FuncEnvironment for ZkasmFuncEnvironment<'zkasm_environ
index: FuncIndex,
) -> WasmResult<ir::FuncRef> {
let sigidx = self.mod_info.functions[index].entity;
// A real implementation would probably add a `vmctx` argument.
// And maybe attempt some signature de-duplication.
let signature = func.import_signature(self.vmctx_sig(sigidx));
let signature = func.import_signature(self.func_sig(sigidx));
let name =
ir::ExternalName::User(func.declare_imported_user_function(ir::UserExternalName {
namespace: 0,
Expand Down Expand Up @@ -446,12 +435,6 @@ impl<'zkasm_environment> FuncEnvironment for ZkasmFuncEnvironment<'zkasm_environ
callee: ir::Value,
call_args: &[ir::Value],
) -> WasmResult<ir::Inst> {
// Pass the current function's vmctx parameter on to the callee.
let vmctx = builder
.func
.special_param(ir::ArgumentPurpose::VMContext)
.expect("Missing vmctx parameter");

// The `callee` value is an index into a table of function pointers.
// Apparently, that table is stored at absolute address 0 in this zkasm environment.
// TODO: Generate bounds checking code.
Expand All @@ -465,12 +448,10 @@ impl<'zkasm_environment> FuncEnvironment for ZkasmFuncEnvironment<'zkasm_environ
let mflags = ir::MemFlags::trusted();
let func_ptr = builder.ins().load(ptr, mflags, callee_offset, 0);

// Build a value list for the indirect call instruction containing the callee, call_args,
// and the vmctx parameter.
// Build a value list for the indirect call instruction containing the callee and call_args.
let mut args = ir::ValueList::default();
args.push(func_ptr, &mut builder.func.dfg.value_lists);
args.extend(call_args.iter().cloned(), &mut builder.func.dfg.value_lists);
args.push(vmctx, &mut builder.func.dfg.value_lists);

Ok(builder
.ins()
Expand Down Expand Up @@ -508,17 +489,9 @@ impl<'zkasm_environment> FuncEnvironment for ZkasmFuncEnvironment<'zkasm_environ
callee: ir::FuncRef,
call_args: &[ir::Value],
) -> WasmResult<ir::Inst> {
// Pass the current function's vmctx parameter on to the callee.
let vmctx = builder
.func
.special_param(ir::ArgumentPurpose::VMContext)
.expect("Missing vmctx parameter");

// Build a value list for the call instruction containing the call_args and the vmctx
// parameter.
// Build a value list for the call instruction containing the call_args.
let mut args = ir::ValueList::default();
args.extend(call_args.iter().cloned(), &mut builder.func.dfg.value_lists);
args.push(vmctx, &mut builder.func.dfg.value_lists);

Ok(builder
.ins()
Expand Down Expand Up @@ -948,7 +921,7 @@ impl<'data> ModuleEnvironment<'data> for ZkasmEnvironment {
let func_index =
FuncIndex::new(self.get_num_func_imports() + self.info.function_bodies.len());

let sig = func_environ.vmctx_sig(self.get_func_type(func_index));
let sig = func_environ.func_sig(self.get_func_type(func_index));
let mut func =
ir::Function::with_name_signature(UserFuncName::user(0, func_index.as_u32()), sig);

Expand Down
58 changes: 29 additions & 29 deletions cranelift/zkasm_data/benchmarks/fibonacci/generated/from_rust.zkasm
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,60 +10,60 @@ function_1:
D :MSTORE(SP - 2)
E :MSTORE(SP - 3)
B :MSTORE(SP - 4)
1n => C ;; LoadConst64(1)
1n => D ;; LoadConst64(1)
0n => B ;; LoadConst64(0)
42949672960000n => A ;; LoadConst32(10000)
A => E
C => A
A :MSTORE(SP)
D => A
:JMP(label_1_1)
label_1_1:
$ => C :ADD
$ => E :ADD
A => B
C => A
$ => C :ADD
E => A
$ => E :ADD
A => B
C => A
$ => C :ADD
E => A
$ => E :ADD
A => B
C => A
$ => C :ADD
E => A
$ => E :ADD
A => B
C => A
$ => C :ADD
E => A
$ => E :ADD
A => B
C => A
$ => C :ADD
E => A
$ => E :ADD
A => B
C => A
$ => C :ADD
E => A
$ => E :ADD
A => B
C => A
$ => C :ADD
E => A
$ => E :ADD
A => B
C => A
$ => C :ADD
E => A
$ => E :ADD
A => B
C => A
$ => D :ADD
A :MSTORE(SP)
18446744030759878656n => B ;; LoadConst32(4294967286)
E => A
$ => E :ADD
A => D
18446744030759878656n => B ;; LoadConst32(4294967286)
$ => A :MLOAD(SP)
$ => A :ADD
A => C
A => E
A :MSTORE(SP)
C => B
0 => A
$ => A :EQ
A :JMPZ(label_1_2)
:JMP(label_1_3)
label_1_2:
C => E
$ => B :MLOAD(SP)
D => A
D => B
C :MSTORE(SP)
E => A
:JMP(label_1_1)
label_1_3:
15574651946073070043n => B ;; LoadConst64(15574651946073070043)
$ => A :MLOAD(SP)
D => A
B :ASSERT
$ => C :MLOAD(SP - 1)
$ => D :MLOAD(SP - 2)
Expand Down
Loading

0 comments on commit 6c9a9b2

Please sign in to comment.