Compiling Rust and other high-level languages to Miden VM

[greenhat]

Hi! As a side project, I've been working on a compiler framework for zero-knowledge VMs. It's called OmniZK, and the idea behind it is to help build compilers from high-level languages to various ZK VMs. Its design resembles the MLIR (LLVM) architecture, where IR transformations are implemented generically and can be reused with different custom IR dialects.

I'm working on adding Miden VM support to OmniZK and looking for your feedback.
So far, I have implemented Wasm local and global variables to direct memory access ops and I/O via extern functions in Rust 'pub_input/pub_output' (Wasm imports) with the idea of using the Wasm component model and WASI in the future. I plan to tackle control flow conversion next.

To give you an idea of what it'd be capable of, check out the Rust to Triton VM compilation example I've implemented for Triton VM.

I've implemented a small subset of instructions in the Wasm frontend (parser, IR dialect, etc.) and Triton VM backend (IR dialect, codegen, etc.) with Wasm -> Triton VM conversion.

Check out how the following Fibonacci example in Rust:

use ozk_stdlib::*;

pub fn fib_seq() {
    let n = pub_input() as u32;
    let mut a: u32 = 0;
    let mut b: u32 = 1;
    for _ in 0..n {
        let c = a + b;
        a = b;
        b = c;
    }
    pub_output(a as u64);
}

is compiled to the following fully executable Triton VM code. Keep in mind that it's in the early stage, there are no optimizations yet so the code is quite lengthy.

[greenhat]

Here is how one can build a compiler from Wasm to Miden VM using OmniZK framework:

pub fn compile(wasm: &[u8]) -> String {
    let frontend_config = WasmFrontendConfig::default();
    let target_config = MidenTargetConfig::default();
    let mut ctx = Context::new();
    frontend_config.register(&mut ctx);
    target_config.register(&mut ctx);
    let wasm_module_op =
        ozk_frontend_wasm::parse_module(&mut ctx, wasm, &frontend_config).unwrap();
    let miden_prog = run_conversion_passes(&mut ctx, wasm_module_op);
    let inst_buf = emit_prog(&ctx, miden_prog, &target_config).unwrap();
    inst_buf.pretty_print()
}

fn run_conversion_passes(ctx: &mut Context, wasm_module: ModuleOp) -> Ptr<Operation> {
    // we need to wrap the wasm in an op because passes cannot replace the root op
    let wrapper_module = builtin::ops::ModuleOp::new(ctx, "wrapper");
    wasm_module
        .get_operation()
        .insert_at_back(wrapper_module.get_body(ctx, 0), ctx);
    let mut pass_manager = PassManager::new();
    pass_manager.add_pass(Box::<WasmToMidenCFLoweringPass>::default());
    pass_manager.add_pass(Box::<WasmToMidenArithLoweringPass>::default());
    pass_manager.add_pass(Box::<WasmToMidenFinalLoweringPass>::default());
    pass_manager
        .run(ctx, wrapper_module.get_operation())
        .unwrap();
    let inner_module = wrapper_module
        .get_body(ctx, 0)
        .deref(ctx)
        .iter(ctx)
        .collect::<Vec<Ptr<Operation>>>()
        .first()
        .cloned()
        .unwrap();
    inner_module
}

https://github.com/greenhat/omnizk/blob/ce2f0ebc7efa7bd82487000a0df2f7733be7304d/crates/codegen-midenvm/tests/sem_tests.rs#L35-L70

You can define your custom transformations as passes and extend IRs with your custom ops.

[bobbinth]

This is very cool! Thank you! (and the whole OmniZK project sounds very interesting!). I haven't looked too deeply yet, but wanted to ask a couple of questions:

• For Miden, we are working on a slightly higher level MidenIR which should make compilation from WASM simpler. Maybe it would make sense to target that rather than MASM directly? (MidenIR text format still needs to be described, but hopefully, we'll have an initial description in a couple of weeks or so).
• One of the things I was thinking about is how to handle field element operations. WASM has only a few primitive types, but it would be super useful to know when we could use field element ops instead of u32 or u64 operations. What's your thinking on this?

[greenhat]

Please see my comments below. I missed that Github's "Quote replay" feature puts the quote into a separate comment.

[greenhat]

@bobbinth Would you consider sponsoring the Rust -> Wasm -> MidenIR compilation pipeline in OmniZK?

I started it as my side project, and I'm working on it in the evenings and weekends. I'm ready to give it more effort and time to speed up its development.

My email is [email protected]

[greenhat]

For Miden, we are working on a slightly higher level MidenIR which should make compilation from WASM simpler. Maybe it would make sense to target that rather than MASM directly? (MidenIR text format still needs to be described, but hopefully, we'll have an initial description in a couple of weeks or so).

From what I saw in the hir module, it seems higher than Wasm IR and has an SSA form (Wasm is stack-based). So it'd take more effort to target MidenIR from Wasm. MidenIR text format is not that important here as I can use MidenIR Rust types in OmniZK.
For now, I'm emitting MASM and looking into Operation. I believe it might be a good(easier) fit as a target in OmniZK.

One of the things I was thinking about is how to handle field element operations. WASM has only a few primitive types, but it would be super useful to know when we could use field element ops instead of u32 or u64 operations. What's your thinking on this?

I'd look into exposing Rust field element ops to Wasm in a way that would allow us to recognize them in Wasm and substitute them with Miden field element ops.
One option is to craft a Felt type wrapping u64 similar to what @bitwalker made in 0xPolygonMiden/compiler#12 (reply in thread) but with all the implementations calling extern C functions (like fn add(u64, u64) -> u64) which should end up as imports in Wasm module, and we substitute them with native field element ops. For non-wasm targets, a regular Rust implementation is used.

EDIT: Something along these lines:

extern "C" {
   fn felt_add(a: u64, b: u64) -> u64;
}

/// A field element
#[derive(Default, Copy, Clone, PartialEq, Eq)]
#[repr(transparent)]
pub struct Felt(u64);
impl Add for Felt {
   type Output = Felt;

   fn add(self, other: Self) -> Self::Output {
       #[cfg(not(target_arch = "wasm32"))]
       {
           const M: u64 = 0xFFFFFFFF00000001;
           // We compute a + b = a - (p - b).
           let (x1, c1) = self.0.overflowing_sub(M - other.0);
           let adj = 0u32.wrapping_sub(c1 as u32);
           Self(x1.wrapping_sub(adj as u64))
       }

       #[cfg(target_arch = "wasm32")]
       Self(unsafe { felt_add(self.0, other.0) })
   }
}

#[no_mangle]
fn run() {
   let a = Felt(1);
   let b = Felt(2);
   let c = a + b;
}

is compiled to the following Wasm:

(module
  (type (;0;) (func (param i64 i64) (result i64)))
  (type (;1;) (func))
  (import "env" "felt_add" (func $felt_add (;0;) (type 0)))
  (func $run (;1;) (type 1)
    i64.const 1
    i64.const 2
    call $felt_add
    drop
  )

where we substitute call $felt_add with add Miden VM op.

Another option, which I like more, is to look into WASI and build something similar to WASI-crypto https://www.fastly.com/blog/standard-cryptography-api-webassembly but for field element operations and all other custom ops like hashing, etc.

[greenhat]

Another option, which I like more, is to look into WASI and build something similar to WASI-crypto https://www.fastly.com/blog/standard-cryptography-api-webassembly but for field element operations and all other custom ops like hashing, etc.

After digging into the Wasm component model, it seems that its world is a viable option to expose Miden cryptographic ops (hash, etc.) and Miden Standard Library functions in a high-level language like Rust. Any Wasm host can define any world with any interfaces they want. The guest writes Rust code using these interfaces and wit-bindgen generates a glue code leaving host-implementable parts as import functions in Wasm. In our case, the host is not a Wasm runtime, but the Wasm to Miden compiler converting those interface calls into Miden Standard Library functions, special ops, etc. We will need to craft WIT interfaces so that underlying Wasm import function signatures are close to Miden Standard Library and ops counterparts. All these WIT interfaces can be hidden away from the end-users in some miden-std crate so that for hashing, they use fn hash(bytes: Vec<u8>) -> [u8; 16]

For details:
https://github.com/bytecodealliance/wit-bindgen#guest-rust
https://github.com/WebAssembly/component-model

[greenhat]

After a deeper dive into MidenIR, I can see its role in your compilation pipeline, and I think transforming Wasm IR into SSA form and targeting MidenIR in OmniZK might be worth the effort.