refactor: WebGPU backend as accelerator hooks of the native provers - #1817
refactor: WebGPU backend as accelerator hooks of the native provers#1817gbotrel wants to merge 10 commits into
Conversation
Pins gnark-crypto to the head of Consensys-Incorporated/gnark-crypto#879, which adds kzg.Committer, OpenWithCommitter and BatchOpenSinglePointWithCommitter. The bump from v0.20.1 to v0.21.x regenerates tinyfield and renames one koalabear test call (MulByNonResidue -> MulByQuadraticNonResidue).
The proving keys get a runtime-only accelerator field (SetAccelerator). Groth16 hooks: MultiExpG1, MultiExpG2, ComputeH. PLONK hooks: MultiExp, ToCanonical, ComputeNumerator(*NumeratorInput). KZG openings go through kzg.OpenWithCommitter / BatchOpenSinglePointWithCommitter. With no accelerator set the code path is unchanged (one nil check per call).
…ng them Deletes the six hand-copied per-curve prover packages. Each curve now has one generated accelerator.go (templates/go) that uploads the key bases once, resolves MSM bases by slice pointer and sends gnark-crypto elements as raw Montgomery limbs (no re-encoding). Prepare attaches it to the native proving key; groth16.Prove / plonk.Prove run unchanged. The wasm runtime and entrypoints are simplified accordingly.
shaders/curves/* are produced from templates/wgsl by the generator, with all field constants computed from gnark-crypto (they matched the previous hand-typed values). Fp and Fr share one field core template. Dead twist and scalar-multiplication helpers removed.
…one submit per op G1/G2, MSM, Groth16/PLONK modules and the two bridges are each one generic implementation. Proving key bases are uploaded once at prepareKey and released with releaseKey. Multi-stage operations record into a single command encoder with one readback. Fixes the G2 combine dispatch that used 64 threads for a 32-thread kernel, releases buffers on error paths and marks the shader bundle as side-effectful. New Go/JS wire protocol is documented at the top of src/curvegpu/bridge.ts.
|
Review the following changes in direct dependencies. Learn more about Socket for GitHub.
|
There was a problem hiding this comment.
🔵 Needs a closer look
It substantially changes the proving execution path across Go/wasm, JS/WebGPU, and generated WGSL, which warrants careful human validation for cryptographic correctness and runtime safety.
Pull request overview
Refactors the experimental browser WebGPU prover so Groth16/PLONK run through the native backend/groth16 and backend/plonk codepaths, with proving keys gaining runtime-only accelerator hooks that offload heavy work (MSMs, quotient work, FFTs) to a generated WebGPU bridge and WGSL kernels.
Changes:
- Add per-prover “accelerator” attachment points on native proving keys and update the WebGPU wasm runtimes to call native provers after
Prepare. - Replace hand-copied per-curve WebGPU code and shaders with generator-driven Go accelerators and WGSL templates; consolidate/modernize the TypeScript runtime and test pages.
- Bump dependencies (notably
gnark-crypto) and do minor small-field/tinyfield and std test adjustments.
File summaries
| File | Description |
|---|---|
| std/internal/fieldextension/koalabear_ext_test.go | Update Koalabear E2 non-residue helper used by the test. |
| internal/smallfields/tinyfield/vector.go | Replace local parallel executor with gnark-crypto/parallel and modernize loops. |
| internal/smallfields/tinyfield/vector_test.go | Modernize loops/bench loops. |
| internal/smallfields/tinyfield/element.go | Use any for SetInterface and add cube-root helpers. |
| internal/generator/backend/template/zkpschemes/plonk/plonk.setup.go.tmpl | Add runtime-only accelerator field to generated PLONK proving keys. |
| internal/generator/backend/template/zkpschemes/groth16/groth16.setup.go.tmpl | Add runtime-only accelerator field to generated Groth16 proving keys. |
| go.mod | Dependency bumps; pin gnark-crypto to a newer pseudo-version. |
| backend/plonk/bw6-761/setup.go | Add runtime-only accelerator field to PLONK proving key. |
| backend/plonk/bn254/setup.go | Add runtime-only accelerator field to PLONK proving key. |
| backend/plonk/bls12-381/setup.go | Add runtime-only accelerator field to PLONK proving key. |
| backend/plonk/bls12-377/setup.go | Add runtime-only accelerator field to PLONK proving key. |
| backend/groth16/bw6-761/setup.go | Add runtime-only accelerator field to Groth16 proving key. |
| backend/groth16/bn254/setup.go | Add runtime-only accelerator field to Groth16 proving key. |
| backend/groth16/bls12-381/setup.go | Add runtime-only accelerator field to Groth16 proving key. |
| backend/groth16/bls12-377/setup.go | Add runtime-only accelerator field to Groth16 proving key. |
| backend/accelerated/webgpu/web/tests/api/src/shared/raw_kernel.ts | New minimal WebGPU plumbing for raw shader exercise pages + profiling. |
| backend/accelerated/webgpu/web/tests/api/src/shared/msm_bench_sources.ts | New shared base-fixture loading logic for MSM benchmarks. |
| backend/accelerated/webgpu/web/tests/api/src/shared/fixtures.ts | New shared fixture path helpers + hex decoding + assertions. |
| backend/accelerated/webgpu/web/tests/api/src/shared/browser_utils.ts | Move adapter info helper to runtime; add scalar generation + diagnostics helpers. |
| backend/accelerated/webgpu/web/tests/api/src/shared/bench_page.ts | New shared benchmark page wrapper (controls/status/log). |
| backend/accelerated/webgpu/web/tests/api/src/g1_scalar_mul_page.ts | Refactor to use shared fixture/assertion utilities. |
| backend/accelerated/webgpu/web/tests/api/src/fr_ops_page.ts | Deleted; replaced by shared field_ops_page.ts. |
| backend/accelerated/webgpu/web/tests/api/src/fr_ntt_page.ts | Refactor to use shared fixture/assertion utilities. |
| backend/accelerated/webgpu/web/tests/api/src/fr_ntt_bench_page.ts | Refactor benchmark page to shared bench chrome + new scalar generator. |
| backend/accelerated/webgpu/web/tests/api/src/fp_ops_page.ts | Deleted; replaced by shared field_ops_page.ts. |
| backend/accelerated/webgpu/web/tests/api/src/field_ops_page.ts | New shared smoke suite for fr/fp ops. |
| backend/accelerated/webgpu/web/tests/api/src/curvegpu_page.ts | Update suite wiring; allow a single page to serve multiple suites via suiteId. |
| backend/accelerated/webgpu/web/src/index.ts | Update public TS API docs/exports for unified modules and proof runtime. |
| backend/accelerated/webgpu/web/src/curvegpu/types.ts | Remove generated shapes table from this file; keep interface/docs. |
| backend/accelerated/webgpu/web/src/curvegpu/msm_shared.ts | Switch to packed-scalar words and add base index offset support. |
| backend/accelerated/webgpu/web/src/curvegpu/kernels.ts | Deleted old kernel loader (superseded by new organization). |
| backend/accelerated/webgpu/web/src/curvegpu/context.ts | Move adapter info helper here and extend context limits/diagnostics. |
| backend/accelerated/webgpu/web/package.json | Mark generated shader bundle as side-effectful for bundlers. |
| backend/accelerated/webgpu/web/eslint.config.js | Tighten ignores and add test JS globals configuration. |
| backend/accelerated/webgpu/shaders/curves/bn254/g2_io.wgsl | Generated shader header + unrolled loads/stores. |
| backend/accelerated/webgpu/shaders/curves/bn254/g2_arith.wgsl | Generated shader header + cleanup of twist helper and more comments. |
| backend/accelerated/webgpu/shaders/curves/bn254/g1_io.wgsl | Generated shader header + reordered helpers. |
| backend/accelerated/webgpu/shaders/curves/bn254/fr_vector.wgsl | Generated shader header + expanded field core helpers. |
| backend/accelerated/webgpu/shaders/curves/bn254/fr_plonk_quotient.wgsl | Generated shader header. |
| backend/accelerated/webgpu/shaders/curves/bn254/fr_ntt.wgsl | Generated shader header + expanded field core helpers. |
| backend/accelerated/webgpu/shaders/curves/bn254/fr_arith.wgsl | Generated shader header + dedicated square helper. |
| backend/accelerated/webgpu/shaders/curves/bn254/fp_arith.wgsl | Generated shader header + inversion loop var naming cleanup. |
| backend/accelerated/webgpu/shaders/curves/bls12_381/g2_io.wgsl | Generated shader header + unrolled loads/stores. |
| backend/accelerated/webgpu/shaders/curves/bls12_381/g2_arith.wgsl | Generated shader header + remove unused twist helper and add comments. |
| backend/accelerated/webgpu/shaders/curves/bls12_381/g1_io.wgsl | Generated shader header + unrolled loads/stores. |
| backend/accelerated/webgpu/shaders/curves/bls12_381/fr_vector.wgsl | Generated shader header + expanded field core helpers. |
| backend/accelerated/webgpu/shaders/curves/bls12_381/fr_plonk_quotient.wgsl | Generated shader header. |
| backend/accelerated/webgpu/shaders/curves/bls12_381/fr_ntt.wgsl | Generated shader header + expanded field core helpers. |
| backend/accelerated/webgpu/shaders/curves/bls12_381/fr_arith.wgsl | Generated shader header + dedicated square helper. |
| backend/accelerated/webgpu/shaders/curves/bls12_377/g2_io.wgsl | Generated shader header + unrolled loads/stores. |
| backend/accelerated/webgpu/shaders/curves/bls12_377/g2_arith.wgsl | Generated shader header + remove non-residue inv helper and add comments. |
| backend/accelerated/webgpu/shaders/curves/bls12_377/g1_io.wgsl | Generated shader header + unrolled loads/stores. |
| backend/accelerated/webgpu/shaders/curves/bls12_377/fr_vector.wgsl | Generated shader header + expanded field core helpers. |
| backend/accelerated/webgpu/shaders/curves/bls12_377/fr_plonk_quotient.wgsl | Generated shader header. |
| backend/accelerated/webgpu/shaders/curves/bls12_377/fr_ntt.wgsl | Generated shader header + expanded field core helpers. |
| backend/accelerated/webgpu/shaders/curves/bls12_377/fr_arith.wgsl | Generated shader header + dedicated square helper. |
| backend/accelerated/webgpu/shaders/common/g1_core.wgsl | Remove unused affine small scalar mul helper. |
| backend/accelerated/webgpu/README.md | Rewrite docs around native-prover accelerator model and generation. |
| backend/accelerated/webgpu/plonk/plonk.go | Replace forked prover wrappers with Prepare/Release that attach accelerators to native keys. |
| backend/accelerated/webgpu/plonk/internal/wasmruntime/webgpu/main.go | Update wasm entrypoint to use native key factories + Prepare hook, then call native Prove/Verify. |
| backend/accelerated/webgpu/plonk/internal/wasmruntime/native/main.go | Simplify native wasm entrypoint wiring and keep program alive. |
| backend/accelerated/webgpu/plonk/internal/bridge/bridge.go | Deleted curve-specific bridge wrapper (centralized in shared bridge). |
| backend/accelerated/webgpu/plonk/doc.go | Update package docs to match accelerator attachment design. |
| backend/accelerated/webgpu/plonk/bn254/provingkey.go | Deleted old per-curve proving-key wrapper implementation. |
| backend/accelerated/webgpu/plonk/bls12-381/provingkey.go | Deleted old per-curve proving-key wrapper implementation. |
| backend/accelerated/webgpu/plonk/bls12-377/provingkey.go | Deleted old per-curve proving-key wrapper implementation. |
| backend/accelerated/webgpu/internal/generator/templates/wgsl/g2_io.wgsl.tmpl | New WGSL template for generated G2 IO helpers. |
| backend/accelerated/webgpu/internal/generator/templates/wgsl/g1_io.wgsl.tmpl | New WGSL template for generated G1 IO helpers. |
| backend/accelerated/webgpu/internal/generator/templates/wgsl/fr_vector.wgsl.tmpl | New WGSL template for generated Fr vector ops kernel. |
| backend/accelerated/webgpu/internal/generator/templates/wgsl/fr_ntt.wgsl.tmpl | New WGSL template for generated Fr NTT stage kernel. |
| backend/accelerated/webgpu/internal/generator/templates/wgsl/field_arith.wgsl.tmpl | New WGSL template for generated standalone field arithmetic kernels. |
| backend/accelerated/webgpu/internal/generator/main.go | Ensure WGSL and Go accelerators are generated in generator main. |
| backend/accelerated/webgpu/internal/generator/goaccel.go | New generator for per-curve Go accelerator packages. |
| backend/accelerated/webgpu/internal/bridge/basis.go | New basis registry for mapping proving-key subslices to GPU buffers. |
| backend/accelerated/webgpu/groth16/internal/wasmruntime/webgpu/main.go | Update wasm entrypoint to use native key factories + Prepare hook, then call native Prove/Verify. |
| backend/accelerated/webgpu/groth16/internal/wasmruntime/native/main.go | Simplify native wasm entrypoint wiring and keep program alive. |
| backend/accelerated/webgpu/groth16/internal/common/filter_indices.go | Deleted old helper used by forked Groth16 proving key wrapper. |
| backend/accelerated/webgpu/groth16/internal/bridge/bridge.go | Deleted old Groth16-specific bridge wrapper (centralized in shared bridge). |
| backend/accelerated/webgpu/groth16/groth16.go | Replace forked Groth16 Prove wrapper with Prepare/Release attaching accelerators to native keys. |
| backend/accelerated/webgpu/groth16/doc.go | Update package docs to match accelerator attachment design. |
| backend/accelerated/webgpu/groth16/bn254/serialize.go | Deleted old serialization/packing helpers for forked Groth16 wrapper. |
| backend/accelerated/webgpu/groth16/bn254/provingkey.go | Deleted old per-curve proving-key wrapper implementation. |
| backend/accelerated/webgpu/groth16/bls12-381/provingkey.go | Deleted old per-curve proving-key wrapper implementation. |
| backend/accelerated/webgpu/groth16/bls12-377/provingkey.go | Deleted old per-curve proving-key wrapper implementation. |
Review details
- Files reviewed: 88/160 changed files
- Comments generated: 5
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| const count = inputA.byteLength / elementBytes; | ||
| const dataBytes = inputA.byteLength; | ||
| const totalStart = performance.now(); | ||
| const inputABuffer = device.createBuffer({ label: "input-a", size: dataBytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST }); | ||
| const inputBBuffer = device.createBuffer({ label: "input-b", size: dataBytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST }); | ||
| const outputBuffer = device.createBuffer({ label: "output", size: dataBytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC }); | ||
| const stagingBuffer = device.createBuffer({ label: "staging", size: dataBytes, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ }); |
| throw new Error(`fr NTT vectors unavailable for curve ${module.id}`); | ||
| } | ||
| log(`=== ${config.title} ===`); | ||
| export async function runSuite(module: CurveModule, log: Log): Promise<SuiteResult> { |
| throw new Error(`g1 scalar-mul vectors unavailable for curve ${module.id}`); | ||
| } | ||
| log(`=== ${config.title} ===`); | ||
| export async function runSuite(module: CurveModule, log: Log): Promise<SuiteResult> { |
| const kernelStart = performance.now(); | ||
| const encoder = device.createCommandEncoder({ label: "encoder" }); | ||
| const pass = encoder.beginComputePass({ label: "pass" }); | ||
| pass.setPipeline(kernel.pipeline); | ||
| pass.setBindGroup(0, bindGroup); | ||
| pass.dispatchWorkgroups(Math.ceil(count / WORKGROUP_SIZE)); | ||
| pass.end(); | ||
| encoder.copyBufferToBuffer(outputBuffer, 0, stagingBuffer, 0, dataBytes); | ||
| device.queue.submit([encoder.finish()]); | ||
| const kernelMs = performance.now() - kernelStart; | ||
|
|
||
| const readbackStart = performance.now(); | ||
| await stagingBuffer.mapAsync(GPUMapMode.READ); | ||
| const out = new Uint8Array(stagingBuffer.getMappedRange().slice(0)); | ||
| stagingBuffer.unmap(); | ||
| const readbackMs = performance.now() - readbackStart; |
| // Cbrt z = ∛x (mod q) | ||
| // if the cube root doesn't exist (x is not a cube mod q) | ||
| // Cbrt leaves z unchanged and returns nil | ||
| func (z *Element) Cbrt(x *Element) *Element { |
| return out, nil | ||
| } | ||
|
|
||
| func (r *Runtime[PK, VK, Proof]) verify(args []js.Value) (js.Value, error) { |
There was a problem hiding this comment.
Wasm release skips GPU cleanup
Medium Severity
Runtime release only deletes the Go handle maps. It never calls groth16.Release / plonk.Release, and Config has no cleanup hook, so prepared GPU keys stay in gnarkGroth16WebGPU / gnarkPlonkWebGPU until the page is torn down.
Additional Locations (2)
Reviewed by Cursor Bugbot for commit ac94f14. Configure here.
web/scripts/e2e.mjs runs the API suites and one verified proof per system and curve with Playwright (npm run test:e2e), or the prover benchmark matrix (npm run bench:e2e). .github/workflows/webgpu.yml runs it on macOS.
… for the MSM The signed-digit bucket assignment moves from JavaScript to shaders/common/msm_sort.wgsl (count, scan, scatter, chunks). Buckets are reduced with per-thread running sums instead of a scalar multiplication per bucket, heavy buckets are folded first, and the combine stage gets one point per window. Bases are stored on the GPU as two affine coordinates. G2 kernels are kept small and read each point from a fixed buffer: larger fused G2 kernels were miscompiled by Metal. One command encoder per MSM.
…fused NTT stages Field add, sub, neg, double and normalize use adc/sbb chains and select instead of a limb-by-limb compare and branch. The Montgomery reduction loop is peeled (outer unrolling is a generator knob, off by default: 3-7% faster but 3x larger shaders and 2x slower cold init). The NTT performs up to nine radix-2 stages per dispatch in workgroup memory, so a 2^20 transform is three dispatches instead of twenty; twiddles and coset tables come from cached GPU buffers. The NTT module exposes a recording API (recordForward/recordInverse/recordMulVector) for callers that batch several passes into one submission.
… GPU The circuit polynomials are canonicalized once and uploaded once, with the per-coset tables; each proof recomputes their coset evaluations on the GPU in the same submission as the witness transforms and the numerator kernel, which now writes Montgomery output directly in the bit-reversed layout of the large domain. This removes the rho coset copies from the Go heap and the per-proof assembly of a cosetCount x vectorCount buffer in JavaScript. All PLONK bridge inputs are Montgomery, so Go sends raw element memory. Shared key handling and limb decoding move to leaf generics in the bridge package; the generated accelerators shrink accordingly.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
There are 2 total unresolved issues (including 1 from previous review).
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit af21c07. Configure here.
| dynamicEvals, | ||
| staticEvals, | ||
| s.twiddles, | ||
| { buffer: s.denominators, offset: coset * vectorBytes, size: vectorBytes }, |
There was a problem hiding this comment.
Unaligned GPU binds on small domains
Medium Severity
evaluate binds each coset’s scaling and denominators at offset coset * n * 32. WebGPU requires storage-buffer offsets to be multiples of minStorageBufferOffsetAlignment (256). For small PLONK domains (n is 1, 2, or 4 when the system has fewer than 6 constraints), those offsets are 32, 64, or 128, so bind-group creation throws after the first coset and proving fails. The tables are packed tightly with no 256-byte padding.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit af21c07. Configure here.


This sits on top of #1789. It keeps the same idea (Go prover in wasm, heavy math on the GPU) but removes the copy-pasted code, fixes a few things found while reviewing, and makes the GPU side faster.
What changed
Native provers get hooks instead of being forked.
backend/groth16andbackend/plonkproving keys now have anAcceleratorfield (pk.SetAccelerator(...)). Groth16 has 3 hooks (MultiExpG1,MultiExpG2,ComputeH), PLONK has 3 (MultiExp,ToCanonical,ComputeNumerator). When nothing is attached, the code path is the old one behind a nil check. This is done in the templates and regenerated. The KZG openings use the newkzg.OpenWithCommitter/BatchOpenSinglePointWithCommitterfrom Consensys-Incorporated/gnark-crypto#879, so we don't copy the KZG folding code anymore. That PR must be merged first; go.mod points at its commit for now.WebGPU Go code. The 6 hand-copied per-curve prover packages are deleted. Each curve now has one generated
accelerator.go(170 lines Groth16, 460 lines PLONK) that implements the hooks. Bases are uploaded to the GPU once; Go sends field elements and points as their raw in-memory bytes (Montgomery limbs, little endian) instead of re-encoding them. Results coming back from the GPU are checked (reduced limbs, point on curve). For PLONK the circuit polynomials are canonicalized once and kept on the GPU; each proof recomputes their coset evaluations there, so the Go heap no longer holds the coset copies (that was ~300 MB at 2^18 and would not fit at 2^20).Shaders.
shaders/curves/*are generated from templates, with all field constants computed from gnark-crypto (checked to match the old hand-typed ones). Field add/sub are branchless. The MSM does its bucket sorting on the GPU (shaders/common/msm_sort.wgsl), reduces buckets with running sums instead of one scalar multiplication per bucket, and reads affine bases (two coordinates instead of three). The NTT runs up to nine radix-2 stages per dispatch in workgroup memory (a 2^20 transform is 3 dispatches instead of 20). One thing learned the hard way: large fused G2 kernels are miscompiled by Metal in Chrome, so the G2 stages are kept small and identical in shape to G1.TypeScript. G1/G2, MSM, Groth16/PLONK modules and the two bridges are each a single generic implementation. Multi-stage operations use one command encoder and one readback. Fixed a real bug: the G2 combine kernel was dispatched with 64 threads per workgroup but compiled with 32. Buffers are released on error paths. The Go/JS protocol is written down at the top of
web/src/curvegpu/bridge.ts.Tests.
npm run test:e2eruns the API suites and one verified proof per system and curve headlessly with Playwright;npm run bench:e2eruns the prover matrix..github/workflows/webgpu.ymlruns the former on macOS runners (WebGPU needs a GPU; not yet verified on GitHub's runners).Numbers
Lines in
backend/accelerated/webgpuvs #1789: hand-written −9,655, generated output +6,712 (more shader is generated now: the NTT and MSM kernels grew), net −2,943. Native prover templates: +295 / −26.Headless Chrome on an M-series Mac, 3 proofs each, all proofs verified. GPU prove time, #1789 → this PR (CPU wasm prover for scale):
With one BSB22 commitment: Groth16 bn254 2^18 974 → 322 ms, PLONK bn254 2^18 3.93 → 1.64 s. Prepare (one-time per key): Groth16 2^18 1.06 s → 8 ms, PLONK 2^18 7.3 s → 0.41 s.
Known limits
The MSM combine stage is a single-thread chain of 256 doublings per MSM (~7 ms G1, ~20 ms G2 on bn254); it is now the largest fixed cost. The field multiply still uses 16-bit limbs with a carry after every product. The PLONK
preloadre-uploads nothing per proof, but the numerator is computed coset by coset, so the largest GPU buffer is one coset's worth of vectors.How to check