[ProjectTracking]: ZK WASM Stage 2, WASM Interpreter

[aborg-dev]

Goals

Our long-term goal is to use ZK Proofs to scale NEAR protocol to more shards while maintaining high security guarantees.

The goal of Stage 2 is to support proofs for a WASM interpreter (compiled to WASM) that supports WASM MVP features.

Supported workloads:

• WASM interpreter - WASMI (subject to change)
• Running and passing all our spectests and benchmarks on ZK ASM circuit corresponding to WASM interpreter

We will also set an ambitious performance goal - we want the cycle amplification factor from WASM to ZK ASM (number of ZK ASM cycles to implement a WASM opcode) to be less than 5 on average. Right now it is around 30.

Links to external documentation and discussions

The documentation for the project can be found in near/wasmtime.

All meeting notes and design docs are indexed on HackMD.

Estimated effort

The concrete work items for this project are tracked by the milestone.

It is expected that the runtime team will be working on this project. The people actively working on this project are: @akashin , @MCJOHN974 , @mooori , and @nagisa .

Stage 2 is expected to be completed by April 2023.

Assumptions

Assumptions and risks are enumerated in the roadmap.

Pre-requisites

Pre-requisite work will be completed as part of Stage 1

Out of scope

• Auditing the generated code

[MCJOHN974]

Just some notes.

(number of ZK ASM cycles to implement a WASM opcode) to be less than 5 on average

1) I think 5 is good goal for stage 3 or 4, while for stage 2 I would suggest 12 -- realistic, but still almost 3 times better than current one. 2) Should @mooori be added to list of actively involved persons? 3) I suggest to add work on 32-bit zk-processor to milestone, IMO, it is one of the most effective ways to boost performance. Also, if we planning to create 32-bit zk-processor, we should do it ASAP, because all work with masking for 32-bit we do now, will be useless after 32-bit processor will be created. So, I suggest to set 32-bit zk-processor as second priority after supporting instruction set for wasm interpreter.

[aborg-dev]

Just some notes.

(number of ZK ASM cycles to implement a WASM opcode) to be less than 5 on average

1. I think 5 is good goal for stage 3 or 4, while for stage 2 I would suggest 12 -- realistic, but still almost 3 times better than current one.

I agree that 5 is likely too optimistic. I think we should identify the methodology on how to measure this, as you suggested in the last meeting, and that would help us to come up with a better number/set of numbers to target. I filed an issue for this - https://github.com/near/wasmtime/issues/137

2. Should @mooori be added to list of actively involved persons?

Yes! My bad, thank you for flagging this.

3. I suggest to add work on 32-bit zk-processor to milestone, IMO, it is one of the most effective ways to boost performance. Also, if we planning to create 32-bit zk-processor,  we should do it ASAP, because all work with masking for 32-bit we do now, will be useless after 32-bit processor will be created. So, I suggest to set 32-bit zk-processor as second priority after supporting instruction set for wasm interpreter.

I agree that dedicated support of 32-bit ops in ZK ASM will likely bring a good performance boost. I still think that this should be treated as an implementation detail and I'm not ready to commit to it before we identify other things on our plate and other optimizations that we can work on and weigh them against each other. I suggest we make this call slightly later when we study the results of SHA256 benchmark and do some basic analysis of the "hotness" of each instruction.

[MCJOHN974]

I agree that dedicated support of 32-bit ops in ZK ASM will likely bring a good performance boost. I still think that this should be treated as an implementation detail and I'm not ready to commit to it before we identify other things on our plate and other optimizations that we can work on and weigh them against each other. I suggest we make this call slightly later when we study the results of SHA256 benchmark and do some basic analysis of the "hotness" of each instruction.

My main point here is, that, while we don't have 32-bit type in zkasm, we have to implement some sort of masking. Once we will have 32-bit supported by zkasm, we will instantly move to trash all this masking. So, I just suggest to try make as small effort on masking as possible

About how it will increase performance, totally agree, see results of more benchmarks will be better

[aborg-dev]

We now have a milestone tracking all the work that we want to do. Given that there are a lot of items there, we added some structure by introducing a top-level tracking items that encompass smaller items:

1. Test infrastructure
2. Benchmarking infrastructure
3. WASM Interpreter in ZKASM
4. Supporting WASM opcodes for Stage 2
5. Performance optimizations for Stage 2

Items 1-3 will be worked on as individual projects for the next 4-8 weeks given their well-defined scope and limited context. For items 4-5, I think we will all work on solving them collectively - this will ensure that everyone has the opportunity to contribute to the actual backend implementation and understand WASM and ZKASM on a deeper level. This will also help us to dogfood the end-to-end tooling that we've built for items 1-3.

We have started this stage this week and will aim to complete it by the end of April 2024.

[aborg-dev]

Here is an update on the work that happened in the last two weeks.

We finally fixed SHA256 benchmark and are working now on including code changes into the main branch. The resulting SHA256 ZKASM program takes 27k cycles on 32 bytes of input and 210k cycles on 448 bytes of input (there are some constant startup overheads). This is within 7x of a dedicated circuit for SHA256.

We have completed a round of design reviews:

• Testing infrastructure https://github.com/near/wasmtime/pull/183 by @MCJOHN974
• Benchmarking infrastructure https://github.com/near/wasmtime/pull/182 by @mooori

and will start implementing these designs as the next step.

We also had a productive discussion about the design and features of the ZKASM processor specialized for WASM and agreed to start working on it at the beginning of March.

[aborg-dev]

We've had a lot of progress over the last two weeks:

• We agreed and started building towards the benchmarking MVP for benchmarking infrastructure (by @mooori )

• We agreed and started building towards using cranelift filetest for the test infrastructure (by @MCJOHN974 )

• We have finally submitted all the code necessary to run SHA256 benchmark

• Based on the work for SHA256, we were able to run Keccak256 hash function

  • This is very cool, as it further confirms that we have a correct implementation of tricky memory operations and have implemented enough instructions to support a wide range of benchmarks
  • The performance is within 10x of a highly optimized Polygon implementation (34k steps vs 3k steps), very similar to SHA256 benchmark

• Implemented a prototype of WASM interpreter for Stage 2

  • Overall, this was quite straightforward and we're only missing a few instructions to run it
  • The size is quite large: 450k lines of ZKASM, 570KB of WASM. I hope we can bring it down closer to WAMR size - 85KB binary

The current timeline is:

• Finish infrastructure work and WASM interpreter by the end of February 2024
• Design zkAsm processor with Polygon and migrate to it in March 2024
• Optimize the benchmarks till we reach the performance target in April 2024

[bowenwang1996]

Optimize the benchmarks till we reach the performance target in April 2024

@akashin what is the performance target?

[aborg-dev]

Optimize the benchmarks till we reach the performance target in April 2024

@akashin what is the performance target?

I think in Stage 2 we should aim for:

• Performance of EVM-specific benchmarks (SHA256, Blake2, Elliptic Curves) within 2x of their optimized implementation in PIL
• Overhead of WASM interpreter in zkAsm running benchmarks within 2x of performance of the dedicated circuit for that benchmark

This seems realistic and will help us to keep our work grounded. For the Stage 3 we can start introducing NEAR-specific performance targets.

[aborg-dev]

Actually, I take back my last statement, we can start targeting some NEAR-specific workloads now, but we first need to figure out what is worth pursuing and at what performance it will become valuable for NEAR protocol. We should discuss this in more detail.

[aborg-dev]

Our progress over the last 3 weeks:

Test Infra (by @MCJOHN974)

• Implemented zkAsm code generation for filetests #211
• Implemented JS helper to collect test results #229
• Working on enabling run statements in runtests for zkAsm

Benchmarking Infra (by @mooori)

• Extended zkAsm JS runner to enable benchmarking #226
• Implemented profiling instrumentation during zkAsm codegen #230

WASM Interpreter (by @akashin)

• Fixed Rust std version of WASMI interpreter to fully run in WASM
• Started implementing stubs for floating point operations required to run the interpreter

Misc

• Participated in the workshop about PIL2 with a goal of understanding how we can use it to build a ZK processor for WASM
  • We will be working with Polygon in March to migrate to PIL2
• Investigated compiling NEAR Light client
  • Takes 40M cycles on Risc0, proving time around 15m on M1 Macbook. It should take around 8M cycles on zkAsm
  • This is 1000x slowdown compared to block production
  • We will include this as a benchmark for Stage 2
• Discussed using zkWasm for Stateless Validation with @robin-near
  • Looks like the most promising direction to investigate would be removing the need to send the whole state witness and replace it (or parts of it) with a short zk proof.
• Investigated using a faster JS runtime to interpreter zkAsm
  • We also agreed with Polygon that we need to have Rust interface for executing zkAsm and generating proofs
• Merged 234 commits from upstream wasmtime

[bowenwang1996]

We will be working with Polygon in March to migrate to PIL2

Sounds cool! Where can I read more about this?

This is 1000x slowdown compared to block production

What do you mean by "block production"? Light clients do not verify state transitions, so I am not sure whether it is an apple to apple comparison

[aborg-dev]

We will be working with Polygon in March to migrate to PIL2

Sounds cool! Where can I read more about this?

This is 1000x slowdown compared to block production

What do you mean by "block production"? Light clients do not verify state transitions, so I am not sure whether it is an apple to apple comparison

I will share the design doc for PIL2 migration when it is ready, for now our discussions can be found in https://hackmd.io/MTmmOMOxShaMqMi83bxzRQ. For now, I've shared PIL2 presentation (over email) that we had during the Workshop last week.

Regarding the light client, I meant that we need ~1000 machines producing light client proofs to keep up with the block production rate (every 1 second). I agree that this is not apples to apples comparison, just an attempt to understand what would it take to make ZK practical in this context.

[mooori]

Progress over the last two weeks:

Testing infra (by @MCJOH974)

• Implemented compiling of .clif run commands
• Implemented assembling all parts of zkasm test file together
• Implemented pretty-printing of test results
• Now zkasm filetests finally generate valid zkasm program, run it, check assertions, and print details in case of fail

TODO:

• Still need rework preamble generation to make tests with non-default memory settings work (for example with global variables).
• Start enabling zkasm-filetest pipeline until they will cover everything we had in our old testing infra

Benchmarking infra (by @mooori)

• Setup the analyze-zkasm CLI and added the instrument-inst subcommand that allows compiling WAT to zkASM instrumented for tracing instructions executed at runtime #233
• Finalized tooling to generate instructions traces #240
  • Extended run-tests-zkasm.js to profile instrumented zkASM.
  • Reworked the zkevm-proverjs helper to collect raw traces and write them in chunks.

Above changes enable going WAT -> instrumented zkASM -> trace of executed instructions as described in 240. Started working on a analyze-zkasm subcommand to facilitate and better document that workflow.

Tooling and codegen (by @akashin)

• Fixed an issue with node module path definition that broke tests if executed from certain directories #241.
• Made zkasm_runner public and enabled using it for filetests #243.
• Make a module's main function the entrypoint for generated zkASM and stop abusing the start section for that purpose #244.
  • This helps generating zkASM from Rust, in which case generating a start section would not be trivial.
• Started adding benchmarks written in Rust to the wasmtime repo, compile them them in build scripts and execute them as part of CI #231.

Planning (by @akashin)

• Exploring use cases of zk wasm, for instance in relation to Merkle path compressions and state proofs ref.

---

@akashin I hope we captured your progress correctly. In case something is missing or wrong, please leave a message or (if possible) update this post.

[mooori]

Timeline update regarding benchmarking infra:

• Instruction profiling should be finalized next week.

TBD if afterwards work on benchmarking infra continues or if something else gets prioritized. Possible next steps for benchmarking infra could be:

• Label profiling: Answers the questions which labels were entered how often at runtime. Should be in the ballpark of 1-2 weeks since it is similar to instruction profiling.
• Flame graphs: Would give insights into function call hierarchies and where most time is spent. Further investigations would be required to estimate the workload.

Other tasks that might be prioritized, for now, over further benchmarking infra work:

• Work on the PIL2 processor for zk wasm.
• Work related to evaluating use cases of zk wasm.