[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 - near/wasmtime#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 docs: infra wasmtime#183 by @MCJOHN974
• Benchmarking infrastructure bench: Add design specs for ZK WASM benchmarking infrastructure wasmtime#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?