o1-labs / o1js

TypeScript framework for zk-SNARKs and zkApps
https://docs.minaprotocol.com/en/zkapps/how-to-write-a-zkapp
Apache License 2.0
509 stars 116 forks source link

Swap out Wasm curve arithmetic with faster implementations #914

Open mitschabaude opened 1 year ago

mitschabaude commented 1 year ago

Under the hood, snarkyjs depends on a Rust code base which is based on arkworks for finite field and elliptic curve implementations. These are fairly slow when compiled to Wasm, in part due to some Wasm idiosyncracies, and there are faster options available:

There is some promising work by Parity to swap out the arkworks curve arithmetic (with native implementations) to get major speedups over the Wasm version: https://github.com/paritytech/ark-substrate

stevenpack commented 1 year ago

Looks promising. @mitschabaude can you fill out the estimate for this?

mitschabaude commented 2 months ago

Current progress: I implemented a Wasm multiplication representing field elements as 5x51-bit limbs. Performance is improved by about 30% compared to the fastest Wasm multiplication I ever wrote previously, and should be at least 3x faster than current arkworks.

The improved multiplication relies on an experimental Wasm feature (relaxed SIMD), so currently I'm working on developing a fallback which does 51x5 multiplication with non-experimental opcodes. This seems necessary so that we don't ruin compatibility with browsers like Firefox and Safari which currently only have relaxed SIMD support behind a flag. It looks like the fallback version will have speed at least comparable to arkworks, so we won't have a regression on older platforms.

All of this is currently just done in a plain TS environment which facilitates quick iteration. Once all the details of low-level arithmetic are figured out, the second task is implementing them as Rust crates that can replace what we currently use in Kimchi for low-level Field arithmetic.

I spent about 6 days in total so far and estimate that the entire project will take another ~3-4 eng weeks. If all is done, our Wasm prover should be at least 3x faster than it is now. cc @georgeee

georgeee commented 2 months ago

Relates to https://github.com/mitschabaude/montgomery/pull/19