Tracking issue: Proposal

[TalDerei]

Proposed Work: Implement a general-purpose, open-source WebGPU-based Groth16 prover to accelerate client-side proof generation for the Penumbra Protocol.

Phases: The end-to-end WebGPU POC will comprise three distinct, high-level phases described in greater detail below. Generally, phase (ii) scopes the open-source gluing / integration of the Penumbra web-code with a WebGPU backend, and phase (iii) scopes the closed-source zprize work based on a sparse-matrix MSM design (cuZK). Since (iii) is a single, sub-component in the larger system (ii) that can be built entirely independently and without conflict, then (ii) would be optimally structured to ease the eventual merging with (iii). For instance, hot-swapping the reference MSM algorithm in (ii) with the optimized MSM algorithm in (iii) using an API function call. Phases (ii) and (iii) would be concurrently built (in the open and closed-source), and remain independent until the codebases can be merged.

Phase (I): Architecture Code and Performance Review

• Description: Conduct general architecture, code-review, and comprehensive performance pass comprising the scoped penumbra crates (arkworks groth16 prover, crypto-modules, penumbra-web-assembly.
• Success: Measured by successfully understanding the system architecture design (e.g. transaction planning, proof generation, storage, etc.) and properly identifying the performance bottlenecks.
  • [x] #3
  • [x] #4 (link to spreadsheet)

Phase (II): WebGPU Backend for Penumbra Web Code + Web Improvements

• Description: Integration of the naive (unoptimized) WebGPU components with Penumbra WASM and web-client. The reference WebGPU MSM repository from Demox-Labs will serve as the naive WebGPU implementation. Additionally, implement necessary changes to (a) Penumbra WASM crate and (b) integration strategies with web-code related to accelerating the existing web-client.
• Success: Measured by successfully generating a client-side proof (2^16 constraints) on the GPU.
  • [x] #1
  • [x] #6, https://github.com/penumbra-zone/webgpu/pull/8
  • [x] https://github.com/penumbra-zone/penumbra/issues/3236 https://github.com/penumbra-zone/penumbra/pull/3289
  • [x] https://github.com/penumbra-zone/penumbra/pull/3337 https://github.com/penumbra-zone/penumbra/pull/3411/
  • [x] https://github.com/penumbra-zone/webgpu/pull/9,
  • [x] https://github.com/penumbra-zone/web/pull/303
  • [x] https://github.com/penumbra-zone/web/issues/324
  • [x] Optimize crypto primitives to improve performance in the WASM context https://github.com/penumbra-zone/penumbra/issues/3526

Phase (III): ZPrize Competition Related-Work

• Description: Integration and adaptation of optimized WebGPU MSM with Penumbra WASM and web clients. As part of the Zprize submission collaboration, there will be a closed-source WebGPU MSM which will be open-sourced upon completion of the Zprize contest submission deadline (Early February). Upon open-sourcing, the optimized WebGPU MSM will be adapted and integrated into the Penumbra prover from phase (ii), swapping the new optimized MSM for the unoptimized MSM from the demox-labs repository. The MSM will be rewritten in WGSL (WebGPU shader language) and utilize web-workers for multi-threading, rather than using the existing variable-base MSM from Arkworks in the Penumbra monorepo.

• Success: Measured by achieving sub-second MSM execution times.

  • [x] Open-Source codebases (https://github.com/td-kwj-zp2023)
  • [ ] #11

[TalDerei]

Phase III is currently a post-mainnet task