Skip to content

Commit bf52a3a

Browse files
committed
Add GPU-accelerated PLONK prover (bn254) via ICICLE
Ports the Gnark_plonk_on_GPU work onto current gnark master: - New backend/accelerated/icicle/plonk/ package: full bn254 PLONK prover on GPU (NTT, MSM, quotient, linearization, KZG openings via ICICLE), selected with -tags=icicle; falls back to the CPU prover without the tag. Stage-timing breakdown under ICICLE_STEP_PROFILE. - Solver output caching (constraint/bn254): save/load solved witness vectors and LRO solutions to skip redundant gnark solving runs. - Optional blinding: USE_BLINDINGS env gates blinding-polynomial generation in the bn254 PLONK prover (off = deterministic proofs). - More descriptive PLONK verifier errors; witness-creation logging. - Depends on the extended icicle-gnark fork (batch inverse, poly eval, permutation supports, shard merge - see ingonyama-zk/icicle-gnark#4) via a local replace directive to ../icicle-gnark-extended.
1 parent 2ea1515 commit bf52a3a

24 files changed

Lines changed: 8968 additions & 94 deletions

AGENTS.md

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,3 +56,25 @@ For PRs, include clear purpose, linked issues when relevant, and test evidence.
5656
## Security & Reporting
5757

5858
Do not open public issues for vulnerabilities. Follow `SECURITY.md` for responsible disclosure.
59+
60+
## GPU acceleration rules (this fork)
61+
62+
This fork implements GPU-accelerated proving components.
63+
64+
Main goals:
65+
66+
- correctness first
67+
- Prefer GPU operations over CPU operations
68+
69+
The codebase may include:
70+
71+
- A library called icicle to wrap GPU kernels over go bindings
72+
73+
General rules:
74+
75+
- Always take into account the layout and Basis for any polynomial used
76+
- Always warn the user about CPU pre-computations required for the given GPU operation
77+
- Always prefer GPU wrapped functions over any CPU work
78+
- Try to never copy any polynomial from GPU back to CPU unless absolutely required
79+
- Always write multithreaded code on CPU, parallelize any loop that can be parallelized
80+
- Prefer async code over synchronous. Wait for the channels where the results are required, not earlier

README.md

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,58 @@ go generate ./...
157157

158158
See [CHANGELOG.md](CHANGELOG.md).
159159

160+
## Solver Output Cache
161+
162+
When proving the same circuit repeatedly (e.g. during development), the
163+
constraint-system solver is the single largest bottleneck. This fork adds a
164+
**raw solver-values cache** that dumps the solver's wire-value array to disk
165+
after the first run and reloads it on subsequent runs, skipping the solver
166+
entirely.
167+
168+
### How it works
169+
170+
1. **First run** -- the solver runs normally and writes every wire value
171+
(Montgomery form, `[]fr.Element`) to a binary file. BSB22 commitment
172+
polynomials are saved alongside.
173+
2. **Subsequent runs** -- the cache file is memory-mapped back, the L/R/O
174+
Lagrange evaluations are derived via `evaluateLROSmallDomain`, and the
175+
BSB22 commitment is recomputed from the cached committed-wire values.
176+
The solver is never invoked.
177+
178+
### Usage
179+
180+
Set the `GNARK_RAW_SOLVER_CACHE` environment variable to a file path
181+
(ideally on a tmpfs / RAM-disk such as `/dev/shm`):
182+
183+
```bash
184+
export GNARK_RAW_SOLVER_CACHE=/dev/shm/raw_solver.bin
185+
```
186+
187+
The first proving run creates the file; every subsequent run loads it.
188+
Delete the file whenever the witness or circuit changes.
189+
190+
A slower but self-contained alternative caches the full `SparseR1CSSolution`
191+
(L, R, O + BSB22 data) via the `WithSolutionCachePath` prover option or the
192+
`GNARK_SOLUTION_CACHE` environment variable.
193+
194+
### Performance (sha256 circuit, RTX 4090, ICICLE PLONK BN254)
195+
196+
Measured with `ICICLE_STEP_PROFILE=1`. Blinding disabled (`USE_BLINDINGS`
197+
not set).
198+
199+
| Step | No Cache | With Cache | Saved |
200+
|------|----------|------------|-------|
201+
| **Solve constraints** | **4,243 ms** | **164 ms** | **4,079 ms** |
202+
| Commit L, R, O | 446 | 445 | -- |
203+
| Build ratio copy constraint | 462 | 445 | -- |
204+
| Commit Z | 155 | 166 | -- |
205+
| Compute quotient (total) | 5,794 | 4,919 | 875 ms |
206+
| Open Z | 896 | 807 | 89 ms |
207+
| Linearized polynomial | 1,230 | 1,287 | -- |
208+
| **Total prover** | **22,889 ms** | **18,271 ms** | **4,618 ms (20%)** |
209+
210+
Cache file sizes: `raw_solver.bin` ~163 MB, `bsb22_commit_0.bin` ~257 MB.
211+
160212
## Citing
161213

162214
If you use `gnark` in research, please cite the latest release:

backend/accelerated/icicle/groth16/bn254/icicle.go

Lines changed: 26 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
1+
# Cleanup Plan: ICICLE PLONK Prover (`icicle.go`, 6280 lines)
2+
3+
## Context
4+
5+
The ICICLE-accelerated PLONK prover (`backend/accelerated/icicle/plonk/bn254/icicle.go`) is functional but has accumulated significant technical debt: 48 `panic()` calls instead of error returns, 33 unchecked ICICLE return codes that silently discard errors, ~100 lines of dead code, ~350 lines of profiling boilerplate, duplicated validation blocks, and sprawling function signatures. This cleanup improves stability, readability, and maintainability without changing any proving logic.
6+
7+
---
8+
9+
## Resource Ownership Principle
10+
11+
When converting panics to error returns, follow this rule at every new early-return point:
12+
13+
> **Whoever allocates a `DeviceSlice` is responsible for freeing it on the error path.** Use `defer` for slices whose lifetime spans the function; use explicit `Free()` or `putTempDeviceSlice()` before each `return err`. When ownership transfers (e.g. a slice is returned to the caller), document the transfer with a comment.
14+
15+
This prevents GPU memory leaks that would otherwise be introduced by the new error-return paths.
16+
17+
---
18+
19+
## TIER 1 — HIGH PRIORITY (Stability / Correctness)
20+
21+
### 1A. Convert panics in constraint evaluation functions to error returns
22+
**Effort:** Medium (2-3h) | **Risk:** Low
23+
24+
30 panics across 6 functions that are already called from `RunOnDevice` closures with `done chan error`:
25+
26+
| Function | Line | Panic count |
27+
|---|---|---|
28+
| `computeGateConstraint` | 3373 | 4 |
29+
| `computeOrderingConstraint` | 3425 | 12 |
30+
| `computeLocalConstraint` | 3586 | 1 |
31+
| `combineConstraints` | 3638 | 3 |
32+
| `scaleSVectorsByBeta` | 3338 | 3 |
33+
| `multiplyMontgomerySlices` | 3181 | 2 |
34+
35+
**Change:** Return `(icicle_core.DeviceSlice, error)` instead of just `icicle_core.DeviceSlice`. The single caller (`gpuEvaluateConstraints` at line 3790) already uses `done <- err` patterns.
36+
37+
### 1B. Check unchecked ICICLE return codes in constraint evaluation
38+
**Effort:** Small-Medium (1-2h) | **Risk:** Low (same functions as 1A)
39+
40+
**This is worse than the panics** — these silently discard errors, producing invalid proofs:
41+
42+
| Call type | Unchecked count | Location |
43+
|---|---|---|
44+
| `icicle_vecops.VecOp(...)` (return discarded) | 31 | constraint eval functions (3305-3881) |
45+
| `_ = icicle_ntt.Ntt(...)` (explicit discard) | 2 | `gpuNTTInverseScaleForwardOnDevice` (2666, 2694) |
46+
| `icicle_runtime.SynchronizeStream(...)` (unchecked) | 1 | same function (2696) |
47+
48+
**Do together with 1A** since they are in the same functions. Every `VecOp` / `Ntt` / `SynchronizeStream` call must have its return code checked and propagated.
49+
50+
### 1C. Convert panics in `allocDeviceUninitialized` + `gpuMemoryPool.Get`
51+
**Effort:** Large (4-6h) | **Risk:** Medium
52+
53+
`allocDeviceUninitialized` (line 2943) is called from ~50 sites and panics on malloc failure. `gpuMemoryPool.Get` (line 2912) panics after retry.
54+
55+
**Strategy:**
56+
1. Create `mustAllocDeviceUninitialized(n int) icicle_core.DeviceSlice` (preserves panic behavior)
57+
2. Make `allocDeviceUninitialized` return `(icicle_core.DeviceSlice, error)`
58+
3. Delete `tryAllocDeviceUninitialized` (line 2862) — it exists solely to `recover()` from the panic, which becomes unnecessary
59+
4. Convert call sites incrementally, starting with those already inside error-returning code paths
60+
61+
### 1D. Convert panics in Montgomery/standard form helpers
62+
**Effort:** Small (1h) | **Risk:** Low
63+
64+
4 functions, each with 2 panic sites:
65+
- `toStandardFormInPlace` (3143)
66+
- `toStandardFormInPlaceWithCfg` (3148)
67+
- `toMontgomeryFormInPlaceWithCfg` (3159)
68+
- Plus upload helpers: `uploadScalarStd` (3047), `uploadScalarStdOnCurrentDevice` (3059), etc.
69+
70+
Do after 1A since these are called from the same code paths.
71+
72+
---
73+
74+
## TIER 2 — MEDIUM PRIORITY (Maintainability)
75+
76+
### 2A. Remove dead code: `saveSolveCache` + `loadSolveCache`
77+
**Effort:** Tiny (10 min) | **Risk:** None
78+
79+
Lines 510-613 (~103 lines). Neither function is called anywhere — `solveConstraints()` uses `cs.LoadRawSolverValues` / `cs.SaveRawSolverValues` instead (lines 424-478). Delete both functions.
80+
81+
### 2B. Extract profiling helper to reduce `isProfileMode` boilerplate
82+
**Effort:** Small (1-2h) | **Risk:** Very low
83+
84+
58 occurrences of the pattern (each ~6 lines, ~350 lines total):
85+
```go
86+
var startFoo time.Time
87+
if isProfileMode { startFoo = time.Now() }
88+
// ... work ...
89+
if isProfileMode {
90+
l := logger.Logger()
91+
l.Debug().Dur("took", time.Since(startFoo)).Msg("foo")
92+
}
93+
```
94+
95+
**Replace with:**
96+
```go
97+
func profileStep(msg string) func() {
98+
if !isProfileMode { return func() {} }
99+
start := time.Now()
100+
return func() {
101+
logger.Logger().Debug().Dur("took", time.Since(start)).Msg(msg)
102+
}
103+
}
104+
```
105+
106+
Usage: `done := profileStep("commitToLRO"); defer done()` — saves ~120 lines for the ~20 simple cases. Leave the ~9 complex cases (with extra fields like `.Int("idx", i)`) for now.
107+
108+
**Scope rule:** This is a purely mechanical extraction. Do not change message text, log levels, or field names. Each converted block must produce byte-identical log output.
109+
110+
### 2C. Extract `RunOnDevice` boilerplate into helper
111+
**Effort:** Medium (2-3h) | **Risk:** Low-medium (closure capture semantics)
112+
113+
~18 instances of the `done := make(chan error, 1)` + `RunOnDevice` + `createAsyncVecOpsConfig` + `finish` pattern (~15 lines each, ~270 lines total).
114+
115+
**Extract two helpers** — error-only and value+error:
116+
```go
117+
// runOnDeviceWithStream runs fn inside a RunOnDevice closure with automatic
118+
// stream creation, sync, and destroy. A panic inside fn is recovered and
119+
// returned as an error so the caller never deadlocks on the done channel.
120+
func (s *instance) runOnDeviceWithStream(label string, fn func(cfg icicle_core.VecOpsConfig) error) error {
121+
done := make(chan error, 1)
122+
icicle_runtime.RunOnDevice(&s.device, func(args ...any) {
123+
defer func() {
124+
if r := recover(); r != nil {
125+
done <- fmt.Errorf("%s: panic: %v", label, r)
126+
}
127+
}()
128+
cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label)
129+
if cfgErr != nil { done <- cfgErr; return }
130+
runErr := fn(cfg)
131+
if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil {
132+
runErr = syncErr
133+
}
134+
done <- runErr
135+
})
136+
return <-done
137+
}
138+
139+
// runOnDeviceWithStreamValue is like runOnDeviceWithStream but returns a value.
140+
func runOnDeviceWithStreamValue[T any](s *instance, label string, fn func(cfg icicle_core.VecOpsConfig) (T, error)) (T, error) {
141+
// same pattern with result channel
142+
}
143+
```
144+
145+
The `recover()` prevents deadlocks when a panic occurs inside the callback (currently possible since many helpers still panic). ~12-15 of the 18 patterns fit the error-only mold.
146+
147+
### 2D. Deduplicate GPU state prerequisite checks
148+
**Effort:** Small (30 min) | **Risk:** Very low
149+
150+
`computeLinearizedPolynomial` (1303-1323) and `batchOpening` (1530-1559) have ~30 lines of near-identical validation. Extract:
151+
```go
152+
func (s *instance) getValidatedLinearizedEvalState(caller string) (*gpuPolysState, error)
153+
```
154+
155+
### 2E. Collapse `commitToLRO` branches
156+
**Effort:** Small (30-45 min) | **Risk:** Low
157+
158+
Lines 651-761: 4 branches (sequential/concurrent x blinding/no-blinding). Extract a `commitFn` closure that selects blinding vs direct, then the sequential/concurrent decision only appears once. Reduces from ~110 lines to ~50 lines.
159+
160+
### 2F. Consolidate scalar upload helper variants
161+
**Effort:** Small (1h) | **Risk:** Low
162+
163+
8 upload helpers come in sync/async pairs differing by ~2 lines each (lines 3023-3139). Make sync variants delegate to the async variants with `DefaultVecOpsConfig()`.
164+
165+
---
166+
167+
## TIER 3 — LOW PRIORITY (Nice-to-have)
168+
169+
### 3A. Split `computeNumeratorLoopContext` into config + iteration state
170+
Lines 1721-1741: 17-field struct mixes immutable config with mutable iteration state.
171+
172+
### 3B. Group `instance` struct fields into logical sub-structs
173+
33+ fields mixing proof state, GPU handles, channels, polynomials, memory pools.
174+
175+
### 3C. Decouple `multiplyMontgomerySlices` from `gpuConstraintEvalState`
176+
Accept `getTempSlice`/`putTempSlice` functions directly instead of the full state struct.
177+
178+
### 3D. Simplify `ICICLE_LRO_COMMIT_SEQUENTIAL` double-configuration
179+
Lines 656-659: domain-size threshold AND env var override is confusing.
180+
181+
### 3E. Remove `getDeviceSlice` function pointer from `gpuConstraintEvalState`
182+
Only used for BSB gate lookups; pass BSB device slices as a separate parameter instead.
183+
184+
---
185+
186+
## Recommended Implementation Order
187+
188+
| Phase | Items | Time | Definition of Done |
189+
|-------|-------|------|-------------------|
190+
| 1 | 2A | 10 min | Dead code removed; build passes |
191+
| 2 | 1A + 1B | 3-4h | 0 panics and 0 unchecked ICICLE returns in constraint eval functions; `grep -c 'panic(' icicle.go` drops by ≥25 |
192+
| 3 | 1C | 4-6h | `allocDeviceUninitialized` returns error; `tryAllocDeviceUninitialized` deleted; `grep -c 'panic(' icicle.go` drops by ≥5 |
193+
| 4 | 1D | 1h | Montgomery/upload helpers return errors; `grep -c 'panic(' icicle.go` drops to ≤5 (only `warmUpDevice` and `init`-level panics remain) |
194+
| 5 | 2B, 2C, 2D, 2E, 2F | 4-6h | Structure/dedup cleanup; no functional change; `grep -c 'isProfileMode' icicle.go` drops by ≥30 |
195+
| 6 | 3A-3E | As time permits | Optional nice-to-haves |
196+
197+
## Verification
198+
199+
**Per-phase checks:**
200+
- Build: `go build -tags=icicle ./backend/accelerated/icicle/plonk/...`
201+
- Tests: `go test -tags=icicle ./backend/accelerated/icicle/plonk/...` (requires CUDA GPU)
202+
- Each phase should be a separate commit with its own build+test pass
203+
204+
**Performance guardrail (before Phase 2, after Phase 5):**
205+
- Record baseline prove wall-time: run sha256 e2e test twice (cold-start + steady-state)
206+
- After all phases complete, re-run the same test and compare
207+
- Acceptable delta: < 5% regression (these changes should not affect hot-path logic)
208+
- Command: `ICICLE_STEP_PROFILE=1 bash scripts/test_sha256_gnark.sh` from `/root/neon/openvm`
209+
210+
**Final completion check:**
211+
- Integration: full sha256 e2e test passes on GPU machine
212+
- `grep -c 'panic(' icicle.go` ≤ 5
213+
- `grep -cP 'icicle_vecops\.\w+\([^)]*\)$' icicle.go` = 0 (no unchecked ICICLE calls)

0 commit comments

Comments
 (0)