[ProjectTracking]: ZK WASM Stage 1, WASM Spec Test

[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 1 is to support proofs for individual simple WASM programs and pass WASM MVP Spec test.

Supported workloads:

• WASM MVP Spec tests
• Integers ops, Control flow ops, Memory ops from WASM Instruction set

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 , and @nagisa .

Stage 1 is expected to be completed by the end of November 2023.

Assumptions

Assumptions and risks are enumerated in the roadmap.

Pre-requisites

All preparatory work has been done in August-September 2023 and can be found here:

• https://github.com/near/wasmtime/issues/28
• https://github.com/near/wasmtime/issues/27
• https://github.com/near/wasmtime/issues/25

Out of scope

• Floating point operations support
• Performance optimizations for bulk memory operations (memory.fill and memory.copy)

[aborg-dev]

We are nearing the completion of this milestone. Over the last two weeks, we have:

• Imported i32 and i64 WASM Spec tests
• Supported new opcodes for i32 and i64: mul, ne, or, shl, shr_u, xor, and
• Built infrastructure to measure continuously test success rates and performance of all test programs
• Introduced two benchmarks that will be within the scope of Stage 1:
  • Fibonacci
  • SHA256

We aim to complete Stage 1 by the end of next week by:

• Passing tests in 2 benchmarks above
• Passing i32 and i64 spec tests for instructions involved in these benchmarks

We still need to implement the following instructions:

• i32.rotl, i64.extend_i32_u, i32.load, i32.store8, i32.lt_u

---

As of now, the spec test pass rates are as follows:

• i32: 155 out of 364 tests
• i64: 161 out of 374 tests

We don't expect a 100% pass rate for Stage 1, as this involves obscure instructions that we don't need (e.g. popcount). Right now we fully support following instructions:

• i32: add, eq, lt_s, mul, ne, or, shl, shr_u, xor, and
• i64: add, eq, lt_s, mul, ne, or, shl, shr_u, xor, and

The up-to-date data about pass rates can be found at https://github.com/near/wasmtime/blob/main/docs/zkasm/test_summary.csv

[aborg-dev]

Quick summary of where we are now:

• All necessary i32/i64 instructions are supported and pass spec tests (Great job, @Viktar Makouski )
• We have Fibonacci and SHA256 WAT files imported to wasmtime and run them as part of the test suite
• Fibonacci works and we have meaningful comparison with multiple other implementations
• SHA256 compiles to ZK ASM after a small chain of PRs that fix memory operations and runs to completion after a much hackier set of changes

The missing pieces right now are:

• Global variables - these are declared in WAT file and live at specific memory offsets. We can replace them with ZK ASM Global variables and would need to detect the access to these offsets and replace them with loads/store from globals
• Data segment (e.g. "hello world" string in SHA256 benchmark) - these are very similar to global variables - reads from a very large offset. But this time it's from a contiguous slice of memory, rather than discrete slots of globals - one way to handle these is to dedicate a memory region and fill it with data at the start of the program
• Loads from frame pointer + offset - I'm not sure what these are, likely part of the function ABI and would have to use stack or registers instead of the frame. Right now these are commented out

As a next step, we will take relevant spec tests and make sure they work:

• https://github.com/WebAssembly/spec/blob/master/test/core/global.wast
• https://github.com/WebAssembly/spec/blob/master/test/core/data.wast
• https://github.com/WebAssembly/spec/blob/master/test/core/load.wast
• https://github.com/WebAssembly/spec/blob/master/test/core/store.wast

Overall, it's been great progress, but it will take us a bit longer to complete Stage 1.

[akhi3030]

Thanks for the summary, great work all around!

[aborg-dev]

Here is the latest update:

• Changes necessary to support memory operations have been submitted
• Plain memory loads and stores from register + offset are working now
• Support for global variables is in progress

We also have implemented a new approach to storing i32 integers - this yields smaller and faster code for most of the instructions. In particular, this cuts more than half (30k) of lines off ZKASM generated file for SHA256 benchmark. This should hopefully make the debugging easier.

We are also getting a pretty good coverage for i32/i64 instructions now. Here is the improvement from the last time:

• i32: 155 -> 287 out of 364 tests
• i64: 161 -> 291 out of 374 tests
• (new) type conversions: 0 -> 12 out of 24 tests

We still need more work on the implementation of memory-related ops mentioned in https://github.com/near/near-one-project-tracking/issues/2#issuecomment-1838085253 and on introducing relevant spectests.

[aborg-dev]

Since the last time, we have finished implementing global variables, data segments, and frame pointers. The SHA256 benchmark that we were blocked on now runs till completion within 30k cycles, but it still yields incorrect results.

Even though we still need to fully fix SHA256 benchmark, we decided that it does not make sense to block the start of Stage 2 on that result, as the work on SHA256 is quite tricky to parallelize.

Given that we completed all our other goals, we will declare Stage 1 as complete and will continue working on SHA256 correctness in Stage 2.