From 0917b4577ac45523f020d2033f82be154dd1f5c5 Mon Sep 17 00:00:00 2001 From: Martun Karapetyan Date: Mon, 27 Jul 2026 16:06:54 +0400 Subject: [PATCH 1/5] Add GPU-accelerated PLONK prover via ICICLE (bn254, bls12-377, bls12-381, bw6-761) Ports the Gnark_plonk_on_GPU work onto current gnark master: - New backend/accelerated/icicle/plonk/ package: full PLONK prover on GPU (NTT, MSM, quotient, linearization, KZG openings via ICICLE), selected with -tags=icicle. Like the accelerated Groth16 backend, the per-curve packages are generated from templates by backend/accelerated/icicle/internal/generator for all four supported curves. Stage-timing breakdown under ICICLE_STEP_PROFILE. - Solver output caching for the accelerated backends: raw wire-value cache for the ICICLE PLONK prover (GNARK_RAW_SOLVER_CACHE) and a WithSolutionCachePath prover option for ICICLE Groth16 (now part of the groth16 template, all curves). Cached data is validated against the circuit's wire count before use. Helpers are generated into constraint//solution_cache.go; no gnark-generator files are modified. - Blinding is on by default, matching the native prover. Setting GNARK_DISABLE_BLINDING trades zero-knowledge for a faster, deterministic GPU prover. - Tests: KZG-MSM and NTT GPU/CPU parity, proving-key marshal roundtrip, accelerated setup, and a 2^12-constraint end-to-end prove/verify with a BSB22 commitment on every supported curve (all behind -tags=icicle). - 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. --- README.md | 53 + .../icicle/groth16/bls12-377/icicle.go | 35 +- .../icicle/groth16/bls12-381/icicle.go | 35 +- .../icicle/groth16/bn254/icicle.go | 35 +- .../icicle/groth16/bw6-761/icicle.go | 35 +- .../icicle/internal/generator/main.go | 29 +- .../constraint.solution_cache.go.tmpl | 247 + .../templates/groth16.icicle.go.tmpl | 35 +- .../templates/plonk.icicle.doc.go.tmpl | 2 + .../generator/templates/plonk.icicle.go.tmpl | 7086 ++++++++++++++++ .../templates/plonk.icicle.provingkey.go.tmpl | 93 + .../accelerated/icicle/plonk/bls12-377/doc.go | 7 + .../icicle/plonk/bls12-377/icicle.go | 7093 +++++++++++++++++ .../icicle/plonk/bls12-377/provingkey.go | 100 + .../accelerated/icicle/plonk/bls12-381/doc.go | 7 + .../icicle/plonk/bls12-381/icicle.go | 7093 +++++++++++++++++ .../icicle/plonk/bls12-381/provingkey.go | 100 + backend/accelerated/icicle/plonk/bn254/doc.go | 7 + .../accelerated/icicle/plonk/bn254/icicle.go | 7093 +++++++++++++++++ .../icicle/plonk/bn254/icicle_msm_kzg_test.go | 374 + .../icicle/plonk/bn254/icicle_ntt_test.go | 125 + .../icicle/plonk/bn254/provingkey.go | 100 + .../accelerated/icicle/plonk/bw6-761/doc.go | 7 + .../icicle/plonk/bw6-761/icicle.go | 7093 +++++++++++++++++ .../icicle/plonk/bw6-761/provingkey.go | 100 + backend/accelerated/icicle/plonk/e2e_test.go | 88 + .../accelerated/icicle/plonk/fallback_test.go | 51 + .../accelerated/icicle/plonk/marshal_test.go | 105 + backend/accelerated/icicle/plonk/plonk_all.go | 30 + .../accelerated/icicle/plonk/plonk_icicle.go | 209 + .../icicle/plonk/plonk_noicicle.go | 30 + backend/backend.go | 27 +- backend/plonk/bench_impl_icicle_test.go | 19 + backend/plonk/bench_impl_native_test.go | 18 + backend/plonk/plonk_test.go | 4 +- constraint/bls12-377/solution_cache.go | 254 + constraint/bls12-381/solution_cache.go | 254 + constraint/bn254/solution_cache.go | 254 + constraint/bw6-761/solution_cache.go | 254 + go.mod | 7 + 40 files changed, 38558 insertions(+), 30 deletions(-) create mode 100644 backend/accelerated/icicle/internal/generator/templates/constraint.solution_cache.go.tmpl create mode 100644 backend/accelerated/icicle/internal/generator/templates/plonk.icicle.doc.go.tmpl create mode 100644 backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl create mode 100644 backend/accelerated/icicle/internal/generator/templates/plonk.icicle.provingkey.go.tmpl create mode 100644 backend/accelerated/icicle/plonk/bls12-377/doc.go create mode 100644 backend/accelerated/icicle/plonk/bls12-377/icicle.go create mode 100644 backend/accelerated/icicle/plonk/bls12-377/provingkey.go create mode 100644 backend/accelerated/icicle/plonk/bls12-381/doc.go create mode 100644 backend/accelerated/icicle/plonk/bls12-381/icicle.go create mode 100644 backend/accelerated/icicle/plonk/bls12-381/provingkey.go create mode 100644 backend/accelerated/icicle/plonk/bn254/doc.go create mode 100644 backend/accelerated/icicle/plonk/bn254/icicle.go create mode 100644 backend/accelerated/icicle/plonk/bn254/icicle_msm_kzg_test.go create mode 100644 backend/accelerated/icicle/plonk/bn254/icicle_ntt_test.go create mode 100644 backend/accelerated/icicle/plonk/bn254/provingkey.go create mode 100644 backend/accelerated/icicle/plonk/bw6-761/doc.go create mode 100644 backend/accelerated/icicle/plonk/bw6-761/icicle.go create mode 100644 backend/accelerated/icicle/plonk/bw6-761/provingkey.go create mode 100644 backend/accelerated/icicle/plonk/e2e_test.go create mode 100644 backend/accelerated/icicle/plonk/fallback_test.go create mode 100644 backend/accelerated/icicle/plonk/marshal_test.go create mode 100644 backend/accelerated/icicle/plonk/plonk_all.go create mode 100644 backend/accelerated/icicle/plonk/plonk_icicle.go create mode 100644 backend/accelerated/icicle/plonk/plonk_noicicle.go create mode 100644 backend/plonk/bench_impl_icicle_test.go create mode 100644 backend/plonk/bench_impl_native_test.go create mode 100644 constraint/bls12-377/solution_cache.go create mode 100644 constraint/bls12-381/solution_cache.go create mode 100644 constraint/bn254/solution_cache.go create mode 100644 constraint/bw6-761/solution_cache.go diff --git a/README.md b/README.md index 50af91482e..ea89b3364c 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,59 @@ go generate ./... See [CHANGELOG.md](CHANGELOG.md). +## Solver Output Cache + +When proving the same circuit repeatedly (e.g. during development), the +constraint-system solver is the single largest bottleneck. The ICICLE +backends support a **raw solver-values cache** that dumps the solver's +wire-value array to disk +after the first run and reloads it on subsequent runs, skipping the solver +entirely. + +### How it works + +1. **First run** -- the solver runs normally and writes every wire value + (Montgomery form, `[]fr.Element`) to a binary file. BSB22 commitment + polynomials are saved alongside. +2. **Subsequent runs** -- the cache file is memory-mapped back, the L/R/O + Lagrange evaluations are derived via `evaluateLROSmallDomain`, and the + BSB22 commitment is recomputed from the cached committed-wire values. + The solver is never invoked. + +### Usage + +Set the `GNARK_RAW_SOLVER_CACHE` environment variable to a file path +(ideally on a tmpfs / RAM-disk such as `/dev/shm`): + +```bash +export GNARK_RAW_SOLVER_CACHE=/dev/shm/raw_solver.bin +``` + +The first proving run creates the file; every subsequent run loads it. +Delete the file whenever the witness or circuit changes. + +For the ICICLE Groth16 backend, the `WithSolutionCachePath` prover option +caches the full `R1CSSolution` instead. + +### Performance (sha256 circuit, RTX 4090, ICICLE PLONK BN254) + +Measured with `ICICLE_STEP_PROFILE=1` and blinding disabled +(`GNARK_DISABLE_BLINDING` set; the default is blinding on, matching the +native prover). + +| Step | No Cache | With Cache | Saved | +|------|----------|------------|-------| +| **Solve constraints** | **4,243 ms** | **164 ms** | **4,079 ms** | +| Commit L, R, O | 446 | 445 | -- | +| Build ratio copy constraint | 462 | 445 | -- | +| Commit Z | 155 | 166 | -- | +| Compute quotient (total) | 5,794 | 4,919 | 875 ms | +| Open Z | 896 | 807 | 89 ms | +| Linearized polynomial | 1,230 | 1,287 | -- | +| **Total prover** | **22,889 ms** | **18,271 ms** | **4,618 ms (20%)** | + +Cache file sizes: `raw_solver.bin` ~163 MB, `bsb22_commit_0.bin` ~257 MB. + ## Citing If you use `gnark` in research, please cite the latest release: diff --git a/backend/accelerated/icicle/groth16/bls12-377/icicle.go b/backend/accelerated/icicle/groth16/bls12-377/icicle.go index 39c8ecd1b9..cb6946b098 100644 --- a/backend/accelerated/icicle/groth16/bls12-377/icicle.go +++ b/backend/accelerated/icicle/groth16/bls12-377/icicle.go @@ -888,12 +888,39 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic return nil })) - _solution, err := r1cs.Solve(fullWitness, solverOpts...) - if err != nil { - return nil, err + cachePath := opt.SolutionCachePath + canCache := cachePath != "" && len(commitmentInfo) == 0 + + var solution *cs.R1CSSolution + + if canCache { + if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() + if len(cached.W) != expectedWires { + log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). + Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + } else { + solution = cached + log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver") + } + } } - solution := _solution.(*cs.R1CSSolution) + if solution == nil { + _solution, err := r1cs.Solve(fullWitness, solverOpts...) + if err != nil { + return nil, err + } + solution = _solution.(*cs.R1CSSolution) + + if canCache { + if err := cs.SaveR1CSSolution(cachePath, solution); err != nil { + log.Warn().Err(err).Msg("failed to cache Groth16 solution") + } else { + log.Debug().Str("file", cachePath).Msg("cached Groth16 solution to disk") + } + } + } wireValues := []fr.Element(solution.W) start := time.Now() diff --git a/backend/accelerated/icicle/groth16/bls12-381/icicle.go b/backend/accelerated/icicle/groth16/bls12-381/icicle.go index 5dda9ebf14..f385fd1718 100644 --- a/backend/accelerated/icicle/groth16/bls12-381/icicle.go +++ b/backend/accelerated/icicle/groth16/bls12-381/icicle.go @@ -888,12 +888,39 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic return nil })) - _solution, err := r1cs.Solve(fullWitness, solverOpts...) - if err != nil { - return nil, err + cachePath := opt.SolutionCachePath + canCache := cachePath != "" && len(commitmentInfo) == 0 + + var solution *cs.R1CSSolution + + if canCache { + if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() + if len(cached.W) != expectedWires { + log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). + Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + } else { + solution = cached + log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver") + } + } } - solution := _solution.(*cs.R1CSSolution) + if solution == nil { + _solution, err := r1cs.Solve(fullWitness, solverOpts...) + if err != nil { + return nil, err + } + solution = _solution.(*cs.R1CSSolution) + + if canCache { + if err := cs.SaveR1CSSolution(cachePath, solution); err != nil { + log.Warn().Err(err).Msg("failed to cache Groth16 solution") + } else { + log.Debug().Str("file", cachePath).Msg("cached Groth16 solution to disk") + } + } + } wireValues := []fr.Element(solution.W) start := time.Now() diff --git a/backend/accelerated/icicle/groth16/bn254/icicle.go b/backend/accelerated/icicle/groth16/bn254/icicle.go index f0c33bc76c..9b36fb8bcb 100644 --- a/backend/accelerated/icicle/groth16/bn254/icicle.go +++ b/backend/accelerated/icicle/groth16/bn254/icicle.go @@ -888,12 +888,39 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic return nil })) - _solution, err := r1cs.Solve(fullWitness, solverOpts...) - if err != nil { - return nil, err + cachePath := opt.SolutionCachePath + canCache := cachePath != "" && len(commitmentInfo) == 0 + + var solution *cs.R1CSSolution + + if canCache { + if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() + if len(cached.W) != expectedWires { + log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). + Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + } else { + solution = cached + log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver") + } + } } - solution := _solution.(*cs.R1CSSolution) + if solution == nil { + _solution, err := r1cs.Solve(fullWitness, solverOpts...) + if err != nil { + return nil, err + } + solution = _solution.(*cs.R1CSSolution) + + if canCache { + if err := cs.SaveR1CSSolution(cachePath, solution); err != nil { + log.Warn().Err(err).Msg("failed to cache Groth16 solution") + } else { + log.Debug().Str("file", cachePath).Msg("cached Groth16 solution to disk") + } + } + } wireValues := []fr.Element(solution.W) start := time.Now() diff --git a/backend/accelerated/icicle/groth16/bw6-761/icicle.go b/backend/accelerated/icicle/groth16/bw6-761/icicle.go index e557fee455..05b2bb3eeb 100644 --- a/backend/accelerated/icicle/groth16/bw6-761/icicle.go +++ b/backend/accelerated/icicle/groth16/bw6-761/icicle.go @@ -877,12 +877,39 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic return nil })) - _solution, err := r1cs.Solve(fullWitness, solverOpts...) - if err != nil { - return nil, err + cachePath := opt.SolutionCachePath + canCache := cachePath != "" && len(commitmentInfo) == 0 + + var solution *cs.R1CSSolution + + if canCache { + if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() + if len(cached.W) != expectedWires { + log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). + Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + } else { + solution = cached + log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver") + } + } } - solution := _solution.(*cs.R1CSSolution) + if solution == nil { + _solution, err := r1cs.Solve(fullWitness, solverOpts...) + if err != nil { + return nil, err + } + solution = _solution.(*cs.R1CSSolution) + + if canCache { + if err := cs.SaveR1CSSolution(cachePath, solution); err != nil { + log.Warn().Err(err).Msg("failed to cache Groth16 solution") + } else { + log.Debug().Str("file", cachePath).Msg("cached Groth16 solution to disk") + } + } + } wireValues := []fr.Element(solution.W) start := time.Now() diff --git a/backend/accelerated/icicle/internal/generator/main.go b/backend/accelerated/icicle/internal/generator/main.go index 786855436c..f0f828643f 100644 --- a/backend/accelerated/icicle/internal/generator/main.go +++ b/backend/accelerated/icicle/internal/generator/main.go @@ -63,9 +63,32 @@ func main() { if err := bgen.Generate(d, d.CurvePkg, "./templates/", entries...); err != nil { panic(err) } + + plonkRoot := strings.Replace(d.RootPath, "groth16", "plonk", 1) + plonkEntries := []bavard.Entry{ + {File: filepath.Join(plonkRoot, "doc.go"), Templates: []string{"plonk.icicle.doc.go.tmpl"}}, + {File: filepath.Join(plonkRoot, "icicle.go"), Templates: []string{"plonk.icicle.go.tmpl"}}, + {File: filepath.Join(plonkRoot, "provingkey.go"), Templates: []string{"plonk.icicle.provingkey.go.tmpl"}}, + } + if err := bgen.Generate(d, d.CurvePkg, "./templates/", plonkEntries...); err != nil { + panic(err) + } + + // solver-output cache helpers used by the accelerated backends; these + // live in the per-curve constraint package to access unexported types. + cacheEntry := bavard.Entry{ + File: filepath.Join("../../../../../constraint", strings.ToLower(d.Curve), "solution_cache.go"), + Templates: []string{"constraint.solution_cache.go.tmpl"}, + } + if err := bgen.Generate(d, "cs", "./templates/", cacheEntry); err != nil { + panic(err) + } } - runCmd("gofmt", "-w", "../../groth16") + runCmd("gofmt", "-w", "../../groth16", "../../plonk") + for _, d := range data { + runCmd("gofmt", "-w", filepath.Join("../../../../../constraint", strings.ToLower(d.Curve), "solution_cache.go")) + } runGoImports() } @@ -80,8 +103,8 @@ func runCmd(name string, arg ...string) { } func runGoImports() { - fmt.Println("go tool goimports", "-w", "../../groth16") - cmd := exec.Command("go", "tool", "goimports", "-w", "../../groth16") + fmt.Println("go tool goimports", "-w", "../../groth16", "../../plonk") + cmd := exec.Command("go", "tool", "goimports", "-w", "../../groth16", "../../plonk") cmd.Stdout = os.Stdout cmd.Stderr = os.Stderr if err := cmd.Run(); err != nil { diff --git a/backend/accelerated/icicle/internal/generator/templates/constraint.solution_cache.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/constraint.solution_cache.go.tmpl new file mode 100644 index 0000000000..e9388dae74 --- /dev/null +++ b/backend/accelerated/icicle/internal/generator/templates/constraint.solution_cache.go.tmpl @@ -0,0 +1,247 @@ +// Package-level helpers for caching solver output to disk, used by the +// ICICLE-accelerated backends (backend/accelerated/icicle/...) to skip +// redundant solver runs when proving the same (circuit, witness) repeatedly. + +import ( + "encoding/binary" + "fmt" + "io" + "os" + "path/filepath" + "time" + + "github.com/consensys/gnark-crypto/ecc" + "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}/fr" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" + csolver "github.com/consensys/gnark/constraint/solver" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + "unsafe" +) + +// evaluateLROSmallDomain extracts the solver l, r, o, and returns it in lagrange form. +// solver = [ public | secret | internal ] +func evaluateLROSmallDomain(cs *system, solution []fr.Element) ([]fr.Element, []fr.Element, []fr.Element) { + s := cs.GetNbConstraints() + len(cs.Public) // len(spr.Public) is for the placeholder constraints + s = int(ecc.NextPowerOfTwo(uint64(s))) + + var l, r, o []fr.Element + l = make([]fr.Element, s, s+4) // +4 to leave room for the blinding in plonk + r = make([]fr.Element, s, s+4) + o = make([]fr.Element, s, s+4) + s0 := solution[0] + + utils.Parallelize(len(cs.Public), func(start, end int) { + for i := start; i < end; i++ { // placeholders + l[i] = solution[i] + r[i] = s0 + o[i] = s0 + } + }) + offset := len(cs.Public) + nbConstraints := cs.GetNbConstraints() + + var sparseR1C constraint.SparseR1C + j := 0 + for _, inst := range cs.Instructions { + blueprint := cs.Blueprints[inst.BlueprintID] + if bc, ok := blueprint.(constraint.BlueprintSparseR1C); ok { + bc.DecompressSparseR1C(&sparseR1C, inst.Unpack(&cs.System)) + + l[offset+j] = solution[sparseR1C.XA] + r[offset+j] = solution[sparseR1C.XB] + o[offset+j] = solution[sparseR1C.XC] + j++ + } + } + + offset += nbConstraints + + padding := s - offset + utils.Parallelize(padding, func(start, end int) { + for i := start; i < end; i++ { // offset to reach 2**n constraints (where the id of l,r,o is 0, so we assign solver[0]) + j := offset + i + l[j] = s0 + r[j] = s0 + o[j] = s0 + } + }) + + return l, r, o +} + +// EvaluateLROSmallDomainFromValues derives the L, R, O Lagrange evaluations from a +// pre-computed full wire-value vector (as produced by SolveAndSaveRawValues / +// LoadRawSolverValues). It validates that the vector length matches the system. +func (cs *system) EvaluateLROSmallDomainFromValues(values []fr.Element) ([]fr.Element, []fr.Element, []fr.Element, error) { + expected := cs.GetNbPublicVariables() + cs.GetNbSecretVariables() + cs.GetNbInternalVariables() + if len(values) != expected { + return nil, nil, nil, fmt.Errorf("wire-value vector length mismatch: got %d, expected %d (stale cache?)", len(values), expected) + } + l, r, o := evaluateLROSmallDomain(cs, values) + return l, r, o, nil +} + +// SolveAndSaveRawValues behaves like Solve but, when rawCachePath is non-empty, +// additionally writes the solver's full wire-value vector to that path so that +// subsequent runs can skip the solver via LoadRawSolverValues + +// EvaluateLROSmallDomainFromValues. The caller is responsible for invalidating +// the cache when the circuit or witness changes. +func (cs *system) SolveAndSaveRawValues(witness witness.Witness, rawCachePath string, opts ...csolver.Option) (any, error) { + log := logger.Logger().With().Int("nbConstraints", cs.GetNbConstraints()).Logger() + start := time.Now() + + v := witness.Vector().(fr.Vector) + + // init the solver + solver, err := newSolver(cs, v, opts...) + if err != nil { + log.Err(err).Send() + return nil, err + } + + // reset the stateful blueprints + for i := range cs.Blueprints { + if b, ok := cs.Blueprints[i].(constraint.BlueprintStateful[constraint.U64]); ok { + b.Reset() + } + } + + // defer log printing once all solver.values are computed + // (or sooner, if a constraint is not satisfied) + defer solver.printLogs(cs.Logs) + + // run it. + if err := solver.run(); err != nil { + log.Err(err).Send() + return nil, err + } + + log.Debug().Dur("took", time.Since(start)).Msg("constraint system solver done") + + if rawCachePath != "" { + if err := SaveRawSolverValues(rawCachePath, solver.values); err != nil { + log.Warn().Err(err).Msg("failed to save raw solver values cache") + } else { + log.Debug().Str("file", rawCachePath).Int("wires", len(solver.values)).Msg("saved raw solver values cache") + } + } + + // format the solution + if cs.Type == constraint.SystemR1CS { + var res R1CSSolution + res.W = solver.values + res.A = solver.a + res.B = solver.b + res.C = solver.c + return &res, nil + } else { + // sparse R1CS: solver fills L,R,O during solving. + var res SparseR1CSSolution + res.L = solver.l + res.R = solver.r + res.O = solver.o + return &res, nil + } +} + +// LoadR1CSSolution reads a cached R1CSSolution (Groth16) from disk. +func LoadR1CSSolution(path string) (*R1CSSolution, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + var sol R1CSSolution + if _, err := sol.ReadFrom(f); err != nil { + return nil, fmt.Errorf("read cached Groth16 solution: %w", err) + } + return &sol, nil +} + +// SaveR1CSSolution writes a R1CSSolution (Groth16) to disk atomically. +func SaveR1CSSolution(path string, solution *R1CSSolution) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".groth16_solution_cache_*") + if err != nil { + return err + } + tmpName := tmp.Name() + if _, err := solution.WriteTo(tmp); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, +// no per-element conversion) with a little-endian uint64 length prefix. +// The format is not portable across architectures with different endianness. +func SaveRawSolverValues(path string, values []fr.Element) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".raw_solver_*") + if err != nil { + return err + } + tmpName := tmp.Name() + + nWires := uint64(len(values)) + if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if nWires > 0 { + // Bulk write: cast []fr.Element to []byte + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := tmp.Write(byteSlice); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. +// The declared length is validated against the file size before allocating. +func LoadRawSolverValues(path string) ([]fr.Element, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + fi, err := f.Stat() + if err != nil { + return nil, err + } + + var nWires uint64 + if err := binary.Read(f, binary.LittleEndian, &nWires); err != nil { + return nil, fmt.Errorf("read wire count: %w", err) + } + expectedSize := int64(8) + int64(nWires)*int64(fr.Bytes) + if fi.Size() != expectedSize { + return nil, fmt.Errorf("raw solver cache size mismatch: file is %d bytes, header declares %d wires (%d bytes)", + fi.Size(), nWires, expectedSize) + } + values := make([]fr.Element, nWires) + if nWires > 0 { + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := io.ReadFull(f, byteSlice); err != nil { + return nil, fmt.Errorf("read wires: %w", err) + } + } + return values, nil +} diff --git a/backend/accelerated/icicle/internal/generator/templates/groth16.icicle.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/groth16.icicle.go.tmpl index 559ea3b626..b4d2a80a07 100644 --- a/backend/accelerated/icicle/internal/generator/templates/groth16.icicle.go.tmpl +++ b/backend/accelerated/icicle/internal/generator/templates/groth16.icicle.go.tmpl @@ -902,12 +902,39 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic return nil })) - _solution, err := r1cs.Solve(fullWitness, solverOpts...) - if err != nil { - return nil, err + cachePath := opt.SolutionCachePath + canCache := cachePath != "" && len(commitmentInfo) == 0 + + var solution *cs.R1CSSolution + + if canCache { + if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() + if len(cached.W) != expectedWires { + log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). + Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + } else { + solution = cached + log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver") + } + } } - solution := _solution.(*cs.R1CSSolution) + if solution == nil { + _solution, err := r1cs.Solve(fullWitness, solverOpts...) + if err != nil { + return nil, err + } + solution = _solution.(*cs.R1CSSolution) + + if canCache { + if err := cs.SaveR1CSSolution(cachePath, solution); err != nil { + log.Warn().Err(err).Msg("failed to cache Groth16 solution") + } else { + log.Debug().Str("file", cachePath).Msg("cached Groth16 solution to disk") + } + } + } wireValues := []fr.Element(solution.W) start := time.Now() diff --git a/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.doc.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.doc.go.tmpl new file mode 100644 index 0000000000..f2fef22247 --- /dev/null +++ b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.doc.go.tmpl @@ -0,0 +1,2 @@ +// Package {{ .CurvePkg }} implements ICICLE acceleration for {{ .Curve }} PLONK backend. +package {{ .CurvePkg }} \ No newline at end of file diff --git a/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl new file mode 100644 index 0000000000..04a8b4b7b3 --- /dev/null +++ b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl @@ -0,0 +1,7086 @@ +//go:build icicle + +import ( + "context" + "errors" + "fmt" + "hash" + "io" + "math/big" + "math/bits" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + "unsafe" + + "golang.org/x/sync/errgroup" + + "github.com/consensys/gnark/backend" + plonk_{{ .CurvePkg }} "github.com/consensys/gnark/backend/plonk/{{ toLower .Curve }}" + "github.com/consensys/gnark/backend/witness" + constraint "github.com/consensys/gnark/constraint" + cs "github.com/consensys/gnark/constraint/{{ toLower .Curve }}" + "github.com/consensys/gnark/constraint/solver" + fcs "github.com/consensys/gnark/frontend/cs" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + + "github.com/consensys/gnark-crypto/ecc" + curve "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}" + "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}/fp" + "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}/fr" + "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}/fr/fft" + "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}/fr/hash_to_field" + "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}/fr/iop" + "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}/kzg" + fiatshamir "github.com/consensys/gnark-crypto/fiat-shamir" + + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" + icicle_{{ .CurvePkg }} "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/{{ .CurvePkg }}" + icicle_msm "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/{{ .CurvePkg }}/msm" + icicle_ntt "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/{{ .CurvePkg }}/ntt" + icicle_vecops "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/{{ .CurvePkg }}/vecOps" + icicle_runtime "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime" + "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime/config_extension" +) + +const HasIcicle = true + +var isProfileMode bool + +var useBlinding bool + +var isNttTrace bool + +func init() { + _, isProfileMode = os.LookupEnv("ICICLE_STEP_PROFILE") + // Blinding polynomials (zero-knowledge) are enabled by default, matching the + // native prover. Set GNARK_DISABLE_BLINDING to trade zero-knowledge for a + // faster, deterministic prover (e.g. when the witness is not secret). + _, disableBlinding := os.LookupEnv("GNARK_DISABLE_BLINDING") + useBlinding = !disableBlinding + isNttTrace = envEnabled("ICICLE_NTT_TRACE", false) +} + +// profileStep returns a function that, when called, logs the elapsed time since +// profileStep was invoked. If profiling is disabled, it returns a no-op. +// Usage: done := profileStep("label"); defer done() +func profileStep(msg string) func() { + if !isProfileMode { + return func() {} + } + start := time.Now() + return func() { + l := logger.Logger() + l.Debug().Dur("took", time.Since(start)).Msg(msg) + } +} + +// stageTiming is a single recorded prover stage and its wall-clock duration. +type stageTiming struct { + name string + dur time.Duration +} + +// stageTimings is a concurrency-safe, ordered recorder of prover stage +// durations. The PLONK prover runs its stages as concurrent goroutines, so the +// recorded durations OVERLAP and do not sum to the total — the printed table +// flags this. +type stageTimings struct { + mu sync.Mutex + entries []stageTiming +} + +// record appends a (stage, duration) entry. Safe to call from any goroutine and +// safe on a nil receiver (records nothing). +func (t *stageTimings) record(name string, d time.Duration) { + if t == nil { + return + } + t.mu.Lock() + t.entries = append(t.entries, stageTiming{name: name, dur: d}) + t.mu.Unlock() +} + +// printTable writes an aligned breakdown of the recorded stages to w, sorted by +// duration (largest first), followed by the overall prover total. Stages run +// concurrently, so the rows overlap and intentionally do not sum to the total. +func (t *stageTimings) printTable(w io.Writer, total time.Duration) { + if t == nil { + return + } + t.mu.Lock() + rows := make([]stageTiming, len(t.entries)) + copy(rows, t.entries) + t.mu.Unlock() + + sort.SliceStable(rows, func(i, j int) bool { return rows[i].dur > rows[j].dur }) + + nameW := len("TOTAL (prover done)") + for _, r := range rows { + if len(r.name) > nameW { + nameW = len(r.name) + } + } + + fmt.Fprintln(w, "") + fmt.Fprintln(w, "================ gnark PLONK prove breakdown (GPU) ================") + fmt.Fprintln(w, "(stages run concurrently — durations overlap and do not sum to TOTAL)") + fmt.Fprintf(w, " %-*s %12s %6s\n", nameW, "STAGE", "TIME", "%TOTAL") + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + for _, r := range rows { + pct := 0.0 + if total > 0 { + pct = 100 * float64(r.dur) / float64(total) + } + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, r.name, r.dur.Round(time.Millisecond), pct) + } + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, "TOTAL (prover done)", total.Round(time.Millisecond), 100.0) + fmt.Fprintln(w, "===================================================================") + fmt.Fprintln(w, "") +} + +func envEnabled(key string, defaultVal bool) bool { + v, ok := os.LookupEnv(key) + if !ok { + return defaultVal + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "0", "false", "off", "no": + return false + default: + return true + } +} + +func nttAlgorithmFromEnv(key string, fallback icicle_core.NttAlgorithm) icicle_core.NttAlgorithm { + v, ok := os.LookupEnv(key) + if !ok { + return fallback + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "auto", "0": + return icicle_core.Auto + case "radix2", "radix-2", "r2", "1": + return icicle_core.Radix2 + case "mixed", "mixedradix", "mixed-radix", "2": + return icicle_core.MixedRadix + default: + return fallback + } +} + +const ( + id_L int = iota + id_R + id_O + id_Z + id_ZS + id_Ql + id_Qr + id_Qm + id_Qo + id_Qk + id_S1 + id_S2 + id_S3 + id_Qci // [ .. , Qc_i, Pi_i, ...] +) + +// blinding factors +const ( + id_Bl int = iota + id_Br + id_Bo + id_Bz + nb_blinding_polynomials +) + +// blinding orders (-1 to deactivate) +const ( + order_blinding_L = 1 + order_blinding_R = 1 + order_blinding_O = 1 + order_blinding_Z = 2 +) + +// Prove generates a PLONK proof. When the accelerator option is not set to +// "icicle", we delegate to the native prover. Otherwise, we run a local copy +// of the CPU prover logic to enable incremental GPU adaptation. +func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts ...backend.ProverOption) (*plonk_{{ .CurvePkg }}.Proof, error) { + opt, err := backend.NewProverConfig(opts...) + if err != nil { + return nil, err + } + + log := logger.Logger().With(). + Str("curve", spr.CurveID().String()). + Int("nbConstraints", spr.GetNbConstraints()). + Str("backend", "plonk").Logger() + + // parse the options + if opt.HashToFieldFn == nil { + opt.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + + // When blinding is disabled (GNARK_DISABLE_BLINDING), also disable StatisticalZK, it makes no sense + // to use statistical zero knowledge when we don't use blinding. + if !useBlinding { + opt.StatisticalZK = false + } + + start := time.Now() + + // Initialize device and preload KZG bases once per proving key + device := icicle_runtime.CreateDevice("CUDA", 0) + if pk.deviceInfo == nil { + if err := pk.setupDevicePointers(&device); err != nil { + return nil, err + } + } + + // init instance + g, ctx := errgroup.WithContext(context.Background()) + instance, err := newInstance(ctx, spr, pk, fullWitness, &opt) + if err != nil { + return nil, fmt.Errorf("new instance: %w", err) + } + // attach device to instance for GPU calls + instance.device = device + instance.initSharedGPUState() + defer instance.releaseTempGPUMemoryPool() + defer instance.releaseSharedGPUState() + defer instance.releaseLinearizedEvalGPUState() + + // solve constraints + g.Go(instance.solveConstraints) + + // complete qk + g.Go(instance.completeQk) + + // init blinding polynomials + g.Go(instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + g.Go(instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + g.Go(instance.buildRatioCopyConstraint) + + // compute h + g.Go(instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + g.Go(instance.openZ) + + // linearized polynomial + g.Go(instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + g.Go(func() error { + startBatchOpening := time.Now() + err := instance.batchOpening() + if isProfileMode { + instance.timings.record("batchOpening (folded KZG)", time.Since(startBatchOpening)) + } + return err + }) + + if err := g.Wait(); err != nil { + return nil, err + } + + total := time.Since(start) + log.Debug().Dur("took", total).Msg("prover done") + if isProfileMode { + instance.timings.printTable(os.Stderr, total) + } + return instance.proof, nil +} + +// represents a Prover instance +type instance struct { + ctx context.Context + + pk *ProvingKey + proof *plonk_{{ .CurvePkg }}.Proof + spr *cs.SparseR1CS + opt *backend.ProverConfig + + fs *fiatshamir.Transcript + kzgFoldingHash hash.Hash // for KZG folding + htfFunc hash.Hash // hash to field function + + // polynomials + polyL, polyR, polyO *iop.Polynomial + polyZ, polyZS, polyQk *iop.Polynomial + bp []*iop.Polynomial // blinding polynomials + h *iop.Polynomial // h is the quotient polynomial + hGPU *gpuQuotientPolynomial + polyZLagrangeGPU icicle_core.DeviceSlice + blindedZCanonicalGPU icicle_core.DeviceSlice + quotientShardsRandomizers [2]fr.Element // random elements for blinding the shards of the quotient + + linearizedPolynomial []fr.Element + linearizedPolynomialGPU icicle_core.DeviceSlice + linearizedPolynomialClaim fr.Element + linearizedPolynomialDigest kzg.Digest + + fullWitness witness.Witness + + // bsb22 commitment stuff + commitmentInfo constraint.PlonkCommitments + commitmentVal []fr.Element + cCommitments []*iop.Polynomial + + // challenges + gamma, beta, alpha, zeta fr.Element + + // channel to wait for the steps + chLRO, + chQk, + chbp, + chZ, + chH, + chRestoreLRO, + chZOpening, + chLinearizedPolynomial, + chGammaBeta chan struct{} + + domain0, domain1 *fft.Domain + + trace *plonk_{{ .CurvePkg }}.Trace + + // GPU device handle + device icicle_runtime.Device + + // Shared GPU polynomial context reused across buildRatioCopyConstraint + // and computeQuotient to avoid repeated host<->device uploads. + gpuStateMu sync.Mutex + sharedGPUState *gpuPolysState + // Snapshot of immutable polynomial slices used by computeLinearizedPolynomial + // for zeta evaluations after computeQuotient mutates/frees shared state. + linearizedEvalGPUState *gpuPolysState + + // Reusable temporary GPU memory pool for non-state buffers. + tempGPUMemPool *gpuMemoryPool + + // Per-prove stage-timing recorder (used to print the breakdown table when + // ICICLE_STEP_PROFILE is set). + timings *stageTimings +} + +func newInstance(ctx context.Context, spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts *backend.ProverConfig) (*instance, error) { + if opts.HashToFieldFn == nil { + opts.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + s := instance{ + ctx: ctx, + pk: pk, + proof: &plonk_{{ .CurvePkg }}.Proof{}, + spr: spr, + opt: opts, + fullWitness: fullWitness, + bp: make([]*iop.Polynomial, nb_blinding_polynomials), + fs: fiatshamir.NewTranscript(opts.ChallengeHash, "gamma", "beta", "alpha", "zeta"), + kzgFoldingHash: opts.KZGFoldingHash, + htfFunc: opts.HashToFieldFn, + chLRO: make(chan struct{}, 1), + chQk: make(chan struct{}, 1), + chbp: make(chan struct{}, 1), + chGammaBeta: make(chan struct{}, 1), + chZ: make(chan struct{}, 1), + chH: make(chan struct{}, 1), + chZOpening: make(chan struct{}, 1), + chLinearizedPolynomial: make(chan struct{}, 1), + chRestoreLRO: make(chan struct{}, 1), + tempGPUMemPool: newGPUMemoryPool(), + timings: &stageTimings{}, + } + s.initBSB22Commitments() + + // FFT domains and the PLONK trace are witness-independent and expensive to + // build at large n (NewTrace walks every constraint), so they are cached + // on the proving key and shared read-only across proofs. + setup := pk.hostSetupFor(spr) + s.domain0 = setup.domain0 + s.domain1 = setup.domain1 + s.trace = setup.trace + + // sampling random numbers for blinding the quotient + if opts.StatisticalZK { + s.quotientShardsRandomizers[0].SetRandom() + s.quotientShardsRandomizers[1].SetRandom() + } + + return &s, nil +} + +func (s *instance) initBlindingPolynomials() error { + if !useBlinding { + // When blinding is disabled (GNARK_DISABLE_BLINDING), skip creating blinding polynomials entirely + // Just close the channel to unblock any goroutines waiting on it + close(s.chbp) + return nil + } + + s.bp[id_Bl] = getRandomPolynomial(order_blinding_L) + s.bp[id_Br] = getRandomPolynomial(order_blinding_R) + s.bp[id_Bo] = getRandomPolynomial(order_blinding_O) + s.bp[id_Bz] = getRandomPolynomial(order_blinding_Z) + close(s.chbp) + return nil +} + +func (s *instance) initBSB22Commitments() { + s.commitmentInfo = s.spr.CommitmentInfo.(constraint.PlonkCommitments) + s.commitmentVal = make([]fr.Element, len(s.commitmentInfo)) // TODO @Tabaie get rid of this + s.cCommitments = make([]*iop.Polynomial, len(s.commitmentInfo)) + s.proof.Bsb22Commitments = make([]kzg.Digest, len(s.commitmentInfo)) + + // override the hint for the commitment constraints + bsb22ID := solver.GetHintID(fcs.Bsb22CommitmentComputePlaceholder) + s.opt.SolverOpts = append(s.opt.SolverOpts, solver.OverrideHint(bsb22ID, s.bsb22Hint)) +} + +// Computing and verifying Bsb22 multi-commits explained in https://hackmd.io/x8KsadW3RRyX7YTCFJIkHg +func (s *instance) bsb22Hint(_ *big.Int, ins, outs []*big.Int) error { + var err error + commDepth := int(ins[0].Int64()) + ins = ins[1:] + + res := &s.commitmentVal[commDepth] + + commitmentInfo := s.spr.CommitmentInfo.(constraint.PlonkCommitments)[commDepth] + committedValues := make([]fr.Element, s.domain0.Cardinality) + offset := s.spr.GetNbPublicVariables() + for i := range ins { + committedValues[offset+commitmentInfo.Committed[i]].SetBigInt(ins[i]) + } + if _, err = committedValues[offset+commitmentInfo.CommitmentIndex].SetRandom(); err != nil { // Commitment injection constraint has qcp = 0. Safe to use for blinding. + return err + } + if _, err = committedValues[offset+s.spr.GetNbConstraints()-1].SetRandom(); err != nil { // Last constraint has qcp = 0. Safe to use for blinding + return err + } + s.cCommitments[commDepth] = iop.NewPolynomial(&committedValues, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[commDepth], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[commDepth]); err != nil { + return err + } + + s.htfFunc.Write(s.proof.Bsb22Commitments[commDepth].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + res.SetBytes(hashBts[:nbBuf]) // TODO @Tabaie use CommitmentIndex for this; create a new variable CommitmentConstraintIndex for other uses + res.BigInt(outs[0]) + + return nil +} + +// solveConstraints computes the evaluation of the L, R, O polynomials in Lagrange form. +func (s *instance) solveConstraints() error { + startSolve := time.Now() + log := logger.Logger() + + var solution *cs.SparseR1CSSolution + + // Try to load raw solver values from cache (fastest path) + rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + if rawCachePath != "" { + if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { + // Reconstruct L, R, O from raw values + var sol cs.SparseR1CSSolution + if sol.L, sol.R, sol.O, err = s.spr.EvaluateLROSmallDomainFromValues(rawValues); err != nil { + log.Warn().Err(err).Str("file", rawCachePath).Msg("ignoring raw solver cache") + } else { + log.Debug().Dur("took", time.Since(startSolve)).Int("wires", len(rawValues)).Msg("loaded raw solver values from cache") + solution = &sol + } + + // Load cached BSB22 cCommitments polynomials + cacheDir := filepath.Dir(rawCachePath) + for i := 0; solution != nil && i < len(s.commitmentInfo); i++ { + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + coeffs, err := cs.LoadRawSolverValues(bsb22Path) + if err != nil { + log.Warn().Err(err).Int("i", i).Msg("ignoring raw solver cache: missing BSB22 commitment sidecar") + solution = nil + break + } + coeffSlice := []fr.Element(coeffs) + s.cCommitments[i] = iop.NewPolynomial(&coeffSlice, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[i], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[i]); err != nil { + return err + } + s.htfFunc.Write(s.proof.Bsb22Commitments[i].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + s.commitmentVal[i].SetBytes(hashBts[:nbBuf]) + } + } + } + + if solution == nil { + _solution, err := s.spr.SolveAndSaveRawValues(s.fullWitness, rawCachePath, s.opt.SolverOpts...) + if err != nil { + log.Debug().Dur("took", time.Since(startSolve)).Err(err).Msg("solveConstraints: spr.Solve") + return err + } + log.Debug().Dur("took", time.Since(startSolve)).Msg("solveConstraints: spr.Solve") + if isProfileMode { + s.timings.record("solveConstraints: spr.Solve", time.Since(startSolve)) + } + solution = _solution.(*cs.SparseR1CSSolution) + + // Save cCommitments polynomial coefficients for BSB22 reconstruction + if rawCachePath != "" && len(s.commitmentInfo) > 0 { + cacheDir := filepath.Dir(rawCachePath) + for i := range s.commitmentInfo { + if s.cCommitments[i] != nil { + coeffs := s.cCommitments[i].Coefficients() + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + if err := cs.SaveRawSolverValues(bsb22Path, coeffs); err != nil { + log.Warn().Err(err).Int("i", i).Msg("failed to save BSB22 commitment polynomial") + } + } + } + } + } + + evaluationLDomainSmall := []fr.Element(solution.L) + evaluationRDomainSmall := []fr.Element(solution.R) + evaluationODomainSmall := []fr.Element(solution.O) + var wg sync.WaitGroup + wg.Add(2) + go func() { + s.polyL = iop.NewPolynomial(&evaluationLDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + go func() { + s.polyR = iop.NewPolynomial(&evaluationRDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + + s.polyO = iop.NewPolynomial(&evaluationODomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + + wg.Wait() + if _, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}); err != nil { + return err + } + + // commit to l, r, o and add blinding factors + if err := s.commitToLRO(); err != nil { + return err + } + close(s.chLRO) + return nil +} + +func (s *instance) completeQk() error { + qk := s.trace.Qk.Clone() + qkCoeffs := qk.Coefficients() + + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + copy(qkCoeffs, wWitness[:len(s.spr.Public)]) + + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + for i := range s.commitmentInfo { + qkCoeffs[s.spr.GetNbPublicVariables()+s.commitmentInfo[i].CommitmentIndex] = s.commitmentVal[i] + } + + s.polyQk = qk + close(s.chQk) + + return nil +} + +func (s *instance) commitToLRO() error { + var startCommitLRO time.Time + if isProfileMode { + startCommitLRO = time.Now() + } + sequentialLRO := s.domain0 != nil && s.domain0.Cardinality >= (1<<22) + if _, ok := os.LookupEnv("ICICLE_LRO_COMMIT_SEQUENTIAL"); ok { + sequentialLRO = envEnabled("ICICLE_LRO_COMMIT_SEQUENTIAL", true) + } + + if !useBlinding { + // When blinding is disabled, commit directly without waiting for blinding polynomials + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + if err != nil { + return err + } + } else { + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (no blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (no blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil + } + + // wait for blinding polynomials to be initialized or context to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chbp: + } + + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + if err != nil { + return err + } + } else { + // Run the three commitments concurrently. + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (with blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil +} + +// deriveGammaAndBeta (copy constraint) +func (s *instance) deriveGammaAndBeta() error { + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + if err := bindPublicData(s.fs, "gamma", s.pk.VerifyingKey().(*plonk_{{ .CurvePkg }}.VerifyingKey), wWitness[:len(s.spr.Public)]); err != nil { + return err + } + + // wait for LRO to be committed + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + gamma, err := deriveRandomness(s.fs, "gamma", &s.proof.LRO[0], &s.proof.LRO[1], &s.proof.LRO[2]) + if err != nil { + return err + } + + bbeta, err := s.fs.ComputeChallenge("beta") + if err != nil { + return err + } + s.gamma = gamma + s.beta.SetBytes(bbeta) + + close(s.chGammaBeta) + + return nil +} + +// commitToPolyAndBlinding computes the KZG commitment of a polynomial p +// in Lagrange form (large degree) +// and add the contribution of a blinding polynomial b (small degree) +// /!\ The polynomial p is supposed to be in Lagrange form. +// Only used when blinding is enabled (the default). +func (s *instance) commitToPolyAndBlinding(p, b *iop.Polynomial) (commit curve.G1Affine, err error) { + // Commit over the Lagrange SRS using shared device-resident polynomial data. + gpuCommit, err := s.commitLagrangePolynomialOnGPU(p) + if err != nil { + return curve.G1Affine{}, err + } + + // add CPU blinding contribution (two MSMs on canonical SRS) + n := int(s.domain0.Cardinality) + cb := commitBlindingFactor(n, b, s.pk.Kzg) + gpuCommit.Add(&gpuCommit, &cb) + return gpuCommit, nil +} + +func (s *instance) commitLagrangePolynomialOnGPU(p *iop.Polynomial) (curve.G1Affine, error) { + if p == nil { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: nil polynomial") + } + if p.Basis != iop.Lagrange { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial must be in Lagrange basis, got %v", p.Basis) + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{p}) + if err != nil { + return curve.G1Affine{}, err + } + idx, ok := gpuState.polyToIdx[p] + if !ok || idx < 0 || idx >= len(gpuState.deviceSlices) { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial is missing from shared GPU state") + } + if gpuState.deviceSlices[idx].IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: empty device slice for polynomial") + } + + // Keep the polynomial in its native Lagrange basis; large MSMs are split + // into device-side chunks inside commitOnGPULagrangeDevice. + return commitOnGPULagrangeDevice(gpuState.deviceSlices[idx], &s.device, s.pk) +} + +func (s *instance) deriveAlpha() (err error) { + alphaDeps := make([]*curve.G1Affine, len(s.proof.Bsb22Commitments)+1) + for i := range s.proof.Bsb22Commitments { + alphaDeps[i] = &s.proof.Bsb22Commitments[i] + } + alphaDeps[len(alphaDeps)-1] = &s.proof.Z + s.alpha, err = deriveRandomness(s.fs, "alpha", alphaDeps...) + return err +} + +func (s *instance) deriveZeta() (err error) { + s.zeta, err = deriveRandomness(s.fs, "zeta", &s.proof.H[0], &s.proof.H[1], &s.proof.H[2]) + return +} + +func (s *instance) computeQuotient() (err error) { + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + if isProfileMode { + var startComputeQuotient time.Time + startComputeQuotient = time.Now() + defer func() { + l := logger.Logger() + if err != nil { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Err(err).Msg("computeQuotient: total (with error)") + } else { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Msg("computeQuotient: total") + } + s.timings.record("computeQuotient (total)", time.Since(startComputeQuotient)) + }() + } + + // wait for Z to be committed or context done + doneWaitZ := profileStep("computeQuotient: wait Z commit") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZ: + } + doneWaitZ() + + // derive alpha + if err = s.deriveAlpha(); err != nil { + return err + } + + if err := s.waitForComputeNumeratorQk(); err != nil { + return err + } + if s.polyQk == nil { + return fmt.Errorf("computeQuotient: missing completed Qk polynomial") + } + + doneEnsureGPUState := profileStep("computeQuotient: ensure shared GPU state") + gpuState, err := s.ensurePolysOnSharedGPU(s.buildComputeNumeratorGPUBatch()) + if err != nil { + return err + } + doneEnsureGPUState() + + // compute Z shifted by one for copy-constraint terms. + if s.polyZ == nil { + return fmt.Errorf("computeQuotient: missing Z polynomial") + } + s.polyZS = s.polyZ.ShallowClone().Shift(1) + + var numeratorGPU *gpuNumeratorPolynomial + var quotientGPU *gpuQuotientPolynomial + var e error + doneComputeNumerator := profileStep("computeQuotient: computeNumerator") + numeratorGPU, e = s.computeNumerator(gpuState) + if e != nil { + return e + } + doneComputeNumerator() + + doneDivideByZH := profileStep("computeQuotient: divideByZHOnGPU") + quotientGPU, e = s.divideByZHOnGPU(numeratorGPU, [2]*fft.Domain{s.domain0, s.domain1}) + if e != nil { + return e + } + doneDivideByZH() + s.hGPU = quotientGPU + + // Shared state slices were mutated during numerator coset iterations and are no + // longer needed now; computeLinearizedPolynomial uses the immutable snapshot. + s.releaseSharedGPUState() + close(s.chRestoreLRO) + + doneCommitH := profileStep("computeQuotient: commit H from device") + if err := s.commitToQuotientGPUFromDevice(s.hGPU); err != nil { + return err + } + doneCommitH() + + if err := s.deriveZeta(); err != nil { + return err + } + + donePrepareLinearizedEval := profileStep("computeQuotient: prepare linearized eval GPU state") + if err := s.prepareLinearizedEvalGPUStateFromHost(); err != nil { + return fmt.Errorf("computeQuotient: prepare linearized eval GPU state failed: %w", err) + } + donePrepareLinearizedEval() + + close(s.chH) + + return nil +} + +func (s *instance) buildRatioCopyConstraint() (err error) { + // wait for gamma and beta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chGammaBeta: + } + + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("buildRatioCopyConstraint: missing L/R/O polynomials") + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}) + if err != nil { + return err + } + dL, err := getStateDeviceSlice(gpuState, s.polyL, "L") + if err != nil { + return err + } + dR, err := getStateDeviceSlice(gpuState, s.polyR, "R") + if err != nil { + return err + } + dO, err := getStateDeviceSlice(gpuState, s.polyO, "O") + if err != nil { + return err + } + + var startBuildRatioCopyConstraintIcicle time.Time + if isProfileMode { + startBuildRatioCopyConstraintIcicle = time.Now() + } + s.polyZ, err = s.BuildRatioCopyConstraintIcicle( + []icicle_core.DeviceSlice{dL, dR, dO}, + s.trace.S, + s.beta, + s.gamma, + s.domain0, + gpuState, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startBuildRatioCopyConstraintIcicle)).Msg("buildRatioCopyConstraint: BuildRatioCopyConstraintIcicle") + s.timings.record("buildRatioCopyConstraint (perm Z)", time.Since(startBuildRatioCopyConstraintIcicle)) + } + + dZ, err := getStateDeviceSlice(gpuState, s.polyZ, "Z") + if err != nil { + return err + } + copyDone := make(chan error, 1) + var dPersist icicle_core.DeviceSlice + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + var allocErr error + dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) + if allocErr != nil { + finish(fmt.Errorf("buildRatioCopyConstraint: alloc persist Z failed: %w", allocErr)) + return + } + if e := copyDeviceSliceIntoOnCurrentDevice(dPersist, dZ, cfg); e != icicle_runtime.Success { + _ = dPersist.Free() + dPersist = icicle_core.DeviceSlice{} + finish(fmt.Errorf("buildRatioCopyConstraint: persist Z copy failed: %s", e.AsString())) + return + } + finish(nil) + }) + if err := <-copyDone; err != nil { + return err + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + s.polyZLagrangeGPU = dPersist + + // commit to Z (with or without blinding) + var startCommitZ time.Time + if isProfileMode { + startCommitZ = time.Now() + } + if useBlinding { + s.proof.Z, err = s.commitToPolyAndBlinding(s.polyZ, s.bp[id_Bz]) + } else { + s.proof.Z, err = s.commitLagrangePolynomialOnGPU(s.polyZ) + } + if isProfileMode { + l := logger.Logger() + if useBlinding { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (no blinding)") + } + } + s.freeIdleTempGPUMemoryOnDevice() + + close(s.chZ) + + return +} + +// open Z (blinded) at ωζ +func (s *instance) openZ() (err error) { + // wait for H to be committed and zeta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + + var zetaShifted fr.Element + zetaShifted.Mul(&s.zeta, &s.pk.Vk.Generator) + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("openZ: missing GPU Z polynomial") + } + + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + + dZLagrange := s.polyZLagrangeGPU + if dZLagrange.Len() <= 1 { + return fmt.Errorf("openZ: invalid Z size %d", dZLagrange.Len()) + } + + blindSize := order_blinding_Z + 1 + if useBlinding { + if len(s.bp) <= id_Bz || s.bp[id_Bz] == nil { + return fmt.Errorf("openZ: missing Z blinding polynomial") + } + blindSize = len(s.bp[id_Bz].Coefficients()) + if blindSize == 0 { + return fmt.Errorf("openZ: empty Z blinding polynomial") + } + } + + var dBlindedCanonical icicle_core.DeviceSlice + buildDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + buildDone <- cfgErr + return + } + finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + // Async boundary for canonicalization/blinding before exposing output. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + buildDone <- runErr + } + + n := dZLagrange.Len() + dZCanonical := s.getTempDeviceSlice(n) + if err := copyDeviceSliceIntoOnCurrentDevice(dZCanonical, dZLagrange, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy Z to canonical buffer failed: %s", err.AsString()), dZCanonical, false) + return + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + cfgNtt.Ordering = icicle_core.KNN // regular lagrange -> regular canonical + if err := icicle_ntt.Ntt(dZCanonical, icicle_core.KInverse, &cfgNtt, dZCanonical); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: inverse NTT on Z failed: %s", err.AsString()), dZCanonical, false) + return + } + + dBlindedCanonical = s.getTempDeviceSlice(n + blindSize) + dBlindedPrefix := (&dBlindedCanonical).Range(0, n, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedPrefix, dZCanonical, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy canonical Z into blinded buffer failed: %s", err.AsString()), dZCanonical, true) + return + } + + if useBlinding { + dBp := uploadVector(s.bp[id_Bz].Coefficients()) + dBlindedHead := (&dBlindedPrefix).Range(0, blindSize, false) + if err := icicle_vecops.VecOp(dBlindedHead, dBp, dBlindedHead, cfgVec, icicle_core.Sub); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: subtract Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedTail, dBp, cfgVec); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: append Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + _ = dBp.FreeAsync(stream) + } else { + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := zeroDeviceSliceOnCurrentDevice(dBlindedTail, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: zero-pad non-blinded Z failed: %s", err.AsString()), dZCanonical, true) + return + } + } + + finalize(nil, dZCanonical, false) + }) + if err := <-buildDone; err != nil { + return err + } + s.blindedZCanonicalGPU = dBlindedCanonical + + // open z at zeta*w. + var startKzgOpen time.Time + if isProfileMode { + startKzgOpen = time.Now() + } + s.proof.ZShiftedOpening, err = s.openPolynomialOnGPUCanonicalDevice(s.blindedZCanonicalGPU, zetaShifted) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startKzgOpen)).Msg("openZ: open polynomial on GPU") + s.timings.record("openZ (KZG open on GPU)", time.Since(startKzgOpen)) + } + if err != nil { + return err + } + close(s.chZOpening) + return nil +} + +func (s *instance) h1() []fr.Element { + var h1 []fr.Element + if !s.opt.StatisticalZK { + h1 = s.h.Coefficients()[:s.domain0.Cardinality+2] + } else { + h1 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h1, s.h.Coefficients()[:s.domain0.Cardinality+2]) + h1[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[0]) + } + return h1 +} + +func (s *instance) h2() []fr.Element { + var h2 []fr.Element + if !s.opt.StatisticalZK { + h2 = s.h.Coefficients()[s.domain0.Cardinality+2 : 2*(s.domain0.Cardinality+2)] + } else { + h2 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h2, s.h.Coefficients()[s.domain0.Cardinality+2:2*(s.domain0.Cardinality+2)]) + h2[0].Sub(&h2[0], &s.quotientShardsRandomizers[0]) + h2[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[1]) + } + return h2 +} + +func (s *instance) h3() []fr.Element { + var h3 []fr.Element + if !s.opt.StatisticalZK { + h3 = s.h.Coefficients()[2*(s.domain0.Cardinality+2) : 3*(s.domain0.Cardinality+2)] + } else { + h3 = make([]fr.Element, s.domain0.Cardinality+2) + copy(h3, s.h.Coefficients()[2*(s.domain0.Cardinality+2):3*(s.domain0.Cardinality+2)]) + h3[0].Sub(&h3[0], &s.quotientShardsRandomizers[1]) + } + return h3 +} + +// witnessEvalAtZeta holds the scalar evaluations of witness and constraint +// polynomials at the challenge point zeta, as needed by the linearized +// polynomial computation. +type witnessEvalAtZeta struct { + blzeta, brzeta, bozeta fr.Element + s1zeta, s2zeta fr.Element + qcpzeta []fr.Element +} + +type linearizedSelectorScales struct { + s3, ql, qr, qm, qo, qk fr.Element + qcp []fr.Element +} + +// evaluateWitnessPolynomialsAtZeta evaluates L, R, O (with optional blinding), +// S1, S2, and all Qcp polynomials at the point zeta using the GPU-resident +// polynomial state. +func (s *instance) evaluateWitnessPolynomialsAtZeta( + evalGPUState *gpuPolysState, + zeta fr.Element, +) (witnessEvalAtZeta, error) { + doneTotal := profileStep("evaluateWitnessPolynomialsAtZeta: total") + defer doneTotal() + + var result witnessEvalAtZeta + var err error + + result.qcpzeta = make([]fr.Element, len(s.commitmentInfo)) + var startQcp time.Time + if isProfileMode { + startQcp = time.Now() + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i >= len(s.trace.Qcp) || s.trace.Qcp[i] == nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: missing Qcp polynomial at index %d", i) + } + var startQcpItem time.Time + if isProfileMode { + startQcpItem = time.Now() + } + result.qcpzeta[i], err = s.evalPolynomialInCurrentFormOnGPU(s.trace.Qcp[i], evalGPUState, zeta) + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: qcp[%d] GPU evaluation failed: %w", i, err) + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("idx", i).Dur("took", time.Since(startQcpItem)).Msg("evaluateWitnessPolynomialsAtZeta: qcp eval item") + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startQcp)).Msg("evaluateWitnessPolynomialsAtZeta: qcpZeta evaluate on GPU") + } + + if useBlinding { + result.blzeta, err = s.evaluateBlindedOnGPU(s.polyL, s.bp[id_Bl], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: L evaluate on GPU") + result.blzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyL, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: blzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.brzeta, err = s.evaluateBlindedOnGPU(s.polyR, s.bp[id_Br], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: R evaluate on GPU") + result.brzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyR, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: brzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.bozeta, err = s.evaluateBlindedOnGPU(s.polyO, s.bp[id_Bo], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: O evaluate on GPU") + result.bozeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyO, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: bozeta GPU evaluation failed: %w", err) + } + + doneS1 := profileStep("evaluateWitnessPolynomialsAtZeta: S1 evaluate on GPU") + result.s1zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S1, evalGPUState, zeta) + doneS1() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s1(zeta) GPU evaluation failed: %w", err) + } + doneS2 := profileStep("evaluateWitnessPolynomialsAtZeta: S2 evaluate on GPU") + result.s2zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S2, evalGPUState, zeta) + doneS2() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s2(zeta) GPU evaluation failed: %w", err) + } + + return result, nil +} + +func (s *instance) computeLinearizedPolynomial() error { + + // wait for H to be committed and zeta to be derived (or ctx.Done()) + var startWaitH time.Time + if isProfileMode { + startWaitH = time.Now() + } + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startWaitH)).Msg("computeLinearizedPolynomial: wait H and zeta") + s.timings.record("computeLinearizedPoly (wait H+zeta, overlaps)", time.Since(startWaitH)) + } + if s.opt.StatisticalZK { + return fmt.Errorf("computeLinearizedPolynomial: GPU-only opening path does not support StatisticalZK=true") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil || s.polyZ == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing required polynomials") + } + + // Reuse the immutable snapshot prepared in computeQuotient before numerator + // coset iterations mutate shared state slices. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("computeLinearizedPolynomial: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized eval prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: required polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing blinding polynomials for GPU evaluation") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized blinding prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: blinding polynomial is not on GPU: %w", e) + } + } + } + + doneEvaluate := profileStep("computeLinearizedPolynomial: evaluate witness polynomials") + evals, err := s.evaluateWitnessPolynomialsAtZeta(evalGPUState, s.zeta) + doneEvaluate() + if err != nil { + return err + } + + // wait for Z to be opened at zeta (or ctx.Done()) + doneWaitZOpening := profileStep("computeLinearizedPolynomial: wait Z opening") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZOpening: + } + doneWaitZOpening() + if s.blindedZCanonicalGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing canonical blinded Z on GPU") + } + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing lagrange Z on GPU") + } + defer func() { + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + }() + bzuzeta := s.proof.ZShiftedOpening.ClaimedValue + + if s.hGPU == nil || s.hGPU.coeffs.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing GPU quotient polynomial") + } + + doneBuild := profileStep("computeLinearizedPolynomial: build selector terms on GPU") + dLin, err := s.buildLinearizedSelectorTermsOnGPU(evals, bzuzeta, s.blindedZCanonicalGPU.Len()) + doneBuild() + if err != nil { + return err + } + + doneAddZ := profileStep("computeLinearizedPolynomial: add Z contribution on GPU") + err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) + doneAddZ() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + s.linearizedPolynomialGPU = dLin + + doneEvalClaim := profileStep("computeLinearizedPolynomial: evaluate linearized claim") + claim, err := s.evalDevicePolynomialAtPoint(dLin, s.zeta) + doneEvalClaim() + if err != nil { + return err + } + s.linearizedPolynomialClaim = claim + + // Commit the linearized polynomial over the canonical SRS. + var startMSM time.Time + if isProfileMode { + startMSM = time.Now() + } + s.linearizedPolynomialDigest, err = commitOnGPUCanonicalDevice(dLin, &s.device, s.pk) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startMSM)).Msg("computeLinearizedPolynomial: KZG commit") + s.timings.record("computeLinearizedPoly (KZG commit)", time.Since(startMSM)) + } + close(s.chLinearizedPolynomial) + return nil +} + +func (s *instance) batchOpening() error { + + // wait for linearizedPolynomial to be computed (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLinearizedPolynomial: + } + + defer func() { + freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } + }() + + if s.linearizedPolynomialGPU.IsEmpty() { + return fmt.Errorf("batchOpening: missing GPU linearized polynomial") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("batchOpening: missing L/R/O polynomials") + } + + // Reuse immutable GPU snapshot prepared before quotient iterations. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("batchOpening: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening eval prereq"); e != nil { + return fmt.Errorf("batchOpening: required polynomial is not on GPU: %w", e) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] == nil { + return fmt.Errorf("batchOpening: missing Qcp polynomial at index %d", i) + } + if _, e := getStateDeviceSlice(evalGPUState, s.trace.Qcp[i], "batchOpening qcp prereq"); e != nil { + return fmt.Errorf("batchOpening: Qcp polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("batchOpening: missing blinding polynomials") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening blinding prereq"); e != nil { + return fmt.Errorf("batchOpening: blinding polynomial is not on GPU: %w", e) + } + } + } + + devicePolys, ownedPolys, claimed, err := s.prepareBatchOpeningPolynomialsOnGPU(evalGPUState, s.zeta) + if err != nil { + return err + } + defer func() { + for i := 0; i < len(devicePolys); i++ { + if i < len(ownedPolys) && ownedPolys[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + } + } + s.releaseLinearizedEvalGPUState() + }() + + digestsToOpen := make([]curve.G1Affine, len(s.pk.Vk.Qcp)+6) + copy(digestsToOpen[6:], s.pk.Vk.Qcp) + digestsToOpen[0] = s.linearizedPolynomialDigest + digestsToOpen[1] = s.proof.LRO[0] + digestsToOpen[2] = s.proof.LRO[1] + digestsToOpen[3] = s.proof.LRO[2] + digestsToOpen[4] = s.pk.Vk.S[0] + digestsToOpen[5] = s.pk.Vk.S[1] + if len(claimed) != len(digestsToOpen) { + return fmt.Errorf("batchOpening: claimed size mismatch (%d != %d)", len(claimed), len(digestsToOpen)) + } + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + var dFold icicle_core.DeviceSlice + foldDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening fold", foldDone) + + dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: copy linearized polynomial failed: %s", e.AsString())) + return + } + + gammaPow := gamma + for i := 1; i < len(devicePolys); i++ { + dPoly := devicePolys[i] + if dPoly.IsEmpty() { + gammaPow.Mul(&gammaPow, &gamma) + continue + } + dScaled := s.getTempDeviceSlice(dPoly.Len()) + dGammaStd := uploadScalarStdOnCurrentDevice(gammaPow, cfg) + eMul := icicle_vecops.ScalarMulVec(dGammaStd, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dGammaStd.FreeAsync(cfg.StreamHandle) + } else { + _ = dGammaStd.Free() + } + if eMul != icicle_runtime.Success { + if cfg.IsAsync { + _ = icicle_runtime.SynchronizeStream(cfg.StreamHandle) + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: scale polynomial %d failed: %s", i, eMul.AsString())) + return + } + + dPrefix := (&dFold).Range(0, dPoly.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add) + if cfg.IsAsync { + // dScaled is recycled each iteration; wait before returning to pool. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: synchronize stream failed: %s", eSync.AsString())) + return + } + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + if eAdd != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: fold add polynomial %d failed: %s", i, eAdd.AsString())) + return + } + gammaPow.Mul(&gammaPow, &gamma) + } + finish(nil) + }) + if err := <-foldDone; err != nil { + if !dFold.IsEmpty() { + s.putTempDeviceSlice(dFold, dFold.Len()) + } + return err + } + var dWitness icicle_core.DeviceSlice + divDone := make(chan error, 1) + witnessSize := dFold.Len() - 1 + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dFold, dPoint, dWitness, cfg) + if cfg.IsAsync { + _ = dPoint.FreeAsync(cfg.StreamHandle) + } else { + _ = dPoint.Free() + } + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + finish(fmt.Errorf("batchOpening: divide by (x-zeta) failed: %s", eDiv.AsString())) + return + } + finish(nil) + }) + if err := <-divDone; err != nil { + s.putTempDeviceSlice(dFold, dFold.Len()) + return err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dFold, dFold.Len()) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return err + } + + s.proof.BatchedProof = kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + } + if err := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); err != nil { + l := logger.Logger() + l.Warn().Err(err).Msg("batchOpening: GPU folded opening failed raw-G2 validation; falling back to host fold with CPU KZG commitment") + fallbackProof, fallbackErr := s.batchOpeningHostFoldGPUCommitFromDevicePolys(devicePolys, digestsToOpen) + if fallbackErr != nil { + return fallbackErr + } + s.proof.BatchedProof = fallbackProof + if verifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); verifyErr != nil { + nativeProof, nativeErr := s.batchOpeningNativeCPUFromDevicePolys(devicePolys, digestsToOpen) + if nativeErr != nil { + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed: %v", verifyErr, nativeErr) + } + s.proof.BatchedProof = nativeProof + if nativeVerifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); nativeVerifyErr != nil { + diagnostic, diagnosticErr := s.diagnoseBatchOpeningDevicePolynomials(devicePolys, digestsToOpen, claimed) + if diagnosticErr != nil { + diagnostic = fmt.Sprintf("batch opening diagnostic failed: %v", diagnosticErr) + } + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed raw-G2 validation: %v; %s", verifyErr, nativeVerifyErr, diagnostic) + } + l.Warn().Msg("batchOpening: native CPU KZG fallback produced a valid proof after host-fold fallback failed") + } + } + _ = foldedEval // kept for parity with kzg.BatchOpenSinglePoint flow. + return nil +} + +func (s *instance) verifyBatchOpeningWithRecomputedLines(digestsToOpen []curve.G1Affine) error { + vk := s.pk.Vk.Kzg + vk.Lines[0] = curve.PrecomputeLines(vk.G2[0]) + vk.Lines[1] = curve.PrecomputeLines(vk.G2[1]) + return kzg.BatchVerifySinglePoint( + digestsToOpen, + &s.proof.BatchedProof, + s.zeta, + s.kzgFoldingHash, + vk, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningHostFoldGPUCommitFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningNativeCPUFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningNativeCPUFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return kzg.BatchOpenSinglePoint( + polysToOpen, + digestsToOpen, + s.zeta, + s.kzgFoldingHash, + s.pk.Kzg, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) diagnoseBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, + gpuClaimed []fr.Element, +) (string, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "diagnoseBatchOpeningDevicePolynomials", + ) + if err != nil { + return "", err + } + if len(gpuClaimed) != len(polysToOpen) { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: claimed/polynomial mismatch (%d != %d)", len(gpuClaimed), len(polysToOpen)) + } + + l := logger.Logger() + claimMismatches := make([]string, 0) + commitMismatches := make([]string, 0) + for i := range polysToOpen { + label := batchOpeningPolynomialLabel(i) + cpuClaim := evalCanonicalAtPoint(polysToOpen[i], s.zeta) + if !cpuClaim.Equal(&gpuClaimed[i]) { + claimMismatches = append(claimMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("gpuClaim", frFingerprint(gpuClaimed[i])). + Str("cpuClaim", frFingerprint(cpuClaim)). + Msg("batchOpening diagnostic: GPU claim differs from CPU evaluation") + } + + cpuDigest, err := kzg.Commit(polysToOpen[i], s.pk.Kzg) + if err != nil { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: commit %s: %w", label, err) + } + if !cpuDigest.Equal(&digestsToOpen[i]) { + commitMismatches = append(commitMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("expectedDigest", g1Fingerprint(digestsToOpen[i])). + Str("cpuDigest", g1Fingerprint(cpuDigest)). + Msg("batchOpening diagnostic: CPU commitment differs from proof digest") + } + } + + if len(claimMismatches) == 0 && len(commitMismatches) == 0 { + return "batch opening diagnostic found no per-polynomial claim or commitment mismatch", nil + } + return fmt.Sprintf( + "batch opening diagnostic claim mismatches=[%s] commitment mismatches=[%s]", + strings.Join(claimMismatches, ","), + strings.Join(commitMismatches, ","), + ), nil +} + +func batchOpeningPolynomialLabel(index int) string { + switch index { + case 0: + return "linearized" + case 1: + return "L" + case 2: + return "R" + case 3: + return "O" + case 4: + return "S1" + case 5: + return "S2" + default: + return fmt.Sprintf("Qcp[%d]", index-6) + } +} + +func frFingerprint(v fr.Element) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func g1Fingerprint(v curve.G1Affine) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func (s *instance) downloadBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + expectedDigests int, + label string, +) ([][]fr.Element, error) { + if len(devicePolys) != expectedDigests { + return nil, fmt.Errorf("%s: polynomial/digest mismatch (%d != %d)", label, len(devicePolys), expectedDigests) + } + + polysToOpen := make([][]fr.Element, len(devicePolys)) + for i := range devicePolys { + var err error + polysToOpen[i], err = s.downloadCanonicalDeviceCoefficients( + devicePolys[i], + fmt.Sprintf("%s[%d]", label, i), + ) + if err != nil { + return nil, err + } + } + return polysToOpen, nil +} + +func (s *instance) batchOpeningHostFoldGPUCommit(digestsToOpen []curve.G1Affine) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.batchOpeningHostPolynomials() + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( + polysToOpen [][]fr.Element, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + if len(polysToOpen) != len(digestsToOpen) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: polynomial/digest mismatch (%d != %d)", len(polysToOpen), len(digestsToOpen)) + } + + largestPoly := 0 + for i := range polysToOpen { + if len(polysToOpen[i]) == 0 || len(polysToOpen[i]) > len(s.pk.Kzg.G1) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: invalid polynomial %d size %d", i, len(polysToOpen[i])) + } + if len(polysToOpen[i]) > largestPoly { + largestPoly = len(polysToOpen[i]) + } + } + + claimed := make([]fr.Element, len(polysToOpen)) + utils.Parallelize(len(polysToOpen), func(start, end int) { + for i := start; i < end; i++ { + claimed[i] = evalCanonicalAtPoint(polysToOpen[i], s.zeta) + } + }) + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + foldedPolynomials := make([]fr.Element, largestPoly) + copy(foldedPolynomials, polysToOpen[0]) + gammaPow := gamma + for i := 1; i < len(polysToOpen); i++ { + poly := polysToOpen[i] + scale := gammaPow + utils.Parallelize(len(poly), func(start, end int) { + var term fr.Element + for j := start; j < end; j++ { + term.Mul(&poly[j], &scale) + foldedPolynomials[j].Add(&foldedPolynomials[j], &term) + } + }) + gammaPow.Mul(&gammaPow, &gamma) + } + + hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) + var dWitness icicle_core.DeviceSlice + uploadDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + dWitness = uploadVector(hCoeffs) + close(uploadDone) + }) + <-uploadDone + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + freeSliceOnDevice(&dWitness, &s.device) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + return kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + }, nil +} + +func (s *instance) batchOpeningHostPolynomials() ([][]fr.Element, error) { + total := 6 + len(s.trace.Qcp) + polysToOpen := make([][]fr.Element, total) + + var err error + polysToOpen[0], err = s.downloadCanonicalDeviceCoefficients( + s.linearizedPolynomialGPU, + "batchOpeningHostPolynomials linearized", + ) + if err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare linearized: %w", err) + } + + prepareLRO := func(p, bp *iop.Polynomial, blindingOrder int) ([]fr.Element, error) { + base, err := canonicalRegularCoefficientsCopy(p, s.domain0) + if err != nil { + return nil, err + } + if useBlinding { + if bp == nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomial") + } + blind, err := canonicalRegularCoefficientsCopy(bp, s.domain0) + if err != nil { + return nil, err + } + out := make([]fr.Element, len(base)+len(blind)) + copy(out, base) + copy(out[len(base):], blind) + for i := range blind { + out[i].Sub(&out[i], &blind[i]) + } + return out, nil + } + out := make([]fr.Element, len(base)+blindingOrder+1) + copy(out, base) + return out, nil + } + + var bpL, bpR, bpO *iop.Polynomial + if useBlinding { + if len(s.bp) <= id_Bo { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomials") + } + bpL, bpR, bpO = s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo] + } + + if polysToOpen[1], err = prepareLRO(s.polyL, bpL, order_blinding_L); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare L: %w", err) + } + if polysToOpen[2], err = prepareLRO(s.polyR, bpR, order_blinding_R); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare R: %w", err) + } + if polysToOpen[3], err = prepareLRO(s.polyO, bpO, order_blinding_O); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare O: %w", err) + } + if polysToOpen[4], err = canonicalRegularCoefficientsCopy(s.trace.S1, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S1: %w", err) + } + if polysToOpen[5], err = canonicalRegularCoefficientsCopy(s.trace.S2, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S2: %w", err) + } + for i := range s.trace.Qcp { + if polysToOpen[6+i], err = canonicalRegularCoefficientsCopy(s.trace.Qcp[i], s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare Qcp[%d]: %w", i, err) + } + } + return polysToOpen, nil +} + +func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceSlice, label string) ([]fr.Element, error) { + if dPoly.IsEmpty() { + return nil, fmt.Errorf("%s: empty device polynomial", label) + } + + coeffs := make([]fr.Element, dPoly.Len()) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + }) + if err := <-done; err != nil { + return nil, err + } + return coeffs, nil +} + +func canonicalRegularCoefficientsCopy(p *iop.Polynomial, domain *fft.Domain) ([]fr.Element, error) { + if p == nil { + return nil, fmt.Errorf("nil polynomial") + } + cp := p.Clone() + cp.ToCanonical(domain).ToRegular() + coeffs := cp.Coefficients() + out := make([]fr.Element, len(coeffs)) + copy(out, coeffs) + return out, nil +} + +func dividePolyByXMinusAHost(f []fr.Element, fa, a fr.Element) []fr.Element { + f[0].Sub(&f[0], &fa) + var t fr.Element + for i := len(f) - 2; i >= 0; i-- { + t.Mul(&f[i+1], &a) + f[i].Add(&f[i], &t) + } + return f[1:] +} + +// evaluate the full set of constraints on the GPU-resident polynomial state. +type computeNumeratorLoopContext struct { + n int + rho int + mm uint64 + bn *big.Int + shifters []fr.Element + twiddles0 []fr.Element + dTwiddles0 icicle_core.DeviceSlice + dPrecomputedDenominators *icicle_core.DeviceSlice + scalingVector []fr.Element + scalingVectorRev []fr.Element + gpuState *gpuPolysState + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + one fr.Element + cs fr.Element + css fr.Element + nbBsbGates int + numeratorShards []icicle_core.DeviceSlice +} + +type gpuNumeratorPolynomial struct { + shards []icicle_core.DeviceSlice + n int + rho int + mm uint64 +} + +type gpuQuotientPolynomial struct { + coeffs icicle_core.DeviceSlice + size int +} + +func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolynomial, error) { + twiddles0, err := s.buildComputeNumeratorTwiddles() + if err != nil { + return nil, err + } + if err := s.validateComputeNumeratorGPUState(gpuState); err != nil { + return nil, err + } + + var startComputeNumerator time.Time + if isProfileMode { + startComputeNumerator = time.Now() + } + + n := s.domain0.Cardinality + nbBsbGates := len(s.proof.Bsb22Commitments) + + var cs, css fr.Element + cs.Set(&s.domain1.FrMultiplicativeGen) + css.Square(&cs) + + bn := big.NewInt(int64(n)) + + rho := int(s.domain1.Cardinality / n) + shifters := make([]fr.Element, rho) + shifters[0].Set(&s.domain1.FrMultiplicativeGen) + for i := 1; i < rho; i++ { + shifters[i].Set(&s.domain1.Generator) + } + + cosetTable, err := s.domain0.CosetTable() + if err != nil { + return nil, err + } + + // for the first iteration, the scalingVector is the coset table + scalingVector := cosetTable + scalingVectorRev := make([]fr.Element, len(cosetTable)) + copy(scalingVectorRev, cosetTable) + fft.BitReverse(scalingVectorRev) + + // pre-computed to compute the bit reverse index + // of the result polynomial + m := uint64(s.domain1.Cardinality) + mm := uint64(64 - bits.TrailingZeros64(m)) + + var dPrecomputedDenominators icicle_core.DeviceSlice + defer func() { + if !dPrecomputedDenominators.IsEmpty() { + freeDone := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + freeDone <- dPrecomputedDenominators.Free() + }) + if err := <-freeDone; err != icicle_runtime.Success { + panic(fmt.Sprintf("computeNumerator: failed to free dPrecomputedDenominators: %s", err.AsString())) + } + } + }() + + var coset, cosetExponentiatedToNMinusOne, one fr.Element + coset.SetOne() + one.SetOne() + + dTwiddles0, err := s.uploadComputeNumeratorTwiddles(twiddles0) + if err != nil { + return nil, err + } + + loopCtx := &computeNumeratorLoopContext{ + n: int(n), + rho: rho, + mm: mm, + bn: bn, + shifters: shifters, + twiddles0: twiddles0, + dTwiddles0: dTwiddles0, + dPrecomputedDenominators: &dPrecomputedDenominators, + scalingVector: scalingVector, + scalingVectorRev: scalingVectorRev, + gpuState: gpuState, + coset: coset, + cosetExponentiatedToNMinusOne: cosetExponentiatedToNMinusOne, + one: one, + cs: cs, + css: css, + nbBsbGates: nbBsbGates, + numeratorShards: make([]icicle_core.DeviceSlice, rho), + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: setup before iteration loop") + } + if err := s.executeComputeNumeratorCosetIterations(loopCtx); err != nil { + return nil, err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: main body (post-wait)") + } + + return &gpuNumeratorPolynomial{ + shards: loopCtx.numeratorShards, + n: loopCtx.n, + rho: loopCtx.rho, + mm: loopCtx.mm, + }, nil + +} + +func (s *instance) buildComputeNumeratorTwiddles() ([]fr.Element, error) { + n := s.domain0.Cardinality + var startTwiddles time.Time + if isProfileMode { + startTwiddles = time.Now() + } + twiddles0 := make([]fr.Element, n) + if n == 1 { + // edge case + twiddles0[0].SetOne() + } else { + twiddles, err := s.domain0.Twiddles() + if err != nil { + return nil, err + } + copy(twiddles0, twiddles[0]) + w := twiddles0[1] + for i := len(twiddles[0]); i < len(twiddles0); i++ { + twiddles0[i].Mul(&twiddles0[i-1], &w) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startTwiddles)).Msg("computeNumerator: build twiddles") + } + return twiddles0, nil +} + +func (s *instance) waitForComputeNumeratorQk() error { + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chQk: + } + return nil +} + +func (s *instance) buildLinearizedEvalGPUBatch() []*iop.Polynomial { + baseCap := 5 + len(s.trace.Qcp) + if useBlinding { + baseCap += 3 + } + polys := make([]*iop.Polynomial, 0, baseCap) + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + } + if useBlinding && len(s.bp) > id_Bo { + for _, bpPoly := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if bpPoly != nil { + polys = append(polys, bpPoly) + } + } + } + return polys +} + +func (s *instance) buildComputeNumeratorGPUBatch() []*iop.Polynomial { + polys := make([]*iop.Polynomial, 0, 13+2*len(s.commitmentInfo)) + for _, p := range []*iop.Polynomial{ + s.polyL, s.polyR, s.polyO, s.polyZ, + s.trace.Ql, s.trace.Qr, s.trace.Qm, s.trace.Qo, s.polyQk, + s.trace.S1, s.trace.S2, s.trace.S3, + } { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i < len(s.trace.Qcp) && s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + if i < len(s.cCommitments) && s.cCommitments[i] != nil { + polys = append(polys, s.cCommitments[i]) + } + } + return polys +} + +func (s *instance) validateComputeNumeratorGPUState(state *gpuPolysState) error { + if state == nil { + return fmt.Errorf("computeNumerator: shared GPU state is nil") + } + required := s.buildComputeNumeratorGPUBatch() + if len(required) == 0 { + return fmt.Errorf("computeNumerator: no polynomials prepared for GPU batch") + } + for _, p := range required { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; !ok { + return fmt.Errorf("computeNumerator: polynomial ptr=%p is missing from shared GPU state", p) + } + } + return nil +} + +func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicle_core.DeviceSlice, error) { + var dTwiddles0 icicle_core.DeviceSlice + uploadTwiddlesDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + if s.tempGPUMemPool != nil { + s.tempGPUMemPool.FreeAll() + } + host := icicle_core.HostSliceFromElements(twiddles0) + var allocErr error + dTwiddles0, allocErr = allocDeviceUninitialized(len(twiddles0)) + if allocErr != nil { + uploadTwiddlesDone <- fmt.Errorf("uploadComputeNumeratorTwiddles: %w", allocErr) + return + } + host.CopyToDevice(&dTwiddles0, false) + uploadTwiddlesDone <- nil + }) + if err := <-uploadTwiddlesDone; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dTwiddles0, nil +} + +// executeComputeNumeratorCosetIterations runs the rho coset iterations. +func (s *instance) executeComputeNumeratorCosetIterations(loopCtx *computeNumeratorLoopContext) error { + var startIterLoop time.Time + if isProfileMode { + startIterLoop = time.Now() + } + + for i := 0; i < loopCtx.rho; i++ { + if err := s.computeNumeratorIteration(i, loopCtx); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + freeSliceOnDevice(&loopCtx.dTwiddles0, &s.device) + return err + } + } + + // Free twiddles0 device slice (uploaded once before the loop). + freeTwiddlesDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + loopCtx.dTwiddles0.Free() + close(freeTwiddlesDone) + }) + <-freeTwiddlesDone + + if useBlinding { + var startRestoreBlindingPolys time.Time + if isProfileMode { + startRestoreBlindingPolys = time.Now() + } + csInv := inverseShifterProduct(loopCtx.shifters) + if err := s.restoreBlindingPolynomials(csInv); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startRestoreBlindingPolys)).Msg("computeNumerator: restore blinding polys") + } + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startIterLoop)).Msg("computeNumerator: full iteration loop") + } + return nil +} + +func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoopContext) error { + loopCtx.coset.Mul(&loopCtx.coset, &loopCtx.shifters[i]) + loopCtx.cosetExponentiatedToNMinusOne.Exp(loopCtx.coset, loopCtx.bn). + Sub(&loopCtx.cosetExponentiatedToNMinusOne, &loopCtx.one) + + batchInvertDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) + }) + if err := <-batchInvertDone; err != nil { + return err + } + + s.applyNumeratorBlindingScale(i, loopCtx) + if i == 1 { + // We have to update the scalingVector; instead of scaling by + // cosets we scale by the twiddles of the large domain. + w := s.domain1.Generator + loopCtx.scalingVector = make([]fr.Element, loopCtx.n) + fft.BuildExpTable(w, loopCtx.scalingVector) + + // Reuse memory. + copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) + fft.BitReverse(loopCtx.scalingVectorRev) + } + + // We do **a lot** of FFT here, but on the small domain. + // Note that for all the polynomials in the proving key + // (Ql, Qr, Qm, Qo, S1, S2, S3, Qcp, Qc) and ID, LOne + // we could pre-compute these rho*2 FFTs and store them + // at the cost of a huge memory footprint. + var startGpuInverseScaleForward time.Time + if isProfileMode { + startGpuInverseScaleForward = time.Now() + } + + // Inverse NTT -> Scale -> Forward NTT all on GPU using persistent GPU memory. + if err := s.gpuNTTInverseScaleForwardOnDevice(loopCtx.gpuState, loopCtx.scalingVector, loopCtx.scalingVectorRev, s.pk); err != nil { + return err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startGpuInverseScaleForward)).Msg("computeNumerator: gpuNTTInverseScaleForwardOnDevice") + } + + // Evaluate constraints on GPU. + constraintParams := gpuConstraintEvalParams{ + beta: s.beta, + gamma: s.gamma, + alpha: s.alpha, + coset: loopCtx.coset, + cosetExponentiatedToNMinusOne: loopCtx.cosetExponentiatedToNMinusOne, + cs: loopCtx.cs, + css: loopCtx.css, + cardinalityInv: s.domain0.CardinalityInv, + n: loopCtx.n, + nbBsbGates: loopCtx.nbBsbGates, + } + var startEvalConstraints time.Time + if isProfileMode { + startEvalConstraints = time.Now() + } + dNumeratorShard, err := s.gpuEvaluateConstraints( + loopCtx.gpuState, + constraintParams, + loopCtx.twiddles0, // CPU version for computeBlindingPolynomials + loopCtx.dTwiddles0, // GPU version for computeOrderingConstraint + *loopCtx.dPrecomputedDenominators, + s.bp, + nil, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startEvalConstraints)).Msg("computeNumerator: gpuEvaluateConstraints") + } + loopCtx.numeratorShards[i] = dNumeratorShard + + loopCtx.cosetExponentiatedToNMinusOne. + Inverse(&loopCtx.cosetExponentiatedToNMinusOne) + s.applyNumeratorBlindingUnscale(i, loopCtx) + return nil +} + +func (s *instance) buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx *computeNumeratorLoopContext) error { + if loopCtx == nil || loopCtx.dPrecomputedDenominators == nil { + return fmt.Errorf("computeNumerator: nil denominator device slice") + } + if loopCtx.dTwiddles0.IsEmpty() || loopCtx.dTwiddles0.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid twiddles device slice size %d, expected %d", loopCtx.dTwiddles0.Len(), loopCtx.n) + } + + if loopCtx.dPrecomputedDenominators.IsEmpty() { + dDenominators, err := allocDeviceUninitialized(loopCtx.n) + if err != nil { + return err + } + *loopCtx.dPrecomputedDenominators = dDenominators + } else if loopCtx.dPrecomputedDenominators.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid denominator device slice size %d, expected %d", loopCtx.dPrecomputedDenominators.Len(), loopCtx.n) + } + + cfg := icicle_core.DefaultVecOpsConfig() + + // dTwiddles0 is a Montgomery scalar vector in domain0 regular order. + // ScalarMulVec expects the scalar in standard form and preserves the + // Montgomery representation of the vector result. + dCosetStd := uploadScalarStdOnCurrentDevice(loopCtx.coset, cfg) + defer dCosetStd.Free() + if err := icicle_vecops.ScalarMulVec( + dCosetStd, + loopCtx.dTwiddles0, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators coset*twiddles failed: %s", err.AsString()) + } + + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneMont := uploadScalarMontOnCurrentDevice(minusOne, cfg) + defer dMinusOneMont.Free() + if err := icicle_vecops.ScalarAddVec( + dMinusOneMont, + *loopCtx.dPrecomputedDenominators, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators subtract one failed: %s", err.AsString()) + } + + if err := s.batchInvertOnCurrentDevice(*loopCtx.dPrecomputedDenominators); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: batchInvert failed: %s", err.AsString()) + } + return nil +} + +func (s *instance) applyNumeratorBlindingScale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindScale time.Time + if isProfileMode { + startBlindScale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + acc := loopCtx.cosetExponentiatedToNMinusOne + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &acc) + acc.Mul(&acc, &loopCtx.shifters[i]) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindScale)).Msg("computeNumerator: scale blinding polys") + } +} + +func (s *instance) applyNumeratorBlindingUnscale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindUnscale time.Time + if isProfileMode { + startBlindUnscale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &loopCtx.cosetExponentiatedToNMinusOne) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindUnscale)).Msg("computeNumerator: unscale blinding polys") + } +} + +func (s *instance) restoreBlindingPolynomials(csInv fr.Element) error { + for _, q := range s.bp { + if q == nil { + continue + } + cp := q.Coefficients() + if len(cp) == 0 { + continue + } + var acc fr.Element + acc.SetOne() + for i := 0; i < len(cp); i++ { + cp[i].Mul(&cp[i], &acc) + acc.Mul(&acc, &csInv) + } + } + return nil +} + +func inverseShifterProduct(shifters []fr.Element) fr.Element { + var acc fr.Element + acc.SetOne() + for i := 0; i < len(shifters); i++ { + acc.Mul(&acc, &shifters[i]) + } + acc.Inverse(&acc) + return acc +} + +func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial) (_ *iop.Polynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("downloadNumeratorFromGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("downloadNumeratorFromGPU: invalid dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard count mismatch: got %d, expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + + defer func() { + if err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + } + }() + + for i := 0; i < gpuNumerator.rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard %d is empty", i) + } + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + var dMerged icicle_core.DeviceSlice + + mergeDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + return + } + // Async boundary before merged slice is consumed by host copy. + mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + }) + if mergeErr := <-mergeDone; mergeErr != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, mergeErr + } + + cres := make([]fr.Element, totalSize) + cresHost := icicle_core.HostSliceFromElements(cres) + downloadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + }) + if err := <-downloadDone; err != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, err + } + s.putTempDeviceSlice(dMerged, totalSize) + + s.freeNumeratorShards(gpuNumerator.shards) + return iop.NewPolynomial(&cres, iop.Form{Basis: iop.LagrangeCoset, Layout: iop.BitReverse}), nil +} + +func (s *instance) freeNumeratorShards(shards []icicle_core.DeviceSlice) { + if len(shards) == 0 { + return + } + for i := 0; i < len(shards); i++ { + if !shards[i].IsEmpty() { + s.putTempDeviceSlice(shards[i], shards[i].Len()) + shards[i] = icicle_core.DeviceSlice{} + } + } +} + +func (s *instance) batchInvert(dVec icicle_core.DeviceSlice) { + if dVec.Len() == 0 { + return + } + + done := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + done <- s.batchInvertOnCurrentDevice(dVec) + }) + if err := <-done; err != icicle_runtime.Success { + panic(fmt.Sprintf("batchInvert: BatchInverseVec failed: %s", err.AsString())) + } +} + +// batchInvertOnCurrentDevice assumes caller already runs on the active device thread. +func (s *instance) batchInvertOnCurrentDevice(dVec icicle_core.DeviceSlice) icicle_runtime.EIcicleError { + if dVec.Len() == 0 { + return icicle_runtime.Success + } + err := icicle_{{ .CurvePkg }}.FromMontgomery(dVec) + if err == icicle_runtime.Success { + cfg := icicle_core.DefaultVecOpsConfig() + err = icicle_vecops.BatchInverseVec(dVec, dVec, cfg) + } + if err == icicle_runtime.Success { + err = icicle_{{ .CurvePkg }}.ToMontgomery(dVec) + } + return err +} + +// gpuPolysState holds GPU-resident polynomial data to avoid repeated CPU-GPU transfers. +// Use ensurePolysOnSharedGPU to populate/reuse and freeGPUPolys to release GPU memory. +type gpuPolysState struct { + deviceSlices []icicle_core.DeviceSlice + hostSlices []icicle_core.HostSlice[fr.Element] + polys []*iop.Polynomial + originalForm []iop.Form + polyToIdx map[*iop.Polynomial]int +} + +func (s *instance) sharedGPUStateInitialCap(extra int) int { + base := 16 + len(s.bp) + 2*len(s.commitmentInfo) + if extra > 0 { + base += extra + } + return base +} + +func (s *instance) initSharedGPUState() { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + if s.sharedGPUState != nil { + return + } + initialCap := s.sharedGPUStateInitialCap(0) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } +} + +func (s *instance) releaseSharedGPUState() { + s.gpuStateMu.Lock() + state := s.sharedGPUState + s.sharedGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) releaseLinearizedEvalGPUState() { + s.gpuStateMu.Lock() + state := s.linearizedEvalGPUState + s.linearizedEvalGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) freeIdleTempGPUMemoryOnDevice() { + if s == nil || s.tempGPUMemPool == nil { + return + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +func (s *instance) prepareLinearizedEvalGPUState(source *gpuPolysState) error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.clonePolysOnGPUFromState(source, s.buildLinearizedEvalGPUBatch()) + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) prepareLinearizedEvalGPUStateFromHost() error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.uploadPolysToGPUState(s.buildLinearizedEvalGPUBatch(), "prepareLinearizedEvalGPUStateFromHost") + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop.Polynomial) (*gpuPolysState, error) { + if source == nil { + return nil, fmt.Errorf("clonePolysOnGPUFromState: nil source state") + } + + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("clonePolysOnGPUFromState: empty polynomial batch") + } + + srcSlices := make([]icicle_core.DeviceSlice, len(unique)) + useSource := make([]bool, len(unique)) + for i, p := range unique { + idx, ok := source.polyToIdx[p] + if ok && idx >= 0 && idx < len(source.deviceSlices) && !source.deviceSlices[idx].IsEmpty() { + srcSlices[i] = source.deviceSlices[idx] + useSource[i] = true + } + } + + snapshot := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + snapshot.polys[i] = p + snapshot.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + snapshot.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("clonePolysOnGPUFromState") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for snapshot cloning before handing state to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i := range srcSlices { + if useSource[i] { + dst, allocErr := allocDeviceUninitialized(srcSlices[i].Len()) + if allocErr != nil { + runErr = fmt.Errorf("clonePolysOnGPUFromState: alloc failed at index %d: %w", i, allocErr) + return + } + if err := copyDeviceSliceIntoOnCurrentDevice(dst, srcSlices[i], cfg); err != icicle_runtime.Success { + _ = dst.Free() + runErr = fmt.Errorf("clonePolysOnGPUFromState: device copy failed at index %d: %s", i, err.AsString()) + return + } + snapshot.deviceSlices[i] = dst + continue + } + + coeffs := unique[i].Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("clonePolysOnGPUFromState: empty host coefficients at index %d", i) + return + } + host := icicle_core.HostSliceFromElements(coeffs) + var dst icicle_core.DeviceSlice + host.CopyToDeviceAsync(&dst, cfg.StreamHandle, true) + if dst.IsEmpty() { + runErr = fmt.Errorf("clonePolysOnGPUFromState: host upload failed at index %d", i) + return + } + snapshot.deviceSlices[i] = dst + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(snapshot) + return nil, err + } + return snapshot, nil +} + +func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) (*gpuPolysState, error) { + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("%s: empty polynomial batch", label) + } + + state := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + state.polys[i] = p + state.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + state.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label) + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i, p := range unique { + coeffs := p.Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("%s: empty host coefficients at index %d", label, i) + return + } + state.hostSlices[i] = icicle_core.HostSliceFromElements(coeffs) + state.hostSlices[i].CopyToDeviceAsync(&state.deviceSlices[i], cfg.StreamHandle, true) + if state.deviceSlices[i].IsEmpty() { + runErr = fmt.Errorf("%s: host upload failed at index %d", label, i) + return + } + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(state) + return nil, err + } + return state, nil +} + +func (s *instance) getTempDeviceSlice(n int) icicle_core.DeviceSlice { + if s == nil { + panic("getTempDeviceSlice: nil instance") + } + if s.tempGPUMemPool == nil { + panic("getTempDeviceSlice: temp GPU memory pool is not initialized") + } + return s.tempGPUMemPool.Get(n) +} + +func (s *instance) putTempDeviceSlice(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + if s != nil && s.tempGPUMemPool != nil { + s.tempGPUMemPool.Put(ds, n) + return + } + _ = ds.Free() +} + +func (s *instance) releaseTempGPUMemoryPool() { + if s == nil || s.tempGPUMemPool == nil { + return + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +// ensurePolysOnSharedGPU uploads missing polynomials once and reuses already-uploaded slices. +func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysState, error) { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if s.sharedGPUState == nil { + initialCap := s.sharedGPUStateInitialCap(len(polys)) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } + } + state := s.sharedGPUState + if state == nil { + return nil, fmt.Errorf("ensurePolysOnSharedGPU: shared GPU state is nil") + } + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int, len(polys)) + } + + newIndices := make([]int, 0, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; ok { + continue + } + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, icicle_core.DeviceSlice{}) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + newIndices = append(newIndices, idx) + } + if len(newIndices) == 0 { + return state, nil + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("ensurePolysOnSharedGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for GPU uploads in ensurePolysOnSharedGPU. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + for _, idx := range newIndices { + p := state.polys[idx] + if p == nil { + continue + } + cp := p.Coefficients() + state.hostSlices[idx] = icicle_core.HostSliceFromElements(cp) + state.hostSlices[idx].CopyToDeviceAsync(&state.deviceSlices[idx], cfg.StreamHandle, true) + } + }) + if err := <-done; err != nil { + return nil, err + } + return state, nil +} + +func getStateDeviceSlice(state *gpuPolysState, p *iop.Polynomial, label string) (icicle_core.DeviceSlice, error) { + if state == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil GPU state", label) + } + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil polynomial", label) + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): polynomial is not registered on GPU", label) + } + ds := state.deviceSlices[idx] + if ds.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): empty device slice", label) + } + return ds, nil +} + +// gpuNTTInverseScaleForwardOnDevice performs inverse NTT → scale → forward NTT +// on GPU-resident polynomial data without CPU-GPU transfers for polynomial data. +// The scaling vectors are uploaded each call (they may change between iterations). +func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scalingVector, scalingVectorRev []fr.Element, pk *ProvingKey) error { + if state == nil || len(state.polys) == 0 { + return nil + } + + device := &s.device + var scalingVectorDevice, scalingVectorRevDevice icicle_core.DeviceSlice + + // Upload scaling vectors to GPU + uploadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) + scalingHost := icicle_core.HostSliceFromElements(scalingVector) + scalingHost.CopyToDevice(&scalingVectorDevice, false) + scalingVectorRevDevice = s.getTempDeviceSlice(len(scalingVectorRev)) + scalingRevHost := icicle_core.HostSliceFromElements(scalingVectorRev) + scalingRevHost.CopyToDevice(&scalingVectorRevDevice, false) + + // Convert scaling vectors from Montgomery form to standard form + if err := icicle_{{ .CurvePkg }}.FromMontgomery(scalingVectorDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVector failed: %s", err.AsString()) + return + } + if err := icicle_{{ .CurvePkg }}.FromMontgomery(scalingVectorRevDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVectorRev failed: %s", err.AsString()) + return + } + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return err + } + + doneChans := make([]chan error, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + + done := make(chan error, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Step 1: Inverse NTT + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + return + } + + // Step 2: Scale by vector using GPU vecOps + vecCfg := icicle_core.DefaultVecOpsConfig() + vecCfg.StreamHandle = stream + vecCfg.IsAsync = true + + var scaleDevice icicle_core.DeviceSlice + if layout == iop.Regular { + // After KNR inverse, output is BitReverse → use scalingVectorRev + scaleDevice = scalingVectorRevDevice + } else { + // After KRN inverse, output is Regular → use scalingVector + scaleDevice = scalingVectorDevice + } + + if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + return + } + + // Step 3: Forward NTT to Lagrange + one := icicle_ntt.GetDefaultNttConfig().CosetGen + cfg.CosetGen = one + if layout == iop.Regular { + cfg.Ordering = icicle_core.KRN // BitReverse → Regular + } else { + cfg.Ordering = icicle_core.KNR // Regular → BitReverse + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + return + } + done <- nil + }) + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + if err := <-doneChans[i]; err != nil { + return err + } + } + } + + // Update polynomial metadata: final result is in Lagrange, same layout as original + for _, p := range state.polys { + if p != nil { + p.Basis = iop.Lagrange + } + } + + // Free scaling vectors from device + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + return nil +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +// +// NOTE: The prover hot path should use the state-based API and keep data on device. +// This wrapper exists for compatibility/testing where callers expect host coefficients +// to be materialized after the transform. +func (s *instance) gpuNTTInverseBatch(polys []*iop.Polynomial, pk *ProvingKey) { + if len(polys) == 0 { + return + } + state, err := s.ensurePolysOnSharedGPU(polys) + if err != nil { + panic(fmt.Sprintf("gpuNTTInverseBatch: ensurePolysOnSharedGPU failed: %v", err)) + } + + s.gpuNTTInverseBatchOnState(state, pk) + + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + for _, p := range polys { + if p == nil { + continue + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) || state.deviceSlices[idx].IsEmpty() { + continue + } + cp := p.Coefficients() + host := icicle_core.HostSliceFromElements(cp) + host.CopyFromDevice(&state.deviceSlices[idx]) + copy(cp, ([]fr.Element)(host)) + } + close(done) + }) + <-done +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKey) { + if state == nil || len(state.polys) == 0 { + return + } + + device := &s.device + doneChans := make([]chan struct{}, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + switch p.Basis { + case iop.Canonical: + continue // already in canonical form + case iop.Lagrange, iop.LagrangeCoset: + // Schedule GPU work + done := make(chan struct{}, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Select ordering and coset generator depending on basis and input layout + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + + // Base-domain inverse: + // - Regular input → KNR (output BitReverse) + // - BitReverse input → KRN (output Regular) + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + + // Run NTT inverse directly on the existing device slice (in-place) + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) + } + icicle_runtime.SynchronizeStream(stream) + + // Update metadata inside closure to avoid race + p.Basis = iop.Canonical + if layout == iop.Regular { + p.Layout = iop.BitReverse + } else { + p.Layout = iop.Regular + } + close(done) + }) + default: + panic("unsupported polynomial basis") + } + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + <-doneChans[i] + } + } +} + +// freeGPUPolys releases GPU memory for polynomial data. +func (s *instance) freeGPUPolys(state *gpuPolysState) { + if state == nil { + return + } + + device := &s.device + freeDone := make(chan struct{}) + + icicle_runtime.RunOnDevice(device, func(args ...any) { + for i := range state.polys { + if state.deviceSlices[i].IsEmpty() { + continue + } + _ = state.deviceSlices[i].Free() + } + close(freeDone) + }) + <-freeDone +} + +// gpuMemoryPool manages a pool of reusable device slices to avoid repeated allocations. +// Must be used within RunOnDevice context to ensure thread safety per device. +type gpuMemoryPool struct { + freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice + mu sync.Mutex +} + +// newGPUMemoryPool creates a new GPU memory pool. +func newGPUMemoryPool() *gpuMemoryPool { + return &gpuMemoryPool{ + freeSlices: make(map[int][]icicle_core.DeviceSlice), + } +} + +// Get returns a device slice of the specified size, either from the pool or newly allocated. +func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { + p.mu.Lock() + defer p.mu.Unlock() + + // Check if we have a free slice of this size + if slices, ok := p.freeSlices[n]; ok && len(slices) > 0 { + // Reuse the last slice + slice := slices[len(slices)-1] + p.freeSlices[n] = slices[:len(slices)-1] + return slice + } + + // No free slice available, allocate a new one. + // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + // Free all currently cached (idle) slices to reduce memory pressure. + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + panic(fmt.Sprintf("gpuMemoryPool.Get: allocation failed for size %d after clearing idle cache", n)) +} + +// Put returns a device slice to the pool for reuse instead of freeing it. +func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + + p.mu.Lock() + defer p.mu.Unlock() + + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// FreeAll releases all pooled device slices. +func (p *gpuMemoryPool) FreeAll() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) +} + +// allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. +// Use when the destination is fully overwritten by a kernel. +func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { + var ds icicle_core.DeviceSlice + if _, err := ds.Malloc(int(unsafe.Sizeof(fr.Element{})), n); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("allocDeviceUninitialized: malloc failed for size %d: %s", n, err.AsString()) + } + return ds, nil +} + +// mustAllocDeviceUninitialized is like allocDeviceUninitialized but panics on failure. +// Use only in contexts where error propagation is impractical (e.g. upload helpers). +func mustAllocDeviceUninitialized(n int) icicle_core.DeviceSlice { + ds, err := allocDeviceUninitialized(n) + if err != nil { + panic(err) + } + return ds +} + +// freeDeviceSlice frees a device slice if non-empty and zeroes the pointer. +// Use for directly-allocated slices, NOT pool-allocated ones (use putTempDeviceSlice for those). +func freeDeviceSlice(ds *icicle_core.DeviceSlice) { + if ds != nil && !ds.IsEmpty() { + _ = ds.Free() + *ds = icicle_core.DeviceSlice{} + } +} + +// freeSliceOnDevice frees a device slice on the specified device and blocks +// until complete. Use outside RunOnDevice closures. Zeroes the slice after freeing. +func freeSliceOnDevice(ds *icicle_core.DeviceSlice, device *icicle_runtime.Device) { + if ds == nil || ds.IsEmpty() { + return + } + d := *ds + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + _ = d.Free() + close(done) + }) + <-done + *ds = icicle_core.DeviceSlice{} +} + +// copyDeviceSliceIntoOnCurrentDevice copies src into dst entirely on GPU. +func copyDeviceSliceIntoOnCurrentDevice( + dst, src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if src.IsEmpty() || src.Len() <= 0 || dst.IsEmpty() || dst.Len() < src.Len() { + return icicle_runtime.InvalidArgument + } + src.CheckDevice() + dst.CheckDevice() + + srcElemSize := src.SizeOfElement() + dstElemSize := dst.SizeOfElement() + if srcElemSize <= 0 || dstElemSize <= 0 || srcElemSize != dstElemSize { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(src.Len() * srcElemSize) + if cfg.IsAsync { + return icicle_runtime.CopyAsync(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen, cfg.StreamHandle) + } + _, err := icicle_runtime.Copy(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen) + return err +} + +// zeroDeviceSliceOnCurrentDevice zero-fills dst entirely on GPU. +func zeroDeviceSliceOnCurrentDevice( + dst icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if dst.IsEmpty() || dst.Len() <= 0 { + return icicle_runtime.InvalidArgument + } + dst.CheckDevice() + + elemSize := dst.SizeOfElement() + if elemSize <= 0 { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(dst.Len() * elemSize) + if cfg.IsAsync { + return icicle_runtime.MemSetAsync(dst.AsUnsafePointer(), 0, byteLen, cfg.StreamHandle) + } + return icicle_runtime.MemSet(dst.AsUnsafePointer(), 0, byteLen) +} + +func createAsyncVecOpsConfigOnCurrentDevice(label string) (icicle_core.VecOpsConfig, icicle_runtime.Stream, error) { + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + return icicle_core.VecOpsConfig{}, nil, fmt.Errorf("%s: create stream failed: %s", label, eStream.AsString()) + } + cfg := icicle_core.DefaultVecOpsConfig() + cfg.StreamHandle = stream + cfg.IsAsync = true + return cfg, stream, nil +} + +func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label string) error { + if stream == nil { + return nil + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + _ = icicle_runtime.DestroyStream(stream) + return fmt.Errorf("%s: synchronize stream failed: %s", label, eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success { + return fmt.Errorf("%s: destroy stream failed: %s", label, eDestroy.AsString()) + } + return nil +} + +// makeFinisher returns a closure that synchronizes and destroys the stream, +// then sends the (possibly merged) error to done. Use inside RunOnDevice closures. +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { + return func(runErr error) { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + } +} + +// uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. +// Use this for additions where the vector is already in Montgomery form. +func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + return ds +} + +func uploadScalarMontOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + if cfg.IsAsync { + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&ds, false) + } + return ds +} + +// uploadScalarStd uploads a scalar in STANDARD form (not Montgomery) as a single-element device slice. +// For use with ScalarMulVec: (a*R) * b_std = (a*b)*R +func uploadScalarStd(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + // Convert from Montgomery form to standard form + if err := icicle_{{ .CurvePkg }}.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +func uploadScalarStdOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + ds := uploadScalarMontOnCurrentDevice(scalar, cfg) + // Fallback to sync conversion for compatibility with ICICLE wrappers + // that do not expose *_WithConfig APIs. + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_{{ .CurvePkg }}.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +// uploadVectorStd uploads a vector and converts to standard form. +func uploadVectorStd(vec []fr.Element) icicle_core.DeviceSlice { + ds := mustAllocDeviceUninitialized(len(vec)) + uploadVectorStdInto(&ds, vec) + return ds +} + +// uploadVectorStdInto uploads vec into an existing device slice and converts it to standard form. +// The destination must already be allocated with enough capacity for len(vec) elements. +func uploadVectorStdInto(dst *icicle_core.DeviceSlice, vec []fr.Element) { + cfg := icicle_core.DefaultVecOpsConfig() + uploadVectorStdIntoOnCurrentDevice(dst, vec, cfg) +} + +// uploadVectorStdIntoOnCurrentDevice uploads vec into an existing device slice and converts it +// to standard form while honoring the provided vector-op config/stream. +func uploadVectorStdIntoOnCurrentDevice( + dst *icicle_core.DeviceSlice, + vec []fr.Element, + cfg icicle_core.VecOpsConfig, +) { + host := icicle_core.HostSliceFromElements(vec) + if cfg.IsAsync { + host.CopyToDeviceAsync(dst, cfg.StreamHandle, false) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } else { + host.CopyToDevice(dst, false) + } + // Convert from Montgomery form to standard form + if err := icicle_{{ .CurvePkg }}.FromMontgomery(*dst); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +// uploadVector uploads a vector (keeps Montgomery form for additions). +func uploadVector(vec []fr.Element) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(len(vec)) + host.CopyToDevice(&ds, false) + return ds +} + +// uploadInt64Vector uploads int64 indices to a device slice. +func uploadInt64Vector(vec []int64) icicle_core.DeviceSlice { + cfg := icicle_core.DefaultVecOpsConfig() + return uploadInt64VectorOnCurrentDevice(vec, cfg) +} + +func uploadInt64VectorOnCurrentDevice(vec []int64, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + var ds icicle_core.DeviceSlice + if cfg.IsAsync { + if _, err := ds.MallocAsync(int(unsafe.Sizeof(int64(0))), len(vec), cfg.StreamHandle); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc async failed: %s", err.AsString())) + } + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + return ds + } + if _, err := ds.Malloc(int(unsafe.Sizeof(int64(0))), len(vec)); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc failed: %s", err.AsString())) + } + host.CopyToDevice(&ds, false) + return ds +} + +// toStandardFormInPlace converts a device slice to standard form in-place (modifies the source). +// Use this for temporary vectors that won't be needed in Montgomery form. +func toStandardFormInPlace(src icicle_core.DeviceSlice) { + cfg := icicle_core.DefaultVecOpsConfig() + toStandardFormInPlaceWithCfg(src, cfg) +} + +func toStandardFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_{{ .CurvePkg }}.FromMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +func toMontgomeryFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_{{ .CurvePkg }}.ToMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("ToMontgomery failed: %s", err.AsString())) + } +} + +// multiplyMontgomerySlices multiplies two device slices that are both in Montgomery form. +// It creates a copy of dSlice1Mont, converts the copy to standard form, and then multiplies +// it with dSlice2Mont (which remains in Montgomery form). The result is stored in dResult +// and will be in Montgomery form. +// +// Parameters: +// - dSlice1Mont: first device slice in Montgomery form (not modified) +// - dSlice2Mont: second device slice in Montgomery form (not modified) +// - dResult: destination device slice for the result (must be pre-allocated) +// - state: GPU state with memory pool and vector configuration +// - n: size of the slices +func multiplyMontgomerySlices( + dSlice1Mont, dSlice2Mont icicle_core.DeviceSlice, + dResult icicle_core.DeviceSlice, + state *gpuConstraintEvalState, + vecCfg icicle_core.VecOpsConfig, + n int, +) error { + // Copy dSlice1Mont to standard form + dSlice1Std := state.getTempDeviceSlice(n) + defer state.putTempDeviceSlice(dSlice1Std, n) + + if err := copyDeviceSliceIntoOnCurrentDevice(dSlice1Std, dSlice1Mont, vecCfg); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: device copy failed: %s", err.AsString()) + } + toStandardFormInPlace(dSlice1Std) + + // Multiply: dSlice1Std (standard) * dSlice2Mont (Montgomery) = dResult (Montgomery) + if err := icicle_vecops.VecOp(dSlice1Std, dSlice2Mont, dResult, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: VecOp multiplication failed: %s", err.AsString()) + } + return nil +} + +// gpuConstraintEvalParams holds parameters for GPU constraint evaluation +type gpuConstraintEvalParams struct { + beta fr.Element + gamma fr.Element + alpha fr.Element + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + cs fr.Element // domain1.FrMultiplicativeGen + css fr.Element // cs^2 + cardinalityInv fr.Element + n int + nbBsbGates int +} + +// gpuConstraintEvalState holds intermediate state during constraint evaluation +type gpuConstraintEvalState struct { + // Polynomial device slices (may point to gpuState or allocated buffers) + dL, dR, dO, dZ, dZS icicle_core.DeviceSlice + dQl, dQr, dQm, dQo, dQk icicle_core.DeviceSlice + dS1, dS2, dS3 icicle_core.DeviceSlice + // Intermediate results + dGate, dOrdering, dLocal, dResult icicle_core.DeviceSlice + // Scalar device slices + dGammaScalar icicle_core.DeviceSlice + // Configuration + vecCfg icicle_core.VecOpsConfig + // Helper function to get device slices + getDeviceSlice func(int) icicle_core.DeviceSlice + // Shared prover-level temporary GPU memory pool accessors + getTempDeviceSlice func(int) icicle_core.DeviceSlice + putTempDeviceSlice func(icicle_core.DeviceSlice, int) + // Track allocated polynomial buffers for automatic cleanup + allocatedPolyBuffers []struct { + slice icicle_core.DeviceSlice + size int + } +} + +// allocate allocates a new device slice from the memory pool and tracks it for automatic cleanup. +// Returns the allocated device slice. +func (s *gpuConstraintEvalState) allocate(size int) icicle_core.DeviceSlice { + slice := s.getTempDeviceSlice(size) + s.allocatedPolyBuffers = append(s.allocatedPolyBuffers, struct { + slice icicle_core.DeviceSlice + size int + }{slice, size}) + return slice +} + +// freeAllocatedPolyBuffers returns all allocated polynomial buffers to the memory pool. +// This should be called during cleanup to free all buffers allocated via allocate(). +func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { + for _, buf := range s.allocatedPolyBuffers { + s.putTempDeviceSlice(buf.slice, buf.size) + } + s.allocatedPolyBuffers = nil +} + +// computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. +// Returns device slices for the blinding polynomials. +func computeBlindingPolynomials( + n int, + twiddles0 []fr.Element, + bp []*iop.Polynomial, +) (dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice) { + blindL := make([]fr.Element, n) + blindR := make([]fr.Element, n) + blindO := make([]fr.Element, n) + blindZ := make([]fr.Element, n) + blindZS := make([]fr.Element, n) // ZS uses shifted index + + // TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove this. Since we don't normally use blindings, we will + // not make this optimization. + utils.Parallelize(n, func(start, end int) { + for i := start; i < end; i++ { + blindL[i] = bp[id_Bl].Evaluate(twiddles0[i]) + blindR[i] = bp[id_Br].Evaluate(twiddles0[i]) + blindO[i] = bp[id_Bo].Evaluate(twiddles0[i]) + blindZ[i] = bp[id_Bz].Evaluate(twiddles0[i]) + blindZS[i] = bp[id_Bz].Evaluate(twiddles0[(i+1)%n]) + } + }) + + dBlindL = uploadVector(blindL) + dBlindR = uploadVector(blindR) + dBlindO = uploadVector(blindO) + dBlindZ = uploadVector(blindZ) + dBlindZS = uploadVector(blindZS) + + return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS +} + +// applyBlindingToPolynomials applies blinding to polynomials L, R, O, Z, ZS. +// Allocates new buffers for L, R, O, Z (tracked for cleanup) and modifies ZS in-place. +// The original slices in gpuState remain unchanged. +func applyBlindingToPolynomials( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice, +) error { + // L' = L + blindL (allocate new buffer, tracked for cleanup) + dLBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dL, dBlindL, dLBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add L failed: %s", err.AsString()) + } + state.dL = dLBlinded + + // R' = R + blindR (allocate new buffer, tracked for cleanup) + dRBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dR, dBlindR, dRBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add R failed: %s", err.AsString()) + } + state.dR = dRBlinded + + // O' = O + blindO (allocate new buffer, tracked for cleanup) + dOBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dO, dBlindO, dOBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add O failed: %s", err.AsString()) + } + state.dO = dOBlinded + + // Z' = Z + blindZ (allocate new buffer, tracked for cleanup) + dZBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dZ, dBlindZ, dZBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add Z failed: %s", err.AsString()) + } + state.dZ = dZBlinded + + // ZS' = ZS + blindZS + // Note: dZS is a temporary buffer created inside gpuEvaluateConstraints, + // so it's safe to modify it in-place. + if err := icicle_vecops.VecOp(state.dZS, dBlindZS, state.dZS, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) + } + + // Free blinding vectors - no longer needed after creating blinded polynomials + dBlindL.Free() + dBlindR.Free() + dBlindO.Free() + dBlindZ.Free() + dBlindZS.Free() + return nil +} + +// scaleSVectorsByBeta scales S1, S2, S3 by beta. +// Allocates new buffers for S1, S2, S3 (tracked for cleanup). +func scaleSVectorsByBeta( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, +) error { + // S1' = S1 * beta (need to scale S1, S2, S3 by beta for ordering constraint) + // Use standard form for beta so: S1_mont * beta_std = (S1*beta)_mont + dBetaStd := uploadScalarStd(params.beta) + + // S1' = S1 * beta (allocate new buffer, tracked for cleanup) + dS1Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dS1Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S1 failed: %s", err.AsString()) + } + state.dS1 = dS1Scaled + + // S2' = S2 * beta (allocate new buffer, tracked for cleanup) + dS2Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dS2Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S2 failed: %s", err.AsString()) + } + state.dS2 = dS2Scaled + + // S3' = S3 * beta (allocate new buffer, tracked for cleanup) + dS3Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dS3Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S3 failed: %s", err.AsString()) + } + state.dS3 = dS3Scaled + + // Free dBetaStd - no longer needed after scaling S vectors + dBetaStd.Free() + return nil +} + +// computeGateConstraint computes the gate constraint. +// Returns dGate device slice. +func computeGateConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // gate = Ql*L' + Qr*R' + Qm*L'*R' + Qo*O' + Qk + sum(Qci*Pi) + // We use multiplyMontgomerySlices for all poly×poly multiplications. + // Note: dL, dR, dO are used later in ordering constraint, so we preserve them. + + dGate := state.getTempDeviceSlice(params.n) + dTmp := state.getTempDeviceSlice(params.n) + + // Ql * L' + if err := multiplyMontgomerySlices(state.dQl, state.dL, dGate, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Ql*L: %w", err) + } + + // + Qr * R' + if err := multiplyMontgomerySlices(state.dQr, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qr*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qr*R failed: %s", err.AsString()) + } + + // + Qm * L' * R' (need two multiplications) + // First: Qm * L' = dTmp (Montgomery) + if err := multiplyMontgomerySlices(state.dQm, state.dL, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qm*L: %w", err) + } + // Second: dTmp (Montgomery) * R' (Montgomery) + if err := multiplyMontgomerySlices(dTmp, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: (Qm*L)*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qm*L*R failed: %s", err.AsString()) + } + + // + Qo * O' + if err := multiplyMontgomerySlices(state.dQo, state.dO, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qo*O: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qo*O failed: %s", err.AsString()) + } + + // + Qk + if err := icicle_vecops.VecOp(dGate, state.dQk, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qk failed: %s", err.AsString()) + } + + // + BSB gates: sum(Qci[2*i] * Qci[2*i+1]) + for i := 0; i < params.nbBsbGates; i++ { + origQci0 := state.getDeviceSlice(id_Qci + 2*i) + origQci1 := state.getDeviceSlice(id_Qci + 2*i + 1) + if !origQci0.IsEmpty() && !origQci1.IsEmpty() { + // Use helper to multiply Qci0 * Qci1 without modifying original values + if err := multiplyMontgomerySlices(origQci0, origQci1, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qci[%d]: %w", i, err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qci[%d] failed: %s", i, err.AsString()) + } + } + } + + // Return temporary buffer to pool - no longer needed after Step 3 + state.putTempDeviceSlice(dTmp, params.n) + + return dGate, nil +} + +// computeOrderingConstraint computes the ordering constraint. +// Returns dOrdering device slice. +func computeOrderingConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dTwiddles0 icicle_core.DeviceSlice, // twiddles0 already on GPU + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // This is complex: involves ID computation, gamma, beta, Z, ZS, S1, S2, S3 + // id = twiddles[i] * coset * beta + // a = gamma + L' + id + // b = gamma + R' + id*cs + // c = gamma + O' + id*css + // r = a * b * c * Z' + // + // a2 = gamma + L' + S1*beta + // b2 = gamma + R' + S2*beta + // c2 = gamma + O' + S3*beta + // l = a2 * b2 * c2 * ZS' + // + // ordering = l - r + + // Compute ID vector: twiddles * coset * beta (computed on GPU) + // dTwiddles0 is already on GPU (passed as parameter, don't free it here) + + // Compute coset * beta on CPU, then upload as scalar in standard form + var cosetTimesBeta fr.Element + cosetTimesBeta.Mul(¶ms.coset, ¶ms.beta) + dCosetTimesBetaStd := uploadScalarStd(cosetTimesBeta) + dBetaStd := uploadScalarStd(params.beta) + + // Multiply twiddles0 by cosetTimesBeta on GPU: dID = (cosetTimesBeta * twiddles0) * R (Montgomery form) + dID := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dCosetTimesBetaStd, dTwiddles0, dID, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec coset*beta*twiddles failed: %s", err.AsString()) + } + + // Free temporary device slice (dTwiddles0 is owned by caller, don't free it) + dCosetTimesBetaStd.Free() + + // id * cs - use standard form for cs + dIDcs := state.getTempDeviceSlice(params.n) + dCsStd := uploadScalarStd(params.cs) + if err := icicle_vecops.ScalarMulVec(dCsStd, dID, dIDcs, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*cs failed: %s", err.AsString()) + } + + // id * css - use standard form for css + dIDcss := state.getTempDeviceSlice(params.n) + dCssStd := uploadScalarStd(params.css) + if err := icicle_vecops.ScalarMulVec(dCssStd, dID, dIDcss, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*css failed: %s", err.AsString()) + } + + // a = gamma + L' + id (dL now contains L' after in-place blinding) + dA := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dID, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a+id failed: %s", err.AsString()) + } + + // b = gamma + R' + id*cs (dR now contains R' after in-place blinding) + dB := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dIDcs, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b+id*cs failed: %s", err.AsString()) + } + + // c = gamma + O' + id*css (dO now contains O' after in-place blinding) + dC := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dIDcss, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c+id*css failed: %s", err.AsString()) + } + + // Return to pool: dID, dIDcs, dIDcss - no longer needed after computing a, b, c + state.putTempDeviceSlice(dID, params.n) + state.putTempDeviceSlice(dIDcs, params.n) + state.putTempDeviceSlice(dIDcss, params.n) + dCsStd.Free() + dCssStd.Free() + + // r = a * b * c * Z' (dZ now contains Z' after in-place blinding) + // For chain multiplication, convert operands to std form in-place when possible + dR_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a*b failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dR_ord, dC, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*c failed: %s", err.AsString()) + } + // Convert dR_ord to standard form in-place (temporary result, dZ needs to be preserved) + toStandardFormInPlace(dR_ord) + if err := icicle_vecops.VecOp(dR_ord, state.dZ, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*Z failed: %s", err.AsString()) + } + + // Reuse dA, dB, dC for a2, b2, c2 instead of freeing and reallocating. + // To reduce peak memory, we scale S vectors by beta on-demand through a single temp buffer. + dScaledS := state.getTempDeviceSlice(params.n) + + // a2 = gamma + L' + S1*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L (a2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S1 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dScaledS, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2+S1*beta failed: %s", err.AsString()) + } + + // b2 = gamma + R' + S2*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R (b2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S2 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dScaledS, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b2+S2*beta failed: %s", err.AsString()) + } + + // c2 = gamma + O' + S3*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O (c2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S3 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dScaledS, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c2+S3*beta failed: %s", err.AsString()) + } + + // Free dGammaScalar - no longer needed after computing a2, b2, c2 + // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. + state.dGammaScalar.Free() + state.putTempDeviceSlice(dScaledS, params.n) + dBetaStd.Free() + + // l = a2 * b2 * c2 * ZS' (dZS now contains ZS' after in-place blinding) + dL_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2*b2 failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dL_ord, dC, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*c2 failed: %s", err.AsString()) + } + // Convert dZS to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(state.dZS) + if err := icicle_vecops.VecOp(dL_ord, state.dZS, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*ZS failed: %s", err.AsString()) + } + + // Return temporary buffers to pool - no longer needed after computing l + state.putTempDeviceSlice(dA, params.n) + state.putTempDeviceSlice(dB, params.n) + state.putTempDeviceSlice(dC, params.n) + state.putTempDeviceSlice(state.dZS, params.n) + state.dZS = icicle_core.DeviceSlice{} + + // ordering = l - r, reuse dL_ord as the final ordering vector + if err := icicle_vecops.VecOp(dL_ord, dR_ord, dL_ord, vecCfg, icicle_core.Sub); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l-r failed: %s", err.AsString()) + } + + // Return dR_ord to pool - no longer needed after computing ordering + state.putTempDeviceSlice(dR_ord, params.n) + + // Return dL_ord as ordering (caller is responsible for freeing) + return dL_ord, nil +} + +// computeLocalConstraint computes the local constraint. +// Returns dLocal device slice. +func computeLocalConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dPrecomputedDenominators icicle_core.DeviceSlice, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // local = (Z' - 1) * LagrangeOne + // where LagrangeOne[i] = cosetExpMinusOne * cardinalityInv / (coset*twiddles0[i] - 1) + + if dPrecomputedDenominators.IsEmpty() || dPrecomputedDenominators.Len() < params.n { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: invalid denominator device slice size %d, expected at least %d", dPrecomputedDenominators.Len(), params.n) + } + + // Compute LagrangeOne on device. dPrecomputedDenominators is already in + // Montgomery form after batch inversion; ScalarMulVec expects the scalar in + // standard form and preserves a Montgomery vector result. + var lagrangeCoeff fr.Element + lagrangeCoeff.Mul(¶ms.cosetExponentiatedToNMinusOne, ¶ms.cardinalityInv) + dLagrangeCoeffStd := uploadScalarStdOnCurrentDevice(lagrangeCoeff, vecCfg) + dLagrangeOneStd := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dLagrangeCoeffStd, dPrecomputedDenominators, dLagrangeOneStd, vecCfg); err != icicle_runtime.Success { + dLagrangeCoeffStd.Free() + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarMulVec lagrangeOne failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dLagrangeOneStd, vecCfg) + dLagrangeCoeffStd.Free() + + // Z' - 1 using ScalarAddVec with minus one + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneScalar := uploadScalarMont(minusOne) + + dZMinusOne := state.getTempDeviceSlice(params.n) + // dZ now contains Z' after in-place blinding + if err := icicle_vecops.ScalarAddVec(dMinusOneScalar, state.dZ, dZMinusOne, vecCfg); err != icicle_runtime.Success { + dMinusOneScalar.Free() + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarAddVec Z-1 failed: %s", err.AsString()) + } + + // Free dMinusOneScalar - no longer needed after computing Z' - 1 + // Note: dZ is part of gpuState and will be freed later + dMinusOneScalar.Free() + + // local = (Z' - 1) * LagrangeOne_std + dLocal := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.VecOp(dZMinusOne, dLagrangeOneStd, dLocal, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: VecOp (Z-1)*LagrangeOne failed: %s", err.AsString()) + } + + // Return temporary buffers to pool + state.putTempDeviceSlice(dZMinusOne, params.n) + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + + return dLocal, nil +} + +// createGetDeviceSliceFunc creates a function to get device slices for polynomials. +// It returns a function that maps polynomial indices to their device slices. +func (s *instance) polyByID(polyIdx int) *iop.Polynomial { + switch polyIdx { + case id_L: + return s.polyL + case id_R: + return s.polyR + case id_O: + return s.polyO + case id_Z: + return s.polyZ + case id_ZS: + return s.polyZS + case id_Ql: + return s.trace.Ql + case id_Qr: + return s.trace.Qr + case id_Qm: + return s.trace.Qm + case id_Qo: + return s.trace.Qo + case id_Qk: + return s.polyQk + case id_S1: + return s.trace.S1 + case id_S2: + return s.trace.S2 + case id_S3: + return s.trace.S3 + default: + if polyIdx < id_Qci { + return nil + } + offset := polyIdx - id_Qci + i := offset / 2 + if i < 0 { + return nil + } + if offset%2 == 0 { + if i < len(s.trace.Qcp) { + return s.trace.Qcp[i] + } + return nil + } + if i < len(s.cCommitments) { + return s.cCommitments[i] + } + return nil + } +} + +func createGetDeviceSliceFunc( + gpuState *gpuPolysState, + polyToIdx map[*iop.Polynomial]int, + resolvePoly func(int) *iop.Polynomial, +) func(int) icicle_core.DeviceSlice { + return func(polyIdx int) icicle_core.DeviceSlice { + p := resolvePoly(polyIdx) + if p == nil { + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d is nil", polyIdx)) + } + if idx, ok := polyToIdx[p]; ok { + return gpuState.deviceSlices[idx] + } + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d (ptr=%p) not found in polyToIdx (map has %d entries)", polyIdx, p, len(polyToIdx))) + } +} + +// initializeConstraintEvalState initializes the GPU constraint evaluation state. +// It sets up all device slices. +// The slices in gpuState are treated as read-only; helper functions will allocate +// separate working buffers whenever they need to modify data. +func initializeConstraintEvalState( + getDeviceSlice func(int) icicle_core.DeviceSlice, + getTempDeviceSlice func(int) icicle_core.DeviceSlice, + putTempDeviceSlice func(icicle_core.DeviceSlice, int), +) *gpuConstraintEvalState { + vecCfg := icicle_core.DefaultVecOpsConfig() + + state := &gpuConstraintEvalState{ + dL: getDeviceSlice(id_L), + dR: getDeviceSlice(id_R), + dO: getDeviceSlice(id_O), + dZ: getDeviceSlice(id_Z), + dQl: getDeviceSlice(id_Ql), + dQr: getDeviceSlice(id_Qr), + dQm: getDeviceSlice(id_Qm), + dQo: getDeviceSlice(id_Qo), + dQk: getDeviceSlice(id_Qk), + dS1: getDeviceSlice(id_S1), + dS2: getDeviceSlice(id_S2), + dS3: getDeviceSlice(id_S3), + vecCfg: vecCfg, + getDeviceSlice: getDeviceSlice, + getTempDeviceSlice: getTempDeviceSlice, + putTempDeviceSlice: putTempDeviceSlice, + } + + return state +} + +// gpuEvaluateConstraints evaluates all PLONK constraints on GPU. +// It takes polynomials already on GPU (via gpuState), computes blinding polynomial evaluations, +// and evaluates gate, ordering, and local constraints entirely on GPU. +// If result is non-nil, it downloads into result and returns an empty device slice. +// If result is nil, it returns a persistent device slice with the result. +// TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove 'twiddles0'. Since we don't +// normally use blindings, we will not make this optimization. +func (s *instance) gpuEvaluateConstraints( + gpuState *gpuPolysState, + params gpuConstraintEvalParams, + twiddles0 []fr.Element, // CPU vector for computeBlindingPolynomials + dTwiddles0 icicle_core.DeviceSlice, // GPU vector for computeOrderingConstraint + dPrecomputedDenominators icicle_core.DeviceSlice, + bp []*iop.Polynomial, // blinding polynomials (already scaled for this iteration) + result []fr.Element, +) (icicle_core.DeviceSlice, error) { + if gpuState == nil || len(gpuState.polys) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("gpuState is nil or empty") + } + + n := params.n + device := &s.device + + // Create a map from polynomial to its index in gpuState + polyToIdx := make(map[*iop.Polynomial]int) + for i, p := range gpuState.polys { + if p != nil { + polyToIdx[p] = i + } + } + + // Get device slices for the polynomials we need. + getDeviceSlice := createGetDeviceSliceFunc(gpuState, polyToIdx, s.polyByID) + + done := make(chan error, 1) + var resultOnDevice icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(device, func(args ...any) { + state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). + // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). + state.dGammaScalar = uploadScalarMont(params.gamma) + + state.dZS = state.getTempDeviceSlice(n) + if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + return + } + + // Step 1: Compute and apply blinding polynomial evaluations (if enabled) + if useBlinding { + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + done <- err + return + } + } + + // Step 2-4: Compute gate, ordering, and local constraints sequentially on a + // single synchronous stream, folding them into dResult as + // gate + alpha*ordering + alpha^2*local. Computing one family at a time + // keeps peak device allocation low, and sequential is not a compromise: + // the family kernels are memory-bandwidth-bound and each already saturates + // the device, so the parallel three-stream variant this replaces measured + // identical timings (111ms/iteration at n=2^23) — while racing on the + // shared temp-slice pool and lazily materialized inputs (it corrupted the + // numerator at every circuit size). + seqVecCfg := state.vecCfg + seqVecCfg.IsAsync = false + + // Compute ordering first to minimize peak memory before gate/local allocations. + var err error + state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + return + } + + // dResult = alpha * ordering + state.dResult = state.getTempDeviceSlice(params.n) + dAlphaStd := uploadScalarStd(params.alpha) + if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { + dAlphaStd.Free() + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + return + } + dAlphaStd.Free() + state.putTempDeviceSlice(state.dOrdering, params.n) + + // dResult += alpha^2 * local + state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + return + } + var alphaSquared fr.Element + alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) + dAlphaSquaredStd := uploadScalarStd(alphaSquared) + dTmp := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + return + } + dAlphaSquaredStd.Free() + if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dLocal, params.n) + state.putTempDeviceSlice(dTmp, params.n) + + // dResult += gate + state.dGate, err = computeGateConstraint(state, params, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + return + } + if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dGate, params.n) + + // Step 5: materialize result either on host or as a persistent device slice. + if result != nil { + resultHost := icicle_core.HostSliceFromElements(result) + resultHost.CopyFromDevice(&state.dResult) + } else { + resultOnDevice = s.getTempDeviceSlice(params.n) + if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // Return dResult pool slice after materialization. + state.putTempDeviceSlice(state.dResult, params.n) + + // Return all allocated polynomial buffers to the pool. + // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). + // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are + // returned to pool inside computeGateConstraint. dQk and the original gpuState slices + // are owned by gpuState and will be freed separately. + state.freeAllocatedPolyBuffers() + + done <- nil + }) + + err := <-done + + if err != nil { + if !resultOnDevice.IsEmpty() { + s.putTempDeviceSlice(resultOnDevice, resultOnDevice.Len()) + } + return icicle_core.DeviceSlice{}, err + } + return resultOnDevice, nil +} + +func evaluateBlinded(p, bp *iop.Polynomial, zeta fr.Element) fr.Element { + // Get the size of the polynomial + n := big.NewInt(int64(p.Size())) + + var pEvaluatedAtZeta fr.Element + + // Evaluate the polynomial and blinded polynomial at zeta + chP := make(chan struct{}, 1) + go func() { + pEvaluatedAtZeta = p.Evaluate(zeta) + close(chP) + }() + + bpEvaluatedAtZeta := bp.Evaluate(zeta) + + // Multiply the evaluated blinded polynomial by tempElement + var t fr.Element + one := fr.One() + t.Exp(zeta, n).Sub(&t, &one) + bpEvaluatedAtZeta.Mul(&bpEvaluatedAtZeta, &t) + + // Add the evaluated polynomial and the evaluated blinded polynomial + <-chP + pEvaluatedAtZeta.Add(&pEvaluatedAtZeta, &bpEvaluatedAtZeta) + + // Return the result + return pEvaluatedAtZeta +} + +// /!\ modifies the size +func getBlindedCoefficients(p, bp *iop.Polynomial) []fr.Element { + cp := p.Coefficients() + cbp := bp.Coefficients() + cp = append(cp, cbp...) + for i := 0; i < len(cbp); i++ { + cp[i].Sub(&cp[i], &cbp[i]) + } + return cp +} + +// getNonBlindedCoefficients returns a padded copy of polynomial coefficients +// to match the size they would have with blinding enabled. +// The padding size is blindingOrder+1 (e.g., order 2 → 3 coefficients). +func getNonBlindedCoefficients(p *iop.Polynomial, blindingOrder int) []fr.Element { + cp := p.Coefficients() + padded := make([]fr.Element, len(cp)+blindingOrder+1) + copy(padded, cp) + return padded +} + +// commits to a polynomial of the form b*(Xⁿ-1) where b is of small degree +func commitBlindingFactor(n int, b *iop.Polynomial, key kzg.ProvingKey) curve.G1Affine { + cp := b.Coefficients() + np := b.Size() + + // lo + var tmp curve.G1Affine + tmp.MultiExp(key.G1[:np], cp, ecc.MultiExpConfig{}) + + // hi + var res curve.G1Affine + res.MultiExp(key.G1[n:n+np], cp, ecc.MultiExpConfig{}) + res.Sub(&res, &tmp) + return res +} + +// return a random polynomial of degree n, if n==-1 cancel the blinding +func getRandomPolynomial(n int) *iop.Polynomial { + var a []fr.Element + if n == -1 { + a := make([]fr.Element, 1) + a[0].SetZero() + } else { + a = make([]fr.Element, n+1) + for i := 0; i <= n; i++ { + a[i].SetRandom() + } + } + res := iop.NewPolynomial(&a, iop.Form{ + Basis: iop.Canonical, Layout: iop.Regular}) + return res +} + +func coefficients(p []*iop.Polynomial) [][]fr.Element { + res := make([][]fr.Element, len(p)) + for i, pI := range p { + res[i] = pI.Coefficients() + } + return res +} + +func (s *instance) freeGPUQuotient(quotient *gpuQuotientPolynomial) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return + } + s.putTempDeviceSlice(quotient.coeffs, quotient.coeffs.Len()) + quotient.coeffs = icicle_core.DeviceSlice{} + quotient.size = 0 +} + +// commitToQuotientGPUFromDevice commits H1/H2/H3 directly from device memory. +// For StatisticalZK=true we materialize adjusted device vectors for h1/h2/h3 +// and commit those without downloading quotient coefficients to host. +// prepareStatisticalZKQuotientShards constructs blinded quotient polynomial +// shards h1, h2, h3 on the GPU for the Statistical ZK path. Each shard is +// randomized so that the quotient split h = h1 + X^(n+2)*h2 + X^(2(n+2))*h3 +// hides the original polynomial. +// +// Caller is responsible for returning dH1, dH2, dH3 to the temp pool: +// - dH1 and dH2 have size nPlus2+1 +// - dH3 has size nPlus2 +func (s *instance) prepareStatisticalZKQuotientShards( + h1Device, h2Device, h3Device icicle_core.DeviceSlice, + nPlus2 int, +) (dH1, dH2, dH3 icicle_core.DeviceSlice, err error) { + nPlus3 := nPlus2 + 1 + + prepareDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + dH1 = s.getTempDeviceSlice(nPlus3) + dH2 = s.getTempDeviceSlice(nPlus3) + dH3 = s.getTempDeviceSlice(nPlus2) + + // h1 = base h1 with extra randomizer coefficient at degree n+2. + dH1Prefix := (&dH1).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH1Prefix, h1Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h1 failed: %s", e.AsString()) + return + } + dH1Tail := (&dH1).Range(nPlus2, nPlus3, false) + r0Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[0]}) + r0Host.CopyToDevice(&dH1Tail, false) + + // h2 = base h2 with first coefficient adjusted by -r0 and tail = r1. + dH2Prefix := (&dH2).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH2Prefix, h2Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h2 failed: %s", e.AsString()) + return + } + dH2First := (&dH2).Range(0, 1, false) + var negR0 fr.Element + negR0.Neg(&s.quotientShardsRandomizers[0]) + dNegR0 := uploadScalarMont(negR0) + if e := icicle_vecops.ScalarAddVec(dNegR0, dH2First, dH2First, cfg); e != icicle_runtime.Success { + _ = dNegR0.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h2[0] failed: %s", e.AsString()) + return + } + _ = dNegR0.Free() + dH2Tail := (&dH2).Range(nPlus2, nPlus3, false) + r1Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[1]}) + r1Host.CopyToDevice(&dH2Tail, false) + + // h3 = base h3 with first coefficient adjusted by -r1. + if e := copyDeviceSliceIntoOnCurrentDevice(dH3, h3Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h3 failed: %s", e.AsString()) + return + } + dH3First := (&dH3).Range(0, 1, false) + var negR1 fr.Element + negR1.Neg(&s.quotientShardsRandomizers[1]) + dNegR1 := uploadScalarMont(negR1) + if e := icicle_vecops.ScalarAddVec(dNegR1, dH3First, dH3First, cfg); e != icicle_runtime.Success { + _ = dNegR1.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h3[0] failed: %s", e.AsString()) + return + } + _ = dNegR1.Free() + prepareDone <- nil + }) + if err := <-prepareDone; err != nil { + if !dH1.IsEmpty() { + s.putTempDeviceSlice(dH1, nPlus3) + } + if !dH2.IsEmpty() { + s.putTempDeviceSlice(dH2, nPlus3) + } + if !dH3.IsEmpty() { + s.putTempDeviceSlice(dH3, nPlus2) + } + return icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, err + } + return dH1, dH2, dH3, nil +} + +func (s *instance) commitToQuotientGPUFromDevice(quotient *gpuQuotientPolynomial) error { + if quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("commitToQuotientGPUFromDevice: empty quotient") + } + + nPlus2 := int(s.domain0.Cardinality) + 2 + required := 3 * nPlus2 + if quotient.coeffs.Len() < required { + return fmt.Errorf("commitToQuotientGPUFromDevice: quotient too small: got %d need >= %d", quotient.coeffs.Len(), required) + } + + h1Device := ("ient.coeffs).Range(0, nPlus2, false) + h2Device := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3Device := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + if s.opt.StatisticalZK { + nPlus3 := nPlus2 + 1 + dH1, dH2, dH3, err := s.prepareStatisticalZKQuotientShards(h1Device, h2Device, h3Device, nPlus2) + if err != nil { + return err + } + defer s.putTempDeviceSlice(dH1, nPlus3) + defer s.putTempDeviceSlice(dH2, nPlus3) + defer s.putTempDeviceSlice(dH3, nPlus2) + + c0, err := commitOnGPUCanonicalDevice(dH1, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(dH2, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(dH3, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + return nil + } + + // Commit sequentially to avoid 3-way concurrent MSM memory spikes. + c0, err := commitOnGPUCanonicalDevice(h1Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(h2Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(h3Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + + return nil +} + +func (s *instance) inverseAndMergeShards( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (icicle_core.DeviceSlice, error) { + if gpuNumerator == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: nil numerator") + } + n := gpuNumerator.n + rho := gpuNumerator.rho + if n <= 0 || rho <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: invalid n=%d rho=%d", n, rho) + } + if len(gpuNumerator.shards) != rho { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), rho) + } + for i := 0; i < rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard %d is empty", i) + } + } + + expo := big.NewInt(int64(n)) + + // Per-shard cosets: c_i = c * g^i where c=FrMultiplicativeGen, g=Generator. + cosets := make([]fr.Element, rho) + cosets[0].Set(&domains[1].FrMultiplicativeGen) + for i := 1; i < rho; i++ { + cosets[i].Mul(&cosets[i-1], &domains[1].Generator) + } + invCosets := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + invCosets[i].Inverse(&cosets[i]) + } + + // ν = g^n is a rho-th root, used for the rho-point inverse DFT in combine. + var nu, nuInv fr.Element + nu.Exp(domains[1].Generator, expo) + nuInv.Inverse(&nu) + nuInvPowers := make([]fr.Element, rho) + nuInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + nuInvPowers[i].Mul(&nuInvPowers[i-1], &nuInv) + } + + // cN = c^n. Recover original coefficient blocks by scaling with cN^{-t}. + var cN, cNInv fr.Element + cN.Exp(domains[1].FrMultiplicativeGen, expo) + cNInv.Inverse(&cN) + cNInvPowers := make([]fr.Element, rho) + cNInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + cNInvPowers[i].Mul(&cNInvPowers[i-1], &cNInv) + } + + // Each shard inverse contributes a 1/n factor; apply extra 1/rho. + var rhoFr, invRho fr.Element + rhoFr.SetUint64(uint64(rho)) + invRho.Inverse(&rhoFr) + combineScales := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + combineScales[i].Mul(&invRho, &cNInvPowers[i]) + } + + totalSize := rho * n + done := make(chan error, 1) + var dMerged icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish := makeFinisher(stream, "inverseAndMergeShards", done) + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + // KNR is often faster than KNN; we restore regular output explicitly + // by bit-reversing each shard after the inverse NTT. + cfgNtt.Ordering = icicle_core.KNR + + ext := config_extension.Create() + defer config_extension.Delete(ext) + alg := nttAlgorithmFromEnv("ICICLE_DIVIDE_BY_ZH_NTT_ALGO", icicle_core.MixedRadix) + ext.SetInt(icicle_core.CUDA_NTT_ALGORITHM, int(alg)) + cfgNtt.Ext = ext.AsUnsafePointer() + + // Step 1: inverse NTT each shard without coset, reorder to regular, + // then unscale by (c*g^i)^t to recover the coset-inverse equivalent. + nn := uint64(64 - bits.TrailingZeros64(uint64(n))) + invPowers := make([]fr.Element, n) + for i := 0; i < rho; i++ { + if nttErr := icicle_ntt.Ntt(gpuNumerator.shards[i], icicle_core.KInverse, &cfgNtt, gpuNumerator.shards[i]); nttErr != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: inverse NTT failed at shard %d: %s", i, nttErr.AsString())) + return + } + + // KNR outputs bit-reversed coefficients. Reorder back to regular. + dRegular := s.getTempDeviceSlice(n) + mergeErr := icicle_vecops.MergeShardsBitReverse( + []icicle_core.DeviceSlice{gpuNumerator.shards[i]}, + n, + nn, + dRegular, + cfgVec, + ) + if mergeErr != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: reorder failed at shard %d: %s", i, mergeErr.AsString())) + return + } + // gpuNumerator.shards[i] is returned to pool and replaced; wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(gpuNumerator.shards[i], n) + gpuNumerator.shards[i] = dRegular + + fft.BuildExpTable(invCosets[i], invPowers) + dInvPowers := s.getTempDeviceSlice(n) + uploadVectorStdIntoOnCurrentDevice(&dInvPowers, invPowers, cfgVec) + if e := icicle_vecops.VecOp(gpuNumerator.shards[i], dInvPowers, gpuNumerator.shards[i], cfgVec, icicle_core.Mul); e != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: normalize shard %d failed: %s", i, e.AsString())) + return + } + // dInvPowers is temporary and returned to pool each iteration. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(dInvPowers, n) + } + + // Step 2: combine shard results via a size-rho inverse DFT per coefficient index. + dMerged = s.getTempDeviceSlice(totalSize) + keepMerged := false + defer func() { + if !keepMerged && !dMerged.IsEmpty() { + s.putTempDeviceSlice(dMerged, totalSize) + } + }() + + dTmp := s.getTempDeviceSlice(n) + defer func() { + if !dTmp.IsEmpty() { + s.putTempDeviceSlice(dTmp, n) + } + }() + + dNuWeights := make([]icicle_core.DeviceSlice, rho) + dCombineScales := make([]icicle_core.DeviceSlice, rho) + for i := 0; i < rho; i++ { + dNuWeights[i] = uploadScalarStdOnCurrentDevice(nuInvPowers[i], cfgVec) + dCombineScales[i] = uploadScalarStdOnCurrentDevice(combineScales[i], cfgVec) + } + defer func() { + for i := 0; i < rho; i++ { + if !dNuWeights[i].IsEmpty() { + _ = dNuWeights[i].Free() + } + if !dCombineScales[i].IsEmpty() { + _ = dCombineScales[i].Free() + } + } + }() + + for t := 0; t < rho; t++ { + outT := (&dMerged).Range(t*n, (t+1)*n, false) + if e := copyDeviceSliceIntoOnCurrentDevice(outT, gpuNumerator.shards[0], cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: init out[%d] failed: %s", t, e.AsString())) + return + } + for i := 1; i < rho; i++ { + weightIdx := (i * t) % rho + if weightIdx == 0 { + if e := icicle_vecops.VecOp(outT, gpuNumerator.shards[i], outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: add shard %d to out[%d] failed: %s", i, t, e.AsString())) + return + } + continue + } + if e := icicle_vecops.ScalarMulVec(dNuWeights[weightIdx], gpuNumerator.shards[i], dTmp, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: weight shard %d for out[%d] failed: %s", i, t, e.AsString())) + return + } + if e := icicle_vecops.VecOp(outT, dTmp, outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: accumulate shard %d into out[%d] failed: %s", i, t, e.AsString())) + return + } + } + if e := icicle_vecops.ScalarMulVec(dCombineScales[t], outT, outT, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: scale out[%d] failed: %s", t, e.AsString())) + return + } + } + keepMerged = true + finish(nil) + }) + + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dMerged, nil +} + +func (s *instance) divideByZHOnGPU( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (_ *gpuQuotientPolynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("divideByZHOnGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("divideByZHOnGPU: invalid numerator dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + for i := range gpuNumerator.shards { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("divideByZHOnGPU: shard %d is empty", i) + } + } + + rho := int(domains[1].Cardinality / domains[0].Cardinality) + if rho != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: rho mismatch domains=%d numerator=%d", rho, gpuNumerator.rho) + } + + // Evaluate 1/(X^n-1) over the large-domain coset values used by this quotient. + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + + // In bit-reversed merged layout, each shard maps to a fixed (iRev % rho) bucket. + // So we can divide by Z_H by scaling each shard with its corresponding inverse. + scaleDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + for i := 0; i < gpuNumerator.rho; i++ { + dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) + vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if vecErr != icicle_runtime.Success { + finish(fmt.Errorf("divideByZHOnGPU: shard scaling failed at %d: %s", i, vecErr.AsString())) + return + } + } + finish(nil) + }) + if err := <-scaleDone; err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, err + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + dMerged, splitErr := s.inverseAndMergeShards(gpuNumerator, domains) + if splitErr != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, splitErr + } + // Shards are not needed after split inverse+merge. + s.freeNumeratorShards(gpuNumerator.shards) + return &gpuQuotientPolynomial{coeffs: dMerged, size: totalSize}, nil +} + +func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*iop.Polynomial, error) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return nil, fmt.Errorf("downloadQuotientFromGPU: empty quotient") + } + if quotient.size <= 0 { + return nil, fmt.Errorf("downloadQuotientFromGPU: invalid quotient size %d", quotient.size) + } + + coeffs := make([]fr.Element, quotient.size) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) + // Async boundary for host materialization of quotient coefficients. + done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + }) + if err := <-done; err != nil { + return nil, err + } + + return iop.NewPolynomial(&coeffs, iop.Form{Basis: iop.Canonical, Layout: iop.Regular}), nil +} + +func commitOnGPUWithDeviceBases( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, +) (curve.G1Affine, error) { + return commitOnGPUWithDeviceBasesChunked(scalarsDevice, basesDevice, device, icicleMSMChunkSize()) +} + +func commitOnGPUWithDeviceBasesChunked( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, + chunkSize int, +) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty scalar slice") + } + if basesDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty basis slice") + } + if scalarsDevice.Len() > basesDevice.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: invalid scalar size %d", scalarsDevice.Len()) + } + if chunkSize <= 0 || chunkSize > scalarsDevice.Len() { + chunkSize = scalarsDevice.Len() + } + + var commit curve.G1Affine + var msmErr error + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) + close(done) + }) + <-done + if msmErr != nil { + return curve.G1Affine{}, fmt.Errorf("icicle: MSM commit from device bases failed (%d scalars): %w", scalarsDevice.Len(), msmErr) + } + return commit, nil +} + +func commitOnGPUWithDeviceBasesChunkedOnCurrentDevice( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + chunkSize int, +) (curve.G1Affine, error) { + var commit curve.G1Affine + for start := 0; start < scalarsDevice.Len(); start += chunkSize { + end := start + chunkSize + if end > scalarsDevice.Len() { + end = scalarsDevice.Len() + } + + // Each chunk must pair with exactly bases[start:end]: ICICLE treats a + // bases slice longer than the scalars as a batched MSM (and requires + // divisibility), so the full bases buffer cannot be passed as-is when + // it is longer than the scalar vector. + scalarsChunk := scalarsDevice + if start != 0 || end != scalarsDevice.Len() { + scalarsChunk = (&scalarsDevice).Range(start, end, false) + } + basesChunk := basesDevice + if start != 0 || end != basesDevice.Len() { + basesChunk = (&basesDevice).Range(start, end, false) + } + + res := make(icicle_core.HostSlice[icicle_{{ .CurvePkg }}.Projective], 1) + cfg := icicle_msm.GetDefaultMSMConfig() + cfg.AreBasesMontgomeryForm = true + cfg.AreScalarsMontgomeryForm = true + e := icicle_msm.Msm(scalarsChunk, basesChunk, &cfg, res) + if e != icicle_runtime.Success { + return curve.G1Affine{}, fmt.Errorf("icicle MSM failed for chunk [%d:%d]: %s", start, end, e.AsString()) + } + + chunkCommit, err := projectiveToGnarkAffine(res[0]) + if err != nil { + return curve.G1Affine{}, fmt.Errorf("convert chunk [%d:%d]: %w", start, end, err) + } + commit.Add(&commit, &chunkCommit) + } + return commit, nil +} + +func icicleMSMChunkSize() int { + // Production-sized MSMs still need chunking, but tiny chunks add thousands of + // ICICLE calls. 4M-point chunks passed the gnark replay profile; 8M did not. + const defaultChunkSize = 1 << 22 + v := strings.TrimSpace(os.Getenv("ICICLE_MSM_CHUNK_SIZE")) + if v == "" { + return defaultChunkSize + } + n, err := strconv.Atoi(v) + if err != nil || n <= 0 { + return defaultChunkSize + } + return n +} + +func commitOnGPUCanonicalDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: canonical SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgDevice.G1, device) +} + +func evalCanonicalAtPoint(coeffs []fr.Element, point fr.Element) fr.Element { + var acc fr.Element + if len(coeffs) == 0 { + return acc + } + acc.Set(&coeffs[len(coeffs)-1]) + for i := len(coeffs) - 2; i >= 0; i-- { + acc.Mul(&acc, &point).Add(&acc, &coeffs[i]) + } + return acc +} + +func deriveBatchOpeningGamma( + point fr.Element, + digests []curve.G1Affine, + claimedValues []fr.Element, + hf hash.Hash, + dataTranscript ...[]byte, +) (fr.Element, error) { + fs := fiatshamir.NewTranscript(hf, "gamma") + if err := fs.Bind("gamma", point.Marshal()); err != nil { + return fr.Element{}, err + } + for i := range digests { + if err := fs.Bind("gamma", digests[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := range claimedValues { + if err := fs.Bind("gamma", claimedValues[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := 0; i < len(dataTranscript); i++ { + if err := fs.Bind("gamma", dataTranscript[i]); err != nil { + return fr.Element{}, err + } + } + gammaByte, err := fs.ComputeChallenge("gamma") + if err != nil { + return fr.Element{}, err + } + var gamma fr.Element + gamma.SetBytes(gammaByte) + return gamma, nil +} + +func (s *instance) evalDevicePolynomialAtPointOnCurrentDevice( + coeffsDevice icicle_core.DeviceSlice, + point fr.Element, + useBitReverse bool, + cfg icicle_core.VecOpsConfig, +) (fr.Element, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: empty coefficients") + } + + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(point, cfg) + defer dPoint.Free() + + dOut := s.getTempDeviceSlice(1) + defer s.putTempDeviceSlice(dOut, 1) + + opName := "PolyEvalAt" + var eEval icicle_runtime.EIcicleError + if useBitReverse { + opName = "PolyEvalAtBitReverse" + mm := uint64(64 - bits.TrailingZeros64(uint64(coeffsDevice.Len()))) + eEval = icicle_vecops.PolyEvalAtBitReverse(coeffsDevice, dPoint, mm, dOut, cfg) + } else { + eEval = icicle_vecops.PolyEvalAt(coeffsDevice, dPoint, dOut, cfg) + } + if eEval != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: %s failed: %s", opName, eEval.AsString()) + } + + var out fr.Element + hostOut := icicle_core.HostSliceFromElements([]fr.Element{out}) + if cfg.IsAsync { + hostOut.CopyFromDeviceAsync(&dOut, cfg.StreamHandle) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } else { + hostOut.CopyFromDevice(&dOut) + } + return ([]fr.Element)(hostOut)[0], nil +} + +func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (fr.Element, error) { + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + var runErr error + out, runErr = s.evalDevicePolynomialAtPointOnCurrentDevice(coeffsDevice, point, false, cfg) + done <- runErr + }) + return out, <-done +} + +func (s *instance) copyDeviceSliceOnCurrentDevice( + src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, + label string, +) (icicle_core.DeviceSlice, error) { + if src.IsEmpty() || src.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: empty source slice", label) + } + dst := s.getTempDeviceSlice(src.Len()) + eCopy := copyDeviceSliceIntoOnCurrentDevice(dst, src, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dst, src.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: copy failed: %s", label, eCopy.AsString()) + } + return dst, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromStateOnCurrentDevice( + p *iop.Polynomial, + state *gpuPolysState, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "batchOpening poly") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + dCanon, err := s.copyDeviceSliceOnCurrentDevice(dSrc, cfg, "materializePolynomialCanonicalRegularFromStateOnCurrentDevice") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + // Canonical bit-reverse -> canonical regular. + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + dRegular := s.getTempDeviceSlice(dCanon.Len()) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNN // regular -> regular canonical + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular canonical + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) buildBlindedCanonicalPolynomialOnCurrentDevice( + dBaseCanon, dBlindCanon icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBlindCanon.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: empty input") + } + n := dBaseCanon.Len() + blindLen := dBlindCanon.Len() + dOut := s.getTempDeviceSlice(n + blindLen) + + dPrefix := (&dOut).Range(0, n, false) + eCopyBase := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopyBase != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopyBase.AsString()) + } + + dTail := (&dOut).Range(n, n+blindLen, false) + eCopyBlind := copyDeviceSliceIntoOnCurrentDevice(dTail, dBlindCanon, cfg) + if eCopyBlind != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy tail failed: %s", eCopyBlind.AsString()) + } + + dHead := (&dOut).Range(0, blindLen, false) + if eSub := icicle_vecops.VecOp(dHead, dBlindCanon, dHead, cfg, icicle_core.Sub); eSub != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: subtract blind from head failed: %s", eSub.AsString()) + } + return dOut, nil +} + +func (s *instance) buildPaddedCanonicalPolynomialOnCurrentDevice( + dBaseCanon icicle_core.DeviceSlice, + padLen int, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBaseCanon.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: empty base") + } + if padLen < 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: negative pad length %d", padLen) + } + n := dBaseCanon.Len() + dOut := s.getTempDeviceSlice(n + padLen) + dPrefix := (&dOut).Range(0, n, false) + + eCopy := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopy.AsString()) + } + if padLen == 0 { + return dOut, nil + } + + dTail := (&dOut).Range(n, n+padLen, false) + eZero := zeroDeviceSliceOnCurrentDevice(dTail, cfg) + if eZero != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: zero tail failed: %s", eZero.AsString()) + } + return dOut, nil +} + +func (s *instance) prepareBatchOpeningPolynomialsOnGPU( + state *gpuPolysState, + point fr.Element, +) (devicePolys []icicle_core.DeviceSlice, owned []bool, claimed []fr.Element, err error) { + if state == nil { + return nil, nil, nil, fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: nil GPU state") + } + + total := 6 + len(s.trace.Qcp) + devicePolys = make([]icicle_core.DeviceSlice, total) + owned = make([]bool, total) + claimed = make([]fr.Element, total) + devicePolys[0] = s.linearizedPolynomialGPU + claimed[0] = s.linearizedPolynomialClaim + + prepDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + + cleanupOwned := func(from int) { + for i := from; i < len(devicePolys); i++ { + if owned[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + owned[i] = false + } + } + } + + prepareLRORow := func(dstIdx int, p, bp *iop.Polynomial, padLen int) error { + dBase, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + defer s.putTempDeviceSlice(dBase, dBase.Len()) + + var dFinal icicle_core.DeviceSlice + if useBlinding { + if bp == nil { + return fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: missing blinding polynomial") + } + dBlind, eBlind := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(bp, state, cfg) + if eBlind != nil { + return eBlind + } + defer s.putTempDeviceSlice(dBlind, dBlind.Len()) + dFinal, e = s.buildBlindedCanonicalPolynomialOnCurrentDevice(dBase, dBlind, cfg) + } else { + dFinal, e = s.buildPaddedCanonicalPolynomialOnCurrentDevice(dBase, padLen, cfg) + } + if e != nil { + return e + } + devicePolys[dstIdx] = dFinal + owned[dstIdx] = true + + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dFinal, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareLRORow(1, s.polyL, s.bp[id_Bl], order_blinding_L+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare L failed: %w", e)) + return + } + if e := prepareLRORow(2, s.polyR, s.bp[id_Br], order_blinding_R+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare R failed: %w", e)) + return + } + if e := prepareLRORow(3, s.polyO, s.bp[id_Bo], order_blinding_O+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare O failed: %w", e)) + return + } + + prepareDirect := func(dstIdx int, p *iop.Polynomial, label string) error { + dPoly, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + devicePolys[dstIdx] = dPoly + owned[dstIdx] = true + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dPoly, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareDirect(4, s.trace.S1, "S1"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S1 failed: %w", e)) + return + } + if e := prepareDirect(5, s.trace.S2, "S2"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S2 failed: %w", e)) + return + } + + for i := 0; i < len(s.trace.Qcp); i++ { + idx := 6 + i + if e := prepareDirect(idx, s.trace.Qcp[i], "Qcp"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare Qcp[%d] failed: %w", i, e)) + return + } + } + + finish(nil) + }) + if err := <-prepDone; err != nil { + return nil, nil, nil, err + } + return devicePolys, owned, claimed, nil +} + +type evalPolynomialInputPreparationResult struct { + dEval icicle_core.DeviceSlice + ownedLen int + useBitReverseEval bool +} + +func (s *instance) prepareEvalPolynomialInputOnCurrentDevice( + p *iop.Polynomial, + dSrc icicle_core.DeviceSlice, + cfgVec icicle_core.VecOpsConfig, +) (evalPolynomialInputPreparationResult, error) { + if p == nil { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: nil polynomial") + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: empty source polynomial") + } + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] prepareEvalInput begin n=%d basis=%v layout=%v step_profile=%q ntt_trace=%q ntt_profile_full=%q ntt_profile_arbitrary=%q\n", + dSrc.Len(), + p.Basis, + p.Layout, + os.Getenv("ICICLE_STEP_PROFILE"), + os.Getenv("ICICLE_NTT_TRACE"), + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + + result := evalPolynomialInputPreparationResult{ + dEval: dSrc, + } + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return result, nil + } + + releaseOwned := func() { + if result.ownedLen > 0 && !result.dEval.IsEmpty() { + s.putTempDeviceSlice(result.dEval, result.ownedLen) + result.dEval = icicle_core.DeviceSlice{} + result.ownedLen = 0 + } + } + + dWork := s.getTempDeviceSlice(dSrc.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dWork, dSrc, cfgVec); e != icicle_runtime.Success { + s.putTempDeviceSlice(dWork, dSrc.Len()) + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: copy source polynomial failed: %s", e.AsString()) + } + + result.dEval = dWork + result.ownedLen = dSrc.Len() + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfgVec.IsAsync + cfgNtt.StreamHandle = cfgVec.StreamHandle + + var ownsNttStream bool + destroyOwnedNttStream := func() error { + if !ownsNttStream { + return nil + } + return syncAndDestroyStreamOnCurrentDevice(cfgNtt.StreamHandle, "prepareEvalPolynomialInputOnCurrentDevice") + } + + switch p.Basis { + case iop.Canonical: + // No transform required. + result.useBitReverseEval = p.Layout == iop.BitReverse + case iop.Lagrange, iop.LagrangeCoset: + if cfgNtt.StreamHandle == nil { + nttStream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: create NTT stream failed: %s", eStream.AsString()) + } + cfgNtt.StreamHandle = nttStream + cfgNtt.IsAsync = true + ownsNttStream = true + } + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNR // regular -> bitreverse on inverse + result.useBitReverseEval = true + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular on inverse + result.useBitReverseEval = false + } + startNtt := time.Now() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT launch n=%d ordering=%v has_coset=%t basis=%v layout=%v\n", + dSrc.Len(), + cfgNtt.Ordering, + p.Basis == iop.LagrangeCoset, + p.Basis, + p.Layout, + ) + } + + // Martun: This call to Ntt takes about 3 seconds, because Ntt is reusing NTT domain data that + // gets prepared during InitDomain (once per device), not re-derived every call. + // Inside ICICLE, InitDomain precomputes: domain.twiddles (main roots-of-unity table, N+1) + // internal_twiddles and basic_twiddles for mixed-radix kernels + // if fast mode is on (it is by default here), extra forward+inverse fast twiddle tables (fast_external/internal/basic and _inv) — comment says this costs ~4N extra memory + // CPU-side coset_index map (root -> index), then reused by later Ntt calls + eNtt := icicle_ntt.Ntt(result.dEval, icicle_core.KInverse, &cfgNtt, result.dEval) + nttElapsed := time.Since(startNtt) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT done status=%s took=%s\n", + eNtt.AsString(), + nttElapsed, + ) + } + if eNtt != icicle_runtime.Success { + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + l := logger.Logger() + l.Warn().Err(errDestroy).Msg("prepareEvalPolynomialInputOnCurrentDevice") + } + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + // Async boundary for this helper when it owns the stream. + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + releaseOwned() + return evalPolynomialInputPreparationResult{}, errDestroy + } + default: + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: unsupported basis %v", p.Basis) + } + + return result, nil +} + +// evalPolynomialInCurrentFormOnGPU evaluates a polynomial at a point directly +// from the shared GPU state regardless of its current basis/layout by converting +// a temporary device copy to canonical/regular when needed. +func (s *instance) evalPolynomialInCurrentFormOnGPU( + p *iop.Polynomial, + state *gpuPolysState, + point fr.Element, +) (fr.Element, error) { + if p == nil { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "eval") + if err != nil { + return fr.Element{}, err + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: empty device polynomial") + } + + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + dEval := dSrc + ownedLen := 0 + releaseEval := func() { + if ownedLen > 0 && !dEval.IsEmpty() { + s.putTempDeviceSlice(dEval, ownedLen) + dEval = icicle_core.DeviceSlice{} + ownedLen = 0 + } + } + finish := func(runErr error) { + // Async boundary for eval path before handing result back to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + releaseEval() + done <- runErr + } + + prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) + dEval = prepareResult.dEval + ownedLen = prepareResult.ownedLen + if prepErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", prepErr)) + return + } + + evalOut, evalErr := s.evalDevicePolynomialAtPointOnCurrentDevice(dEval, point, prepareResult.useBitReverseEval, cfgVec) + if evalErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", evalErr)) + return + } + out = evalOut + finish(nil) + }) + return out, <-done +} + +func (s *instance) evaluateBlindedOnGPU( + p, bp *iop.Polynomial, + state *gpuPolysState, + zeta fr.Element, +) (fr.Element, error) { + if p == nil || bp == nil { + return fr.Element{}, fmt.Errorf("evaluateBlindedOnGPU: nil polynomial") + } + pAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(p, state, zeta) + if err != nil { + return fr.Element{}, err + } + bpAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(bp, state, zeta) + if err != nil { + return fr.Element{}, err + } + + var t, one fr.Element + one.SetOne() + t.Exp(zeta, big.NewInt(int64(p.Size()))).Sub(&t, &one) + bpAtZeta.Mul(&bpAtZeta, &t) + pAtZeta.Add(&pAtZeta, &bpAtZeta) + return pAtZeta, nil +} + +func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.Element) (kzg.OpeningProof, error) { + if len(coeffs) < 2 || len(coeffs) > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonical: invalid polynomial size %d", len(coeffs)) + } + claimed := evalCanonicalAtPoint(coeffs, point) + + var dWitness icicle_core.DeviceSlice + witnessSize := len(coeffs) - 1 + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dCoeffs := uploadVector(coeffs) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dCoeffs, dPoint, dWitness, cfg) + _ = dCoeffs.Free() + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonical: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + if err := <-divDone; err != nil { + return kzg.OpeningProof{}, err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (kzg.OpeningProof, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() < 2 || coeffsDevice.Len() > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonicalDevice: invalid polynomial size %d", coeffsDevice.Len()) + } + n := coeffsDevice.Len() + + var startEval time.Time + if isProfileMode { + startEval = time.Now() + } + claimed, err := s.evalDevicePolynomialAtPoint(coeffsDevice, point) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startEval)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint") + } + } + if err != nil { + return kzg.OpeningProof{}, err + } + + var dWitness icicle_core.DeviceSlice + witnessSize := coeffsDevice.Len() - 1 + var startDivideByXMinusA time.Time + if isProfileMode { + startDivideByXMinusA = time.Now() + } + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(coeffsDevice, dPoint, dWitness, cfg) + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonicalDevice: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + divideErr := <-divDone + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startDivideByXMinusA)) + if divideErr != nil { + ev.Err(divideErr).Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA") + } + } + if divideErr != nil { + return kzg.OpeningProof{}, divideErr + } + + var startCommit time.Time + if isProfileMode { + startCommit = time.Now() + } + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startCommit)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice") + } + } + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) linearizedZContributionScale(lZeta, rZeta, oZeta fr.Element) fr.Element { + var s2, tmp fr.Element + var uzeta, uuzeta fr.Element + uzeta.Mul(&s.zeta, &s.pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &s.pk.Vk.CosetShift) + + s2.Mul(&s.beta, &s.zeta).Add(&s2, &lZeta).Add(&s2, &s.gamma) + tmp.Mul(&s.beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp) + tmp.Mul(&s.beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp).Neg(&s2).Mul(&s2, &s.alpha) + + var one, alphaSquareLagrangeZero, den fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&s.zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) + den.Sub(&s.zeta, &one).Inverse(&den) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) + + s2.Add(&s2, &alphaSquareLagrangeZero) + return s2 +} + +func (s *instance) linearizedSelectorScales(evals witnessEvalAtZeta, zu fr.Element) linearizedSelectorScales { + var scales linearizedSelectorScales + + // S3 scale: + // alpha * beta * Z(mu*zeta) * + // (L(zeta) + beta*S1(zeta) + gamma) * + // (R(zeta) + beta*S2(zeta) + gamma) + var tmp fr.Element + scales.s3.Mul(&evals.s1zeta, &s.beta).Add(&scales.s3, &evals.blzeta).Add(&scales.s3, &s.gamma) + tmp.Mul(&evals.s2zeta, &s.beta).Add(&tmp, &evals.brzeta).Add(&tmp, &s.gamma) + scales.s3.Mul(&scales.s3, &tmp).Mul(&scales.s3, &zu).Mul(&scales.s3, &s.beta).Mul(&scales.s3, &s.alpha) + + scales.ql.Set(&evals.blzeta) + scales.qr.Set(&evals.brzeta) + scales.qm.Mul(&evals.brzeta, &evals.blzeta) + scales.qo.Set(&evals.bozeta) + scales.qk.SetOne() + scales.qcp = append(scales.qcp, evals.qcpzeta...) + + return scales +} + +func (s *instance) buildLinearizedSelectorTermsOnGPU( + evals witnessEvalAtZeta, + zu fr.Element, + linearizedLen int, +) (icicle_core.DeviceSlice, error) { + if linearizedLen <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: invalid length %d", linearizedLen) + } + if len(evals.qcpzeta) > len(s.cCommitments) { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: qcp/cCommitments mismatch (%d > %d)", len(evals.qcpzeta), len(s.cCommitments)) + } + + scales := s.linearizedSelectorScales(evals, zu) + + var dLinearized icicle_core.DeviceSlice + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildLinearizedSelectorTermsOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if runErr != nil && !dLinearized.IsEmpty() { + s.putTempDeviceSlice(dLinearized, dLinearized.Len()) + dLinearized = icicle_core.DeviceSlice{} + } + done <- runErr + }() + + dLinearized = s.getTempDeviceSlice(linearizedLen) + if eZero := zeroDeviceSliceOnCurrentDevice(dLinearized, cfg); eZero != icicle_runtime.Success { + runErr = fmt.Errorf("buildLinearizedSelectorTermsOnGPU: zero output failed: %s", eZero.AsString()) + return + } + + addTerm := func(p *iop.Polynomial, scale fr.Element, label string) error { + if p == nil { + return fmt.Errorf("missing polynomial %s", label) + } + if scale.IsZero() { + return nil + } + + start := time.Now() + dPoly, err := s.materializePolynomialCanonicalRegularFromHostOnCurrentDevice(p, cfg) + if err != nil { + return fmt.Errorf("%s: %w", label, err) + } + defer s.putTempDeviceSlice(dPoly, dPoly.Len()) + if dPoly.Len() > dLinearized.Len() { + return fmt.Errorf("%s: polynomial too large (%d > %d)", label, dPoly.Len(), dLinearized.Len()) + } + + dScale := uploadScalarStdOnCurrentDevice(scale, cfg) + dScaled := s.getTempDeviceSlice(dPoly.Len()) + defer s.putTempDeviceSlice(dScaled, dScaled.Len()) + + eScale := icicle_vecops.ScalarMulVec(dScale, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if eScale != icicle_runtime.Success { + return fmt.Errorf("%s: scale failed: %s", label, eScale.AsString()) + } + + dPrefix := (&dLinearized).Range(0, dPoly.Len(), false) + if eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add); eAdd != icicle_runtime.Success { + return fmt.Errorf("%s: add failed: %s", label, eAdd.AsString()) + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Str("term", label).Int("n", dPoly.Len()).Dur("took", time.Since(start)).Msg("computeLinearizedPolynomial: add selector term on GPU") + } + return nil + } + + if runErr = addTerm(s.trace.S3, scales.s3, "S3"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Ql, scales.ql, "Ql"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qm, scales.qm, "Qm"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qr, scales.qr, "Qr"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qo, scales.qo, "Qo"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qk, scales.qk, "Qk"); runErr != nil { + return + } + for i := range scales.qcp { + if runErr = addTerm(s.cCommitments[i], scales.qcp[i], fmt.Sprintf("Qcp[%d]", i)); runErr != nil { + return + } + } + }) + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dLinearized, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromHostOnCurrentDevice( + p *iop.Polynomial, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: nil polynomial") + } + coeffs := p.Coefficients() + if len(coeffs) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: empty polynomial") + } + + dCanon := s.getTempDeviceSlice(len(coeffs)) + host := icicle_core.HostSliceFromElements(coeffs) + if cfg.IsAsync { + host.CopyToDeviceAsync(&dCanon, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&dCanon, false) + } + if dCanon.IsEmpty() { + s.putTempDeviceSlice(dCanon, len(coeffs)) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: host upload failed") + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + dRegular := s.getTempDeviceSlice(dCanon.Len()) + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + switch p.Layout { + case iop.Regular: + cfgNtt.Ordering = icicle_core.KNN + case iop.BitReverse: + cfgNtt.Ordering = icicle_core.KRN + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported layout %v", p.Layout) + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) addZContributionToLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + dBlindedZCanonical icicle_core.DeviceSlice, + lZeta, rZeta, oZeta fr.Element, +) error { + if dLinearized.IsEmpty() || dBlindedZCanonical.IsEmpty() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: empty input") + } + if dLinearized.Len() < dBlindedZCanonical.Len() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: linearized too small (%d < %d)", dLinearized.Len(), dBlindedZCanonical.Len()) + } + + zScale := s.linearizedZContributionScale(lZeta, rZeta, oZeta) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("addZContributionToLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dScale icicle_core.DeviceSlice + var dScaledZ icicle_core.DeviceSlice + defer func() { + // Async boundary before returning temporary buffers to the pool. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dScale) + if !dScaledZ.IsEmpty() { + s.putTempDeviceSlice(dScaledZ, dScaledZ.Len()) + } + done <- runErr + }() + + dScale = uploadScalarStdOnCurrentDevice(zScale, cfg) + dScaledZ = s.getTempDeviceSlice(dBlindedZCanonical.Len()) + eMul := icicle_vecops.ScalarMulVec(dScale, dBlindedZCanonical, dScaledZ, cfg) + if eMul != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: scale Z failed: %s", eMul.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, dBlindedZCanonical.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaledZ, dPrefix, cfg, icicle_core.Add) + if eAdd != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: add scaled Z failed: %s", eAdd.AsString()) + return + } + }) + return <-done +} + +func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + quotient *gpuQuotientPolynomial, +) error { + if dLinearized.IsEmpty() || quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: empty input") + } + nPlus2 := int(s.domain0.Cardinality) + 2 + if dLinearized.Len() < nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: linearized too small") + } + if quotient.coeffs.Len() < 3*nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: quotient too small") + } + + var one fr.Element + one.SetOne() + var zetaN, zetaNPlusTwo, zhZeta fr.Element + zetaN.Exp(s.zeta, big.NewInt(int64(s.domain0.Cardinality))) + zhZeta.Sub(&zetaN, &one) + zetaNPlusTwo.Mul(&zetaN, &s.zeta).Mul(&zetaNPlusTwo, &s.zeta) + + h1 := ("ient.coeffs).Range(0, nPlus2, false) + h2 := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3 := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("subtractQuotientContributionFromLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dAcc icicle_core.DeviceSlice + var dZetaStd icicle_core.DeviceSlice + var dZhStd icicle_core.DeviceSlice + defer func() { + // Async boundary before reusing temporary quotient vectors. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dZhStd) + freeDeviceSlice(&dZetaStd) + if !dAcc.IsEmpty() { + s.putTempDeviceSlice(dAcc, dAcc.Len()) + } + done <- runErr + }() + + dAcc = s.getTempDeviceSlice(nPlus2) + dZetaStd = uploadScalarStdOnCurrentDevice(zetaNPlusTwo, cfg) + eMulH3 := icicle_vecops.ScalarMulVec(dZetaStd, h3, dAcc, cfg) + if eMulH3 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale h3 failed: %s", eMulH3.AsString()) + return + } + if eAddH2 := icicle_vecops.VecOp(dAcc, h2, dAcc, cfg, icicle_core.Add); eAddH2 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h2 failed: %s", eAddH2.AsString()) + return + } + if eMulPow := icicle_vecops.ScalarMulVec(dZetaStd, dAcc, dAcc, cfg); eMulPow != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by zeta^(n+2) failed: %s", eMulPow.AsString()) + return + } + if eAddH1 := icicle_vecops.VecOp(dAcc, h1, dAcc, cfg, icicle_core.Add); eAddH1 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h1 failed: %s", eAddH1.AsString()) + return + } + + dZhStd = uploadScalarStdOnCurrentDevice(zhZeta, cfg) + eScaleZh := icicle_vecops.ScalarMulVec(dZhStd, dAcc, dAcc, cfg) + if eScaleZh != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by Z_H(zeta) failed: %s", eScaleZh.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, nPlus2, false) + eSub := icicle_vecops.VecOp(dPrefix, dAcc, dPrefix, cfg, icicle_core.Sub) + if eSub != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: subtract term failed: %s", eSub.AsString()) + return + } + }) + return <-done +} + +// divideByZH +// The input must be in LagrangeCoset. +// The result is in Canonical Regular. (in place using a) +func (s *instance) divideByZH(a *iop.Polynomial, domains [2]*fft.Domain) (*iop.Polynomial, error) { + + // check that the basis is LagrangeCoset + if a.Basis != iop.LagrangeCoset || a.Layout != iop.BitReverse { + return nil, errors.New("invalid form") + } + + // prepare the evaluations of x^n-1 on the big domain's coset + var startEvaluateXnMinusOne time.Time + if isProfileMode { + startEvaluateXnMinusOne = time.Now() + } + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startEvaluateXnMinusOne)).Msg("divideByZH: evaluateXnMinusOneDomainBigCoset") + } + rho := int(domains[1].Cardinality / domains[0].Cardinality) + + r := a.Coefficients() + n := uint64(len(r)) + nn := uint64(64 - bits.TrailingZeros64(n)) + + var startParallelizeMul time.Time + if isProfileMode { + startParallelizeMul = time.Now() + } + utils.Parallelize(len(r), func(start, end int) { + for i := start; i < end; i++ { + iRev := bits.Reverse64(uint64(i)) >> nn + r[i].Mul(&r[i], &xnMinusOneInverseLagrangeCoset[int(iRev)%rho]) + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallelizeMul)).Msg("divideByZH: parallelize multiply coefficients") + } + + // Replace CPU FFT inverse by ICICLE NTT inverse. + var startGpuNTTInverse time.Time + if isProfileMode { + startGpuNTTInverse = time.Now() + } + // It's faster on CPU. + // s.gpuNTTInverse(a) + a.ToCanonical(domains[1]).ToRegular() + if isProfileMode { + l := logger.Logger() + l.Debug(). + Int("size", a.Size()). + Dur("took", time.Since(startGpuNTTInverse)). + Msg("divideByZH: gpuNTTInverse on input of size n") + } + + return a, nil +} + +// evaluateXnMinusOneDomainBigCoset evaluates Xᵐ-1 on DomainBig coset +func evaluateXnMinusOneDomainBigCoset(domains [2]*fft.Domain) []fr.Element { + + rho := domains[1].Cardinality / domains[0].Cardinality + + res := make([]fr.Element, rho) + + expo := big.NewInt(int64(domains[0].Cardinality)) + res[0].Exp(domains[1].FrMultiplicativeGen, expo) + + var t fr.Element + t.Exp(domains[1].Generator, expo) + + one := fr.One() + + for i := 1; i < int(rho); i++ { + res[i].Mul(&res[i-1], &t) + res[i-1].Sub(&res[i-1], &one) + } + res[len(res)-1].Sub(&res[len(res)-1], &one) + + res = fr.BatchInvert(res) + + return res +} + +// innerComputeLinearizedPoly computes the linearized polynomial in canonical basis. +// The purpose is to commit and open all in one ql, qr, qm, qo, qk. +// * lZeta, rZeta, oZeta are the evaluation of l, r, o at zeta +// * z is the permutation polynomial, zu is Z(μX), the shifted version of Z +// * pk is the proving key: the linearized polynomial is a linear combination of ql, qr, qm, qo, qk. +// +// The Linearized polynomial is: +// +// α²*L₁(ζ)*Z(X) +// + α*( (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*(β*s3(X))*Z(μζ) - Z(X)*(l(ζ)+β*id1(ζ)+γ)*(r(ζ)+β*id2(ζ)+γ)*(o(ζ)+β*id3(ζ)+γ)) +// + l(ζ)*Ql(X) + l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) +// - Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) +// +// /!\ blindedZCanonical is modified +func (s *instance) innerComputeLinearizedPoly( + lZeta, rZeta, oZeta, s1Zeta, s2Zeta, + alpha, beta, gamma, zeta, zu fr.Element, + qcpZeta, blindedZCanonical []fr.Element, + pi2Canonical [][]fr.Element, + pk *ProvingKey, +) []fr.Element { + + // l(ζ)r(ζ) + var rl fr.Element + rl.Mul(&rZeta, &lZeta) + + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + // the linearised polynomial is + // α²*L₁(ζ)*Z(X) + + // s1*s3(X)+s2*Z(X) + l(ζ)*Ql(X) + + // l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) - + // Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + var s1, s2, tmp fr.Element + s1.Mul(&s1Zeta, &beta).Add(&s1, &lZeta).Add(&s1, &gamma) // (l(ζ)+β*s1(ζ)+γ) + tmp.Mul(&s2Zeta, &beta).Add(&tmp, &rZeta).Add(&tmp, &gamma) + s1.Mul(&s1, &tmp).Mul(&s1, &zu).Mul(&s1, &beta).Mul(&s1, &alpha) // (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*β*Z(μζ)*α + + var uzeta, uuzeta fr.Element + uzeta.Mul(&zeta, &pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &pk.Vk.CosetShift) + + s2.Mul(&beta, &zeta).Add(&s2, &lZeta).Add(&s2, &gamma) // (l(ζ)+β*ζ+γ) + tmp.Mul(&beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &gamma) // (r(ζ)+β*u*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ) + tmp.Mul(&beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &gamma) // (o(ζ)+β*u²*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + s2.Neg(&s2).Mul(&s2, &alpha) + + // Z_h(ζ), ζⁿ⁺², L₁(ζ)*α²*Z + var zhZeta, zetaNPlusTwo, alphaSquareLagrangeZero, one, den, frNbElmt fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) // ζⁿ + zetaNPlusTwo.Mul(&alphaSquareLagrangeZero, &zeta).Mul(&zetaNPlusTwo, &zeta) // ζⁿ⁺² + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) // ζⁿ - 1 + zhZeta.Set(&alphaSquareLagrangeZero) // Z_h(ζ) = ζⁿ - 1 + frNbElmt.SetUint64(uint64(nbElmt)) + den.Sub(&zeta, &one).Inverse(&den) // 1/(ζ-1) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). // L₁ = (ζⁿ - 1)/(ζ-1) + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) // α²*L₁(ζ) + + s3canonical := s.trace.S3.Coefficients() + // Qk is prepared in canonical/regular form by computeLinearizedPolynomial. + + // len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 when Statistical ZK is activated + // len(h1)=len(h2)=len(h3)=len(blindedZCanonical)-1 when Statistical ZK is deactivated + h1 := s.h1() + h2 := s.h2() + h3 := s.h3() + + // at this stage we have + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + startParallel := time.Now() + utils.Parallelize(len(blindedZCanonical), func(start, end int) { + + cql := s.trace.Ql.Coefficients() + cqr := s.trace.Qr.Coefficients() + cqm := s.trace.Qm.Coefficients() + cqo := s.trace.Qo.Coefficients() + cqk := s.trace.Qk.Coefficients() + + var t, t0, t1 fr.Element + + for i := start; i < end; i++ { + t.Mul(&blindedZCanonical[i], &s2) // -Z(X)*α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + if i < len(s3canonical) { + t0.Mul(&s3canonical[i], &s1) // α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ)*β*s3(X) + t.Add(&t, &t0) + } + if i < len(cqm) { + t1.Mul(&cqm[i], &rl) // l(ζ)r(ζ)*Qm(X) + t.Add(&t, &t1) // linPol += l(ζ)r(ζ)*Qm(X) + t0.Mul(&cql[i], &lZeta) // l(ζ)Q_l(X) + t.Add(&t, &t0) // linPol += l(ζ)*Ql(X) + t0.Mul(&cqr[i], &rZeta) //r(ζ)*Qr(X) + t.Add(&t, &t0) // linPol += r(ζ)*Qr(X) + t0.Mul(&cqo[i], &oZeta) // o(ζ)*Qo(X) + t.Add(&t, &t0) // linPol += o(ζ)*Qo(X) + t.Add(&t, &cqk[i]) // linPol += Qk(X) + for j := range qcpZeta { // linPol += ∑ᵢQcp_(ζ)Pi_(X) + t0.Mul(&pi2Canonical[j][i], &qcpZeta[j]) + t.Add(&t, &t0) + } + } + + t0.Mul(&blindedZCanonical[i], &alphaSquareLagrangeZero) // α²L₁(ζ)Z(X) + blindedZCanonical[i].Add(&t, &t0) // linPol += α²L₁(ζ)Z(X) + + // if statistical zeroknowledge is deactivated, len(h1)=len(h2)=len(h3)=len(blindedZ)-1. + // Else len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 + if i < len(h3) { + t.Mul(&h3[i], &zetaNPlusTwo). + Add(&t, &h2[i]). + Mul(&t, &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } else { + if s.opt.StatisticalZK { + t.Mul(&h2[i], &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } + } + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallel)).Msg("computeLinearizedPolynomial: inner parallel loop") + } + + return blindedZCanonical +} + +var errContextDone = errors.New("context done") + +// local copies of verification-time helpers used by prover transcript +func bindPublicData(fs *fiatshamir.Transcript, challenge string, vk *plonk_{{ .CurvePkg }}.VerifyingKey, publicInputs []fr.Element) error { + if err := fs.Bind(challenge, vk.S[0].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[1].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[2].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Ql.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qr.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qm.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qo.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qk.Marshal()); err != nil { + return err + } + for i := range vk.Qcp { + if err := fs.Bind(challenge, vk.Qcp[i].Marshal()); err != nil { + return err + } + } + for i := 0; i < len(publicInputs); i++ { + if err := fs.Bind(challenge, publicInputs[i].Marshal()); err != nil { + return err + } + } + return nil +} + +func deriveRandomness(fs *fiatshamir.Transcript, challenge string, points ...*curve.G1Affine) (fr.Element, error) { + var buf [curve.SizeOfG1AffineUncompressed]byte + var r fr.Element + for _, p := range points { + buf = p.RawBytes() + if err := fs.Bind(challenge, buf[:]); err != nil { + return r, err + } + } + b, err := fs.ComputeChallenge(challenge) + if err != nil { + return r, err + } + r.SetBytes(b) + return r, nil +} + +// -------------------- GPU helpers and device setup -------------------- + +func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { + if pk.deviceInfo != nil { + return nil + } + pk.deviceInfo = &deviceInfo{} + + // Initialize ICICLE NTT domain (root of unity) and store coset generator for coset NTTs. + // ICICLE InitDomain expects a primitive root of unity; for coset transforms we use 𝔽ᵣ* generator. + + var gen fr.Element + var err error + if pk.Vk.Size < 6 { + gen, err = fft.Generator(8 * pk.Vk.Size) + if err != nil { + return err + } + } else { + gen, err = fft.Generator(4 * pk.Vk.Size) + if err != nil { + return err + } + } + genBits := gen.Bits() + limbs := icicle_core.ConvertUint64ArrToUint32Arr(genBits[:]) + // Initialize ICICLE NTT domain with root of unity + var rouIcicle icicle_{{ .CurvePkg }}.ScalarField + rouIcicle.FromLimbs(limbs) + + // Store coset generator = generator of 𝔽ᵣ* (matches CPU ToLagrangeCoset) + { + cosetGen := fft.GeneratorFullMultiplicativeGroup() + cosetBits := cosetGen.Bits() + cosetLimbs := icicle_core.ConvertUint64ArrToUint32Arr(cosetBits[:]) + copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) + } + + chInitDomain := make(chan struct{}) + initDomainQueuedAt := time.Now() + icicle_runtime.RunOnDevice(device, func(args ...any) { + initDomainStartedAt := time.Now() + initCfg := icicle_core.GetDefaultNTTInitDomainConfig() + ext := config_extension.Create() + defer config_extension.Delete(ext) + fastTwiddles := envEnabled("ICICLE_NTT_FAST_TWIDDLES", true) + ext.SetBool(icicle_core.CUDA_NTT_FAST_TWIDDLES_MODE, fastTwiddles) + initCfg.Ext = ext.AsUnsafePointer() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain start vk_size=%d fast_twiddles=%t profile_full=%q profile_arbitrary=%q\n", + pk.Vk.Size, + fastTwiddles, + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + e := icicle_ntt.InitDomain(rouIcicle, initCfg) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain end status=%s call_took=%s\n", + e.AsString(), + time.Since(initDomainStartedAt), + ) + } + if e != icicle_runtime.Success { + panic("icicle: InitDomain failed") + } + close(chInitDomain) + }) + + <-chInitDomain + if isNttTrace { + fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) + } + + chLag := make(chan struct{}) + chCan := make(chan struct{}) + + if len(pk.KzgLagrange.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) + g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) + close(chLag) + }) + } else { + close(chLag) + } + + if len(pk.Kzg.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) + g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) + close(chCan) + }) + } else { + close(chCan) + } + + <-chLag + <-chCan + return nil +} + +func projectiveToGnarkAffine(p icicle_{{ .CurvePkg }}.Projective) (curve.G1Affine, error) { + x, err := icicleBaseFieldToGnarkFp(p.X) + if err != nil { + return curve.G1Affine{}, err + } + y, err := icicleBaseFieldToGnarkFp(p.Y) + if err != nil { + return curve.G1Affine{}, err + } + z, err := icicleBaseFieldToGnarkFp(p.Z) + if err != nil { + return curve.G1Affine{}, err + } + if z.IsZero() { + return curve.G1Affine{}, nil + } + + var zInv fp.Element + zInv.Inverse(&z) + x.Mul(&x, &zInv) + y.Mul(&y, &zInv) + return curve.G1Affine{X: x, Y: y}, nil +} + +func icicleBaseFieldToGnarkFp(v icicle_{{ .CurvePkg }}.BaseField) (fp.Element, error) { + bytes := v.ToBytesLittleEndian() + if len(bytes) != fp.Bytes { + return fp.Element{}, fmt.Errorf("invalid ICICLE base field byte length %d", len(bytes)) + } + var buf [fp.Bytes]byte + copy(buf[:], bytes) + return fp.LittleEndian.Element(&buf) +} + +func commitOnGPULagrangeDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgLagrangeDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: lagrange SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgLagrangeDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgLagrangeDevice.G1, device) +} + +func (s *instance) registerDevicePolynomialInSharedState(state *gpuPolysState, p *iop.Polynomial, dSlice icicle_core.DeviceSlice) error { + if state == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil shared state") + } + if p == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil polynomial") + } + if dSlice.IsEmpty() { + return fmt.Errorf("registerDevicePolynomialInSharedState: empty device slice") + } + + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int) + } + if idx, ok := state.polyToIdx[p]; ok { + if idx < 0 || idx >= len(state.deviceSlices) { + return fmt.Errorf("registerDevicePolynomialInSharedState: invalid index %d", idx) + } + state.deviceSlices[idx] = dSlice + state.hostSlices[idx] = nil + state.originalForm[idx] = iop.Form{Basis: p.Basis, Layout: p.Layout} + return nil + } + + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, dSlice) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + return nil +} + +func (s *instance) gpuInclusivePrefixProductOnCurrentDevice( + dVec icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + if dVec.IsEmpty() || dVec.Len() <= 1 { + return nil + } + + n := dVec.Len() + for step := 1; step < n; step <<= 1 { + src := (&dVec).Range(0, n-step, false) + dst := (&dVec).Range(step, n, false) + + tmpStd := s.getTempDeviceSlice(n - step) + if err := copyDeviceSliceIntoOnCurrentDevice(tmpStd, src, cfg); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: copy stage failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(tmpStd, cfg) + if err := icicle_vecops.VecOp(tmpStd, dst, dst, cfg, icicle_core.Mul); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: multiply stage failed: %s", err.AsString()) + } + if cfg.IsAsync { + // tmpStd is returned to pool each stage, so we must wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } + s.putTempDeviceSlice(tmpStd, n-step) + } + return nil +} + +// buildPermutationGatherIndices prepares the subset of permutation indices that +// are consumed by the copy-constraint ratio loop (only rows [0, n-1) per copy). +func buildPermutationGatherIndices(permutation []int64, nbPolynomials, n, supportLen int) ([]int64, error) { + if n <= 1 { + return nil, nil + } + total := nbPolynomials * (n - 1) + indices := make([]int64, total) + + var permBuildErr error + var permBuildErrOnce sync.Once + utils.Parallelize(total, func(start, end int) { + for k := start; k < end; k++ { + j := k / (n - 1) + i := k % (n - 1) + base := j * n + permIdx := permutation[base+i] + if permIdx < 0 || int(permIdx) >= supportLen { + jj, ii, bad := j, i, permIdx + permBuildErrOnce.Do(func() { + permBuildErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: permutation index out of range at (%d,%d): %d", jj, ii, bad) + }) + continue + } + indices[k] = permIdx + } + }) + if permBuildErr != nil { + return nil, permBuildErr + } + return indices, nil +} + +func (s *instance) prepareCopyConstraintSupportsOnCurrentDevice( + n, nbPolynomials int, + domain *fft.Domain, + permGatherIndices []int64, + cfg icicle_core.VecOpsConfig, +) (dSupportFlat, dPermFlat icicle_core.DeviceSlice, err error) { + defer func() { + if err == nil { + return + } + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + }() + + if len(permGatherIndices) == 0 { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: empty permutation gather indices") + return + } + + dOmegaStd := uploadScalarStdOnCurrentDevice(domain.Generator, cfg) + defer dOmegaStd.Free() + dShiftStd := uploadScalarStdOnCurrentDevice(domain.FrMultiplicativeGen, cfg) + defer dShiftStd.Free() + + dSupportFlat, err = allocDeviceUninitialized(nbPolynomials * n) + if err != nil { + return + } + if e := icicle_vecops.SupportIdentity(dOmegaStd, dShiftStd, n, nbPolynomials, dSupportFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: generate identity support on GPU failed: %s", e.AsString()) + return + } + toMontgomeryFormInPlaceWithCfg(dSupportFlat, cfg) + + dPermIndicesDevice := uploadInt64VectorOnCurrentDevice(permGatherIndices, cfg) + defer dPermIndicesDevice.Free() + + dPermFlat, err = allocDeviceUninitialized(len(permGatherIndices)) + if err != nil { + return + } + if e := icicle_vecops.GatherByIndices(dSupportFlat, dPermIndicesDevice, dPermFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: gather permutation support on GPU failed: %s", e.AsString()) + return + } + if cfg.IsAsync { + // Ensure temporary support/index slices are safe to free on return. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + return + } + } + + return +} + +func (s *instance) accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail icicle_core.DeviceSlice, + dBetaStd, dGammaMont icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + nMinusOne := dEntryTail.Len() + dScaled := s.getTempDeviceSlice(nMinusOne) + dTerm := s.getTempDeviceSlice(nMinusOne) + defer func() { + s.putTempDeviceSlice(dScaled, nMinusOne) + s.putTempDeviceSlice(dTerm, nMinusOne) + }() + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dID, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale identity support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add beta*id failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dNumTail, dNumTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply numerator term failed: %s", err.AsString()) + } + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dPerm, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale permutation support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add beta*sigma failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dDenTail, dDenTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply denominator term failed: %s", err.AsString()) + } + if cfg.IsAsync { + // Temp vectors are released at function exit, so ensure queued work is complete. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + } + } + + return nil +} + +// validateDeviceEntries checks that all entries are non-empty and have consistent length. +// Returns the common length n. +func validateDeviceEntries(entries []icicle_core.DeviceSlice, label string) (int, error) { + if len(entries) == 0 { + return 0, fmt.Errorf("%s: no entries", label) + } + n := entries[0].Len() + if n == 0 { + return 0, fmt.Errorf("%s: empty device entry 0", label) + } + for i := range entries { + if entries[i].IsEmpty() { + return 0, fmt.Errorf("%s: empty device entry %d", label, i) + } + if entries[i].Len() != n { + return 0, fmt.Errorf("%s: inconsistent device entry size at %d (%d != %d)", label, i, entries[i].Len(), n) + } + } + return n, nil +} + +// BuildRatioCopyConstraintIcicle builds the accumulating ratio polynomial to prove that +// [P₁ ∥ .. ∥ P_{n—1}] is invariant by the permutation \sigma. +// Namely it returns the polynomial Z whose evaluation on the j-th root of unity is +// Z(ω^j) = Π_{i 1 { + dNumTail := (&dNum).Range(1, n, false) + dDenTail := (&dDen).Range(1, n, false) + var supportErr error + dSupportFlat, dPermFlat, supportErr = s.prepareCopyConstraintSupportsOnCurrentDevice(n, nbPolynomials, domain, permGatherIndices, cfg) + if supportErr != nil { + runErr = supportErr + return + } + + dBetaStd := uploadScalarStdOnCurrentDevice(beta, cfg) + dGammaMont := uploadScalarMontOnCurrentDevice(gamma, cfg) + + for j := 0; j < nbPolynomials; j++ { + dEntryTail := (&entriesDevice[j]).Range(0, n-1, false) + baseID := j * n + dID := (&dSupportFlat).Range(baseID, baseID+n-1, false) + basePerm := j * (n - 1) + dPerm := (&dPermFlat).Range(basePerm, basePerm+(n-1), false) + if err := s.accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail, dBetaStd, dGammaMont, cfg, + ); err != nil { + runErr = err + return + } + } + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize before releasing copy-constraint workspace failed: %s", eSync.AsString()) + return + } + } + _ = dBetaStd.Free() + _ = dGammaMont.Free() + + // Support vectors and loop temps are only needed for term accumulation. + // Free them before prefix products and batch inversion, whose ICICLE + // kernels allocate additional full-domain workspace internally. + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + s.tempGPUMemPool.FreeAll() + } + + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dNum, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dDen, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + + if invErr := s.batchInvertOnCurrentDevice(dDen); invErr != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: GPU batch inversion failed: %s", invErr.AsString()) + return + } + + toStandardFormInPlace(dDen) + if err := icicle_vecops.VecOp(dDen, dNum, dNum, cfg, icicle_core.Mul); err != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: final numerator*denominatorInv multiplication failed: %s", err.AsString()) + return + } + + dResult = dNum + dNum = icicle_core.DeviceSlice{} // transfer ownership to dResult + }) + if err := <-buildDone; err != nil { + return nil, err + } + + hostMirror := make([]fr.Element, n) + if len(hostMirror) > 0 { + hostMirror[0].SetOne() + } + res := iop.NewPolynomial(&hostMirror, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if err := s.registerDevicePolynomialInSharedState(gpuState, res, dResult); err != nil { + freeSliceOnDevice(&dResult, &s.device) + return nil, err + } + + return res, nil +} diff --git a/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.provingkey.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.provingkey.go.tmpl new file mode 100644 index 0000000000..febb50e11f --- /dev/null +++ b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.provingkey.go.tmpl @@ -0,0 +1,93 @@ +//go:build icicle + +import ( + "sync" + "time" + + "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}/fr" + "github.com/consensys/gnark-crypto/ecc/{{ toLower .Curve }}/fr/fft" + plonk_{{ .CurvePkg }} "github.com/consensys/gnark/backend/plonk/{{ toLower .Curve }}" + cs "github.com/consensys/gnark/constraint/{{ toLower .Curve }}" + "github.com/consensys/gnark/logger" + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" +) + +// deviceInfo holds device-resident buffers for GPU acceleration. +type deviceInfo struct { + CosetGenerator [fr.Limbs * 2]uint32 + KzgDevice struct { + G1 icicle_core.DeviceSlice + } + KzgLagrangeDevice struct { + G1 icicle_core.DeviceSlice + } +} + +// hostSetup holds host-side, witness-independent prover state derived from the +// constraint system: the FFT domains and the PLONK trace (selector + +// permutation polynomials). Building the trace walks every constraint (~3s at +// 23M constraints), so it is computed once per proving key and shared across +// proofs. Everything here is read-only during proving: the prover clones Qk +// before patching public inputs into it, and every basis conversion of a trace +// polynomial copies first (see canonicalRegularCoefficientsCopy). +type hostSetup struct { + sizeSystem uint64 + domain0 *fft.Domain + domain1 *fft.Domain + trace *plonk_{{ .CurvePkg }}.Trace +} + +// ProvingKey wraps the native PLONK proving key with device-resident state +// (KZG bases, NTT domains, cached trace) that is uploaded once and reused +// across Prove calls. +// +// Concurrency: Prove calls sharing the same ProvingKey must be serialized by +// the caller. The device state hangs off the key and proofs share a single +// GPU; concurrent proves against the same key are not safe. +type ProvingKey struct { + plonk_{{ .CurvePkg }}.ProvingKey + *deviceInfo + hostSetupOnce sync.Once + hostSetup *hostSetup +} + +func buildHostSetup(spr *cs.SparseR1CS, sizeSystem uint64) *hostSetup { + domain0 := fft.NewDomain(sizeSystem) + + // h, the quotient polynomial is of degree 3(n+1)+2, so it's in a 3(n+2) dim + // vector space, the domain is the next power of 2 superior to 3(n+2). + // 4*domainNum is enough in all cases except when n<6. + var domain1 *fft.Domain + if sizeSystem < 6 { + domain1 = fft.NewDomain(8*sizeSystem, fft.WithoutPrecompute()) + } else { + domain1 = fft.NewDomain(4*sizeSystem, fft.WithoutPrecompute()) + } + + return &hostSetup{ + sizeSystem: sizeSystem, + domain0: domain0, + domain1: domain1, + trace: plonk_{{ .CurvePkg }}.NewTrace(spr, domain0), + } +} + +// hostSetupFor returns the FFT domains and trace for spr, building them on +// first use and caching them on the proving key. A PLONK proving key is bound +// to exactly one constraint system, so per-key caching is sound; as a +// defensive measure a system-size mismatch falls back to an uncached build +// rather than ever serving another circuit's trace. +func (pk *ProvingKey) hostSetupFor(spr *cs.SparseR1CS) *hostSetup { + nbConstraints := spr.GetNbConstraints() + sizeSystem := uint64(nbConstraints + len(spr.Public)) // len(spr.Public) is for the placeholder constraints + pk.hostSetupOnce.Do(func() { + start := time.Now() + pk.hostSetup = buildHostSetup(spr, sizeSystem) + log := logger.Logger() + log.Debug().Dur("took", time.Since(start)).Msg("built prover host setup (fft domains + trace)") + }) + if pk.hostSetup.sizeSystem != sizeSystem { + return buildHostSetup(spr, sizeSystem) + } + return pk.hostSetup +} diff --git a/backend/accelerated/icicle/plonk/bls12-377/doc.go b/backend/accelerated/icicle/plonk/bls12-377/doc.go new file mode 100644 index 0000000000..6b8b35d66b --- /dev/null +++ b/backend/accelerated/icicle/plonk/bls12-377/doc.go @@ -0,0 +1,7 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +// Package bls12377 implements ICICLE acceleration for BLS12-377 PLONK backend. +package bls12377 diff --git a/backend/accelerated/icicle/plonk/bls12-377/icicle.go b/backend/accelerated/icicle/plonk/bls12-377/icicle.go new file mode 100644 index 0000000000..b448e58a43 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bls12-377/icicle.go @@ -0,0 +1,7093 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +//go:build icicle + +package bls12377 + +import ( + "context" + "errors" + "fmt" + "hash" + "io" + "math/big" + "math/bits" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + "unsafe" + + "golang.org/x/sync/errgroup" + + "github.com/consensys/gnark/backend" + plonk_bls12377 "github.com/consensys/gnark/backend/plonk/bls12-377" + "github.com/consensys/gnark/backend/witness" + constraint "github.com/consensys/gnark/constraint" + cs "github.com/consensys/gnark/constraint/bls12-377" + "github.com/consensys/gnark/constraint/solver" + fcs "github.com/consensys/gnark/frontend/cs" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + + "github.com/consensys/gnark-crypto/ecc" + curve "github.com/consensys/gnark-crypto/ecc/bls12-377" + "github.com/consensys/gnark-crypto/ecc/bls12-377/fp" + "github.com/consensys/gnark-crypto/ecc/bls12-377/fr" + "github.com/consensys/gnark-crypto/ecc/bls12-377/fr/fft" + "github.com/consensys/gnark-crypto/ecc/bls12-377/fr/hash_to_field" + "github.com/consensys/gnark-crypto/ecc/bls12-377/fr/iop" + "github.com/consensys/gnark-crypto/ecc/bls12-377/kzg" + fiatshamir "github.com/consensys/gnark-crypto/fiat-shamir" + + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" + icicle_bls12377 "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bls12377" + icicle_msm "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bls12377/msm" + icicle_ntt "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bls12377/ntt" + icicle_vecops "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bls12377/vecOps" + icicle_runtime "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime" + "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime/config_extension" +) + +const HasIcicle = true + +var isProfileMode bool + +var useBlinding bool + +var isNttTrace bool + +func init() { + _, isProfileMode = os.LookupEnv("ICICLE_STEP_PROFILE") + // Blinding polynomials (zero-knowledge) are enabled by default, matching the + // native prover. Set GNARK_DISABLE_BLINDING to trade zero-knowledge for a + // faster, deterministic prover (e.g. when the witness is not secret). + _, disableBlinding := os.LookupEnv("GNARK_DISABLE_BLINDING") + useBlinding = !disableBlinding + isNttTrace = envEnabled("ICICLE_NTT_TRACE", false) +} + +// profileStep returns a function that, when called, logs the elapsed time since +// profileStep was invoked. If profiling is disabled, it returns a no-op. +// Usage: done := profileStep("label"); defer done() +func profileStep(msg string) func() { + if !isProfileMode { + return func() {} + } + start := time.Now() + return func() { + l := logger.Logger() + l.Debug().Dur("took", time.Since(start)).Msg(msg) + } +} + +// stageTiming is a single recorded prover stage and its wall-clock duration. +type stageTiming struct { + name string + dur time.Duration +} + +// stageTimings is a concurrency-safe, ordered recorder of prover stage +// durations. The PLONK prover runs its stages as concurrent goroutines, so the +// recorded durations OVERLAP and do not sum to the total — the printed table +// flags this. +type stageTimings struct { + mu sync.Mutex + entries []stageTiming +} + +// record appends a (stage, duration) entry. Safe to call from any goroutine and +// safe on a nil receiver (records nothing). +func (t *stageTimings) record(name string, d time.Duration) { + if t == nil { + return + } + t.mu.Lock() + t.entries = append(t.entries, stageTiming{name: name, dur: d}) + t.mu.Unlock() +} + +// printTable writes an aligned breakdown of the recorded stages to w, sorted by +// duration (largest first), followed by the overall prover total. Stages run +// concurrently, so the rows overlap and intentionally do not sum to the total. +func (t *stageTimings) printTable(w io.Writer, total time.Duration) { + if t == nil { + return + } + t.mu.Lock() + rows := make([]stageTiming, len(t.entries)) + copy(rows, t.entries) + t.mu.Unlock() + + sort.SliceStable(rows, func(i, j int) bool { return rows[i].dur > rows[j].dur }) + + nameW := len("TOTAL (prover done)") + for _, r := range rows { + if len(r.name) > nameW { + nameW = len(r.name) + } + } + + fmt.Fprintln(w, "") + fmt.Fprintln(w, "================ gnark PLONK prove breakdown (GPU) ================") + fmt.Fprintln(w, "(stages run concurrently — durations overlap and do not sum to TOTAL)") + fmt.Fprintf(w, " %-*s %12s %6s\n", nameW, "STAGE", "TIME", "%TOTAL") + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + for _, r := range rows { + pct := 0.0 + if total > 0 { + pct = 100 * float64(r.dur) / float64(total) + } + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, r.name, r.dur.Round(time.Millisecond), pct) + } + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, "TOTAL (prover done)", total.Round(time.Millisecond), 100.0) + fmt.Fprintln(w, "===================================================================") + fmt.Fprintln(w, "") +} + +func envEnabled(key string, defaultVal bool) bool { + v, ok := os.LookupEnv(key) + if !ok { + return defaultVal + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "0", "false", "off", "no": + return false + default: + return true + } +} + +func nttAlgorithmFromEnv(key string, fallback icicle_core.NttAlgorithm) icicle_core.NttAlgorithm { + v, ok := os.LookupEnv(key) + if !ok { + return fallback + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "auto", "0": + return icicle_core.Auto + case "radix2", "radix-2", "r2", "1": + return icicle_core.Radix2 + case "mixed", "mixedradix", "mixed-radix", "2": + return icicle_core.MixedRadix + default: + return fallback + } +} + +const ( + id_L int = iota + id_R + id_O + id_Z + id_ZS + id_Ql + id_Qr + id_Qm + id_Qo + id_Qk + id_S1 + id_S2 + id_S3 + id_Qci // [ .. , Qc_i, Pi_i, ...] +) + +// blinding factors +const ( + id_Bl int = iota + id_Br + id_Bo + id_Bz + nb_blinding_polynomials +) + +// blinding orders (-1 to deactivate) +const ( + order_blinding_L = 1 + order_blinding_R = 1 + order_blinding_O = 1 + order_blinding_Z = 2 +) + +// Prove generates a PLONK proof. When the accelerator option is not set to +// "icicle", we delegate to the native prover. Otherwise, we run a local copy +// of the CPU prover logic to enable incremental GPU adaptation. +func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts ...backend.ProverOption) (*plonk_bls12377.Proof, error) { + opt, err := backend.NewProverConfig(opts...) + if err != nil { + return nil, err + } + + log := logger.Logger().With(). + Str("curve", spr.CurveID().String()). + Int("nbConstraints", spr.GetNbConstraints()). + Str("backend", "plonk").Logger() + + // parse the options + if opt.HashToFieldFn == nil { + opt.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + + // When blinding is disabled (GNARK_DISABLE_BLINDING), also disable StatisticalZK, it makes no sense + // to use statistical zero knowledge when we don't use blinding. + if !useBlinding { + opt.StatisticalZK = false + } + + start := time.Now() + + // Initialize device and preload KZG bases once per proving key + device := icicle_runtime.CreateDevice("CUDA", 0) + if pk.deviceInfo == nil { + if err := pk.setupDevicePointers(&device); err != nil { + return nil, err + } + } + + // init instance + g, ctx := errgroup.WithContext(context.Background()) + instance, err := newInstance(ctx, spr, pk, fullWitness, &opt) + if err != nil { + return nil, fmt.Errorf("new instance: %w", err) + } + // attach device to instance for GPU calls + instance.device = device + instance.initSharedGPUState() + defer instance.releaseTempGPUMemoryPool() + defer instance.releaseSharedGPUState() + defer instance.releaseLinearizedEvalGPUState() + + // solve constraints + g.Go(instance.solveConstraints) + + // complete qk + g.Go(instance.completeQk) + + // init blinding polynomials + g.Go(instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + g.Go(instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + g.Go(instance.buildRatioCopyConstraint) + + // compute h + g.Go(instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + g.Go(instance.openZ) + + // linearized polynomial + g.Go(instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + g.Go(func() error { + startBatchOpening := time.Now() + err := instance.batchOpening() + if isProfileMode { + instance.timings.record("batchOpening (folded KZG)", time.Since(startBatchOpening)) + } + return err + }) + + if err := g.Wait(); err != nil { + return nil, err + } + + total := time.Since(start) + log.Debug().Dur("took", total).Msg("prover done") + if isProfileMode { + instance.timings.printTable(os.Stderr, total) + } + return instance.proof, nil +} + +// represents a Prover instance +type instance struct { + ctx context.Context + + pk *ProvingKey + proof *plonk_bls12377.Proof + spr *cs.SparseR1CS + opt *backend.ProverConfig + + fs *fiatshamir.Transcript + kzgFoldingHash hash.Hash // for KZG folding + htfFunc hash.Hash // hash to field function + + // polynomials + polyL, polyR, polyO *iop.Polynomial + polyZ, polyZS, polyQk *iop.Polynomial + bp []*iop.Polynomial // blinding polynomials + h *iop.Polynomial // h is the quotient polynomial + hGPU *gpuQuotientPolynomial + polyZLagrangeGPU icicle_core.DeviceSlice + blindedZCanonicalGPU icicle_core.DeviceSlice + quotientShardsRandomizers [2]fr.Element // random elements for blinding the shards of the quotient + + linearizedPolynomial []fr.Element + linearizedPolynomialGPU icicle_core.DeviceSlice + linearizedPolynomialClaim fr.Element + linearizedPolynomialDigest kzg.Digest + + fullWitness witness.Witness + + // bsb22 commitment stuff + commitmentInfo constraint.PlonkCommitments + commitmentVal []fr.Element + cCommitments []*iop.Polynomial + + // challenges + gamma, beta, alpha, zeta fr.Element + + // channel to wait for the steps + chLRO, + chQk, + chbp, + chZ, + chH, + chRestoreLRO, + chZOpening, + chLinearizedPolynomial, + chGammaBeta chan struct{} + + domain0, domain1 *fft.Domain + + trace *plonk_bls12377.Trace + + // GPU device handle + device icicle_runtime.Device + + // Shared GPU polynomial context reused across buildRatioCopyConstraint + // and computeQuotient to avoid repeated host<->device uploads. + gpuStateMu sync.Mutex + sharedGPUState *gpuPolysState + // Snapshot of immutable polynomial slices used by computeLinearizedPolynomial + // for zeta evaluations after computeQuotient mutates/frees shared state. + linearizedEvalGPUState *gpuPolysState + + // Reusable temporary GPU memory pool for non-state buffers. + tempGPUMemPool *gpuMemoryPool + + // Per-prove stage-timing recorder (used to print the breakdown table when + // ICICLE_STEP_PROFILE is set). + timings *stageTimings +} + +func newInstance(ctx context.Context, spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts *backend.ProverConfig) (*instance, error) { + if opts.HashToFieldFn == nil { + opts.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + s := instance{ + ctx: ctx, + pk: pk, + proof: &plonk_bls12377.Proof{}, + spr: spr, + opt: opts, + fullWitness: fullWitness, + bp: make([]*iop.Polynomial, nb_blinding_polynomials), + fs: fiatshamir.NewTranscript(opts.ChallengeHash, "gamma", "beta", "alpha", "zeta"), + kzgFoldingHash: opts.KZGFoldingHash, + htfFunc: opts.HashToFieldFn, + chLRO: make(chan struct{}, 1), + chQk: make(chan struct{}, 1), + chbp: make(chan struct{}, 1), + chGammaBeta: make(chan struct{}, 1), + chZ: make(chan struct{}, 1), + chH: make(chan struct{}, 1), + chZOpening: make(chan struct{}, 1), + chLinearizedPolynomial: make(chan struct{}, 1), + chRestoreLRO: make(chan struct{}, 1), + tempGPUMemPool: newGPUMemoryPool(), + timings: &stageTimings{}, + } + s.initBSB22Commitments() + + // FFT domains and the PLONK trace are witness-independent and expensive to + // build at large n (NewTrace walks every constraint), so they are cached + // on the proving key and shared read-only across proofs. + setup := pk.hostSetupFor(spr) + s.domain0 = setup.domain0 + s.domain1 = setup.domain1 + s.trace = setup.trace + + // sampling random numbers for blinding the quotient + if opts.StatisticalZK { + s.quotientShardsRandomizers[0].SetRandom() + s.quotientShardsRandomizers[1].SetRandom() + } + + return &s, nil +} + +func (s *instance) initBlindingPolynomials() error { + if !useBlinding { + // When blinding is disabled (GNARK_DISABLE_BLINDING), skip creating blinding polynomials entirely + // Just close the channel to unblock any goroutines waiting on it + close(s.chbp) + return nil + } + + s.bp[id_Bl] = getRandomPolynomial(order_blinding_L) + s.bp[id_Br] = getRandomPolynomial(order_blinding_R) + s.bp[id_Bo] = getRandomPolynomial(order_blinding_O) + s.bp[id_Bz] = getRandomPolynomial(order_blinding_Z) + close(s.chbp) + return nil +} + +func (s *instance) initBSB22Commitments() { + s.commitmentInfo = s.spr.CommitmentInfo.(constraint.PlonkCommitments) + s.commitmentVal = make([]fr.Element, len(s.commitmentInfo)) // TODO @Tabaie get rid of this + s.cCommitments = make([]*iop.Polynomial, len(s.commitmentInfo)) + s.proof.Bsb22Commitments = make([]kzg.Digest, len(s.commitmentInfo)) + + // override the hint for the commitment constraints + bsb22ID := solver.GetHintID(fcs.Bsb22CommitmentComputePlaceholder) + s.opt.SolverOpts = append(s.opt.SolverOpts, solver.OverrideHint(bsb22ID, s.bsb22Hint)) +} + +// Computing and verifying Bsb22 multi-commits explained in https://hackmd.io/x8KsadW3RRyX7YTCFJIkHg +func (s *instance) bsb22Hint(_ *big.Int, ins, outs []*big.Int) error { + var err error + commDepth := int(ins[0].Int64()) + ins = ins[1:] + + res := &s.commitmentVal[commDepth] + + commitmentInfo := s.spr.CommitmentInfo.(constraint.PlonkCommitments)[commDepth] + committedValues := make([]fr.Element, s.domain0.Cardinality) + offset := s.spr.GetNbPublicVariables() + for i := range ins { + committedValues[offset+commitmentInfo.Committed[i]].SetBigInt(ins[i]) + } + if _, err = committedValues[offset+commitmentInfo.CommitmentIndex].SetRandom(); err != nil { // Commitment injection constraint has qcp = 0. Safe to use for blinding. + return err + } + if _, err = committedValues[offset+s.spr.GetNbConstraints()-1].SetRandom(); err != nil { // Last constraint has qcp = 0. Safe to use for blinding + return err + } + s.cCommitments[commDepth] = iop.NewPolynomial(&committedValues, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[commDepth], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[commDepth]); err != nil { + return err + } + + s.htfFunc.Write(s.proof.Bsb22Commitments[commDepth].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + res.SetBytes(hashBts[:nbBuf]) // TODO @Tabaie use CommitmentIndex for this; create a new variable CommitmentConstraintIndex for other uses + res.BigInt(outs[0]) + + return nil +} + +// solveConstraints computes the evaluation of the L, R, O polynomials in Lagrange form. +func (s *instance) solveConstraints() error { + startSolve := time.Now() + log := logger.Logger() + + var solution *cs.SparseR1CSSolution + + // Try to load raw solver values from cache (fastest path) + rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + if rawCachePath != "" { + if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { + // Reconstruct L, R, O from raw values + var sol cs.SparseR1CSSolution + if sol.L, sol.R, sol.O, err = s.spr.EvaluateLROSmallDomainFromValues(rawValues); err != nil { + log.Warn().Err(err).Str("file", rawCachePath).Msg("ignoring raw solver cache") + } else { + log.Debug().Dur("took", time.Since(startSolve)).Int("wires", len(rawValues)).Msg("loaded raw solver values from cache") + solution = &sol + } + + // Load cached BSB22 cCommitments polynomials + cacheDir := filepath.Dir(rawCachePath) + for i := 0; solution != nil && i < len(s.commitmentInfo); i++ { + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + coeffs, err := cs.LoadRawSolverValues(bsb22Path) + if err != nil { + log.Warn().Err(err).Int("i", i).Msg("ignoring raw solver cache: missing BSB22 commitment sidecar") + solution = nil + break + } + coeffSlice := []fr.Element(coeffs) + s.cCommitments[i] = iop.NewPolynomial(&coeffSlice, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[i], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[i]); err != nil { + return err + } + s.htfFunc.Write(s.proof.Bsb22Commitments[i].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + s.commitmentVal[i].SetBytes(hashBts[:nbBuf]) + } + } + } + + if solution == nil { + _solution, err := s.spr.SolveAndSaveRawValues(s.fullWitness, rawCachePath, s.opt.SolverOpts...) + if err != nil { + log.Debug().Dur("took", time.Since(startSolve)).Err(err).Msg("solveConstraints: spr.Solve") + return err + } + log.Debug().Dur("took", time.Since(startSolve)).Msg("solveConstraints: spr.Solve") + if isProfileMode { + s.timings.record("solveConstraints: spr.Solve", time.Since(startSolve)) + } + solution = _solution.(*cs.SparseR1CSSolution) + + // Save cCommitments polynomial coefficients for BSB22 reconstruction + if rawCachePath != "" && len(s.commitmentInfo) > 0 { + cacheDir := filepath.Dir(rawCachePath) + for i := range s.commitmentInfo { + if s.cCommitments[i] != nil { + coeffs := s.cCommitments[i].Coefficients() + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + if err := cs.SaveRawSolverValues(bsb22Path, coeffs); err != nil { + log.Warn().Err(err).Int("i", i).Msg("failed to save BSB22 commitment polynomial") + } + } + } + } + } + + evaluationLDomainSmall := []fr.Element(solution.L) + evaluationRDomainSmall := []fr.Element(solution.R) + evaluationODomainSmall := []fr.Element(solution.O) + var wg sync.WaitGroup + wg.Add(2) + go func() { + s.polyL = iop.NewPolynomial(&evaluationLDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + go func() { + s.polyR = iop.NewPolynomial(&evaluationRDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + + s.polyO = iop.NewPolynomial(&evaluationODomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + + wg.Wait() + if _, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}); err != nil { + return err + } + + // commit to l, r, o and add blinding factors + if err := s.commitToLRO(); err != nil { + return err + } + close(s.chLRO) + return nil +} + +func (s *instance) completeQk() error { + qk := s.trace.Qk.Clone() + qkCoeffs := qk.Coefficients() + + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + copy(qkCoeffs, wWitness[:len(s.spr.Public)]) + + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + for i := range s.commitmentInfo { + qkCoeffs[s.spr.GetNbPublicVariables()+s.commitmentInfo[i].CommitmentIndex] = s.commitmentVal[i] + } + + s.polyQk = qk + close(s.chQk) + + return nil +} + +func (s *instance) commitToLRO() error { + var startCommitLRO time.Time + if isProfileMode { + startCommitLRO = time.Now() + } + sequentialLRO := s.domain0 != nil && s.domain0.Cardinality >= (1<<22) + if _, ok := os.LookupEnv("ICICLE_LRO_COMMIT_SEQUENTIAL"); ok { + sequentialLRO = envEnabled("ICICLE_LRO_COMMIT_SEQUENTIAL", true) + } + + if !useBlinding { + // When blinding is disabled, commit directly without waiting for blinding polynomials + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + if err != nil { + return err + } + } else { + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (no blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (no blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil + } + + // wait for blinding polynomials to be initialized or context to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chbp: + } + + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + if err != nil { + return err + } + } else { + // Run the three commitments concurrently. + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (with blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil +} + +// deriveGammaAndBeta (copy constraint) +func (s *instance) deriveGammaAndBeta() error { + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + if err := bindPublicData(s.fs, "gamma", s.pk.VerifyingKey().(*plonk_bls12377.VerifyingKey), wWitness[:len(s.spr.Public)]); err != nil { + return err + } + + // wait for LRO to be committed + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + gamma, err := deriveRandomness(s.fs, "gamma", &s.proof.LRO[0], &s.proof.LRO[1], &s.proof.LRO[2]) + if err != nil { + return err + } + + bbeta, err := s.fs.ComputeChallenge("beta") + if err != nil { + return err + } + s.gamma = gamma + s.beta.SetBytes(bbeta) + + close(s.chGammaBeta) + + return nil +} + +// commitToPolyAndBlinding computes the KZG commitment of a polynomial p +// in Lagrange form (large degree) +// and add the contribution of a blinding polynomial b (small degree) +// /!\ The polynomial p is supposed to be in Lagrange form. +// Only used when blinding is enabled (the default). +func (s *instance) commitToPolyAndBlinding(p, b *iop.Polynomial) (commit curve.G1Affine, err error) { + // Commit over the Lagrange SRS using shared device-resident polynomial data. + gpuCommit, err := s.commitLagrangePolynomialOnGPU(p) + if err != nil { + return curve.G1Affine{}, err + } + + // add CPU blinding contribution (two MSMs on canonical SRS) + n := int(s.domain0.Cardinality) + cb := commitBlindingFactor(n, b, s.pk.Kzg) + gpuCommit.Add(&gpuCommit, &cb) + return gpuCommit, nil +} + +func (s *instance) commitLagrangePolynomialOnGPU(p *iop.Polynomial) (curve.G1Affine, error) { + if p == nil { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: nil polynomial") + } + if p.Basis != iop.Lagrange { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial must be in Lagrange basis, got %v", p.Basis) + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{p}) + if err != nil { + return curve.G1Affine{}, err + } + idx, ok := gpuState.polyToIdx[p] + if !ok || idx < 0 || idx >= len(gpuState.deviceSlices) { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial is missing from shared GPU state") + } + if gpuState.deviceSlices[idx].IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: empty device slice for polynomial") + } + + // Keep the polynomial in its native Lagrange basis; large MSMs are split + // into device-side chunks inside commitOnGPULagrangeDevice. + return commitOnGPULagrangeDevice(gpuState.deviceSlices[idx], &s.device, s.pk) +} + +func (s *instance) deriveAlpha() (err error) { + alphaDeps := make([]*curve.G1Affine, len(s.proof.Bsb22Commitments)+1) + for i := range s.proof.Bsb22Commitments { + alphaDeps[i] = &s.proof.Bsb22Commitments[i] + } + alphaDeps[len(alphaDeps)-1] = &s.proof.Z + s.alpha, err = deriveRandomness(s.fs, "alpha", alphaDeps...) + return err +} + +func (s *instance) deriveZeta() (err error) { + s.zeta, err = deriveRandomness(s.fs, "zeta", &s.proof.H[0], &s.proof.H[1], &s.proof.H[2]) + return +} + +func (s *instance) computeQuotient() (err error) { + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + if isProfileMode { + var startComputeQuotient time.Time + startComputeQuotient = time.Now() + defer func() { + l := logger.Logger() + if err != nil { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Err(err).Msg("computeQuotient: total (with error)") + } else { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Msg("computeQuotient: total") + } + s.timings.record("computeQuotient (total)", time.Since(startComputeQuotient)) + }() + } + + // wait for Z to be committed or context done + doneWaitZ := profileStep("computeQuotient: wait Z commit") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZ: + } + doneWaitZ() + + // derive alpha + if err = s.deriveAlpha(); err != nil { + return err + } + + if err := s.waitForComputeNumeratorQk(); err != nil { + return err + } + if s.polyQk == nil { + return fmt.Errorf("computeQuotient: missing completed Qk polynomial") + } + + doneEnsureGPUState := profileStep("computeQuotient: ensure shared GPU state") + gpuState, err := s.ensurePolysOnSharedGPU(s.buildComputeNumeratorGPUBatch()) + if err != nil { + return err + } + doneEnsureGPUState() + + // compute Z shifted by one for copy-constraint terms. + if s.polyZ == nil { + return fmt.Errorf("computeQuotient: missing Z polynomial") + } + s.polyZS = s.polyZ.ShallowClone().Shift(1) + + var numeratorGPU *gpuNumeratorPolynomial + var quotientGPU *gpuQuotientPolynomial + var e error + doneComputeNumerator := profileStep("computeQuotient: computeNumerator") + numeratorGPU, e = s.computeNumerator(gpuState) + if e != nil { + return e + } + doneComputeNumerator() + + doneDivideByZH := profileStep("computeQuotient: divideByZHOnGPU") + quotientGPU, e = s.divideByZHOnGPU(numeratorGPU, [2]*fft.Domain{s.domain0, s.domain1}) + if e != nil { + return e + } + doneDivideByZH() + s.hGPU = quotientGPU + + // Shared state slices were mutated during numerator coset iterations and are no + // longer needed now; computeLinearizedPolynomial uses the immutable snapshot. + s.releaseSharedGPUState() + close(s.chRestoreLRO) + + doneCommitH := profileStep("computeQuotient: commit H from device") + if err := s.commitToQuotientGPUFromDevice(s.hGPU); err != nil { + return err + } + doneCommitH() + + if err := s.deriveZeta(); err != nil { + return err + } + + donePrepareLinearizedEval := profileStep("computeQuotient: prepare linearized eval GPU state") + if err := s.prepareLinearizedEvalGPUStateFromHost(); err != nil { + return fmt.Errorf("computeQuotient: prepare linearized eval GPU state failed: %w", err) + } + donePrepareLinearizedEval() + + close(s.chH) + + return nil +} + +func (s *instance) buildRatioCopyConstraint() (err error) { + // wait for gamma and beta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chGammaBeta: + } + + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("buildRatioCopyConstraint: missing L/R/O polynomials") + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}) + if err != nil { + return err + } + dL, err := getStateDeviceSlice(gpuState, s.polyL, "L") + if err != nil { + return err + } + dR, err := getStateDeviceSlice(gpuState, s.polyR, "R") + if err != nil { + return err + } + dO, err := getStateDeviceSlice(gpuState, s.polyO, "O") + if err != nil { + return err + } + + var startBuildRatioCopyConstraintIcicle time.Time + if isProfileMode { + startBuildRatioCopyConstraintIcicle = time.Now() + } + s.polyZ, err = s.BuildRatioCopyConstraintIcicle( + []icicle_core.DeviceSlice{dL, dR, dO}, + s.trace.S, + s.beta, + s.gamma, + s.domain0, + gpuState, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startBuildRatioCopyConstraintIcicle)).Msg("buildRatioCopyConstraint: BuildRatioCopyConstraintIcicle") + s.timings.record("buildRatioCopyConstraint (perm Z)", time.Since(startBuildRatioCopyConstraintIcicle)) + } + + dZ, err := getStateDeviceSlice(gpuState, s.polyZ, "Z") + if err != nil { + return err + } + copyDone := make(chan error, 1) + var dPersist icicle_core.DeviceSlice + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + var allocErr error + dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) + if allocErr != nil { + finish(fmt.Errorf("buildRatioCopyConstraint: alloc persist Z failed: %w", allocErr)) + return + } + if e := copyDeviceSliceIntoOnCurrentDevice(dPersist, dZ, cfg); e != icicle_runtime.Success { + _ = dPersist.Free() + dPersist = icicle_core.DeviceSlice{} + finish(fmt.Errorf("buildRatioCopyConstraint: persist Z copy failed: %s", e.AsString())) + return + } + finish(nil) + }) + if err := <-copyDone; err != nil { + return err + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + s.polyZLagrangeGPU = dPersist + + // commit to Z (with or without blinding) + var startCommitZ time.Time + if isProfileMode { + startCommitZ = time.Now() + } + if useBlinding { + s.proof.Z, err = s.commitToPolyAndBlinding(s.polyZ, s.bp[id_Bz]) + } else { + s.proof.Z, err = s.commitLagrangePolynomialOnGPU(s.polyZ) + } + if isProfileMode { + l := logger.Logger() + if useBlinding { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (no blinding)") + } + } + s.freeIdleTempGPUMemoryOnDevice() + + close(s.chZ) + + return +} + +// open Z (blinded) at ωζ +func (s *instance) openZ() (err error) { + // wait for H to be committed and zeta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + + var zetaShifted fr.Element + zetaShifted.Mul(&s.zeta, &s.pk.Vk.Generator) + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("openZ: missing GPU Z polynomial") + } + + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + + dZLagrange := s.polyZLagrangeGPU + if dZLagrange.Len() <= 1 { + return fmt.Errorf("openZ: invalid Z size %d", dZLagrange.Len()) + } + + blindSize := order_blinding_Z + 1 + if useBlinding { + if len(s.bp) <= id_Bz || s.bp[id_Bz] == nil { + return fmt.Errorf("openZ: missing Z blinding polynomial") + } + blindSize = len(s.bp[id_Bz].Coefficients()) + if blindSize == 0 { + return fmt.Errorf("openZ: empty Z blinding polynomial") + } + } + + var dBlindedCanonical icicle_core.DeviceSlice + buildDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + buildDone <- cfgErr + return + } + finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + // Async boundary for canonicalization/blinding before exposing output. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + buildDone <- runErr + } + + n := dZLagrange.Len() + dZCanonical := s.getTempDeviceSlice(n) + if err := copyDeviceSliceIntoOnCurrentDevice(dZCanonical, dZLagrange, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy Z to canonical buffer failed: %s", err.AsString()), dZCanonical, false) + return + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + cfgNtt.Ordering = icicle_core.KNN // regular lagrange -> regular canonical + if err := icicle_ntt.Ntt(dZCanonical, icicle_core.KInverse, &cfgNtt, dZCanonical); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: inverse NTT on Z failed: %s", err.AsString()), dZCanonical, false) + return + } + + dBlindedCanonical = s.getTempDeviceSlice(n + blindSize) + dBlindedPrefix := (&dBlindedCanonical).Range(0, n, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedPrefix, dZCanonical, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy canonical Z into blinded buffer failed: %s", err.AsString()), dZCanonical, true) + return + } + + if useBlinding { + dBp := uploadVector(s.bp[id_Bz].Coefficients()) + dBlindedHead := (&dBlindedPrefix).Range(0, blindSize, false) + if err := icicle_vecops.VecOp(dBlindedHead, dBp, dBlindedHead, cfgVec, icicle_core.Sub); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: subtract Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedTail, dBp, cfgVec); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: append Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + _ = dBp.FreeAsync(stream) + } else { + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := zeroDeviceSliceOnCurrentDevice(dBlindedTail, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: zero-pad non-blinded Z failed: %s", err.AsString()), dZCanonical, true) + return + } + } + + finalize(nil, dZCanonical, false) + }) + if err := <-buildDone; err != nil { + return err + } + s.blindedZCanonicalGPU = dBlindedCanonical + + // open z at zeta*w. + var startKzgOpen time.Time + if isProfileMode { + startKzgOpen = time.Now() + } + s.proof.ZShiftedOpening, err = s.openPolynomialOnGPUCanonicalDevice(s.blindedZCanonicalGPU, zetaShifted) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startKzgOpen)).Msg("openZ: open polynomial on GPU") + s.timings.record("openZ (KZG open on GPU)", time.Since(startKzgOpen)) + } + if err != nil { + return err + } + close(s.chZOpening) + return nil +} + +func (s *instance) h1() []fr.Element { + var h1 []fr.Element + if !s.opt.StatisticalZK { + h1 = s.h.Coefficients()[:s.domain0.Cardinality+2] + } else { + h1 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h1, s.h.Coefficients()[:s.domain0.Cardinality+2]) + h1[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[0]) + } + return h1 +} + +func (s *instance) h2() []fr.Element { + var h2 []fr.Element + if !s.opt.StatisticalZK { + h2 = s.h.Coefficients()[s.domain0.Cardinality+2 : 2*(s.domain0.Cardinality+2)] + } else { + h2 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h2, s.h.Coefficients()[s.domain0.Cardinality+2:2*(s.domain0.Cardinality+2)]) + h2[0].Sub(&h2[0], &s.quotientShardsRandomizers[0]) + h2[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[1]) + } + return h2 +} + +func (s *instance) h3() []fr.Element { + var h3 []fr.Element + if !s.opt.StatisticalZK { + h3 = s.h.Coefficients()[2*(s.domain0.Cardinality+2) : 3*(s.domain0.Cardinality+2)] + } else { + h3 = make([]fr.Element, s.domain0.Cardinality+2) + copy(h3, s.h.Coefficients()[2*(s.domain0.Cardinality+2):3*(s.domain0.Cardinality+2)]) + h3[0].Sub(&h3[0], &s.quotientShardsRandomizers[1]) + } + return h3 +} + +// witnessEvalAtZeta holds the scalar evaluations of witness and constraint +// polynomials at the challenge point zeta, as needed by the linearized +// polynomial computation. +type witnessEvalAtZeta struct { + blzeta, brzeta, bozeta fr.Element + s1zeta, s2zeta fr.Element + qcpzeta []fr.Element +} + +type linearizedSelectorScales struct { + s3, ql, qr, qm, qo, qk fr.Element + qcp []fr.Element +} + +// evaluateWitnessPolynomialsAtZeta evaluates L, R, O (with optional blinding), +// S1, S2, and all Qcp polynomials at the point zeta using the GPU-resident +// polynomial state. +func (s *instance) evaluateWitnessPolynomialsAtZeta( + evalGPUState *gpuPolysState, + zeta fr.Element, +) (witnessEvalAtZeta, error) { + doneTotal := profileStep("evaluateWitnessPolynomialsAtZeta: total") + defer doneTotal() + + var result witnessEvalAtZeta + var err error + + result.qcpzeta = make([]fr.Element, len(s.commitmentInfo)) + var startQcp time.Time + if isProfileMode { + startQcp = time.Now() + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i >= len(s.trace.Qcp) || s.trace.Qcp[i] == nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: missing Qcp polynomial at index %d", i) + } + var startQcpItem time.Time + if isProfileMode { + startQcpItem = time.Now() + } + result.qcpzeta[i], err = s.evalPolynomialInCurrentFormOnGPU(s.trace.Qcp[i], evalGPUState, zeta) + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: qcp[%d] GPU evaluation failed: %w", i, err) + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("idx", i).Dur("took", time.Since(startQcpItem)).Msg("evaluateWitnessPolynomialsAtZeta: qcp eval item") + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startQcp)).Msg("evaluateWitnessPolynomialsAtZeta: qcpZeta evaluate on GPU") + } + + if useBlinding { + result.blzeta, err = s.evaluateBlindedOnGPU(s.polyL, s.bp[id_Bl], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: L evaluate on GPU") + result.blzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyL, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: blzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.brzeta, err = s.evaluateBlindedOnGPU(s.polyR, s.bp[id_Br], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: R evaluate on GPU") + result.brzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyR, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: brzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.bozeta, err = s.evaluateBlindedOnGPU(s.polyO, s.bp[id_Bo], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: O evaluate on GPU") + result.bozeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyO, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: bozeta GPU evaluation failed: %w", err) + } + + doneS1 := profileStep("evaluateWitnessPolynomialsAtZeta: S1 evaluate on GPU") + result.s1zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S1, evalGPUState, zeta) + doneS1() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s1(zeta) GPU evaluation failed: %w", err) + } + doneS2 := profileStep("evaluateWitnessPolynomialsAtZeta: S2 evaluate on GPU") + result.s2zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S2, evalGPUState, zeta) + doneS2() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s2(zeta) GPU evaluation failed: %w", err) + } + + return result, nil +} + +func (s *instance) computeLinearizedPolynomial() error { + + // wait for H to be committed and zeta to be derived (or ctx.Done()) + var startWaitH time.Time + if isProfileMode { + startWaitH = time.Now() + } + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startWaitH)).Msg("computeLinearizedPolynomial: wait H and zeta") + s.timings.record("computeLinearizedPoly (wait H+zeta, overlaps)", time.Since(startWaitH)) + } + if s.opt.StatisticalZK { + return fmt.Errorf("computeLinearizedPolynomial: GPU-only opening path does not support StatisticalZK=true") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil || s.polyZ == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing required polynomials") + } + + // Reuse the immutable snapshot prepared in computeQuotient before numerator + // coset iterations mutate shared state slices. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("computeLinearizedPolynomial: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized eval prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: required polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing blinding polynomials for GPU evaluation") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized blinding prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: blinding polynomial is not on GPU: %w", e) + } + } + } + + doneEvaluate := profileStep("computeLinearizedPolynomial: evaluate witness polynomials") + evals, err := s.evaluateWitnessPolynomialsAtZeta(evalGPUState, s.zeta) + doneEvaluate() + if err != nil { + return err + } + + // wait for Z to be opened at zeta (or ctx.Done()) + doneWaitZOpening := profileStep("computeLinearizedPolynomial: wait Z opening") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZOpening: + } + doneWaitZOpening() + if s.blindedZCanonicalGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing canonical blinded Z on GPU") + } + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing lagrange Z on GPU") + } + defer func() { + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + }() + bzuzeta := s.proof.ZShiftedOpening.ClaimedValue + + if s.hGPU == nil || s.hGPU.coeffs.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing GPU quotient polynomial") + } + + doneBuild := profileStep("computeLinearizedPolynomial: build selector terms on GPU") + dLin, err := s.buildLinearizedSelectorTermsOnGPU(evals, bzuzeta, s.blindedZCanonicalGPU.Len()) + doneBuild() + if err != nil { + return err + } + + doneAddZ := profileStep("computeLinearizedPolynomial: add Z contribution on GPU") + err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) + doneAddZ() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + s.linearizedPolynomialGPU = dLin + + doneEvalClaim := profileStep("computeLinearizedPolynomial: evaluate linearized claim") + claim, err := s.evalDevicePolynomialAtPoint(dLin, s.zeta) + doneEvalClaim() + if err != nil { + return err + } + s.linearizedPolynomialClaim = claim + + // Commit the linearized polynomial over the canonical SRS. + var startMSM time.Time + if isProfileMode { + startMSM = time.Now() + } + s.linearizedPolynomialDigest, err = commitOnGPUCanonicalDevice(dLin, &s.device, s.pk) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startMSM)).Msg("computeLinearizedPolynomial: KZG commit") + s.timings.record("computeLinearizedPoly (KZG commit)", time.Since(startMSM)) + } + close(s.chLinearizedPolynomial) + return nil +} + +func (s *instance) batchOpening() error { + + // wait for linearizedPolynomial to be computed (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLinearizedPolynomial: + } + + defer func() { + freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } + }() + + if s.linearizedPolynomialGPU.IsEmpty() { + return fmt.Errorf("batchOpening: missing GPU linearized polynomial") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("batchOpening: missing L/R/O polynomials") + } + + // Reuse immutable GPU snapshot prepared before quotient iterations. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("batchOpening: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening eval prereq"); e != nil { + return fmt.Errorf("batchOpening: required polynomial is not on GPU: %w", e) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] == nil { + return fmt.Errorf("batchOpening: missing Qcp polynomial at index %d", i) + } + if _, e := getStateDeviceSlice(evalGPUState, s.trace.Qcp[i], "batchOpening qcp prereq"); e != nil { + return fmt.Errorf("batchOpening: Qcp polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("batchOpening: missing blinding polynomials") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening blinding prereq"); e != nil { + return fmt.Errorf("batchOpening: blinding polynomial is not on GPU: %w", e) + } + } + } + + devicePolys, ownedPolys, claimed, err := s.prepareBatchOpeningPolynomialsOnGPU(evalGPUState, s.zeta) + if err != nil { + return err + } + defer func() { + for i := 0; i < len(devicePolys); i++ { + if i < len(ownedPolys) && ownedPolys[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + } + } + s.releaseLinearizedEvalGPUState() + }() + + digestsToOpen := make([]curve.G1Affine, len(s.pk.Vk.Qcp)+6) + copy(digestsToOpen[6:], s.pk.Vk.Qcp) + digestsToOpen[0] = s.linearizedPolynomialDigest + digestsToOpen[1] = s.proof.LRO[0] + digestsToOpen[2] = s.proof.LRO[1] + digestsToOpen[3] = s.proof.LRO[2] + digestsToOpen[4] = s.pk.Vk.S[0] + digestsToOpen[5] = s.pk.Vk.S[1] + if len(claimed) != len(digestsToOpen) { + return fmt.Errorf("batchOpening: claimed size mismatch (%d != %d)", len(claimed), len(digestsToOpen)) + } + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + var dFold icicle_core.DeviceSlice + foldDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening fold", foldDone) + + dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: copy linearized polynomial failed: %s", e.AsString())) + return + } + + gammaPow := gamma + for i := 1; i < len(devicePolys); i++ { + dPoly := devicePolys[i] + if dPoly.IsEmpty() { + gammaPow.Mul(&gammaPow, &gamma) + continue + } + dScaled := s.getTempDeviceSlice(dPoly.Len()) + dGammaStd := uploadScalarStdOnCurrentDevice(gammaPow, cfg) + eMul := icicle_vecops.ScalarMulVec(dGammaStd, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dGammaStd.FreeAsync(cfg.StreamHandle) + } else { + _ = dGammaStd.Free() + } + if eMul != icicle_runtime.Success { + if cfg.IsAsync { + _ = icicle_runtime.SynchronizeStream(cfg.StreamHandle) + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: scale polynomial %d failed: %s", i, eMul.AsString())) + return + } + + dPrefix := (&dFold).Range(0, dPoly.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add) + if cfg.IsAsync { + // dScaled is recycled each iteration; wait before returning to pool. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: synchronize stream failed: %s", eSync.AsString())) + return + } + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + if eAdd != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: fold add polynomial %d failed: %s", i, eAdd.AsString())) + return + } + gammaPow.Mul(&gammaPow, &gamma) + } + finish(nil) + }) + if err := <-foldDone; err != nil { + if !dFold.IsEmpty() { + s.putTempDeviceSlice(dFold, dFold.Len()) + } + return err + } + var dWitness icicle_core.DeviceSlice + divDone := make(chan error, 1) + witnessSize := dFold.Len() - 1 + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dFold, dPoint, dWitness, cfg) + if cfg.IsAsync { + _ = dPoint.FreeAsync(cfg.StreamHandle) + } else { + _ = dPoint.Free() + } + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + finish(fmt.Errorf("batchOpening: divide by (x-zeta) failed: %s", eDiv.AsString())) + return + } + finish(nil) + }) + if err := <-divDone; err != nil { + s.putTempDeviceSlice(dFold, dFold.Len()) + return err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dFold, dFold.Len()) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return err + } + + s.proof.BatchedProof = kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + } + if err := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); err != nil { + l := logger.Logger() + l.Warn().Err(err).Msg("batchOpening: GPU folded opening failed raw-G2 validation; falling back to host fold with CPU KZG commitment") + fallbackProof, fallbackErr := s.batchOpeningHostFoldGPUCommitFromDevicePolys(devicePolys, digestsToOpen) + if fallbackErr != nil { + return fallbackErr + } + s.proof.BatchedProof = fallbackProof + if verifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); verifyErr != nil { + nativeProof, nativeErr := s.batchOpeningNativeCPUFromDevicePolys(devicePolys, digestsToOpen) + if nativeErr != nil { + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed: %v", verifyErr, nativeErr) + } + s.proof.BatchedProof = nativeProof + if nativeVerifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); nativeVerifyErr != nil { + diagnostic, diagnosticErr := s.diagnoseBatchOpeningDevicePolynomials(devicePolys, digestsToOpen, claimed) + if diagnosticErr != nil { + diagnostic = fmt.Sprintf("batch opening diagnostic failed: %v", diagnosticErr) + } + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed raw-G2 validation: %v; %s", verifyErr, nativeVerifyErr, diagnostic) + } + l.Warn().Msg("batchOpening: native CPU KZG fallback produced a valid proof after host-fold fallback failed") + } + } + _ = foldedEval // kept for parity with kzg.BatchOpenSinglePoint flow. + return nil +} + +func (s *instance) verifyBatchOpeningWithRecomputedLines(digestsToOpen []curve.G1Affine) error { + vk := s.pk.Vk.Kzg + vk.Lines[0] = curve.PrecomputeLines(vk.G2[0]) + vk.Lines[1] = curve.PrecomputeLines(vk.G2[1]) + return kzg.BatchVerifySinglePoint( + digestsToOpen, + &s.proof.BatchedProof, + s.zeta, + s.kzgFoldingHash, + vk, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningHostFoldGPUCommitFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningNativeCPUFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningNativeCPUFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return kzg.BatchOpenSinglePoint( + polysToOpen, + digestsToOpen, + s.zeta, + s.kzgFoldingHash, + s.pk.Kzg, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) diagnoseBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, + gpuClaimed []fr.Element, +) (string, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "diagnoseBatchOpeningDevicePolynomials", + ) + if err != nil { + return "", err + } + if len(gpuClaimed) != len(polysToOpen) { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: claimed/polynomial mismatch (%d != %d)", len(gpuClaimed), len(polysToOpen)) + } + + l := logger.Logger() + claimMismatches := make([]string, 0) + commitMismatches := make([]string, 0) + for i := range polysToOpen { + label := batchOpeningPolynomialLabel(i) + cpuClaim := evalCanonicalAtPoint(polysToOpen[i], s.zeta) + if !cpuClaim.Equal(&gpuClaimed[i]) { + claimMismatches = append(claimMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("gpuClaim", frFingerprint(gpuClaimed[i])). + Str("cpuClaim", frFingerprint(cpuClaim)). + Msg("batchOpening diagnostic: GPU claim differs from CPU evaluation") + } + + cpuDigest, err := kzg.Commit(polysToOpen[i], s.pk.Kzg) + if err != nil { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: commit %s: %w", label, err) + } + if !cpuDigest.Equal(&digestsToOpen[i]) { + commitMismatches = append(commitMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("expectedDigest", g1Fingerprint(digestsToOpen[i])). + Str("cpuDigest", g1Fingerprint(cpuDigest)). + Msg("batchOpening diagnostic: CPU commitment differs from proof digest") + } + } + + if len(claimMismatches) == 0 && len(commitMismatches) == 0 { + return "batch opening diagnostic found no per-polynomial claim or commitment mismatch", nil + } + return fmt.Sprintf( + "batch opening diagnostic claim mismatches=[%s] commitment mismatches=[%s]", + strings.Join(claimMismatches, ","), + strings.Join(commitMismatches, ","), + ), nil +} + +func batchOpeningPolynomialLabel(index int) string { + switch index { + case 0: + return "linearized" + case 1: + return "L" + case 2: + return "R" + case 3: + return "O" + case 4: + return "S1" + case 5: + return "S2" + default: + return fmt.Sprintf("Qcp[%d]", index-6) + } +} + +func frFingerprint(v fr.Element) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func g1Fingerprint(v curve.G1Affine) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func (s *instance) downloadBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + expectedDigests int, + label string, +) ([][]fr.Element, error) { + if len(devicePolys) != expectedDigests { + return nil, fmt.Errorf("%s: polynomial/digest mismatch (%d != %d)", label, len(devicePolys), expectedDigests) + } + + polysToOpen := make([][]fr.Element, len(devicePolys)) + for i := range devicePolys { + var err error + polysToOpen[i], err = s.downloadCanonicalDeviceCoefficients( + devicePolys[i], + fmt.Sprintf("%s[%d]", label, i), + ) + if err != nil { + return nil, err + } + } + return polysToOpen, nil +} + +func (s *instance) batchOpeningHostFoldGPUCommit(digestsToOpen []curve.G1Affine) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.batchOpeningHostPolynomials() + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( + polysToOpen [][]fr.Element, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + if len(polysToOpen) != len(digestsToOpen) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: polynomial/digest mismatch (%d != %d)", len(polysToOpen), len(digestsToOpen)) + } + + largestPoly := 0 + for i := range polysToOpen { + if len(polysToOpen[i]) == 0 || len(polysToOpen[i]) > len(s.pk.Kzg.G1) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: invalid polynomial %d size %d", i, len(polysToOpen[i])) + } + if len(polysToOpen[i]) > largestPoly { + largestPoly = len(polysToOpen[i]) + } + } + + claimed := make([]fr.Element, len(polysToOpen)) + utils.Parallelize(len(polysToOpen), func(start, end int) { + for i := start; i < end; i++ { + claimed[i] = evalCanonicalAtPoint(polysToOpen[i], s.zeta) + } + }) + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + foldedPolynomials := make([]fr.Element, largestPoly) + copy(foldedPolynomials, polysToOpen[0]) + gammaPow := gamma + for i := 1; i < len(polysToOpen); i++ { + poly := polysToOpen[i] + scale := gammaPow + utils.Parallelize(len(poly), func(start, end int) { + var term fr.Element + for j := start; j < end; j++ { + term.Mul(&poly[j], &scale) + foldedPolynomials[j].Add(&foldedPolynomials[j], &term) + } + }) + gammaPow.Mul(&gammaPow, &gamma) + } + + hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) + var dWitness icicle_core.DeviceSlice + uploadDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + dWitness = uploadVector(hCoeffs) + close(uploadDone) + }) + <-uploadDone + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + freeSliceOnDevice(&dWitness, &s.device) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + return kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + }, nil +} + +func (s *instance) batchOpeningHostPolynomials() ([][]fr.Element, error) { + total := 6 + len(s.trace.Qcp) + polysToOpen := make([][]fr.Element, total) + + var err error + polysToOpen[0], err = s.downloadCanonicalDeviceCoefficients( + s.linearizedPolynomialGPU, + "batchOpeningHostPolynomials linearized", + ) + if err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare linearized: %w", err) + } + + prepareLRO := func(p, bp *iop.Polynomial, blindingOrder int) ([]fr.Element, error) { + base, err := canonicalRegularCoefficientsCopy(p, s.domain0) + if err != nil { + return nil, err + } + if useBlinding { + if bp == nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomial") + } + blind, err := canonicalRegularCoefficientsCopy(bp, s.domain0) + if err != nil { + return nil, err + } + out := make([]fr.Element, len(base)+len(blind)) + copy(out, base) + copy(out[len(base):], blind) + for i := range blind { + out[i].Sub(&out[i], &blind[i]) + } + return out, nil + } + out := make([]fr.Element, len(base)+blindingOrder+1) + copy(out, base) + return out, nil + } + + var bpL, bpR, bpO *iop.Polynomial + if useBlinding { + if len(s.bp) <= id_Bo { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomials") + } + bpL, bpR, bpO = s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo] + } + + if polysToOpen[1], err = prepareLRO(s.polyL, bpL, order_blinding_L); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare L: %w", err) + } + if polysToOpen[2], err = prepareLRO(s.polyR, bpR, order_blinding_R); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare R: %w", err) + } + if polysToOpen[3], err = prepareLRO(s.polyO, bpO, order_blinding_O); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare O: %w", err) + } + if polysToOpen[4], err = canonicalRegularCoefficientsCopy(s.trace.S1, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S1: %w", err) + } + if polysToOpen[5], err = canonicalRegularCoefficientsCopy(s.trace.S2, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S2: %w", err) + } + for i := range s.trace.Qcp { + if polysToOpen[6+i], err = canonicalRegularCoefficientsCopy(s.trace.Qcp[i], s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare Qcp[%d]: %w", i, err) + } + } + return polysToOpen, nil +} + +func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceSlice, label string) ([]fr.Element, error) { + if dPoly.IsEmpty() { + return nil, fmt.Errorf("%s: empty device polynomial", label) + } + + coeffs := make([]fr.Element, dPoly.Len()) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + }) + if err := <-done; err != nil { + return nil, err + } + return coeffs, nil +} + +func canonicalRegularCoefficientsCopy(p *iop.Polynomial, domain *fft.Domain) ([]fr.Element, error) { + if p == nil { + return nil, fmt.Errorf("nil polynomial") + } + cp := p.Clone() + cp.ToCanonical(domain).ToRegular() + coeffs := cp.Coefficients() + out := make([]fr.Element, len(coeffs)) + copy(out, coeffs) + return out, nil +} + +func dividePolyByXMinusAHost(f []fr.Element, fa, a fr.Element) []fr.Element { + f[0].Sub(&f[0], &fa) + var t fr.Element + for i := len(f) - 2; i >= 0; i-- { + t.Mul(&f[i+1], &a) + f[i].Add(&f[i], &t) + } + return f[1:] +} + +// evaluate the full set of constraints on the GPU-resident polynomial state. +type computeNumeratorLoopContext struct { + n int + rho int + mm uint64 + bn *big.Int + shifters []fr.Element + twiddles0 []fr.Element + dTwiddles0 icicle_core.DeviceSlice + dPrecomputedDenominators *icicle_core.DeviceSlice + scalingVector []fr.Element + scalingVectorRev []fr.Element + gpuState *gpuPolysState + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + one fr.Element + cs fr.Element + css fr.Element + nbBsbGates int + numeratorShards []icicle_core.DeviceSlice +} + +type gpuNumeratorPolynomial struct { + shards []icicle_core.DeviceSlice + n int + rho int + mm uint64 +} + +type gpuQuotientPolynomial struct { + coeffs icicle_core.DeviceSlice + size int +} + +func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolynomial, error) { + twiddles0, err := s.buildComputeNumeratorTwiddles() + if err != nil { + return nil, err + } + if err := s.validateComputeNumeratorGPUState(gpuState); err != nil { + return nil, err + } + + var startComputeNumerator time.Time + if isProfileMode { + startComputeNumerator = time.Now() + } + + n := s.domain0.Cardinality + nbBsbGates := len(s.proof.Bsb22Commitments) + + var cs, css fr.Element + cs.Set(&s.domain1.FrMultiplicativeGen) + css.Square(&cs) + + bn := big.NewInt(int64(n)) + + rho := int(s.domain1.Cardinality / n) + shifters := make([]fr.Element, rho) + shifters[0].Set(&s.domain1.FrMultiplicativeGen) + for i := 1; i < rho; i++ { + shifters[i].Set(&s.domain1.Generator) + } + + cosetTable, err := s.domain0.CosetTable() + if err != nil { + return nil, err + } + + // for the first iteration, the scalingVector is the coset table + scalingVector := cosetTable + scalingVectorRev := make([]fr.Element, len(cosetTable)) + copy(scalingVectorRev, cosetTable) + fft.BitReverse(scalingVectorRev) + + // pre-computed to compute the bit reverse index + // of the result polynomial + m := uint64(s.domain1.Cardinality) + mm := uint64(64 - bits.TrailingZeros64(m)) + + var dPrecomputedDenominators icicle_core.DeviceSlice + defer func() { + if !dPrecomputedDenominators.IsEmpty() { + freeDone := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + freeDone <- dPrecomputedDenominators.Free() + }) + if err := <-freeDone; err != icicle_runtime.Success { + panic(fmt.Sprintf("computeNumerator: failed to free dPrecomputedDenominators: %s", err.AsString())) + } + } + }() + + var coset, cosetExponentiatedToNMinusOne, one fr.Element + coset.SetOne() + one.SetOne() + + dTwiddles0, err := s.uploadComputeNumeratorTwiddles(twiddles0) + if err != nil { + return nil, err + } + + loopCtx := &computeNumeratorLoopContext{ + n: int(n), + rho: rho, + mm: mm, + bn: bn, + shifters: shifters, + twiddles0: twiddles0, + dTwiddles0: dTwiddles0, + dPrecomputedDenominators: &dPrecomputedDenominators, + scalingVector: scalingVector, + scalingVectorRev: scalingVectorRev, + gpuState: gpuState, + coset: coset, + cosetExponentiatedToNMinusOne: cosetExponentiatedToNMinusOne, + one: one, + cs: cs, + css: css, + nbBsbGates: nbBsbGates, + numeratorShards: make([]icicle_core.DeviceSlice, rho), + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: setup before iteration loop") + } + if err := s.executeComputeNumeratorCosetIterations(loopCtx); err != nil { + return nil, err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: main body (post-wait)") + } + + return &gpuNumeratorPolynomial{ + shards: loopCtx.numeratorShards, + n: loopCtx.n, + rho: loopCtx.rho, + mm: loopCtx.mm, + }, nil + +} + +func (s *instance) buildComputeNumeratorTwiddles() ([]fr.Element, error) { + n := s.domain0.Cardinality + var startTwiddles time.Time + if isProfileMode { + startTwiddles = time.Now() + } + twiddles0 := make([]fr.Element, n) + if n == 1 { + // edge case + twiddles0[0].SetOne() + } else { + twiddles, err := s.domain0.Twiddles() + if err != nil { + return nil, err + } + copy(twiddles0, twiddles[0]) + w := twiddles0[1] + for i := len(twiddles[0]); i < len(twiddles0); i++ { + twiddles0[i].Mul(&twiddles0[i-1], &w) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startTwiddles)).Msg("computeNumerator: build twiddles") + } + return twiddles0, nil +} + +func (s *instance) waitForComputeNumeratorQk() error { + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chQk: + } + return nil +} + +func (s *instance) buildLinearizedEvalGPUBatch() []*iop.Polynomial { + baseCap := 5 + len(s.trace.Qcp) + if useBlinding { + baseCap += 3 + } + polys := make([]*iop.Polynomial, 0, baseCap) + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + } + if useBlinding && len(s.bp) > id_Bo { + for _, bpPoly := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if bpPoly != nil { + polys = append(polys, bpPoly) + } + } + } + return polys +} + +func (s *instance) buildComputeNumeratorGPUBatch() []*iop.Polynomial { + polys := make([]*iop.Polynomial, 0, 13+2*len(s.commitmentInfo)) + for _, p := range []*iop.Polynomial{ + s.polyL, s.polyR, s.polyO, s.polyZ, + s.trace.Ql, s.trace.Qr, s.trace.Qm, s.trace.Qo, s.polyQk, + s.trace.S1, s.trace.S2, s.trace.S3, + } { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i < len(s.trace.Qcp) && s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + if i < len(s.cCommitments) && s.cCommitments[i] != nil { + polys = append(polys, s.cCommitments[i]) + } + } + return polys +} + +func (s *instance) validateComputeNumeratorGPUState(state *gpuPolysState) error { + if state == nil { + return fmt.Errorf("computeNumerator: shared GPU state is nil") + } + required := s.buildComputeNumeratorGPUBatch() + if len(required) == 0 { + return fmt.Errorf("computeNumerator: no polynomials prepared for GPU batch") + } + for _, p := range required { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; !ok { + return fmt.Errorf("computeNumerator: polynomial ptr=%p is missing from shared GPU state", p) + } + } + return nil +} + +func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicle_core.DeviceSlice, error) { + var dTwiddles0 icicle_core.DeviceSlice + uploadTwiddlesDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + if s.tempGPUMemPool != nil { + s.tempGPUMemPool.FreeAll() + } + host := icicle_core.HostSliceFromElements(twiddles0) + var allocErr error + dTwiddles0, allocErr = allocDeviceUninitialized(len(twiddles0)) + if allocErr != nil { + uploadTwiddlesDone <- fmt.Errorf("uploadComputeNumeratorTwiddles: %w", allocErr) + return + } + host.CopyToDevice(&dTwiddles0, false) + uploadTwiddlesDone <- nil + }) + if err := <-uploadTwiddlesDone; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dTwiddles0, nil +} + +// executeComputeNumeratorCosetIterations runs the rho coset iterations. +func (s *instance) executeComputeNumeratorCosetIterations(loopCtx *computeNumeratorLoopContext) error { + var startIterLoop time.Time + if isProfileMode { + startIterLoop = time.Now() + } + + for i := 0; i < loopCtx.rho; i++ { + if err := s.computeNumeratorIteration(i, loopCtx); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + freeSliceOnDevice(&loopCtx.dTwiddles0, &s.device) + return err + } + } + + // Free twiddles0 device slice (uploaded once before the loop). + freeTwiddlesDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + loopCtx.dTwiddles0.Free() + close(freeTwiddlesDone) + }) + <-freeTwiddlesDone + + if useBlinding { + var startRestoreBlindingPolys time.Time + if isProfileMode { + startRestoreBlindingPolys = time.Now() + } + csInv := inverseShifterProduct(loopCtx.shifters) + if err := s.restoreBlindingPolynomials(csInv); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startRestoreBlindingPolys)).Msg("computeNumerator: restore blinding polys") + } + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startIterLoop)).Msg("computeNumerator: full iteration loop") + } + return nil +} + +func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoopContext) error { + loopCtx.coset.Mul(&loopCtx.coset, &loopCtx.shifters[i]) + loopCtx.cosetExponentiatedToNMinusOne.Exp(loopCtx.coset, loopCtx.bn). + Sub(&loopCtx.cosetExponentiatedToNMinusOne, &loopCtx.one) + + batchInvertDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) + }) + if err := <-batchInvertDone; err != nil { + return err + } + + s.applyNumeratorBlindingScale(i, loopCtx) + if i == 1 { + // We have to update the scalingVector; instead of scaling by + // cosets we scale by the twiddles of the large domain. + w := s.domain1.Generator + loopCtx.scalingVector = make([]fr.Element, loopCtx.n) + fft.BuildExpTable(w, loopCtx.scalingVector) + + // Reuse memory. + copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) + fft.BitReverse(loopCtx.scalingVectorRev) + } + + // We do **a lot** of FFT here, but on the small domain. + // Note that for all the polynomials in the proving key + // (Ql, Qr, Qm, Qo, S1, S2, S3, Qcp, Qc) and ID, LOne + // we could pre-compute these rho*2 FFTs and store them + // at the cost of a huge memory footprint. + var startGpuInverseScaleForward time.Time + if isProfileMode { + startGpuInverseScaleForward = time.Now() + } + + // Inverse NTT -> Scale -> Forward NTT all on GPU using persistent GPU memory. + if err := s.gpuNTTInverseScaleForwardOnDevice(loopCtx.gpuState, loopCtx.scalingVector, loopCtx.scalingVectorRev, s.pk); err != nil { + return err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startGpuInverseScaleForward)).Msg("computeNumerator: gpuNTTInverseScaleForwardOnDevice") + } + + // Evaluate constraints on GPU. + constraintParams := gpuConstraintEvalParams{ + beta: s.beta, + gamma: s.gamma, + alpha: s.alpha, + coset: loopCtx.coset, + cosetExponentiatedToNMinusOne: loopCtx.cosetExponentiatedToNMinusOne, + cs: loopCtx.cs, + css: loopCtx.css, + cardinalityInv: s.domain0.CardinalityInv, + n: loopCtx.n, + nbBsbGates: loopCtx.nbBsbGates, + } + var startEvalConstraints time.Time + if isProfileMode { + startEvalConstraints = time.Now() + } + dNumeratorShard, err := s.gpuEvaluateConstraints( + loopCtx.gpuState, + constraintParams, + loopCtx.twiddles0, // CPU version for computeBlindingPolynomials + loopCtx.dTwiddles0, // GPU version for computeOrderingConstraint + *loopCtx.dPrecomputedDenominators, + s.bp, + nil, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startEvalConstraints)).Msg("computeNumerator: gpuEvaluateConstraints") + } + loopCtx.numeratorShards[i] = dNumeratorShard + + loopCtx.cosetExponentiatedToNMinusOne. + Inverse(&loopCtx.cosetExponentiatedToNMinusOne) + s.applyNumeratorBlindingUnscale(i, loopCtx) + return nil +} + +func (s *instance) buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx *computeNumeratorLoopContext) error { + if loopCtx == nil || loopCtx.dPrecomputedDenominators == nil { + return fmt.Errorf("computeNumerator: nil denominator device slice") + } + if loopCtx.dTwiddles0.IsEmpty() || loopCtx.dTwiddles0.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid twiddles device slice size %d, expected %d", loopCtx.dTwiddles0.Len(), loopCtx.n) + } + + if loopCtx.dPrecomputedDenominators.IsEmpty() { + dDenominators, err := allocDeviceUninitialized(loopCtx.n) + if err != nil { + return err + } + *loopCtx.dPrecomputedDenominators = dDenominators + } else if loopCtx.dPrecomputedDenominators.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid denominator device slice size %d, expected %d", loopCtx.dPrecomputedDenominators.Len(), loopCtx.n) + } + + cfg := icicle_core.DefaultVecOpsConfig() + + // dTwiddles0 is a Montgomery scalar vector in domain0 regular order. + // ScalarMulVec expects the scalar in standard form and preserves the + // Montgomery representation of the vector result. + dCosetStd := uploadScalarStdOnCurrentDevice(loopCtx.coset, cfg) + defer dCosetStd.Free() + if err := icicle_vecops.ScalarMulVec( + dCosetStd, + loopCtx.dTwiddles0, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators coset*twiddles failed: %s", err.AsString()) + } + + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneMont := uploadScalarMontOnCurrentDevice(minusOne, cfg) + defer dMinusOneMont.Free() + if err := icicle_vecops.ScalarAddVec( + dMinusOneMont, + *loopCtx.dPrecomputedDenominators, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators subtract one failed: %s", err.AsString()) + } + + if err := s.batchInvertOnCurrentDevice(*loopCtx.dPrecomputedDenominators); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: batchInvert failed: %s", err.AsString()) + } + return nil +} + +func (s *instance) applyNumeratorBlindingScale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindScale time.Time + if isProfileMode { + startBlindScale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + acc := loopCtx.cosetExponentiatedToNMinusOne + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &acc) + acc.Mul(&acc, &loopCtx.shifters[i]) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindScale)).Msg("computeNumerator: scale blinding polys") + } +} + +func (s *instance) applyNumeratorBlindingUnscale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindUnscale time.Time + if isProfileMode { + startBlindUnscale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &loopCtx.cosetExponentiatedToNMinusOne) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindUnscale)).Msg("computeNumerator: unscale blinding polys") + } +} + +func (s *instance) restoreBlindingPolynomials(csInv fr.Element) error { + for _, q := range s.bp { + if q == nil { + continue + } + cp := q.Coefficients() + if len(cp) == 0 { + continue + } + var acc fr.Element + acc.SetOne() + for i := 0; i < len(cp); i++ { + cp[i].Mul(&cp[i], &acc) + acc.Mul(&acc, &csInv) + } + } + return nil +} + +func inverseShifterProduct(shifters []fr.Element) fr.Element { + var acc fr.Element + acc.SetOne() + for i := 0; i < len(shifters); i++ { + acc.Mul(&acc, &shifters[i]) + } + acc.Inverse(&acc) + return acc +} + +func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial) (_ *iop.Polynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("downloadNumeratorFromGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("downloadNumeratorFromGPU: invalid dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard count mismatch: got %d, expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + + defer func() { + if err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + } + }() + + for i := 0; i < gpuNumerator.rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard %d is empty", i) + } + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + var dMerged icicle_core.DeviceSlice + + mergeDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + return + } + // Async boundary before merged slice is consumed by host copy. + mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + }) + if mergeErr := <-mergeDone; mergeErr != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, mergeErr + } + + cres := make([]fr.Element, totalSize) + cresHost := icicle_core.HostSliceFromElements(cres) + downloadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + }) + if err := <-downloadDone; err != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, err + } + s.putTempDeviceSlice(dMerged, totalSize) + + s.freeNumeratorShards(gpuNumerator.shards) + return iop.NewPolynomial(&cres, iop.Form{Basis: iop.LagrangeCoset, Layout: iop.BitReverse}), nil +} + +func (s *instance) freeNumeratorShards(shards []icicle_core.DeviceSlice) { + if len(shards) == 0 { + return + } + for i := 0; i < len(shards); i++ { + if !shards[i].IsEmpty() { + s.putTempDeviceSlice(shards[i], shards[i].Len()) + shards[i] = icicle_core.DeviceSlice{} + } + } +} + +func (s *instance) batchInvert(dVec icicle_core.DeviceSlice) { + if dVec.Len() == 0 { + return + } + + done := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + done <- s.batchInvertOnCurrentDevice(dVec) + }) + if err := <-done; err != icicle_runtime.Success { + panic(fmt.Sprintf("batchInvert: BatchInverseVec failed: %s", err.AsString())) + } +} + +// batchInvertOnCurrentDevice assumes caller already runs on the active device thread. +func (s *instance) batchInvertOnCurrentDevice(dVec icicle_core.DeviceSlice) icicle_runtime.EIcicleError { + if dVec.Len() == 0 { + return icicle_runtime.Success + } + err := icicle_bls12377.FromMontgomery(dVec) + if err == icicle_runtime.Success { + cfg := icicle_core.DefaultVecOpsConfig() + err = icicle_vecops.BatchInverseVec(dVec, dVec, cfg) + } + if err == icicle_runtime.Success { + err = icicle_bls12377.ToMontgomery(dVec) + } + return err +} + +// gpuPolysState holds GPU-resident polynomial data to avoid repeated CPU-GPU transfers. +// Use ensurePolysOnSharedGPU to populate/reuse and freeGPUPolys to release GPU memory. +type gpuPolysState struct { + deviceSlices []icicle_core.DeviceSlice + hostSlices []icicle_core.HostSlice[fr.Element] + polys []*iop.Polynomial + originalForm []iop.Form + polyToIdx map[*iop.Polynomial]int +} + +func (s *instance) sharedGPUStateInitialCap(extra int) int { + base := 16 + len(s.bp) + 2*len(s.commitmentInfo) + if extra > 0 { + base += extra + } + return base +} + +func (s *instance) initSharedGPUState() { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + if s.sharedGPUState != nil { + return + } + initialCap := s.sharedGPUStateInitialCap(0) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } +} + +func (s *instance) releaseSharedGPUState() { + s.gpuStateMu.Lock() + state := s.sharedGPUState + s.sharedGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) releaseLinearizedEvalGPUState() { + s.gpuStateMu.Lock() + state := s.linearizedEvalGPUState + s.linearizedEvalGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) freeIdleTempGPUMemoryOnDevice() { + if s == nil || s.tempGPUMemPool == nil { + return + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +func (s *instance) prepareLinearizedEvalGPUState(source *gpuPolysState) error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.clonePolysOnGPUFromState(source, s.buildLinearizedEvalGPUBatch()) + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) prepareLinearizedEvalGPUStateFromHost() error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.uploadPolysToGPUState(s.buildLinearizedEvalGPUBatch(), "prepareLinearizedEvalGPUStateFromHost") + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop.Polynomial) (*gpuPolysState, error) { + if source == nil { + return nil, fmt.Errorf("clonePolysOnGPUFromState: nil source state") + } + + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("clonePolysOnGPUFromState: empty polynomial batch") + } + + srcSlices := make([]icicle_core.DeviceSlice, len(unique)) + useSource := make([]bool, len(unique)) + for i, p := range unique { + idx, ok := source.polyToIdx[p] + if ok && idx >= 0 && idx < len(source.deviceSlices) && !source.deviceSlices[idx].IsEmpty() { + srcSlices[i] = source.deviceSlices[idx] + useSource[i] = true + } + } + + snapshot := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + snapshot.polys[i] = p + snapshot.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + snapshot.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("clonePolysOnGPUFromState") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for snapshot cloning before handing state to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i := range srcSlices { + if useSource[i] { + dst, allocErr := allocDeviceUninitialized(srcSlices[i].Len()) + if allocErr != nil { + runErr = fmt.Errorf("clonePolysOnGPUFromState: alloc failed at index %d: %w", i, allocErr) + return + } + if err := copyDeviceSliceIntoOnCurrentDevice(dst, srcSlices[i], cfg); err != icicle_runtime.Success { + _ = dst.Free() + runErr = fmt.Errorf("clonePolysOnGPUFromState: device copy failed at index %d: %s", i, err.AsString()) + return + } + snapshot.deviceSlices[i] = dst + continue + } + + coeffs := unique[i].Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("clonePolysOnGPUFromState: empty host coefficients at index %d", i) + return + } + host := icicle_core.HostSliceFromElements(coeffs) + var dst icicle_core.DeviceSlice + host.CopyToDeviceAsync(&dst, cfg.StreamHandle, true) + if dst.IsEmpty() { + runErr = fmt.Errorf("clonePolysOnGPUFromState: host upload failed at index %d", i) + return + } + snapshot.deviceSlices[i] = dst + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(snapshot) + return nil, err + } + return snapshot, nil +} + +func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) (*gpuPolysState, error) { + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("%s: empty polynomial batch", label) + } + + state := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + state.polys[i] = p + state.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + state.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label) + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i, p := range unique { + coeffs := p.Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("%s: empty host coefficients at index %d", label, i) + return + } + state.hostSlices[i] = icicle_core.HostSliceFromElements(coeffs) + state.hostSlices[i].CopyToDeviceAsync(&state.deviceSlices[i], cfg.StreamHandle, true) + if state.deviceSlices[i].IsEmpty() { + runErr = fmt.Errorf("%s: host upload failed at index %d", label, i) + return + } + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(state) + return nil, err + } + return state, nil +} + +func (s *instance) getTempDeviceSlice(n int) icicle_core.DeviceSlice { + if s == nil { + panic("getTempDeviceSlice: nil instance") + } + if s.tempGPUMemPool == nil { + panic("getTempDeviceSlice: temp GPU memory pool is not initialized") + } + return s.tempGPUMemPool.Get(n) +} + +func (s *instance) putTempDeviceSlice(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + if s != nil && s.tempGPUMemPool != nil { + s.tempGPUMemPool.Put(ds, n) + return + } + _ = ds.Free() +} + +func (s *instance) releaseTempGPUMemoryPool() { + if s == nil || s.tempGPUMemPool == nil { + return + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +// ensurePolysOnSharedGPU uploads missing polynomials once and reuses already-uploaded slices. +func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysState, error) { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if s.sharedGPUState == nil { + initialCap := s.sharedGPUStateInitialCap(len(polys)) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } + } + state := s.sharedGPUState + if state == nil { + return nil, fmt.Errorf("ensurePolysOnSharedGPU: shared GPU state is nil") + } + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int, len(polys)) + } + + newIndices := make([]int, 0, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; ok { + continue + } + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, icicle_core.DeviceSlice{}) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + newIndices = append(newIndices, idx) + } + if len(newIndices) == 0 { + return state, nil + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("ensurePolysOnSharedGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for GPU uploads in ensurePolysOnSharedGPU. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + for _, idx := range newIndices { + p := state.polys[idx] + if p == nil { + continue + } + cp := p.Coefficients() + state.hostSlices[idx] = icicle_core.HostSliceFromElements(cp) + state.hostSlices[idx].CopyToDeviceAsync(&state.deviceSlices[idx], cfg.StreamHandle, true) + } + }) + if err := <-done; err != nil { + return nil, err + } + return state, nil +} + +func getStateDeviceSlice(state *gpuPolysState, p *iop.Polynomial, label string) (icicle_core.DeviceSlice, error) { + if state == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil GPU state", label) + } + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil polynomial", label) + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): polynomial is not registered on GPU", label) + } + ds := state.deviceSlices[idx] + if ds.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): empty device slice", label) + } + return ds, nil +} + +// gpuNTTInverseScaleForwardOnDevice performs inverse NTT → scale → forward NTT +// on GPU-resident polynomial data without CPU-GPU transfers for polynomial data. +// The scaling vectors are uploaded each call (they may change between iterations). +func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scalingVector, scalingVectorRev []fr.Element, pk *ProvingKey) error { + if state == nil || len(state.polys) == 0 { + return nil + } + + device := &s.device + var scalingVectorDevice, scalingVectorRevDevice icicle_core.DeviceSlice + + // Upload scaling vectors to GPU + uploadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) + scalingHost := icicle_core.HostSliceFromElements(scalingVector) + scalingHost.CopyToDevice(&scalingVectorDevice, false) + scalingVectorRevDevice = s.getTempDeviceSlice(len(scalingVectorRev)) + scalingRevHost := icicle_core.HostSliceFromElements(scalingVectorRev) + scalingRevHost.CopyToDevice(&scalingVectorRevDevice, false) + + // Convert scaling vectors from Montgomery form to standard form + if err := icicle_bls12377.FromMontgomery(scalingVectorDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVector failed: %s", err.AsString()) + return + } + if err := icicle_bls12377.FromMontgomery(scalingVectorRevDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVectorRev failed: %s", err.AsString()) + return + } + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return err + } + + doneChans := make([]chan error, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + + done := make(chan error, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Step 1: Inverse NTT + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + return + } + + // Step 2: Scale by vector using GPU vecOps + vecCfg := icicle_core.DefaultVecOpsConfig() + vecCfg.StreamHandle = stream + vecCfg.IsAsync = true + + var scaleDevice icicle_core.DeviceSlice + if layout == iop.Regular { + // After KNR inverse, output is BitReverse → use scalingVectorRev + scaleDevice = scalingVectorRevDevice + } else { + // After KRN inverse, output is Regular → use scalingVector + scaleDevice = scalingVectorDevice + } + + if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + return + } + + // Step 3: Forward NTT to Lagrange + one := icicle_ntt.GetDefaultNttConfig().CosetGen + cfg.CosetGen = one + if layout == iop.Regular { + cfg.Ordering = icicle_core.KRN // BitReverse → Regular + } else { + cfg.Ordering = icicle_core.KNR // Regular → BitReverse + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + return + } + done <- nil + }) + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + if err := <-doneChans[i]; err != nil { + return err + } + } + } + + // Update polynomial metadata: final result is in Lagrange, same layout as original + for _, p := range state.polys { + if p != nil { + p.Basis = iop.Lagrange + } + } + + // Free scaling vectors from device + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + return nil +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +// +// NOTE: The prover hot path should use the state-based API and keep data on device. +// This wrapper exists for compatibility/testing where callers expect host coefficients +// to be materialized after the transform. +func (s *instance) gpuNTTInverseBatch(polys []*iop.Polynomial, pk *ProvingKey) { + if len(polys) == 0 { + return + } + state, err := s.ensurePolysOnSharedGPU(polys) + if err != nil { + panic(fmt.Sprintf("gpuNTTInverseBatch: ensurePolysOnSharedGPU failed: %v", err)) + } + + s.gpuNTTInverseBatchOnState(state, pk) + + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + for _, p := range polys { + if p == nil { + continue + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) || state.deviceSlices[idx].IsEmpty() { + continue + } + cp := p.Coefficients() + host := icicle_core.HostSliceFromElements(cp) + host.CopyFromDevice(&state.deviceSlices[idx]) + copy(cp, ([]fr.Element)(host)) + } + close(done) + }) + <-done +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKey) { + if state == nil || len(state.polys) == 0 { + return + } + + device := &s.device + doneChans := make([]chan struct{}, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + switch p.Basis { + case iop.Canonical: + continue // already in canonical form + case iop.Lagrange, iop.LagrangeCoset: + // Schedule GPU work + done := make(chan struct{}, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Select ordering and coset generator depending on basis and input layout + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + + // Base-domain inverse: + // - Regular input → KNR (output BitReverse) + // - BitReverse input → KRN (output Regular) + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + + // Run NTT inverse directly on the existing device slice (in-place) + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) + } + icicle_runtime.SynchronizeStream(stream) + + // Update metadata inside closure to avoid race + p.Basis = iop.Canonical + if layout == iop.Regular { + p.Layout = iop.BitReverse + } else { + p.Layout = iop.Regular + } + close(done) + }) + default: + panic("unsupported polynomial basis") + } + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + <-doneChans[i] + } + } +} + +// freeGPUPolys releases GPU memory for polynomial data. +func (s *instance) freeGPUPolys(state *gpuPolysState) { + if state == nil { + return + } + + device := &s.device + freeDone := make(chan struct{}) + + icicle_runtime.RunOnDevice(device, func(args ...any) { + for i := range state.polys { + if state.deviceSlices[i].IsEmpty() { + continue + } + _ = state.deviceSlices[i].Free() + } + close(freeDone) + }) + <-freeDone +} + +// gpuMemoryPool manages a pool of reusable device slices to avoid repeated allocations. +// Must be used within RunOnDevice context to ensure thread safety per device. +type gpuMemoryPool struct { + freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice + mu sync.Mutex +} + +// newGPUMemoryPool creates a new GPU memory pool. +func newGPUMemoryPool() *gpuMemoryPool { + return &gpuMemoryPool{ + freeSlices: make(map[int][]icicle_core.DeviceSlice), + } +} + +// Get returns a device slice of the specified size, either from the pool or newly allocated. +func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { + p.mu.Lock() + defer p.mu.Unlock() + + // Check if we have a free slice of this size + if slices, ok := p.freeSlices[n]; ok && len(slices) > 0 { + // Reuse the last slice + slice := slices[len(slices)-1] + p.freeSlices[n] = slices[:len(slices)-1] + return slice + } + + // No free slice available, allocate a new one. + // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + // Free all currently cached (idle) slices to reduce memory pressure. + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + panic(fmt.Sprintf("gpuMemoryPool.Get: allocation failed for size %d after clearing idle cache", n)) +} + +// Put returns a device slice to the pool for reuse instead of freeing it. +func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + + p.mu.Lock() + defer p.mu.Unlock() + + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// FreeAll releases all pooled device slices. +func (p *gpuMemoryPool) FreeAll() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) +} + +// allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. +// Use when the destination is fully overwritten by a kernel. +func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { + var ds icicle_core.DeviceSlice + if _, err := ds.Malloc(int(unsafe.Sizeof(fr.Element{})), n); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("allocDeviceUninitialized: malloc failed for size %d: %s", n, err.AsString()) + } + return ds, nil +} + +// mustAllocDeviceUninitialized is like allocDeviceUninitialized but panics on failure. +// Use only in contexts where error propagation is impractical (e.g. upload helpers). +func mustAllocDeviceUninitialized(n int) icicle_core.DeviceSlice { + ds, err := allocDeviceUninitialized(n) + if err != nil { + panic(err) + } + return ds +} + +// freeDeviceSlice frees a device slice if non-empty and zeroes the pointer. +// Use for directly-allocated slices, NOT pool-allocated ones (use putTempDeviceSlice for those). +func freeDeviceSlice(ds *icicle_core.DeviceSlice) { + if ds != nil && !ds.IsEmpty() { + _ = ds.Free() + *ds = icicle_core.DeviceSlice{} + } +} + +// freeSliceOnDevice frees a device slice on the specified device and blocks +// until complete. Use outside RunOnDevice closures. Zeroes the slice after freeing. +func freeSliceOnDevice(ds *icicle_core.DeviceSlice, device *icicle_runtime.Device) { + if ds == nil || ds.IsEmpty() { + return + } + d := *ds + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + _ = d.Free() + close(done) + }) + <-done + *ds = icicle_core.DeviceSlice{} +} + +// copyDeviceSliceIntoOnCurrentDevice copies src into dst entirely on GPU. +func copyDeviceSliceIntoOnCurrentDevice( + dst, src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if src.IsEmpty() || src.Len() <= 0 || dst.IsEmpty() || dst.Len() < src.Len() { + return icicle_runtime.InvalidArgument + } + src.CheckDevice() + dst.CheckDevice() + + srcElemSize := src.SizeOfElement() + dstElemSize := dst.SizeOfElement() + if srcElemSize <= 0 || dstElemSize <= 0 || srcElemSize != dstElemSize { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(src.Len() * srcElemSize) + if cfg.IsAsync { + return icicle_runtime.CopyAsync(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen, cfg.StreamHandle) + } + _, err := icicle_runtime.Copy(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen) + return err +} + +// zeroDeviceSliceOnCurrentDevice zero-fills dst entirely on GPU. +func zeroDeviceSliceOnCurrentDevice( + dst icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if dst.IsEmpty() || dst.Len() <= 0 { + return icicle_runtime.InvalidArgument + } + dst.CheckDevice() + + elemSize := dst.SizeOfElement() + if elemSize <= 0 { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(dst.Len() * elemSize) + if cfg.IsAsync { + return icicle_runtime.MemSetAsync(dst.AsUnsafePointer(), 0, byteLen, cfg.StreamHandle) + } + return icicle_runtime.MemSet(dst.AsUnsafePointer(), 0, byteLen) +} + +func createAsyncVecOpsConfigOnCurrentDevice(label string) (icicle_core.VecOpsConfig, icicle_runtime.Stream, error) { + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + return icicle_core.VecOpsConfig{}, nil, fmt.Errorf("%s: create stream failed: %s", label, eStream.AsString()) + } + cfg := icicle_core.DefaultVecOpsConfig() + cfg.StreamHandle = stream + cfg.IsAsync = true + return cfg, stream, nil +} + +func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label string) error { + if stream == nil { + return nil + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + _ = icicle_runtime.DestroyStream(stream) + return fmt.Errorf("%s: synchronize stream failed: %s", label, eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success { + return fmt.Errorf("%s: destroy stream failed: %s", label, eDestroy.AsString()) + } + return nil +} + +// makeFinisher returns a closure that synchronizes and destroys the stream, +// then sends the (possibly merged) error to done. Use inside RunOnDevice closures. +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { + return func(runErr error) { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + } +} + +// uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. +// Use this for additions where the vector is already in Montgomery form. +func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + return ds +} + +func uploadScalarMontOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + if cfg.IsAsync { + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&ds, false) + } + return ds +} + +// uploadScalarStd uploads a scalar in STANDARD form (not Montgomery) as a single-element device slice. +// For use with ScalarMulVec: (a*R) * b_std = (a*b)*R +func uploadScalarStd(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + // Convert from Montgomery form to standard form + if err := icicle_bls12377.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +func uploadScalarStdOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + ds := uploadScalarMontOnCurrentDevice(scalar, cfg) + // Fallback to sync conversion for compatibility with ICICLE wrappers + // that do not expose *_WithConfig APIs. + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bls12377.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +// uploadVectorStd uploads a vector and converts to standard form. +func uploadVectorStd(vec []fr.Element) icicle_core.DeviceSlice { + ds := mustAllocDeviceUninitialized(len(vec)) + uploadVectorStdInto(&ds, vec) + return ds +} + +// uploadVectorStdInto uploads vec into an existing device slice and converts it to standard form. +// The destination must already be allocated with enough capacity for len(vec) elements. +func uploadVectorStdInto(dst *icicle_core.DeviceSlice, vec []fr.Element) { + cfg := icicle_core.DefaultVecOpsConfig() + uploadVectorStdIntoOnCurrentDevice(dst, vec, cfg) +} + +// uploadVectorStdIntoOnCurrentDevice uploads vec into an existing device slice and converts it +// to standard form while honoring the provided vector-op config/stream. +func uploadVectorStdIntoOnCurrentDevice( + dst *icicle_core.DeviceSlice, + vec []fr.Element, + cfg icicle_core.VecOpsConfig, +) { + host := icicle_core.HostSliceFromElements(vec) + if cfg.IsAsync { + host.CopyToDeviceAsync(dst, cfg.StreamHandle, false) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } else { + host.CopyToDevice(dst, false) + } + // Convert from Montgomery form to standard form + if err := icicle_bls12377.FromMontgomery(*dst); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +// uploadVector uploads a vector (keeps Montgomery form for additions). +func uploadVector(vec []fr.Element) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(len(vec)) + host.CopyToDevice(&ds, false) + return ds +} + +// uploadInt64Vector uploads int64 indices to a device slice. +func uploadInt64Vector(vec []int64) icicle_core.DeviceSlice { + cfg := icicle_core.DefaultVecOpsConfig() + return uploadInt64VectorOnCurrentDevice(vec, cfg) +} + +func uploadInt64VectorOnCurrentDevice(vec []int64, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + var ds icicle_core.DeviceSlice + if cfg.IsAsync { + if _, err := ds.MallocAsync(int(unsafe.Sizeof(int64(0))), len(vec), cfg.StreamHandle); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc async failed: %s", err.AsString())) + } + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + return ds + } + if _, err := ds.Malloc(int(unsafe.Sizeof(int64(0))), len(vec)); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc failed: %s", err.AsString())) + } + host.CopyToDevice(&ds, false) + return ds +} + +// toStandardFormInPlace converts a device slice to standard form in-place (modifies the source). +// Use this for temporary vectors that won't be needed in Montgomery form. +func toStandardFormInPlace(src icicle_core.DeviceSlice) { + cfg := icicle_core.DefaultVecOpsConfig() + toStandardFormInPlaceWithCfg(src, cfg) +} + +func toStandardFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bls12377.FromMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +func toMontgomeryFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bls12377.ToMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("ToMontgomery failed: %s", err.AsString())) + } +} + +// multiplyMontgomerySlices multiplies two device slices that are both in Montgomery form. +// It creates a copy of dSlice1Mont, converts the copy to standard form, and then multiplies +// it with dSlice2Mont (which remains in Montgomery form). The result is stored in dResult +// and will be in Montgomery form. +// +// Parameters: +// - dSlice1Mont: first device slice in Montgomery form (not modified) +// - dSlice2Mont: second device slice in Montgomery form (not modified) +// - dResult: destination device slice for the result (must be pre-allocated) +// - state: GPU state with memory pool and vector configuration +// - n: size of the slices +func multiplyMontgomerySlices( + dSlice1Mont, dSlice2Mont icicle_core.DeviceSlice, + dResult icicle_core.DeviceSlice, + state *gpuConstraintEvalState, + vecCfg icicle_core.VecOpsConfig, + n int, +) error { + // Copy dSlice1Mont to standard form + dSlice1Std := state.getTempDeviceSlice(n) + defer state.putTempDeviceSlice(dSlice1Std, n) + + if err := copyDeviceSliceIntoOnCurrentDevice(dSlice1Std, dSlice1Mont, vecCfg); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: device copy failed: %s", err.AsString()) + } + toStandardFormInPlace(dSlice1Std) + + // Multiply: dSlice1Std (standard) * dSlice2Mont (Montgomery) = dResult (Montgomery) + if err := icicle_vecops.VecOp(dSlice1Std, dSlice2Mont, dResult, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: VecOp multiplication failed: %s", err.AsString()) + } + return nil +} + +// gpuConstraintEvalParams holds parameters for GPU constraint evaluation +type gpuConstraintEvalParams struct { + beta fr.Element + gamma fr.Element + alpha fr.Element + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + cs fr.Element // domain1.FrMultiplicativeGen + css fr.Element // cs^2 + cardinalityInv fr.Element + n int + nbBsbGates int +} + +// gpuConstraintEvalState holds intermediate state during constraint evaluation +type gpuConstraintEvalState struct { + // Polynomial device slices (may point to gpuState or allocated buffers) + dL, dR, dO, dZ, dZS icicle_core.DeviceSlice + dQl, dQr, dQm, dQo, dQk icicle_core.DeviceSlice + dS1, dS2, dS3 icicle_core.DeviceSlice + // Intermediate results + dGate, dOrdering, dLocal, dResult icicle_core.DeviceSlice + // Scalar device slices + dGammaScalar icicle_core.DeviceSlice + // Configuration + vecCfg icicle_core.VecOpsConfig + // Helper function to get device slices + getDeviceSlice func(int) icicle_core.DeviceSlice + // Shared prover-level temporary GPU memory pool accessors + getTempDeviceSlice func(int) icicle_core.DeviceSlice + putTempDeviceSlice func(icicle_core.DeviceSlice, int) + // Track allocated polynomial buffers for automatic cleanup + allocatedPolyBuffers []struct { + slice icicle_core.DeviceSlice + size int + } +} + +// allocate allocates a new device slice from the memory pool and tracks it for automatic cleanup. +// Returns the allocated device slice. +func (s *gpuConstraintEvalState) allocate(size int) icicle_core.DeviceSlice { + slice := s.getTempDeviceSlice(size) + s.allocatedPolyBuffers = append(s.allocatedPolyBuffers, struct { + slice icicle_core.DeviceSlice + size int + }{slice, size}) + return slice +} + +// freeAllocatedPolyBuffers returns all allocated polynomial buffers to the memory pool. +// This should be called during cleanup to free all buffers allocated via allocate(). +func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { + for _, buf := range s.allocatedPolyBuffers { + s.putTempDeviceSlice(buf.slice, buf.size) + } + s.allocatedPolyBuffers = nil +} + +// computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. +// Returns device slices for the blinding polynomials. +func computeBlindingPolynomials( + n int, + twiddles0 []fr.Element, + bp []*iop.Polynomial, +) (dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice) { + blindL := make([]fr.Element, n) + blindR := make([]fr.Element, n) + blindO := make([]fr.Element, n) + blindZ := make([]fr.Element, n) + blindZS := make([]fr.Element, n) // ZS uses shifted index + + // TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove this. Since we don't normally use blindings, we will + // not make this optimization. + utils.Parallelize(n, func(start, end int) { + for i := start; i < end; i++ { + blindL[i] = bp[id_Bl].Evaluate(twiddles0[i]) + blindR[i] = bp[id_Br].Evaluate(twiddles0[i]) + blindO[i] = bp[id_Bo].Evaluate(twiddles0[i]) + blindZ[i] = bp[id_Bz].Evaluate(twiddles0[i]) + blindZS[i] = bp[id_Bz].Evaluate(twiddles0[(i+1)%n]) + } + }) + + dBlindL = uploadVector(blindL) + dBlindR = uploadVector(blindR) + dBlindO = uploadVector(blindO) + dBlindZ = uploadVector(blindZ) + dBlindZS = uploadVector(blindZS) + + return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS +} + +// applyBlindingToPolynomials applies blinding to polynomials L, R, O, Z, ZS. +// Allocates new buffers for L, R, O, Z (tracked for cleanup) and modifies ZS in-place. +// The original slices in gpuState remain unchanged. +func applyBlindingToPolynomials( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice, +) error { + // L' = L + blindL (allocate new buffer, tracked for cleanup) + dLBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dL, dBlindL, dLBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add L failed: %s", err.AsString()) + } + state.dL = dLBlinded + + // R' = R + blindR (allocate new buffer, tracked for cleanup) + dRBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dR, dBlindR, dRBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add R failed: %s", err.AsString()) + } + state.dR = dRBlinded + + // O' = O + blindO (allocate new buffer, tracked for cleanup) + dOBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dO, dBlindO, dOBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add O failed: %s", err.AsString()) + } + state.dO = dOBlinded + + // Z' = Z + blindZ (allocate new buffer, tracked for cleanup) + dZBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dZ, dBlindZ, dZBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add Z failed: %s", err.AsString()) + } + state.dZ = dZBlinded + + // ZS' = ZS + blindZS + // Note: dZS is a temporary buffer created inside gpuEvaluateConstraints, + // so it's safe to modify it in-place. + if err := icicle_vecops.VecOp(state.dZS, dBlindZS, state.dZS, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) + } + + // Free blinding vectors - no longer needed after creating blinded polynomials + dBlindL.Free() + dBlindR.Free() + dBlindO.Free() + dBlindZ.Free() + dBlindZS.Free() + return nil +} + +// scaleSVectorsByBeta scales S1, S2, S3 by beta. +// Allocates new buffers for S1, S2, S3 (tracked for cleanup). +func scaleSVectorsByBeta( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, +) error { + // S1' = S1 * beta (need to scale S1, S2, S3 by beta for ordering constraint) + // Use standard form for beta so: S1_mont * beta_std = (S1*beta)_mont + dBetaStd := uploadScalarStd(params.beta) + + // S1' = S1 * beta (allocate new buffer, tracked for cleanup) + dS1Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dS1Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S1 failed: %s", err.AsString()) + } + state.dS1 = dS1Scaled + + // S2' = S2 * beta (allocate new buffer, tracked for cleanup) + dS2Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dS2Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S2 failed: %s", err.AsString()) + } + state.dS2 = dS2Scaled + + // S3' = S3 * beta (allocate new buffer, tracked for cleanup) + dS3Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dS3Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S3 failed: %s", err.AsString()) + } + state.dS3 = dS3Scaled + + // Free dBetaStd - no longer needed after scaling S vectors + dBetaStd.Free() + return nil +} + +// computeGateConstraint computes the gate constraint. +// Returns dGate device slice. +func computeGateConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // gate = Ql*L' + Qr*R' + Qm*L'*R' + Qo*O' + Qk + sum(Qci*Pi) + // We use multiplyMontgomerySlices for all poly×poly multiplications. + // Note: dL, dR, dO are used later in ordering constraint, so we preserve them. + + dGate := state.getTempDeviceSlice(params.n) + dTmp := state.getTempDeviceSlice(params.n) + + // Ql * L' + if err := multiplyMontgomerySlices(state.dQl, state.dL, dGate, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Ql*L: %w", err) + } + + // + Qr * R' + if err := multiplyMontgomerySlices(state.dQr, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qr*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qr*R failed: %s", err.AsString()) + } + + // + Qm * L' * R' (need two multiplications) + // First: Qm * L' = dTmp (Montgomery) + if err := multiplyMontgomerySlices(state.dQm, state.dL, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qm*L: %w", err) + } + // Second: dTmp (Montgomery) * R' (Montgomery) + if err := multiplyMontgomerySlices(dTmp, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: (Qm*L)*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qm*L*R failed: %s", err.AsString()) + } + + // + Qo * O' + if err := multiplyMontgomerySlices(state.dQo, state.dO, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qo*O: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qo*O failed: %s", err.AsString()) + } + + // + Qk + if err := icicle_vecops.VecOp(dGate, state.dQk, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qk failed: %s", err.AsString()) + } + + // + BSB gates: sum(Qci[2*i] * Qci[2*i+1]) + for i := 0; i < params.nbBsbGates; i++ { + origQci0 := state.getDeviceSlice(id_Qci + 2*i) + origQci1 := state.getDeviceSlice(id_Qci + 2*i + 1) + if !origQci0.IsEmpty() && !origQci1.IsEmpty() { + // Use helper to multiply Qci0 * Qci1 without modifying original values + if err := multiplyMontgomerySlices(origQci0, origQci1, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qci[%d]: %w", i, err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qci[%d] failed: %s", i, err.AsString()) + } + } + } + + // Return temporary buffer to pool - no longer needed after Step 3 + state.putTempDeviceSlice(dTmp, params.n) + + return dGate, nil +} + +// computeOrderingConstraint computes the ordering constraint. +// Returns dOrdering device slice. +func computeOrderingConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dTwiddles0 icicle_core.DeviceSlice, // twiddles0 already on GPU + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // This is complex: involves ID computation, gamma, beta, Z, ZS, S1, S2, S3 + // id = twiddles[i] * coset * beta + // a = gamma + L' + id + // b = gamma + R' + id*cs + // c = gamma + O' + id*css + // r = a * b * c * Z' + // + // a2 = gamma + L' + S1*beta + // b2 = gamma + R' + S2*beta + // c2 = gamma + O' + S3*beta + // l = a2 * b2 * c2 * ZS' + // + // ordering = l - r + + // Compute ID vector: twiddles * coset * beta (computed on GPU) + // dTwiddles0 is already on GPU (passed as parameter, don't free it here) + + // Compute coset * beta on CPU, then upload as scalar in standard form + var cosetTimesBeta fr.Element + cosetTimesBeta.Mul(¶ms.coset, ¶ms.beta) + dCosetTimesBetaStd := uploadScalarStd(cosetTimesBeta) + dBetaStd := uploadScalarStd(params.beta) + + // Multiply twiddles0 by cosetTimesBeta on GPU: dID = (cosetTimesBeta * twiddles0) * R (Montgomery form) + dID := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dCosetTimesBetaStd, dTwiddles0, dID, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec coset*beta*twiddles failed: %s", err.AsString()) + } + + // Free temporary device slice (dTwiddles0 is owned by caller, don't free it) + dCosetTimesBetaStd.Free() + + // id * cs - use standard form for cs + dIDcs := state.getTempDeviceSlice(params.n) + dCsStd := uploadScalarStd(params.cs) + if err := icicle_vecops.ScalarMulVec(dCsStd, dID, dIDcs, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*cs failed: %s", err.AsString()) + } + + // id * css - use standard form for css + dIDcss := state.getTempDeviceSlice(params.n) + dCssStd := uploadScalarStd(params.css) + if err := icicle_vecops.ScalarMulVec(dCssStd, dID, dIDcss, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*css failed: %s", err.AsString()) + } + + // a = gamma + L' + id (dL now contains L' after in-place blinding) + dA := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dID, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a+id failed: %s", err.AsString()) + } + + // b = gamma + R' + id*cs (dR now contains R' after in-place blinding) + dB := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dIDcs, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b+id*cs failed: %s", err.AsString()) + } + + // c = gamma + O' + id*css (dO now contains O' after in-place blinding) + dC := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dIDcss, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c+id*css failed: %s", err.AsString()) + } + + // Return to pool: dID, dIDcs, dIDcss - no longer needed after computing a, b, c + state.putTempDeviceSlice(dID, params.n) + state.putTempDeviceSlice(dIDcs, params.n) + state.putTempDeviceSlice(dIDcss, params.n) + dCsStd.Free() + dCssStd.Free() + + // r = a * b * c * Z' (dZ now contains Z' after in-place blinding) + // For chain multiplication, convert operands to std form in-place when possible + dR_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a*b failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dR_ord, dC, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*c failed: %s", err.AsString()) + } + // Convert dR_ord to standard form in-place (temporary result, dZ needs to be preserved) + toStandardFormInPlace(dR_ord) + if err := icicle_vecops.VecOp(dR_ord, state.dZ, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*Z failed: %s", err.AsString()) + } + + // Reuse dA, dB, dC for a2, b2, c2 instead of freeing and reallocating. + // To reduce peak memory, we scale S vectors by beta on-demand through a single temp buffer. + dScaledS := state.getTempDeviceSlice(params.n) + + // a2 = gamma + L' + S1*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L (a2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S1 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dScaledS, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2+S1*beta failed: %s", err.AsString()) + } + + // b2 = gamma + R' + S2*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R (b2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S2 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dScaledS, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b2+S2*beta failed: %s", err.AsString()) + } + + // c2 = gamma + O' + S3*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O (c2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S3 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dScaledS, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c2+S3*beta failed: %s", err.AsString()) + } + + // Free dGammaScalar - no longer needed after computing a2, b2, c2 + // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. + state.dGammaScalar.Free() + state.putTempDeviceSlice(dScaledS, params.n) + dBetaStd.Free() + + // l = a2 * b2 * c2 * ZS' (dZS now contains ZS' after in-place blinding) + dL_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2*b2 failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dL_ord, dC, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*c2 failed: %s", err.AsString()) + } + // Convert dZS to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(state.dZS) + if err := icicle_vecops.VecOp(dL_ord, state.dZS, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*ZS failed: %s", err.AsString()) + } + + // Return temporary buffers to pool - no longer needed after computing l + state.putTempDeviceSlice(dA, params.n) + state.putTempDeviceSlice(dB, params.n) + state.putTempDeviceSlice(dC, params.n) + state.putTempDeviceSlice(state.dZS, params.n) + state.dZS = icicle_core.DeviceSlice{} + + // ordering = l - r, reuse dL_ord as the final ordering vector + if err := icicle_vecops.VecOp(dL_ord, dR_ord, dL_ord, vecCfg, icicle_core.Sub); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l-r failed: %s", err.AsString()) + } + + // Return dR_ord to pool - no longer needed after computing ordering + state.putTempDeviceSlice(dR_ord, params.n) + + // Return dL_ord as ordering (caller is responsible for freeing) + return dL_ord, nil +} + +// computeLocalConstraint computes the local constraint. +// Returns dLocal device slice. +func computeLocalConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dPrecomputedDenominators icicle_core.DeviceSlice, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // local = (Z' - 1) * LagrangeOne + // where LagrangeOne[i] = cosetExpMinusOne * cardinalityInv / (coset*twiddles0[i] - 1) + + if dPrecomputedDenominators.IsEmpty() || dPrecomputedDenominators.Len() < params.n { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: invalid denominator device slice size %d, expected at least %d", dPrecomputedDenominators.Len(), params.n) + } + + // Compute LagrangeOne on device. dPrecomputedDenominators is already in + // Montgomery form after batch inversion; ScalarMulVec expects the scalar in + // standard form and preserves a Montgomery vector result. + var lagrangeCoeff fr.Element + lagrangeCoeff.Mul(¶ms.cosetExponentiatedToNMinusOne, ¶ms.cardinalityInv) + dLagrangeCoeffStd := uploadScalarStdOnCurrentDevice(lagrangeCoeff, vecCfg) + dLagrangeOneStd := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dLagrangeCoeffStd, dPrecomputedDenominators, dLagrangeOneStd, vecCfg); err != icicle_runtime.Success { + dLagrangeCoeffStd.Free() + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarMulVec lagrangeOne failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dLagrangeOneStd, vecCfg) + dLagrangeCoeffStd.Free() + + // Z' - 1 using ScalarAddVec with minus one + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneScalar := uploadScalarMont(minusOne) + + dZMinusOne := state.getTempDeviceSlice(params.n) + // dZ now contains Z' after in-place blinding + if err := icicle_vecops.ScalarAddVec(dMinusOneScalar, state.dZ, dZMinusOne, vecCfg); err != icicle_runtime.Success { + dMinusOneScalar.Free() + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarAddVec Z-1 failed: %s", err.AsString()) + } + + // Free dMinusOneScalar - no longer needed after computing Z' - 1 + // Note: dZ is part of gpuState and will be freed later + dMinusOneScalar.Free() + + // local = (Z' - 1) * LagrangeOne_std + dLocal := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.VecOp(dZMinusOne, dLagrangeOneStd, dLocal, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: VecOp (Z-1)*LagrangeOne failed: %s", err.AsString()) + } + + // Return temporary buffers to pool + state.putTempDeviceSlice(dZMinusOne, params.n) + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + + return dLocal, nil +} + +// createGetDeviceSliceFunc creates a function to get device slices for polynomials. +// It returns a function that maps polynomial indices to their device slices. +func (s *instance) polyByID(polyIdx int) *iop.Polynomial { + switch polyIdx { + case id_L: + return s.polyL + case id_R: + return s.polyR + case id_O: + return s.polyO + case id_Z: + return s.polyZ + case id_ZS: + return s.polyZS + case id_Ql: + return s.trace.Ql + case id_Qr: + return s.trace.Qr + case id_Qm: + return s.trace.Qm + case id_Qo: + return s.trace.Qo + case id_Qk: + return s.polyQk + case id_S1: + return s.trace.S1 + case id_S2: + return s.trace.S2 + case id_S3: + return s.trace.S3 + default: + if polyIdx < id_Qci { + return nil + } + offset := polyIdx - id_Qci + i := offset / 2 + if i < 0 { + return nil + } + if offset%2 == 0 { + if i < len(s.trace.Qcp) { + return s.trace.Qcp[i] + } + return nil + } + if i < len(s.cCommitments) { + return s.cCommitments[i] + } + return nil + } +} + +func createGetDeviceSliceFunc( + gpuState *gpuPolysState, + polyToIdx map[*iop.Polynomial]int, + resolvePoly func(int) *iop.Polynomial, +) func(int) icicle_core.DeviceSlice { + return func(polyIdx int) icicle_core.DeviceSlice { + p := resolvePoly(polyIdx) + if p == nil { + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d is nil", polyIdx)) + } + if idx, ok := polyToIdx[p]; ok { + return gpuState.deviceSlices[idx] + } + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d (ptr=%p) not found in polyToIdx (map has %d entries)", polyIdx, p, len(polyToIdx))) + } +} + +// initializeConstraintEvalState initializes the GPU constraint evaluation state. +// It sets up all device slices. +// The slices in gpuState are treated as read-only; helper functions will allocate +// separate working buffers whenever they need to modify data. +func initializeConstraintEvalState( + getDeviceSlice func(int) icicle_core.DeviceSlice, + getTempDeviceSlice func(int) icicle_core.DeviceSlice, + putTempDeviceSlice func(icicle_core.DeviceSlice, int), +) *gpuConstraintEvalState { + vecCfg := icicle_core.DefaultVecOpsConfig() + + state := &gpuConstraintEvalState{ + dL: getDeviceSlice(id_L), + dR: getDeviceSlice(id_R), + dO: getDeviceSlice(id_O), + dZ: getDeviceSlice(id_Z), + dQl: getDeviceSlice(id_Ql), + dQr: getDeviceSlice(id_Qr), + dQm: getDeviceSlice(id_Qm), + dQo: getDeviceSlice(id_Qo), + dQk: getDeviceSlice(id_Qk), + dS1: getDeviceSlice(id_S1), + dS2: getDeviceSlice(id_S2), + dS3: getDeviceSlice(id_S3), + vecCfg: vecCfg, + getDeviceSlice: getDeviceSlice, + getTempDeviceSlice: getTempDeviceSlice, + putTempDeviceSlice: putTempDeviceSlice, + } + + return state +} + +// gpuEvaluateConstraints evaluates all PLONK constraints on GPU. +// It takes polynomials already on GPU (via gpuState), computes blinding polynomial evaluations, +// and evaluates gate, ordering, and local constraints entirely on GPU. +// If result is non-nil, it downloads into result and returns an empty device slice. +// If result is nil, it returns a persistent device slice with the result. +// TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove 'twiddles0'. Since we don't +// normally use blindings, we will not make this optimization. +func (s *instance) gpuEvaluateConstraints( + gpuState *gpuPolysState, + params gpuConstraintEvalParams, + twiddles0 []fr.Element, // CPU vector for computeBlindingPolynomials + dTwiddles0 icicle_core.DeviceSlice, // GPU vector for computeOrderingConstraint + dPrecomputedDenominators icicle_core.DeviceSlice, + bp []*iop.Polynomial, // blinding polynomials (already scaled for this iteration) + result []fr.Element, +) (icicle_core.DeviceSlice, error) { + if gpuState == nil || len(gpuState.polys) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("gpuState is nil or empty") + } + + n := params.n + device := &s.device + + // Create a map from polynomial to its index in gpuState + polyToIdx := make(map[*iop.Polynomial]int) + for i, p := range gpuState.polys { + if p != nil { + polyToIdx[p] = i + } + } + + // Get device slices for the polynomials we need. + getDeviceSlice := createGetDeviceSliceFunc(gpuState, polyToIdx, s.polyByID) + + done := make(chan error, 1) + var resultOnDevice icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(device, func(args ...any) { + state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). + // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). + state.dGammaScalar = uploadScalarMont(params.gamma) + + state.dZS = state.getTempDeviceSlice(n) + if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + return + } + + // Step 1: Compute and apply blinding polynomial evaluations (if enabled) + if useBlinding { + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + done <- err + return + } + } + + // Step 2-4: Compute gate, ordering, and local constraints sequentially on a + // single synchronous stream, folding them into dResult as + // gate + alpha*ordering + alpha^2*local. Computing one family at a time + // keeps peak device allocation low, and sequential is not a compromise: + // the family kernels are memory-bandwidth-bound and each already saturates + // the device, so the parallel three-stream variant this replaces measured + // identical timings (111ms/iteration at n=2^23) — while racing on the + // shared temp-slice pool and lazily materialized inputs (it corrupted the + // numerator at every circuit size). + seqVecCfg := state.vecCfg + seqVecCfg.IsAsync = false + + // Compute ordering first to minimize peak memory before gate/local allocations. + var err error + state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + return + } + + // dResult = alpha * ordering + state.dResult = state.getTempDeviceSlice(params.n) + dAlphaStd := uploadScalarStd(params.alpha) + if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { + dAlphaStd.Free() + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + return + } + dAlphaStd.Free() + state.putTempDeviceSlice(state.dOrdering, params.n) + + // dResult += alpha^2 * local + state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + return + } + var alphaSquared fr.Element + alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) + dAlphaSquaredStd := uploadScalarStd(alphaSquared) + dTmp := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + return + } + dAlphaSquaredStd.Free() + if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dLocal, params.n) + state.putTempDeviceSlice(dTmp, params.n) + + // dResult += gate + state.dGate, err = computeGateConstraint(state, params, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + return + } + if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dGate, params.n) + + // Step 5: materialize result either on host or as a persistent device slice. + if result != nil { + resultHost := icicle_core.HostSliceFromElements(result) + resultHost.CopyFromDevice(&state.dResult) + } else { + resultOnDevice = s.getTempDeviceSlice(params.n) + if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // Return dResult pool slice after materialization. + state.putTempDeviceSlice(state.dResult, params.n) + + // Return all allocated polynomial buffers to the pool. + // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). + // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are + // returned to pool inside computeGateConstraint. dQk and the original gpuState slices + // are owned by gpuState and will be freed separately. + state.freeAllocatedPolyBuffers() + + done <- nil + }) + + err := <-done + + if err != nil { + if !resultOnDevice.IsEmpty() { + s.putTempDeviceSlice(resultOnDevice, resultOnDevice.Len()) + } + return icicle_core.DeviceSlice{}, err + } + return resultOnDevice, nil +} + +func evaluateBlinded(p, bp *iop.Polynomial, zeta fr.Element) fr.Element { + // Get the size of the polynomial + n := big.NewInt(int64(p.Size())) + + var pEvaluatedAtZeta fr.Element + + // Evaluate the polynomial and blinded polynomial at zeta + chP := make(chan struct{}, 1) + go func() { + pEvaluatedAtZeta = p.Evaluate(zeta) + close(chP) + }() + + bpEvaluatedAtZeta := bp.Evaluate(zeta) + + // Multiply the evaluated blinded polynomial by tempElement + var t fr.Element + one := fr.One() + t.Exp(zeta, n).Sub(&t, &one) + bpEvaluatedAtZeta.Mul(&bpEvaluatedAtZeta, &t) + + // Add the evaluated polynomial and the evaluated blinded polynomial + <-chP + pEvaluatedAtZeta.Add(&pEvaluatedAtZeta, &bpEvaluatedAtZeta) + + // Return the result + return pEvaluatedAtZeta +} + +// /!\ modifies the size +func getBlindedCoefficients(p, bp *iop.Polynomial) []fr.Element { + cp := p.Coefficients() + cbp := bp.Coefficients() + cp = append(cp, cbp...) + for i := 0; i < len(cbp); i++ { + cp[i].Sub(&cp[i], &cbp[i]) + } + return cp +} + +// getNonBlindedCoefficients returns a padded copy of polynomial coefficients +// to match the size they would have with blinding enabled. +// The padding size is blindingOrder+1 (e.g., order 2 → 3 coefficients). +func getNonBlindedCoefficients(p *iop.Polynomial, blindingOrder int) []fr.Element { + cp := p.Coefficients() + padded := make([]fr.Element, len(cp)+blindingOrder+1) + copy(padded, cp) + return padded +} + +// commits to a polynomial of the form b*(Xⁿ-1) where b is of small degree +func commitBlindingFactor(n int, b *iop.Polynomial, key kzg.ProvingKey) curve.G1Affine { + cp := b.Coefficients() + np := b.Size() + + // lo + var tmp curve.G1Affine + tmp.MultiExp(key.G1[:np], cp, ecc.MultiExpConfig{}) + + // hi + var res curve.G1Affine + res.MultiExp(key.G1[n:n+np], cp, ecc.MultiExpConfig{}) + res.Sub(&res, &tmp) + return res +} + +// return a random polynomial of degree n, if n==-1 cancel the blinding +func getRandomPolynomial(n int) *iop.Polynomial { + var a []fr.Element + if n == -1 { + a := make([]fr.Element, 1) + a[0].SetZero() + } else { + a = make([]fr.Element, n+1) + for i := 0; i <= n; i++ { + a[i].SetRandom() + } + } + res := iop.NewPolynomial(&a, iop.Form{ + Basis: iop.Canonical, Layout: iop.Regular}) + return res +} + +func coefficients(p []*iop.Polynomial) [][]fr.Element { + res := make([][]fr.Element, len(p)) + for i, pI := range p { + res[i] = pI.Coefficients() + } + return res +} + +func (s *instance) freeGPUQuotient(quotient *gpuQuotientPolynomial) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return + } + s.putTempDeviceSlice(quotient.coeffs, quotient.coeffs.Len()) + quotient.coeffs = icicle_core.DeviceSlice{} + quotient.size = 0 +} + +// commitToQuotientGPUFromDevice commits H1/H2/H3 directly from device memory. +// For StatisticalZK=true we materialize adjusted device vectors for h1/h2/h3 +// and commit those without downloading quotient coefficients to host. +// prepareStatisticalZKQuotientShards constructs blinded quotient polynomial +// shards h1, h2, h3 on the GPU for the Statistical ZK path. Each shard is +// randomized so that the quotient split h = h1 + X^(n+2)*h2 + X^(2(n+2))*h3 +// hides the original polynomial. +// +// Caller is responsible for returning dH1, dH2, dH3 to the temp pool: +// - dH1 and dH2 have size nPlus2+1 +// - dH3 has size nPlus2 +func (s *instance) prepareStatisticalZKQuotientShards( + h1Device, h2Device, h3Device icicle_core.DeviceSlice, + nPlus2 int, +) (dH1, dH2, dH3 icicle_core.DeviceSlice, err error) { + nPlus3 := nPlus2 + 1 + + prepareDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + dH1 = s.getTempDeviceSlice(nPlus3) + dH2 = s.getTempDeviceSlice(nPlus3) + dH3 = s.getTempDeviceSlice(nPlus2) + + // h1 = base h1 with extra randomizer coefficient at degree n+2. + dH1Prefix := (&dH1).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH1Prefix, h1Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h1 failed: %s", e.AsString()) + return + } + dH1Tail := (&dH1).Range(nPlus2, nPlus3, false) + r0Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[0]}) + r0Host.CopyToDevice(&dH1Tail, false) + + // h2 = base h2 with first coefficient adjusted by -r0 and tail = r1. + dH2Prefix := (&dH2).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH2Prefix, h2Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h2 failed: %s", e.AsString()) + return + } + dH2First := (&dH2).Range(0, 1, false) + var negR0 fr.Element + negR0.Neg(&s.quotientShardsRandomizers[0]) + dNegR0 := uploadScalarMont(negR0) + if e := icicle_vecops.ScalarAddVec(dNegR0, dH2First, dH2First, cfg); e != icicle_runtime.Success { + _ = dNegR0.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h2[0] failed: %s", e.AsString()) + return + } + _ = dNegR0.Free() + dH2Tail := (&dH2).Range(nPlus2, nPlus3, false) + r1Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[1]}) + r1Host.CopyToDevice(&dH2Tail, false) + + // h3 = base h3 with first coefficient adjusted by -r1. + if e := copyDeviceSliceIntoOnCurrentDevice(dH3, h3Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h3 failed: %s", e.AsString()) + return + } + dH3First := (&dH3).Range(0, 1, false) + var negR1 fr.Element + negR1.Neg(&s.quotientShardsRandomizers[1]) + dNegR1 := uploadScalarMont(negR1) + if e := icicle_vecops.ScalarAddVec(dNegR1, dH3First, dH3First, cfg); e != icicle_runtime.Success { + _ = dNegR1.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h3[0] failed: %s", e.AsString()) + return + } + _ = dNegR1.Free() + prepareDone <- nil + }) + if err := <-prepareDone; err != nil { + if !dH1.IsEmpty() { + s.putTempDeviceSlice(dH1, nPlus3) + } + if !dH2.IsEmpty() { + s.putTempDeviceSlice(dH2, nPlus3) + } + if !dH3.IsEmpty() { + s.putTempDeviceSlice(dH3, nPlus2) + } + return icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, err + } + return dH1, dH2, dH3, nil +} + +func (s *instance) commitToQuotientGPUFromDevice(quotient *gpuQuotientPolynomial) error { + if quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("commitToQuotientGPUFromDevice: empty quotient") + } + + nPlus2 := int(s.domain0.Cardinality) + 2 + required := 3 * nPlus2 + if quotient.coeffs.Len() < required { + return fmt.Errorf("commitToQuotientGPUFromDevice: quotient too small: got %d need >= %d", quotient.coeffs.Len(), required) + } + + h1Device := ("ient.coeffs).Range(0, nPlus2, false) + h2Device := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3Device := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + if s.opt.StatisticalZK { + nPlus3 := nPlus2 + 1 + dH1, dH2, dH3, err := s.prepareStatisticalZKQuotientShards(h1Device, h2Device, h3Device, nPlus2) + if err != nil { + return err + } + defer s.putTempDeviceSlice(dH1, nPlus3) + defer s.putTempDeviceSlice(dH2, nPlus3) + defer s.putTempDeviceSlice(dH3, nPlus2) + + c0, err := commitOnGPUCanonicalDevice(dH1, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(dH2, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(dH3, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + return nil + } + + // Commit sequentially to avoid 3-way concurrent MSM memory spikes. + c0, err := commitOnGPUCanonicalDevice(h1Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(h2Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(h3Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + + return nil +} + +func (s *instance) inverseAndMergeShards( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (icicle_core.DeviceSlice, error) { + if gpuNumerator == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: nil numerator") + } + n := gpuNumerator.n + rho := gpuNumerator.rho + if n <= 0 || rho <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: invalid n=%d rho=%d", n, rho) + } + if len(gpuNumerator.shards) != rho { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), rho) + } + for i := 0; i < rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard %d is empty", i) + } + } + + expo := big.NewInt(int64(n)) + + // Per-shard cosets: c_i = c * g^i where c=FrMultiplicativeGen, g=Generator. + cosets := make([]fr.Element, rho) + cosets[0].Set(&domains[1].FrMultiplicativeGen) + for i := 1; i < rho; i++ { + cosets[i].Mul(&cosets[i-1], &domains[1].Generator) + } + invCosets := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + invCosets[i].Inverse(&cosets[i]) + } + + // ν = g^n is a rho-th root, used for the rho-point inverse DFT in combine. + var nu, nuInv fr.Element + nu.Exp(domains[1].Generator, expo) + nuInv.Inverse(&nu) + nuInvPowers := make([]fr.Element, rho) + nuInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + nuInvPowers[i].Mul(&nuInvPowers[i-1], &nuInv) + } + + // cN = c^n. Recover original coefficient blocks by scaling with cN^{-t}. + var cN, cNInv fr.Element + cN.Exp(domains[1].FrMultiplicativeGen, expo) + cNInv.Inverse(&cN) + cNInvPowers := make([]fr.Element, rho) + cNInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + cNInvPowers[i].Mul(&cNInvPowers[i-1], &cNInv) + } + + // Each shard inverse contributes a 1/n factor; apply extra 1/rho. + var rhoFr, invRho fr.Element + rhoFr.SetUint64(uint64(rho)) + invRho.Inverse(&rhoFr) + combineScales := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + combineScales[i].Mul(&invRho, &cNInvPowers[i]) + } + + totalSize := rho * n + done := make(chan error, 1) + var dMerged icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish := makeFinisher(stream, "inverseAndMergeShards", done) + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + // KNR is often faster than KNN; we restore regular output explicitly + // by bit-reversing each shard after the inverse NTT. + cfgNtt.Ordering = icicle_core.KNR + + ext := config_extension.Create() + defer config_extension.Delete(ext) + alg := nttAlgorithmFromEnv("ICICLE_DIVIDE_BY_ZH_NTT_ALGO", icicle_core.MixedRadix) + ext.SetInt(icicle_core.CUDA_NTT_ALGORITHM, int(alg)) + cfgNtt.Ext = ext.AsUnsafePointer() + + // Step 1: inverse NTT each shard without coset, reorder to regular, + // then unscale by (c*g^i)^t to recover the coset-inverse equivalent. + nn := uint64(64 - bits.TrailingZeros64(uint64(n))) + invPowers := make([]fr.Element, n) + for i := 0; i < rho; i++ { + if nttErr := icicle_ntt.Ntt(gpuNumerator.shards[i], icicle_core.KInverse, &cfgNtt, gpuNumerator.shards[i]); nttErr != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: inverse NTT failed at shard %d: %s", i, nttErr.AsString())) + return + } + + // KNR outputs bit-reversed coefficients. Reorder back to regular. + dRegular := s.getTempDeviceSlice(n) + mergeErr := icicle_vecops.MergeShardsBitReverse( + []icicle_core.DeviceSlice{gpuNumerator.shards[i]}, + n, + nn, + dRegular, + cfgVec, + ) + if mergeErr != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: reorder failed at shard %d: %s", i, mergeErr.AsString())) + return + } + // gpuNumerator.shards[i] is returned to pool and replaced; wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(gpuNumerator.shards[i], n) + gpuNumerator.shards[i] = dRegular + + fft.BuildExpTable(invCosets[i], invPowers) + dInvPowers := s.getTempDeviceSlice(n) + uploadVectorStdIntoOnCurrentDevice(&dInvPowers, invPowers, cfgVec) + if e := icicle_vecops.VecOp(gpuNumerator.shards[i], dInvPowers, gpuNumerator.shards[i], cfgVec, icicle_core.Mul); e != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: normalize shard %d failed: %s", i, e.AsString())) + return + } + // dInvPowers is temporary and returned to pool each iteration. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(dInvPowers, n) + } + + // Step 2: combine shard results via a size-rho inverse DFT per coefficient index. + dMerged = s.getTempDeviceSlice(totalSize) + keepMerged := false + defer func() { + if !keepMerged && !dMerged.IsEmpty() { + s.putTempDeviceSlice(dMerged, totalSize) + } + }() + + dTmp := s.getTempDeviceSlice(n) + defer func() { + if !dTmp.IsEmpty() { + s.putTempDeviceSlice(dTmp, n) + } + }() + + dNuWeights := make([]icicle_core.DeviceSlice, rho) + dCombineScales := make([]icicle_core.DeviceSlice, rho) + for i := 0; i < rho; i++ { + dNuWeights[i] = uploadScalarStdOnCurrentDevice(nuInvPowers[i], cfgVec) + dCombineScales[i] = uploadScalarStdOnCurrentDevice(combineScales[i], cfgVec) + } + defer func() { + for i := 0; i < rho; i++ { + if !dNuWeights[i].IsEmpty() { + _ = dNuWeights[i].Free() + } + if !dCombineScales[i].IsEmpty() { + _ = dCombineScales[i].Free() + } + } + }() + + for t := 0; t < rho; t++ { + outT := (&dMerged).Range(t*n, (t+1)*n, false) + if e := copyDeviceSliceIntoOnCurrentDevice(outT, gpuNumerator.shards[0], cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: init out[%d] failed: %s", t, e.AsString())) + return + } + for i := 1; i < rho; i++ { + weightIdx := (i * t) % rho + if weightIdx == 0 { + if e := icicle_vecops.VecOp(outT, gpuNumerator.shards[i], outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: add shard %d to out[%d] failed: %s", i, t, e.AsString())) + return + } + continue + } + if e := icicle_vecops.ScalarMulVec(dNuWeights[weightIdx], gpuNumerator.shards[i], dTmp, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: weight shard %d for out[%d] failed: %s", i, t, e.AsString())) + return + } + if e := icicle_vecops.VecOp(outT, dTmp, outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: accumulate shard %d into out[%d] failed: %s", i, t, e.AsString())) + return + } + } + if e := icicle_vecops.ScalarMulVec(dCombineScales[t], outT, outT, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: scale out[%d] failed: %s", t, e.AsString())) + return + } + } + keepMerged = true + finish(nil) + }) + + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dMerged, nil +} + +func (s *instance) divideByZHOnGPU( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (_ *gpuQuotientPolynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("divideByZHOnGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("divideByZHOnGPU: invalid numerator dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + for i := range gpuNumerator.shards { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("divideByZHOnGPU: shard %d is empty", i) + } + } + + rho := int(domains[1].Cardinality / domains[0].Cardinality) + if rho != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: rho mismatch domains=%d numerator=%d", rho, gpuNumerator.rho) + } + + // Evaluate 1/(X^n-1) over the large-domain coset values used by this quotient. + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + + // In bit-reversed merged layout, each shard maps to a fixed (iRev % rho) bucket. + // So we can divide by Z_H by scaling each shard with its corresponding inverse. + scaleDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + for i := 0; i < gpuNumerator.rho; i++ { + dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) + vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if vecErr != icicle_runtime.Success { + finish(fmt.Errorf("divideByZHOnGPU: shard scaling failed at %d: %s", i, vecErr.AsString())) + return + } + } + finish(nil) + }) + if err := <-scaleDone; err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, err + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + dMerged, splitErr := s.inverseAndMergeShards(gpuNumerator, domains) + if splitErr != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, splitErr + } + // Shards are not needed after split inverse+merge. + s.freeNumeratorShards(gpuNumerator.shards) + return &gpuQuotientPolynomial{coeffs: dMerged, size: totalSize}, nil +} + +func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*iop.Polynomial, error) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return nil, fmt.Errorf("downloadQuotientFromGPU: empty quotient") + } + if quotient.size <= 0 { + return nil, fmt.Errorf("downloadQuotientFromGPU: invalid quotient size %d", quotient.size) + } + + coeffs := make([]fr.Element, quotient.size) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) + // Async boundary for host materialization of quotient coefficients. + done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + }) + if err := <-done; err != nil { + return nil, err + } + + return iop.NewPolynomial(&coeffs, iop.Form{Basis: iop.Canonical, Layout: iop.Regular}), nil +} + +func commitOnGPUWithDeviceBases( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, +) (curve.G1Affine, error) { + return commitOnGPUWithDeviceBasesChunked(scalarsDevice, basesDevice, device, icicleMSMChunkSize()) +} + +func commitOnGPUWithDeviceBasesChunked( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, + chunkSize int, +) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty scalar slice") + } + if basesDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty basis slice") + } + if scalarsDevice.Len() > basesDevice.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: invalid scalar size %d", scalarsDevice.Len()) + } + if chunkSize <= 0 || chunkSize > scalarsDevice.Len() { + chunkSize = scalarsDevice.Len() + } + + var commit curve.G1Affine + var msmErr error + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) + close(done) + }) + <-done + if msmErr != nil { + return curve.G1Affine{}, fmt.Errorf("icicle: MSM commit from device bases failed (%d scalars): %w", scalarsDevice.Len(), msmErr) + } + return commit, nil +} + +func commitOnGPUWithDeviceBasesChunkedOnCurrentDevice( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + chunkSize int, +) (curve.G1Affine, error) { + var commit curve.G1Affine + for start := 0; start < scalarsDevice.Len(); start += chunkSize { + end := start + chunkSize + if end > scalarsDevice.Len() { + end = scalarsDevice.Len() + } + + // Each chunk must pair with exactly bases[start:end]: ICICLE treats a + // bases slice longer than the scalars as a batched MSM (and requires + // divisibility), so the full bases buffer cannot be passed as-is when + // it is longer than the scalar vector. + scalarsChunk := scalarsDevice + if start != 0 || end != scalarsDevice.Len() { + scalarsChunk = (&scalarsDevice).Range(start, end, false) + } + basesChunk := basesDevice + if start != 0 || end != basesDevice.Len() { + basesChunk = (&basesDevice).Range(start, end, false) + } + + res := make(icicle_core.HostSlice[icicle_bls12377.Projective], 1) + cfg := icicle_msm.GetDefaultMSMConfig() + cfg.AreBasesMontgomeryForm = true + cfg.AreScalarsMontgomeryForm = true + e := icicle_msm.Msm(scalarsChunk, basesChunk, &cfg, res) + if e != icicle_runtime.Success { + return curve.G1Affine{}, fmt.Errorf("icicle MSM failed for chunk [%d:%d]: %s", start, end, e.AsString()) + } + + chunkCommit, err := projectiveToGnarkAffine(res[0]) + if err != nil { + return curve.G1Affine{}, fmt.Errorf("convert chunk [%d:%d]: %w", start, end, err) + } + commit.Add(&commit, &chunkCommit) + } + return commit, nil +} + +func icicleMSMChunkSize() int { + // Production-sized MSMs still need chunking, but tiny chunks add thousands of + // ICICLE calls. 4M-point chunks passed the gnark replay profile; 8M did not. + const defaultChunkSize = 1 << 22 + v := strings.TrimSpace(os.Getenv("ICICLE_MSM_CHUNK_SIZE")) + if v == "" { + return defaultChunkSize + } + n, err := strconv.Atoi(v) + if err != nil || n <= 0 { + return defaultChunkSize + } + return n +} + +func commitOnGPUCanonicalDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: canonical SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgDevice.G1, device) +} + +func evalCanonicalAtPoint(coeffs []fr.Element, point fr.Element) fr.Element { + var acc fr.Element + if len(coeffs) == 0 { + return acc + } + acc.Set(&coeffs[len(coeffs)-1]) + for i := len(coeffs) - 2; i >= 0; i-- { + acc.Mul(&acc, &point).Add(&acc, &coeffs[i]) + } + return acc +} + +func deriveBatchOpeningGamma( + point fr.Element, + digests []curve.G1Affine, + claimedValues []fr.Element, + hf hash.Hash, + dataTranscript ...[]byte, +) (fr.Element, error) { + fs := fiatshamir.NewTranscript(hf, "gamma") + if err := fs.Bind("gamma", point.Marshal()); err != nil { + return fr.Element{}, err + } + for i := range digests { + if err := fs.Bind("gamma", digests[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := range claimedValues { + if err := fs.Bind("gamma", claimedValues[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := 0; i < len(dataTranscript); i++ { + if err := fs.Bind("gamma", dataTranscript[i]); err != nil { + return fr.Element{}, err + } + } + gammaByte, err := fs.ComputeChallenge("gamma") + if err != nil { + return fr.Element{}, err + } + var gamma fr.Element + gamma.SetBytes(gammaByte) + return gamma, nil +} + +func (s *instance) evalDevicePolynomialAtPointOnCurrentDevice( + coeffsDevice icicle_core.DeviceSlice, + point fr.Element, + useBitReverse bool, + cfg icicle_core.VecOpsConfig, +) (fr.Element, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: empty coefficients") + } + + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(point, cfg) + defer dPoint.Free() + + dOut := s.getTempDeviceSlice(1) + defer s.putTempDeviceSlice(dOut, 1) + + opName := "PolyEvalAt" + var eEval icicle_runtime.EIcicleError + if useBitReverse { + opName = "PolyEvalAtBitReverse" + mm := uint64(64 - bits.TrailingZeros64(uint64(coeffsDevice.Len()))) + eEval = icicle_vecops.PolyEvalAtBitReverse(coeffsDevice, dPoint, mm, dOut, cfg) + } else { + eEval = icicle_vecops.PolyEvalAt(coeffsDevice, dPoint, dOut, cfg) + } + if eEval != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: %s failed: %s", opName, eEval.AsString()) + } + + var out fr.Element + hostOut := icicle_core.HostSliceFromElements([]fr.Element{out}) + if cfg.IsAsync { + hostOut.CopyFromDeviceAsync(&dOut, cfg.StreamHandle) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } else { + hostOut.CopyFromDevice(&dOut) + } + return ([]fr.Element)(hostOut)[0], nil +} + +func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (fr.Element, error) { + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + var runErr error + out, runErr = s.evalDevicePolynomialAtPointOnCurrentDevice(coeffsDevice, point, false, cfg) + done <- runErr + }) + return out, <-done +} + +func (s *instance) copyDeviceSliceOnCurrentDevice( + src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, + label string, +) (icicle_core.DeviceSlice, error) { + if src.IsEmpty() || src.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: empty source slice", label) + } + dst := s.getTempDeviceSlice(src.Len()) + eCopy := copyDeviceSliceIntoOnCurrentDevice(dst, src, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dst, src.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: copy failed: %s", label, eCopy.AsString()) + } + return dst, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromStateOnCurrentDevice( + p *iop.Polynomial, + state *gpuPolysState, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "batchOpening poly") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + dCanon, err := s.copyDeviceSliceOnCurrentDevice(dSrc, cfg, "materializePolynomialCanonicalRegularFromStateOnCurrentDevice") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + // Canonical bit-reverse -> canonical regular. + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + dRegular := s.getTempDeviceSlice(dCanon.Len()) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNN // regular -> regular canonical + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular canonical + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) buildBlindedCanonicalPolynomialOnCurrentDevice( + dBaseCanon, dBlindCanon icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBlindCanon.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: empty input") + } + n := dBaseCanon.Len() + blindLen := dBlindCanon.Len() + dOut := s.getTempDeviceSlice(n + blindLen) + + dPrefix := (&dOut).Range(0, n, false) + eCopyBase := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopyBase != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopyBase.AsString()) + } + + dTail := (&dOut).Range(n, n+blindLen, false) + eCopyBlind := copyDeviceSliceIntoOnCurrentDevice(dTail, dBlindCanon, cfg) + if eCopyBlind != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy tail failed: %s", eCopyBlind.AsString()) + } + + dHead := (&dOut).Range(0, blindLen, false) + if eSub := icicle_vecops.VecOp(dHead, dBlindCanon, dHead, cfg, icicle_core.Sub); eSub != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: subtract blind from head failed: %s", eSub.AsString()) + } + return dOut, nil +} + +func (s *instance) buildPaddedCanonicalPolynomialOnCurrentDevice( + dBaseCanon icicle_core.DeviceSlice, + padLen int, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBaseCanon.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: empty base") + } + if padLen < 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: negative pad length %d", padLen) + } + n := dBaseCanon.Len() + dOut := s.getTempDeviceSlice(n + padLen) + dPrefix := (&dOut).Range(0, n, false) + + eCopy := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopy.AsString()) + } + if padLen == 0 { + return dOut, nil + } + + dTail := (&dOut).Range(n, n+padLen, false) + eZero := zeroDeviceSliceOnCurrentDevice(dTail, cfg) + if eZero != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: zero tail failed: %s", eZero.AsString()) + } + return dOut, nil +} + +func (s *instance) prepareBatchOpeningPolynomialsOnGPU( + state *gpuPolysState, + point fr.Element, +) (devicePolys []icicle_core.DeviceSlice, owned []bool, claimed []fr.Element, err error) { + if state == nil { + return nil, nil, nil, fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: nil GPU state") + } + + total := 6 + len(s.trace.Qcp) + devicePolys = make([]icicle_core.DeviceSlice, total) + owned = make([]bool, total) + claimed = make([]fr.Element, total) + devicePolys[0] = s.linearizedPolynomialGPU + claimed[0] = s.linearizedPolynomialClaim + + prepDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + + cleanupOwned := func(from int) { + for i := from; i < len(devicePolys); i++ { + if owned[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + owned[i] = false + } + } + } + + prepareLRORow := func(dstIdx int, p, bp *iop.Polynomial, padLen int) error { + dBase, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + defer s.putTempDeviceSlice(dBase, dBase.Len()) + + var dFinal icicle_core.DeviceSlice + if useBlinding { + if bp == nil { + return fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: missing blinding polynomial") + } + dBlind, eBlind := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(bp, state, cfg) + if eBlind != nil { + return eBlind + } + defer s.putTempDeviceSlice(dBlind, dBlind.Len()) + dFinal, e = s.buildBlindedCanonicalPolynomialOnCurrentDevice(dBase, dBlind, cfg) + } else { + dFinal, e = s.buildPaddedCanonicalPolynomialOnCurrentDevice(dBase, padLen, cfg) + } + if e != nil { + return e + } + devicePolys[dstIdx] = dFinal + owned[dstIdx] = true + + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dFinal, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareLRORow(1, s.polyL, s.bp[id_Bl], order_blinding_L+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare L failed: %w", e)) + return + } + if e := prepareLRORow(2, s.polyR, s.bp[id_Br], order_blinding_R+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare R failed: %w", e)) + return + } + if e := prepareLRORow(3, s.polyO, s.bp[id_Bo], order_blinding_O+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare O failed: %w", e)) + return + } + + prepareDirect := func(dstIdx int, p *iop.Polynomial, label string) error { + dPoly, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + devicePolys[dstIdx] = dPoly + owned[dstIdx] = true + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dPoly, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareDirect(4, s.trace.S1, "S1"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S1 failed: %w", e)) + return + } + if e := prepareDirect(5, s.trace.S2, "S2"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S2 failed: %w", e)) + return + } + + for i := 0; i < len(s.trace.Qcp); i++ { + idx := 6 + i + if e := prepareDirect(idx, s.trace.Qcp[i], "Qcp"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare Qcp[%d] failed: %w", i, e)) + return + } + } + + finish(nil) + }) + if err := <-prepDone; err != nil { + return nil, nil, nil, err + } + return devicePolys, owned, claimed, nil +} + +type evalPolynomialInputPreparationResult struct { + dEval icicle_core.DeviceSlice + ownedLen int + useBitReverseEval bool +} + +func (s *instance) prepareEvalPolynomialInputOnCurrentDevice( + p *iop.Polynomial, + dSrc icicle_core.DeviceSlice, + cfgVec icicle_core.VecOpsConfig, +) (evalPolynomialInputPreparationResult, error) { + if p == nil { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: nil polynomial") + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: empty source polynomial") + } + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] prepareEvalInput begin n=%d basis=%v layout=%v step_profile=%q ntt_trace=%q ntt_profile_full=%q ntt_profile_arbitrary=%q\n", + dSrc.Len(), + p.Basis, + p.Layout, + os.Getenv("ICICLE_STEP_PROFILE"), + os.Getenv("ICICLE_NTT_TRACE"), + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + + result := evalPolynomialInputPreparationResult{ + dEval: dSrc, + } + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return result, nil + } + + releaseOwned := func() { + if result.ownedLen > 0 && !result.dEval.IsEmpty() { + s.putTempDeviceSlice(result.dEval, result.ownedLen) + result.dEval = icicle_core.DeviceSlice{} + result.ownedLen = 0 + } + } + + dWork := s.getTempDeviceSlice(dSrc.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dWork, dSrc, cfgVec); e != icicle_runtime.Success { + s.putTempDeviceSlice(dWork, dSrc.Len()) + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: copy source polynomial failed: %s", e.AsString()) + } + + result.dEval = dWork + result.ownedLen = dSrc.Len() + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfgVec.IsAsync + cfgNtt.StreamHandle = cfgVec.StreamHandle + + var ownsNttStream bool + destroyOwnedNttStream := func() error { + if !ownsNttStream { + return nil + } + return syncAndDestroyStreamOnCurrentDevice(cfgNtt.StreamHandle, "prepareEvalPolynomialInputOnCurrentDevice") + } + + switch p.Basis { + case iop.Canonical: + // No transform required. + result.useBitReverseEval = p.Layout == iop.BitReverse + case iop.Lagrange, iop.LagrangeCoset: + if cfgNtt.StreamHandle == nil { + nttStream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: create NTT stream failed: %s", eStream.AsString()) + } + cfgNtt.StreamHandle = nttStream + cfgNtt.IsAsync = true + ownsNttStream = true + } + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNR // regular -> bitreverse on inverse + result.useBitReverseEval = true + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular on inverse + result.useBitReverseEval = false + } + startNtt := time.Now() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT launch n=%d ordering=%v has_coset=%t basis=%v layout=%v\n", + dSrc.Len(), + cfgNtt.Ordering, + p.Basis == iop.LagrangeCoset, + p.Basis, + p.Layout, + ) + } + + // Martun: This call to Ntt takes about 3 seconds, because Ntt is reusing NTT domain data that + // gets prepared during InitDomain (once per device), not re-derived every call. + // Inside ICICLE, InitDomain precomputes: domain.twiddles (main roots-of-unity table, N+1) + // internal_twiddles and basic_twiddles for mixed-radix kernels + // if fast mode is on (it is by default here), extra forward+inverse fast twiddle tables (fast_external/internal/basic and _inv) — comment says this costs ~4N extra memory + // CPU-side coset_index map (root -> index), then reused by later Ntt calls + eNtt := icicle_ntt.Ntt(result.dEval, icicle_core.KInverse, &cfgNtt, result.dEval) + nttElapsed := time.Since(startNtt) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT done status=%s took=%s\n", + eNtt.AsString(), + nttElapsed, + ) + } + if eNtt != icicle_runtime.Success { + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + l := logger.Logger() + l.Warn().Err(errDestroy).Msg("prepareEvalPolynomialInputOnCurrentDevice") + } + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + // Async boundary for this helper when it owns the stream. + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + releaseOwned() + return evalPolynomialInputPreparationResult{}, errDestroy + } + default: + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: unsupported basis %v", p.Basis) + } + + return result, nil +} + +// evalPolynomialInCurrentFormOnGPU evaluates a polynomial at a point directly +// from the shared GPU state regardless of its current basis/layout by converting +// a temporary device copy to canonical/regular when needed. +func (s *instance) evalPolynomialInCurrentFormOnGPU( + p *iop.Polynomial, + state *gpuPolysState, + point fr.Element, +) (fr.Element, error) { + if p == nil { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "eval") + if err != nil { + return fr.Element{}, err + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: empty device polynomial") + } + + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + dEval := dSrc + ownedLen := 0 + releaseEval := func() { + if ownedLen > 0 && !dEval.IsEmpty() { + s.putTempDeviceSlice(dEval, ownedLen) + dEval = icicle_core.DeviceSlice{} + ownedLen = 0 + } + } + finish := func(runErr error) { + // Async boundary for eval path before handing result back to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + releaseEval() + done <- runErr + } + + prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) + dEval = prepareResult.dEval + ownedLen = prepareResult.ownedLen + if prepErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", prepErr)) + return + } + + evalOut, evalErr := s.evalDevicePolynomialAtPointOnCurrentDevice(dEval, point, prepareResult.useBitReverseEval, cfgVec) + if evalErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", evalErr)) + return + } + out = evalOut + finish(nil) + }) + return out, <-done +} + +func (s *instance) evaluateBlindedOnGPU( + p, bp *iop.Polynomial, + state *gpuPolysState, + zeta fr.Element, +) (fr.Element, error) { + if p == nil || bp == nil { + return fr.Element{}, fmt.Errorf("evaluateBlindedOnGPU: nil polynomial") + } + pAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(p, state, zeta) + if err != nil { + return fr.Element{}, err + } + bpAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(bp, state, zeta) + if err != nil { + return fr.Element{}, err + } + + var t, one fr.Element + one.SetOne() + t.Exp(zeta, big.NewInt(int64(p.Size()))).Sub(&t, &one) + bpAtZeta.Mul(&bpAtZeta, &t) + pAtZeta.Add(&pAtZeta, &bpAtZeta) + return pAtZeta, nil +} + +func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.Element) (kzg.OpeningProof, error) { + if len(coeffs) < 2 || len(coeffs) > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonical: invalid polynomial size %d", len(coeffs)) + } + claimed := evalCanonicalAtPoint(coeffs, point) + + var dWitness icicle_core.DeviceSlice + witnessSize := len(coeffs) - 1 + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dCoeffs := uploadVector(coeffs) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dCoeffs, dPoint, dWitness, cfg) + _ = dCoeffs.Free() + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonical: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + if err := <-divDone; err != nil { + return kzg.OpeningProof{}, err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (kzg.OpeningProof, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() < 2 || coeffsDevice.Len() > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonicalDevice: invalid polynomial size %d", coeffsDevice.Len()) + } + n := coeffsDevice.Len() + + var startEval time.Time + if isProfileMode { + startEval = time.Now() + } + claimed, err := s.evalDevicePolynomialAtPoint(coeffsDevice, point) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startEval)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint") + } + } + if err != nil { + return kzg.OpeningProof{}, err + } + + var dWitness icicle_core.DeviceSlice + witnessSize := coeffsDevice.Len() - 1 + var startDivideByXMinusA time.Time + if isProfileMode { + startDivideByXMinusA = time.Now() + } + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(coeffsDevice, dPoint, dWitness, cfg) + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonicalDevice: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + divideErr := <-divDone + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startDivideByXMinusA)) + if divideErr != nil { + ev.Err(divideErr).Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA") + } + } + if divideErr != nil { + return kzg.OpeningProof{}, divideErr + } + + var startCommit time.Time + if isProfileMode { + startCommit = time.Now() + } + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startCommit)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice") + } + } + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) linearizedZContributionScale(lZeta, rZeta, oZeta fr.Element) fr.Element { + var s2, tmp fr.Element + var uzeta, uuzeta fr.Element + uzeta.Mul(&s.zeta, &s.pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &s.pk.Vk.CosetShift) + + s2.Mul(&s.beta, &s.zeta).Add(&s2, &lZeta).Add(&s2, &s.gamma) + tmp.Mul(&s.beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp) + tmp.Mul(&s.beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp).Neg(&s2).Mul(&s2, &s.alpha) + + var one, alphaSquareLagrangeZero, den fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&s.zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) + den.Sub(&s.zeta, &one).Inverse(&den) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) + + s2.Add(&s2, &alphaSquareLagrangeZero) + return s2 +} + +func (s *instance) linearizedSelectorScales(evals witnessEvalAtZeta, zu fr.Element) linearizedSelectorScales { + var scales linearizedSelectorScales + + // S3 scale: + // alpha * beta * Z(mu*zeta) * + // (L(zeta) + beta*S1(zeta) + gamma) * + // (R(zeta) + beta*S2(zeta) + gamma) + var tmp fr.Element + scales.s3.Mul(&evals.s1zeta, &s.beta).Add(&scales.s3, &evals.blzeta).Add(&scales.s3, &s.gamma) + tmp.Mul(&evals.s2zeta, &s.beta).Add(&tmp, &evals.brzeta).Add(&tmp, &s.gamma) + scales.s3.Mul(&scales.s3, &tmp).Mul(&scales.s3, &zu).Mul(&scales.s3, &s.beta).Mul(&scales.s3, &s.alpha) + + scales.ql.Set(&evals.blzeta) + scales.qr.Set(&evals.brzeta) + scales.qm.Mul(&evals.brzeta, &evals.blzeta) + scales.qo.Set(&evals.bozeta) + scales.qk.SetOne() + scales.qcp = append(scales.qcp, evals.qcpzeta...) + + return scales +} + +func (s *instance) buildLinearizedSelectorTermsOnGPU( + evals witnessEvalAtZeta, + zu fr.Element, + linearizedLen int, +) (icicle_core.DeviceSlice, error) { + if linearizedLen <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: invalid length %d", linearizedLen) + } + if len(evals.qcpzeta) > len(s.cCommitments) { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: qcp/cCommitments mismatch (%d > %d)", len(evals.qcpzeta), len(s.cCommitments)) + } + + scales := s.linearizedSelectorScales(evals, zu) + + var dLinearized icicle_core.DeviceSlice + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildLinearizedSelectorTermsOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if runErr != nil && !dLinearized.IsEmpty() { + s.putTempDeviceSlice(dLinearized, dLinearized.Len()) + dLinearized = icicle_core.DeviceSlice{} + } + done <- runErr + }() + + dLinearized = s.getTempDeviceSlice(linearizedLen) + if eZero := zeroDeviceSliceOnCurrentDevice(dLinearized, cfg); eZero != icicle_runtime.Success { + runErr = fmt.Errorf("buildLinearizedSelectorTermsOnGPU: zero output failed: %s", eZero.AsString()) + return + } + + addTerm := func(p *iop.Polynomial, scale fr.Element, label string) error { + if p == nil { + return fmt.Errorf("missing polynomial %s", label) + } + if scale.IsZero() { + return nil + } + + start := time.Now() + dPoly, err := s.materializePolynomialCanonicalRegularFromHostOnCurrentDevice(p, cfg) + if err != nil { + return fmt.Errorf("%s: %w", label, err) + } + defer s.putTempDeviceSlice(dPoly, dPoly.Len()) + if dPoly.Len() > dLinearized.Len() { + return fmt.Errorf("%s: polynomial too large (%d > %d)", label, dPoly.Len(), dLinearized.Len()) + } + + dScale := uploadScalarStdOnCurrentDevice(scale, cfg) + dScaled := s.getTempDeviceSlice(dPoly.Len()) + defer s.putTempDeviceSlice(dScaled, dScaled.Len()) + + eScale := icicle_vecops.ScalarMulVec(dScale, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if eScale != icicle_runtime.Success { + return fmt.Errorf("%s: scale failed: %s", label, eScale.AsString()) + } + + dPrefix := (&dLinearized).Range(0, dPoly.Len(), false) + if eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add); eAdd != icicle_runtime.Success { + return fmt.Errorf("%s: add failed: %s", label, eAdd.AsString()) + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Str("term", label).Int("n", dPoly.Len()).Dur("took", time.Since(start)).Msg("computeLinearizedPolynomial: add selector term on GPU") + } + return nil + } + + if runErr = addTerm(s.trace.S3, scales.s3, "S3"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Ql, scales.ql, "Ql"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qm, scales.qm, "Qm"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qr, scales.qr, "Qr"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qo, scales.qo, "Qo"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qk, scales.qk, "Qk"); runErr != nil { + return + } + for i := range scales.qcp { + if runErr = addTerm(s.cCommitments[i], scales.qcp[i], fmt.Sprintf("Qcp[%d]", i)); runErr != nil { + return + } + } + }) + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dLinearized, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromHostOnCurrentDevice( + p *iop.Polynomial, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: nil polynomial") + } + coeffs := p.Coefficients() + if len(coeffs) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: empty polynomial") + } + + dCanon := s.getTempDeviceSlice(len(coeffs)) + host := icicle_core.HostSliceFromElements(coeffs) + if cfg.IsAsync { + host.CopyToDeviceAsync(&dCanon, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&dCanon, false) + } + if dCanon.IsEmpty() { + s.putTempDeviceSlice(dCanon, len(coeffs)) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: host upload failed") + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + dRegular := s.getTempDeviceSlice(dCanon.Len()) + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + switch p.Layout { + case iop.Regular: + cfgNtt.Ordering = icicle_core.KNN + case iop.BitReverse: + cfgNtt.Ordering = icicle_core.KRN + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported layout %v", p.Layout) + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) addZContributionToLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + dBlindedZCanonical icicle_core.DeviceSlice, + lZeta, rZeta, oZeta fr.Element, +) error { + if dLinearized.IsEmpty() || dBlindedZCanonical.IsEmpty() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: empty input") + } + if dLinearized.Len() < dBlindedZCanonical.Len() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: linearized too small (%d < %d)", dLinearized.Len(), dBlindedZCanonical.Len()) + } + + zScale := s.linearizedZContributionScale(lZeta, rZeta, oZeta) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("addZContributionToLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dScale icicle_core.DeviceSlice + var dScaledZ icicle_core.DeviceSlice + defer func() { + // Async boundary before returning temporary buffers to the pool. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dScale) + if !dScaledZ.IsEmpty() { + s.putTempDeviceSlice(dScaledZ, dScaledZ.Len()) + } + done <- runErr + }() + + dScale = uploadScalarStdOnCurrentDevice(zScale, cfg) + dScaledZ = s.getTempDeviceSlice(dBlindedZCanonical.Len()) + eMul := icicle_vecops.ScalarMulVec(dScale, dBlindedZCanonical, dScaledZ, cfg) + if eMul != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: scale Z failed: %s", eMul.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, dBlindedZCanonical.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaledZ, dPrefix, cfg, icicle_core.Add) + if eAdd != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: add scaled Z failed: %s", eAdd.AsString()) + return + } + }) + return <-done +} + +func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + quotient *gpuQuotientPolynomial, +) error { + if dLinearized.IsEmpty() || quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: empty input") + } + nPlus2 := int(s.domain0.Cardinality) + 2 + if dLinearized.Len() < nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: linearized too small") + } + if quotient.coeffs.Len() < 3*nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: quotient too small") + } + + var one fr.Element + one.SetOne() + var zetaN, zetaNPlusTwo, zhZeta fr.Element + zetaN.Exp(s.zeta, big.NewInt(int64(s.domain0.Cardinality))) + zhZeta.Sub(&zetaN, &one) + zetaNPlusTwo.Mul(&zetaN, &s.zeta).Mul(&zetaNPlusTwo, &s.zeta) + + h1 := ("ient.coeffs).Range(0, nPlus2, false) + h2 := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3 := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("subtractQuotientContributionFromLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dAcc icicle_core.DeviceSlice + var dZetaStd icicle_core.DeviceSlice + var dZhStd icicle_core.DeviceSlice + defer func() { + // Async boundary before reusing temporary quotient vectors. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dZhStd) + freeDeviceSlice(&dZetaStd) + if !dAcc.IsEmpty() { + s.putTempDeviceSlice(dAcc, dAcc.Len()) + } + done <- runErr + }() + + dAcc = s.getTempDeviceSlice(nPlus2) + dZetaStd = uploadScalarStdOnCurrentDevice(zetaNPlusTwo, cfg) + eMulH3 := icicle_vecops.ScalarMulVec(dZetaStd, h3, dAcc, cfg) + if eMulH3 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale h3 failed: %s", eMulH3.AsString()) + return + } + if eAddH2 := icicle_vecops.VecOp(dAcc, h2, dAcc, cfg, icicle_core.Add); eAddH2 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h2 failed: %s", eAddH2.AsString()) + return + } + if eMulPow := icicle_vecops.ScalarMulVec(dZetaStd, dAcc, dAcc, cfg); eMulPow != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by zeta^(n+2) failed: %s", eMulPow.AsString()) + return + } + if eAddH1 := icicle_vecops.VecOp(dAcc, h1, dAcc, cfg, icicle_core.Add); eAddH1 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h1 failed: %s", eAddH1.AsString()) + return + } + + dZhStd = uploadScalarStdOnCurrentDevice(zhZeta, cfg) + eScaleZh := icicle_vecops.ScalarMulVec(dZhStd, dAcc, dAcc, cfg) + if eScaleZh != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by Z_H(zeta) failed: %s", eScaleZh.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, nPlus2, false) + eSub := icicle_vecops.VecOp(dPrefix, dAcc, dPrefix, cfg, icicle_core.Sub) + if eSub != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: subtract term failed: %s", eSub.AsString()) + return + } + }) + return <-done +} + +// divideByZH +// The input must be in LagrangeCoset. +// The result is in Canonical Regular. (in place using a) +func (s *instance) divideByZH(a *iop.Polynomial, domains [2]*fft.Domain) (*iop.Polynomial, error) { + + // check that the basis is LagrangeCoset + if a.Basis != iop.LagrangeCoset || a.Layout != iop.BitReverse { + return nil, errors.New("invalid form") + } + + // prepare the evaluations of x^n-1 on the big domain's coset + var startEvaluateXnMinusOne time.Time + if isProfileMode { + startEvaluateXnMinusOne = time.Now() + } + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startEvaluateXnMinusOne)).Msg("divideByZH: evaluateXnMinusOneDomainBigCoset") + } + rho := int(domains[1].Cardinality / domains[0].Cardinality) + + r := a.Coefficients() + n := uint64(len(r)) + nn := uint64(64 - bits.TrailingZeros64(n)) + + var startParallelizeMul time.Time + if isProfileMode { + startParallelizeMul = time.Now() + } + utils.Parallelize(len(r), func(start, end int) { + for i := start; i < end; i++ { + iRev := bits.Reverse64(uint64(i)) >> nn + r[i].Mul(&r[i], &xnMinusOneInverseLagrangeCoset[int(iRev)%rho]) + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallelizeMul)).Msg("divideByZH: parallelize multiply coefficients") + } + + // Replace CPU FFT inverse by ICICLE NTT inverse. + var startGpuNTTInverse time.Time + if isProfileMode { + startGpuNTTInverse = time.Now() + } + // It's faster on CPU. + // s.gpuNTTInverse(a) + a.ToCanonical(domains[1]).ToRegular() + if isProfileMode { + l := logger.Logger() + l.Debug(). + Int("size", a.Size()). + Dur("took", time.Since(startGpuNTTInverse)). + Msg("divideByZH: gpuNTTInverse on input of size n") + } + + return a, nil +} + +// evaluateXnMinusOneDomainBigCoset evaluates Xᵐ-1 on DomainBig coset +func evaluateXnMinusOneDomainBigCoset(domains [2]*fft.Domain) []fr.Element { + + rho := domains[1].Cardinality / domains[0].Cardinality + + res := make([]fr.Element, rho) + + expo := big.NewInt(int64(domains[0].Cardinality)) + res[0].Exp(domains[1].FrMultiplicativeGen, expo) + + var t fr.Element + t.Exp(domains[1].Generator, expo) + + one := fr.One() + + for i := 1; i < int(rho); i++ { + res[i].Mul(&res[i-1], &t) + res[i-1].Sub(&res[i-1], &one) + } + res[len(res)-1].Sub(&res[len(res)-1], &one) + + res = fr.BatchInvert(res) + + return res +} + +// innerComputeLinearizedPoly computes the linearized polynomial in canonical basis. +// The purpose is to commit and open all in one ql, qr, qm, qo, qk. +// * lZeta, rZeta, oZeta are the evaluation of l, r, o at zeta +// * z is the permutation polynomial, zu is Z(μX), the shifted version of Z +// * pk is the proving key: the linearized polynomial is a linear combination of ql, qr, qm, qo, qk. +// +// The Linearized polynomial is: +// +// α²*L₁(ζ)*Z(X) +// + α*( (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*(β*s3(X))*Z(μζ) - Z(X)*(l(ζ)+β*id1(ζ)+γ)*(r(ζ)+β*id2(ζ)+γ)*(o(ζ)+β*id3(ζ)+γ)) +// + l(ζ)*Ql(X) + l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) +// - Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) +// +// /!\ blindedZCanonical is modified +func (s *instance) innerComputeLinearizedPoly( + lZeta, rZeta, oZeta, s1Zeta, s2Zeta, + alpha, beta, gamma, zeta, zu fr.Element, + qcpZeta, blindedZCanonical []fr.Element, + pi2Canonical [][]fr.Element, + pk *ProvingKey, +) []fr.Element { + + // l(ζ)r(ζ) + var rl fr.Element + rl.Mul(&rZeta, &lZeta) + + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + // the linearised polynomial is + // α²*L₁(ζ)*Z(X) + + // s1*s3(X)+s2*Z(X) + l(ζ)*Ql(X) + + // l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) - + // Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + var s1, s2, tmp fr.Element + s1.Mul(&s1Zeta, &beta).Add(&s1, &lZeta).Add(&s1, &gamma) // (l(ζ)+β*s1(ζ)+γ) + tmp.Mul(&s2Zeta, &beta).Add(&tmp, &rZeta).Add(&tmp, &gamma) + s1.Mul(&s1, &tmp).Mul(&s1, &zu).Mul(&s1, &beta).Mul(&s1, &alpha) // (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*β*Z(μζ)*α + + var uzeta, uuzeta fr.Element + uzeta.Mul(&zeta, &pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &pk.Vk.CosetShift) + + s2.Mul(&beta, &zeta).Add(&s2, &lZeta).Add(&s2, &gamma) // (l(ζ)+β*ζ+γ) + tmp.Mul(&beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &gamma) // (r(ζ)+β*u*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ) + tmp.Mul(&beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &gamma) // (o(ζ)+β*u²*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + s2.Neg(&s2).Mul(&s2, &alpha) + + // Z_h(ζ), ζⁿ⁺², L₁(ζ)*α²*Z + var zhZeta, zetaNPlusTwo, alphaSquareLagrangeZero, one, den, frNbElmt fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) // ζⁿ + zetaNPlusTwo.Mul(&alphaSquareLagrangeZero, &zeta).Mul(&zetaNPlusTwo, &zeta) // ζⁿ⁺² + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) // ζⁿ - 1 + zhZeta.Set(&alphaSquareLagrangeZero) // Z_h(ζ) = ζⁿ - 1 + frNbElmt.SetUint64(uint64(nbElmt)) + den.Sub(&zeta, &one).Inverse(&den) // 1/(ζ-1) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). // L₁ = (ζⁿ - 1)/(ζ-1) + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) // α²*L₁(ζ) + + s3canonical := s.trace.S3.Coefficients() + // Qk is prepared in canonical/regular form by computeLinearizedPolynomial. + + // len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 when Statistical ZK is activated + // len(h1)=len(h2)=len(h3)=len(blindedZCanonical)-1 when Statistical ZK is deactivated + h1 := s.h1() + h2 := s.h2() + h3 := s.h3() + + // at this stage we have + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + startParallel := time.Now() + utils.Parallelize(len(blindedZCanonical), func(start, end int) { + + cql := s.trace.Ql.Coefficients() + cqr := s.trace.Qr.Coefficients() + cqm := s.trace.Qm.Coefficients() + cqo := s.trace.Qo.Coefficients() + cqk := s.trace.Qk.Coefficients() + + var t, t0, t1 fr.Element + + for i := start; i < end; i++ { + t.Mul(&blindedZCanonical[i], &s2) // -Z(X)*α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + if i < len(s3canonical) { + t0.Mul(&s3canonical[i], &s1) // α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ)*β*s3(X) + t.Add(&t, &t0) + } + if i < len(cqm) { + t1.Mul(&cqm[i], &rl) // l(ζ)r(ζ)*Qm(X) + t.Add(&t, &t1) // linPol += l(ζ)r(ζ)*Qm(X) + t0.Mul(&cql[i], &lZeta) // l(ζ)Q_l(X) + t.Add(&t, &t0) // linPol += l(ζ)*Ql(X) + t0.Mul(&cqr[i], &rZeta) //r(ζ)*Qr(X) + t.Add(&t, &t0) // linPol += r(ζ)*Qr(X) + t0.Mul(&cqo[i], &oZeta) // o(ζ)*Qo(X) + t.Add(&t, &t0) // linPol += o(ζ)*Qo(X) + t.Add(&t, &cqk[i]) // linPol += Qk(X) + for j := range qcpZeta { // linPol += ∑ᵢQcp_(ζ)Pi_(X) + t0.Mul(&pi2Canonical[j][i], &qcpZeta[j]) + t.Add(&t, &t0) + } + } + + t0.Mul(&blindedZCanonical[i], &alphaSquareLagrangeZero) // α²L₁(ζ)Z(X) + blindedZCanonical[i].Add(&t, &t0) // linPol += α²L₁(ζ)Z(X) + + // if statistical zeroknowledge is deactivated, len(h1)=len(h2)=len(h3)=len(blindedZ)-1. + // Else len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 + if i < len(h3) { + t.Mul(&h3[i], &zetaNPlusTwo). + Add(&t, &h2[i]). + Mul(&t, &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } else { + if s.opt.StatisticalZK { + t.Mul(&h2[i], &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } + } + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallel)).Msg("computeLinearizedPolynomial: inner parallel loop") + } + + return blindedZCanonical +} + +var errContextDone = errors.New("context done") + +// local copies of verification-time helpers used by prover transcript +func bindPublicData(fs *fiatshamir.Transcript, challenge string, vk *plonk_bls12377.VerifyingKey, publicInputs []fr.Element) error { + if err := fs.Bind(challenge, vk.S[0].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[1].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[2].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Ql.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qr.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qm.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qo.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qk.Marshal()); err != nil { + return err + } + for i := range vk.Qcp { + if err := fs.Bind(challenge, vk.Qcp[i].Marshal()); err != nil { + return err + } + } + for i := 0; i < len(publicInputs); i++ { + if err := fs.Bind(challenge, publicInputs[i].Marshal()); err != nil { + return err + } + } + return nil +} + +func deriveRandomness(fs *fiatshamir.Transcript, challenge string, points ...*curve.G1Affine) (fr.Element, error) { + var buf [curve.SizeOfG1AffineUncompressed]byte + var r fr.Element + for _, p := range points { + buf = p.RawBytes() + if err := fs.Bind(challenge, buf[:]); err != nil { + return r, err + } + } + b, err := fs.ComputeChallenge(challenge) + if err != nil { + return r, err + } + r.SetBytes(b) + return r, nil +} + +// -------------------- GPU helpers and device setup -------------------- + +func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { + if pk.deviceInfo != nil { + return nil + } + pk.deviceInfo = &deviceInfo{} + + // Initialize ICICLE NTT domain (root of unity) and store coset generator for coset NTTs. + // ICICLE InitDomain expects a primitive root of unity; for coset transforms we use 𝔽ᵣ* generator. + + var gen fr.Element + var err error + if pk.Vk.Size < 6 { + gen, err = fft.Generator(8 * pk.Vk.Size) + if err != nil { + return err + } + } else { + gen, err = fft.Generator(4 * pk.Vk.Size) + if err != nil { + return err + } + } + genBits := gen.Bits() + limbs := icicle_core.ConvertUint64ArrToUint32Arr(genBits[:]) + // Initialize ICICLE NTT domain with root of unity + var rouIcicle icicle_bls12377.ScalarField + rouIcicle.FromLimbs(limbs) + + // Store coset generator = generator of 𝔽ᵣ* (matches CPU ToLagrangeCoset) + { + cosetGen := fft.GeneratorFullMultiplicativeGroup() + cosetBits := cosetGen.Bits() + cosetLimbs := icicle_core.ConvertUint64ArrToUint32Arr(cosetBits[:]) + copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) + } + + chInitDomain := make(chan struct{}) + initDomainQueuedAt := time.Now() + icicle_runtime.RunOnDevice(device, func(args ...any) { + initDomainStartedAt := time.Now() + initCfg := icicle_core.GetDefaultNTTInitDomainConfig() + ext := config_extension.Create() + defer config_extension.Delete(ext) + fastTwiddles := envEnabled("ICICLE_NTT_FAST_TWIDDLES", true) + ext.SetBool(icicle_core.CUDA_NTT_FAST_TWIDDLES_MODE, fastTwiddles) + initCfg.Ext = ext.AsUnsafePointer() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain start vk_size=%d fast_twiddles=%t profile_full=%q profile_arbitrary=%q\n", + pk.Vk.Size, + fastTwiddles, + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + e := icicle_ntt.InitDomain(rouIcicle, initCfg) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain end status=%s call_took=%s\n", + e.AsString(), + time.Since(initDomainStartedAt), + ) + } + if e != icicle_runtime.Success { + panic("icicle: InitDomain failed") + } + close(chInitDomain) + }) + + <-chInitDomain + if isNttTrace { + fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) + } + + chLag := make(chan struct{}) + chCan := make(chan struct{}) + + if len(pk.KzgLagrange.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) + g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) + close(chLag) + }) + } else { + close(chLag) + } + + if len(pk.Kzg.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) + g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) + close(chCan) + }) + } else { + close(chCan) + } + + <-chLag + <-chCan + return nil +} + +func projectiveToGnarkAffine(p icicle_bls12377.Projective) (curve.G1Affine, error) { + x, err := icicleBaseFieldToGnarkFp(p.X) + if err != nil { + return curve.G1Affine{}, err + } + y, err := icicleBaseFieldToGnarkFp(p.Y) + if err != nil { + return curve.G1Affine{}, err + } + z, err := icicleBaseFieldToGnarkFp(p.Z) + if err != nil { + return curve.G1Affine{}, err + } + if z.IsZero() { + return curve.G1Affine{}, nil + } + + var zInv fp.Element + zInv.Inverse(&z) + x.Mul(&x, &zInv) + y.Mul(&y, &zInv) + return curve.G1Affine{X: x, Y: y}, nil +} + +func icicleBaseFieldToGnarkFp(v icicle_bls12377.BaseField) (fp.Element, error) { + bytes := v.ToBytesLittleEndian() + if len(bytes) != fp.Bytes { + return fp.Element{}, fmt.Errorf("invalid ICICLE base field byte length %d", len(bytes)) + } + var buf [fp.Bytes]byte + copy(buf[:], bytes) + return fp.LittleEndian.Element(&buf) +} + +func commitOnGPULagrangeDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgLagrangeDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: lagrange SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgLagrangeDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgLagrangeDevice.G1, device) +} + +func (s *instance) registerDevicePolynomialInSharedState(state *gpuPolysState, p *iop.Polynomial, dSlice icicle_core.DeviceSlice) error { + if state == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil shared state") + } + if p == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil polynomial") + } + if dSlice.IsEmpty() { + return fmt.Errorf("registerDevicePolynomialInSharedState: empty device slice") + } + + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int) + } + if idx, ok := state.polyToIdx[p]; ok { + if idx < 0 || idx >= len(state.deviceSlices) { + return fmt.Errorf("registerDevicePolynomialInSharedState: invalid index %d", idx) + } + state.deviceSlices[idx] = dSlice + state.hostSlices[idx] = nil + state.originalForm[idx] = iop.Form{Basis: p.Basis, Layout: p.Layout} + return nil + } + + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, dSlice) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + return nil +} + +func (s *instance) gpuInclusivePrefixProductOnCurrentDevice( + dVec icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + if dVec.IsEmpty() || dVec.Len() <= 1 { + return nil + } + + n := dVec.Len() + for step := 1; step < n; step <<= 1 { + src := (&dVec).Range(0, n-step, false) + dst := (&dVec).Range(step, n, false) + + tmpStd := s.getTempDeviceSlice(n - step) + if err := copyDeviceSliceIntoOnCurrentDevice(tmpStd, src, cfg); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: copy stage failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(tmpStd, cfg) + if err := icicle_vecops.VecOp(tmpStd, dst, dst, cfg, icicle_core.Mul); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: multiply stage failed: %s", err.AsString()) + } + if cfg.IsAsync { + // tmpStd is returned to pool each stage, so we must wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } + s.putTempDeviceSlice(tmpStd, n-step) + } + return nil +} + +// buildPermutationGatherIndices prepares the subset of permutation indices that +// are consumed by the copy-constraint ratio loop (only rows [0, n-1) per copy). +func buildPermutationGatherIndices(permutation []int64, nbPolynomials, n, supportLen int) ([]int64, error) { + if n <= 1 { + return nil, nil + } + total := nbPolynomials * (n - 1) + indices := make([]int64, total) + + var permBuildErr error + var permBuildErrOnce sync.Once + utils.Parallelize(total, func(start, end int) { + for k := start; k < end; k++ { + j := k / (n - 1) + i := k % (n - 1) + base := j * n + permIdx := permutation[base+i] + if permIdx < 0 || int(permIdx) >= supportLen { + jj, ii, bad := j, i, permIdx + permBuildErrOnce.Do(func() { + permBuildErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: permutation index out of range at (%d,%d): %d", jj, ii, bad) + }) + continue + } + indices[k] = permIdx + } + }) + if permBuildErr != nil { + return nil, permBuildErr + } + return indices, nil +} + +func (s *instance) prepareCopyConstraintSupportsOnCurrentDevice( + n, nbPolynomials int, + domain *fft.Domain, + permGatherIndices []int64, + cfg icicle_core.VecOpsConfig, +) (dSupportFlat, dPermFlat icicle_core.DeviceSlice, err error) { + defer func() { + if err == nil { + return + } + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + }() + + if len(permGatherIndices) == 0 { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: empty permutation gather indices") + return + } + + dOmegaStd := uploadScalarStdOnCurrentDevice(domain.Generator, cfg) + defer dOmegaStd.Free() + dShiftStd := uploadScalarStdOnCurrentDevice(domain.FrMultiplicativeGen, cfg) + defer dShiftStd.Free() + + dSupportFlat, err = allocDeviceUninitialized(nbPolynomials * n) + if err != nil { + return + } + if e := icicle_vecops.SupportIdentity(dOmegaStd, dShiftStd, n, nbPolynomials, dSupportFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: generate identity support on GPU failed: %s", e.AsString()) + return + } + toMontgomeryFormInPlaceWithCfg(dSupportFlat, cfg) + + dPermIndicesDevice := uploadInt64VectorOnCurrentDevice(permGatherIndices, cfg) + defer dPermIndicesDevice.Free() + + dPermFlat, err = allocDeviceUninitialized(len(permGatherIndices)) + if err != nil { + return + } + if e := icicle_vecops.GatherByIndices(dSupportFlat, dPermIndicesDevice, dPermFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: gather permutation support on GPU failed: %s", e.AsString()) + return + } + if cfg.IsAsync { + // Ensure temporary support/index slices are safe to free on return. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + return + } + } + + return +} + +func (s *instance) accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail icicle_core.DeviceSlice, + dBetaStd, dGammaMont icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + nMinusOne := dEntryTail.Len() + dScaled := s.getTempDeviceSlice(nMinusOne) + dTerm := s.getTempDeviceSlice(nMinusOne) + defer func() { + s.putTempDeviceSlice(dScaled, nMinusOne) + s.putTempDeviceSlice(dTerm, nMinusOne) + }() + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dID, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale identity support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add beta*id failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dNumTail, dNumTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply numerator term failed: %s", err.AsString()) + } + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dPerm, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale permutation support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add beta*sigma failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dDenTail, dDenTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply denominator term failed: %s", err.AsString()) + } + if cfg.IsAsync { + // Temp vectors are released at function exit, so ensure queued work is complete. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + } + } + + return nil +} + +// validateDeviceEntries checks that all entries are non-empty and have consistent length. +// Returns the common length n. +func validateDeviceEntries(entries []icicle_core.DeviceSlice, label string) (int, error) { + if len(entries) == 0 { + return 0, fmt.Errorf("%s: no entries", label) + } + n := entries[0].Len() + if n == 0 { + return 0, fmt.Errorf("%s: empty device entry 0", label) + } + for i := range entries { + if entries[i].IsEmpty() { + return 0, fmt.Errorf("%s: empty device entry %d", label, i) + } + if entries[i].Len() != n { + return 0, fmt.Errorf("%s: inconsistent device entry size at %d (%d != %d)", label, i, entries[i].Len(), n) + } + } + return n, nil +} + +// BuildRatioCopyConstraintIcicle builds the accumulating ratio polynomial to prove that +// [P₁ ∥ .. ∥ P_{n—1}] is invariant by the permutation \sigma. +// Namely it returns the polynomial Z whose evaluation on the j-th root of unity is +// Z(ω^j) = Π_{i 1 { + dNumTail := (&dNum).Range(1, n, false) + dDenTail := (&dDen).Range(1, n, false) + var supportErr error + dSupportFlat, dPermFlat, supportErr = s.prepareCopyConstraintSupportsOnCurrentDevice(n, nbPolynomials, domain, permGatherIndices, cfg) + if supportErr != nil { + runErr = supportErr + return + } + + dBetaStd := uploadScalarStdOnCurrentDevice(beta, cfg) + dGammaMont := uploadScalarMontOnCurrentDevice(gamma, cfg) + + for j := 0; j < nbPolynomials; j++ { + dEntryTail := (&entriesDevice[j]).Range(0, n-1, false) + baseID := j * n + dID := (&dSupportFlat).Range(baseID, baseID+n-1, false) + basePerm := j * (n - 1) + dPerm := (&dPermFlat).Range(basePerm, basePerm+(n-1), false) + if err := s.accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail, dBetaStd, dGammaMont, cfg, + ); err != nil { + runErr = err + return + } + } + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize before releasing copy-constraint workspace failed: %s", eSync.AsString()) + return + } + } + _ = dBetaStd.Free() + _ = dGammaMont.Free() + + // Support vectors and loop temps are only needed for term accumulation. + // Free them before prefix products and batch inversion, whose ICICLE + // kernels allocate additional full-domain workspace internally. + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + s.tempGPUMemPool.FreeAll() + } + + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dNum, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dDen, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + + if invErr := s.batchInvertOnCurrentDevice(dDen); invErr != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: GPU batch inversion failed: %s", invErr.AsString()) + return + } + + toStandardFormInPlace(dDen) + if err := icicle_vecops.VecOp(dDen, dNum, dNum, cfg, icicle_core.Mul); err != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: final numerator*denominatorInv multiplication failed: %s", err.AsString()) + return + } + + dResult = dNum + dNum = icicle_core.DeviceSlice{} // transfer ownership to dResult + }) + if err := <-buildDone; err != nil { + return nil, err + } + + hostMirror := make([]fr.Element, n) + if len(hostMirror) > 0 { + hostMirror[0].SetOne() + } + res := iop.NewPolynomial(&hostMirror, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if err := s.registerDevicePolynomialInSharedState(gpuState, res, dResult); err != nil { + freeSliceOnDevice(&dResult, &s.device) + return nil, err + } + + return res, nil +} diff --git a/backend/accelerated/icicle/plonk/bls12-377/provingkey.go b/backend/accelerated/icicle/plonk/bls12-377/provingkey.go new file mode 100644 index 0000000000..96aa44b1b9 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bls12-377/provingkey.go @@ -0,0 +1,100 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +//go:build icicle + +package bls12377 + +import ( + "sync" + "time" + + "github.com/consensys/gnark-crypto/ecc/bls12-377/fr" + "github.com/consensys/gnark-crypto/ecc/bls12-377/fr/fft" + plonk_bls12377 "github.com/consensys/gnark/backend/plonk/bls12-377" + cs "github.com/consensys/gnark/constraint/bls12-377" + "github.com/consensys/gnark/logger" + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" +) + +// deviceInfo holds device-resident buffers for GPU acceleration. +type deviceInfo struct { + CosetGenerator [fr.Limbs * 2]uint32 + KzgDevice struct { + G1 icicle_core.DeviceSlice + } + KzgLagrangeDevice struct { + G1 icicle_core.DeviceSlice + } +} + +// hostSetup holds host-side, witness-independent prover state derived from the +// constraint system: the FFT domains and the PLONK trace (selector + +// permutation polynomials). Building the trace walks every constraint (~3s at +// 23M constraints), so it is computed once per proving key and shared across +// proofs. Everything here is read-only during proving: the prover clones Qk +// before patching public inputs into it, and every basis conversion of a trace +// polynomial copies first (see canonicalRegularCoefficientsCopy). +type hostSetup struct { + sizeSystem uint64 + domain0 *fft.Domain + domain1 *fft.Domain + trace *plonk_bls12377.Trace +} + +// ProvingKey wraps the native PLONK proving key with device-resident state +// (KZG bases, NTT domains, cached trace) that is uploaded once and reused +// across Prove calls. +// +// Concurrency: Prove calls sharing the same ProvingKey must be serialized by +// the caller. The device state hangs off the key and proofs share a single +// GPU; concurrent proves against the same key are not safe. +type ProvingKey struct { + plonk_bls12377.ProvingKey + *deviceInfo + hostSetupOnce sync.Once + hostSetup *hostSetup +} + +func buildHostSetup(spr *cs.SparseR1CS, sizeSystem uint64) *hostSetup { + domain0 := fft.NewDomain(sizeSystem) + + // h, the quotient polynomial is of degree 3(n+1)+2, so it's in a 3(n+2) dim + // vector space, the domain is the next power of 2 superior to 3(n+2). + // 4*domainNum is enough in all cases except when n<6. + var domain1 *fft.Domain + if sizeSystem < 6 { + domain1 = fft.NewDomain(8*sizeSystem, fft.WithoutPrecompute()) + } else { + domain1 = fft.NewDomain(4*sizeSystem, fft.WithoutPrecompute()) + } + + return &hostSetup{ + sizeSystem: sizeSystem, + domain0: domain0, + domain1: domain1, + trace: plonk_bls12377.NewTrace(spr, domain0), + } +} + +// hostSetupFor returns the FFT domains and trace for spr, building them on +// first use and caching them on the proving key. A PLONK proving key is bound +// to exactly one constraint system, so per-key caching is sound; as a +// defensive measure a system-size mismatch falls back to an uncached build +// rather than ever serving another circuit's trace. +func (pk *ProvingKey) hostSetupFor(spr *cs.SparseR1CS) *hostSetup { + nbConstraints := spr.GetNbConstraints() + sizeSystem := uint64(nbConstraints + len(spr.Public)) // len(spr.Public) is for the placeholder constraints + pk.hostSetupOnce.Do(func() { + start := time.Now() + pk.hostSetup = buildHostSetup(spr, sizeSystem) + log := logger.Logger() + log.Debug().Dur("took", time.Since(start)).Msg("built prover host setup (fft domains + trace)") + }) + if pk.hostSetup.sizeSystem != sizeSystem { + return buildHostSetup(spr, sizeSystem) + } + return pk.hostSetup +} diff --git a/backend/accelerated/icicle/plonk/bls12-381/doc.go b/backend/accelerated/icicle/plonk/bls12-381/doc.go new file mode 100644 index 0000000000..8492216755 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bls12-381/doc.go @@ -0,0 +1,7 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +// Package bls12381 implements ICICLE acceleration for BLS12-381 PLONK backend. +package bls12381 diff --git a/backend/accelerated/icicle/plonk/bls12-381/icicle.go b/backend/accelerated/icicle/plonk/bls12-381/icicle.go new file mode 100644 index 0000000000..4914aa4051 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bls12-381/icicle.go @@ -0,0 +1,7093 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +//go:build icicle + +package bls12381 + +import ( + "context" + "errors" + "fmt" + "hash" + "io" + "math/big" + "math/bits" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + "unsafe" + + "golang.org/x/sync/errgroup" + + "github.com/consensys/gnark/backend" + plonk_bls12381 "github.com/consensys/gnark/backend/plonk/bls12-381" + "github.com/consensys/gnark/backend/witness" + constraint "github.com/consensys/gnark/constraint" + cs "github.com/consensys/gnark/constraint/bls12-381" + "github.com/consensys/gnark/constraint/solver" + fcs "github.com/consensys/gnark/frontend/cs" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + + "github.com/consensys/gnark-crypto/ecc" + curve "github.com/consensys/gnark-crypto/ecc/bls12-381" + "github.com/consensys/gnark-crypto/ecc/bls12-381/fp" + "github.com/consensys/gnark-crypto/ecc/bls12-381/fr" + "github.com/consensys/gnark-crypto/ecc/bls12-381/fr/fft" + "github.com/consensys/gnark-crypto/ecc/bls12-381/fr/hash_to_field" + "github.com/consensys/gnark-crypto/ecc/bls12-381/fr/iop" + "github.com/consensys/gnark-crypto/ecc/bls12-381/kzg" + fiatshamir "github.com/consensys/gnark-crypto/fiat-shamir" + + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" + icicle_bls12381 "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bls12381" + icicle_msm "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bls12381/msm" + icicle_ntt "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bls12381/ntt" + icicle_vecops "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bls12381/vecOps" + icicle_runtime "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime" + "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime/config_extension" +) + +const HasIcicle = true + +var isProfileMode bool + +var useBlinding bool + +var isNttTrace bool + +func init() { + _, isProfileMode = os.LookupEnv("ICICLE_STEP_PROFILE") + // Blinding polynomials (zero-knowledge) are enabled by default, matching the + // native prover. Set GNARK_DISABLE_BLINDING to trade zero-knowledge for a + // faster, deterministic prover (e.g. when the witness is not secret). + _, disableBlinding := os.LookupEnv("GNARK_DISABLE_BLINDING") + useBlinding = !disableBlinding + isNttTrace = envEnabled("ICICLE_NTT_TRACE", false) +} + +// profileStep returns a function that, when called, logs the elapsed time since +// profileStep was invoked. If profiling is disabled, it returns a no-op. +// Usage: done := profileStep("label"); defer done() +func profileStep(msg string) func() { + if !isProfileMode { + return func() {} + } + start := time.Now() + return func() { + l := logger.Logger() + l.Debug().Dur("took", time.Since(start)).Msg(msg) + } +} + +// stageTiming is a single recorded prover stage and its wall-clock duration. +type stageTiming struct { + name string + dur time.Duration +} + +// stageTimings is a concurrency-safe, ordered recorder of prover stage +// durations. The PLONK prover runs its stages as concurrent goroutines, so the +// recorded durations OVERLAP and do not sum to the total — the printed table +// flags this. +type stageTimings struct { + mu sync.Mutex + entries []stageTiming +} + +// record appends a (stage, duration) entry. Safe to call from any goroutine and +// safe on a nil receiver (records nothing). +func (t *stageTimings) record(name string, d time.Duration) { + if t == nil { + return + } + t.mu.Lock() + t.entries = append(t.entries, stageTiming{name: name, dur: d}) + t.mu.Unlock() +} + +// printTable writes an aligned breakdown of the recorded stages to w, sorted by +// duration (largest first), followed by the overall prover total. Stages run +// concurrently, so the rows overlap and intentionally do not sum to the total. +func (t *stageTimings) printTable(w io.Writer, total time.Duration) { + if t == nil { + return + } + t.mu.Lock() + rows := make([]stageTiming, len(t.entries)) + copy(rows, t.entries) + t.mu.Unlock() + + sort.SliceStable(rows, func(i, j int) bool { return rows[i].dur > rows[j].dur }) + + nameW := len("TOTAL (prover done)") + for _, r := range rows { + if len(r.name) > nameW { + nameW = len(r.name) + } + } + + fmt.Fprintln(w, "") + fmt.Fprintln(w, "================ gnark PLONK prove breakdown (GPU) ================") + fmt.Fprintln(w, "(stages run concurrently — durations overlap and do not sum to TOTAL)") + fmt.Fprintf(w, " %-*s %12s %6s\n", nameW, "STAGE", "TIME", "%TOTAL") + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + for _, r := range rows { + pct := 0.0 + if total > 0 { + pct = 100 * float64(r.dur) / float64(total) + } + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, r.name, r.dur.Round(time.Millisecond), pct) + } + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, "TOTAL (prover done)", total.Round(time.Millisecond), 100.0) + fmt.Fprintln(w, "===================================================================") + fmt.Fprintln(w, "") +} + +func envEnabled(key string, defaultVal bool) bool { + v, ok := os.LookupEnv(key) + if !ok { + return defaultVal + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "0", "false", "off", "no": + return false + default: + return true + } +} + +func nttAlgorithmFromEnv(key string, fallback icicle_core.NttAlgorithm) icicle_core.NttAlgorithm { + v, ok := os.LookupEnv(key) + if !ok { + return fallback + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "auto", "0": + return icicle_core.Auto + case "radix2", "radix-2", "r2", "1": + return icicle_core.Radix2 + case "mixed", "mixedradix", "mixed-radix", "2": + return icicle_core.MixedRadix + default: + return fallback + } +} + +const ( + id_L int = iota + id_R + id_O + id_Z + id_ZS + id_Ql + id_Qr + id_Qm + id_Qo + id_Qk + id_S1 + id_S2 + id_S3 + id_Qci // [ .. , Qc_i, Pi_i, ...] +) + +// blinding factors +const ( + id_Bl int = iota + id_Br + id_Bo + id_Bz + nb_blinding_polynomials +) + +// blinding orders (-1 to deactivate) +const ( + order_blinding_L = 1 + order_blinding_R = 1 + order_blinding_O = 1 + order_blinding_Z = 2 +) + +// Prove generates a PLONK proof. When the accelerator option is not set to +// "icicle", we delegate to the native prover. Otherwise, we run a local copy +// of the CPU prover logic to enable incremental GPU adaptation. +func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts ...backend.ProverOption) (*plonk_bls12381.Proof, error) { + opt, err := backend.NewProverConfig(opts...) + if err != nil { + return nil, err + } + + log := logger.Logger().With(). + Str("curve", spr.CurveID().String()). + Int("nbConstraints", spr.GetNbConstraints()). + Str("backend", "plonk").Logger() + + // parse the options + if opt.HashToFieldFn == nil { + opt.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + + // When blinding is disabled (GNARK_DISABLE_BLINDING), also disable StatisticalZK, it makes no sense + // to use statistical zero knowledge when we don't use blinding. + if !useBlinding { + opt.StatisticalZK = false + } + + start := time.Now() + + // Initialize device and preload KZG bases once per proving key + device := icicle_runtime.CreateDevice("CUDA", 0) + if pk.deviceInfo == nil { + if err := pk.setupDevicePointers(&device); err != nil { + return nil, err + } + } + + // init instance + g, ctx := errgroup.WithContext(context.Background()) + instance, err := newInstance(ctx, spr, pk, fullWitness, &opt) + if err != nil { + return nil, fmt.Errorf("new instance: %w", err) + } + // attach device to instance for GPU calls + instance.device = device + instance.initSharedGPUState() + defer instance.releaseTempGPUMemoryPool() + defer instance.releaseSharedGPUState() + defer instance.releaseLinearizedEvalGPUState() + + // solve constraints + g.Go(instance.solveConstraints) + + // complete qk + g.Go(instance.completeQk) + + // init blinding polynomials + g.Go(instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + g.Go(instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + g.Go(instance.buildRatioCopyConstraint) + + // compute h + g.Go(instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + g.Go(instance.openZ) + + // linearized polynomial + g.Go(instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + g.Go(func() error { + startBatchOpening := time.Now() + err := instance.batchOpening() + if isProfileMode { + instance.timings.record("batchOpening (folded KZG)", time.Since(startBatchOpening)) + } + return err + }) + + if err := g.Wait(); err != nil { + return nil, err + } + + total := time.Since(start) + log.Debug().Dur("took", total).Msg("prover done") + if isProfileMode { + instance.timings.printTable(os.Stderr, total) + } + return instance.proof, nil +} + +// represents a Prover instance +type instance struct { + ctx context.Context + + pk *ProvingKey + proof *plonk_bls12381.Proof + spr *cs.SparseR1CS + opt *backend.ProverConfig + + fs *fiatshamir.Transcript + kzgFoldingHash hash.Hash // for KZG folding + htfFunc hash.Hash // hash to field function + + // polynomials + polyL, polyR, polyO *iop.Polynomial + polyZ, polyZS, polyQk *iop.Polynomial + bp []*iop.Polynomial // blinding polynomials + h *iop.Polynomial // h is the quotient polynomial + hGPU *gpuQuotientPolynomial + polyZLagrangeGPU icicle_core.DeviceSlice + blindedZCanonicalGPU icicle_core.DeviceSlice + quotientShardsRandomizers [2]fr.Element // random elements for blinding the shards of the quotient + + linearizedPolynomial []fr.Element + linearizedPolynomialGPU icicle_core.DeviceSlice + linearizedPolynomialClaim fr.Element + linearizedPolynomialDigest kzg.Digest + + fullWitness witness.Witness + + // bsb22 commitment stuff + commitmentInfo constraint.PlonkCommitments + commitmentVal []fr.Element + cCommitments []*iop.Polynomial + + // challenges + gamma, beta, alpha, zeta fr.Element + + // channel to wait for the steps + chLRO, + chQk, + chbp, + chZ, + chH, + chRestoreLRO, + chZOpening, + chLinearizedPolynomial, + chGammaBeta chan struct{} + + domain0, domain1 *fft.Domain + + trace *plonk_bls12381.Trace + + // GPU device handle + device icicle_runtime.Device + + // Shared GPU polynomial context reused across buildRatioCopyConstraint + // and computeQuotient to avoid repeated host<->device uploads. + gpuStateMu sync.Mutex + sharedGPUState *gpuPolysState + // Snapshot of immutable polynomial slices used by computeLinearizedPolynomial + // for zeta evaluations after computeQuotient mutates/frees shared state. + linearizedEvalGPUState *gpuPolysState + + // Reusable temporary GPU memory pool for non-state buffers. + tempGPUMemPool *gpuMemoryPool + + // Per-prove stage-timing recorder (used to print the breakdown table when + // ICICLE_STEP_PROFILE is set). + timings *stageTimings +} + +func newInstance(ctx context.Context, spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts *backend.ProverConfig) (*instance, error) { + if opts.HashToFieldFn == nil { + opts.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + s := instance{ + ctx: ctx, + pk: pk, + proof: &plonk_bls12381.Proof{}, + spr: spr, + opt: opts, + fullWitness: fullWitness, + bp: make([]*iop.Polynomial, nb_blinding_polynomials), + fs: fiatshamir.NewTranscript(opts.ChallengeHash, "gamma", "beta", "alpha", "zeta"), + kzgFoldingHash: opts.KZGFoldingHash, + htfFunc: opts.HashToFieldFn, + chLRO: make(chan struct{}, 1), + chQk: make(chan struct{}, 1), + chbp: make(chan struct{}, 1), + chGammaBeta: make(chan struct{}, 1), + chZ: make(chan struct{}, 1), + chH: make(chan struct{}, 1), + chZOpening: make(chan struct{}, 1), + chLinearizedPolynomial: make(chan struct{}, 1), + chRestoreLRO: make(chan struct{}, 1), + tempGPUMemPool: newGPUMemoryPool(), + timings: &stageTimings{}, + } + s.initBSB22Commitments() + + // FFT domains and the PLONK trace are witness-independent and expensive to + // build at large n (NewTrace walks every constraint), so they are cached + // on the proving key and shared read-only across proofs. + setup := pk.hostSetupFor(spr) + s.domain0 = setup.domain0 + s.domain1 = setup.domain1 + s.trace = setup.trace + + // sampling random numbers for blinding the quotient + if opts.StatisticalZK { + s.quotientShardsRandomizers[0].SetRandom() + s.quotientShardsRandomizers[1].SetRandom() + } + + return &s, nil +} + +func (s *instance) initBlindingPolynomials() error { + if !useBlinding { + // When blinding is disabled (GNARK_DISABLE_BLINDING), skip creating blinding polynomials entirely + // Just close the channel to unblock any goroutines waiting on it + close(s.chbp) + return nil + } + + s.bp[id_Bl] = getRandomPolynomial(order_blinding_L) + s.bp[id_Br] = getRandomPolynomial(order_blinding_R) + s.bp[id_Bo] = getRandomPolynomial(order_blinding_O) + s.bp[id_Bz] = getRandomPolynomial(order_blinding_Z) + close(s.chbp) + return nil +} + +func (s *instance) initBSB22Commitments() { + s.commitmentInfo = s.spr.CommitmentInfo.(constraint.PlonkCommitments) + s.commitmentVal = make([]fr.Element, len(s.commitmentInfo)) // TODO @Tabaie get rid of this + s.cCommitments = make([]*iop.Polynomial, len(s.commitmentInfo)) + s.proof.Bsb22Commitments = make([]kzg.Digest, len(s.commitmentInfo)) + + // override the hint for the commitment constraints + bsb22ID := solver.GetHintID(fcs.Bsb22CommitmentComputePlaceholder) + s.opt.SolverOpts = append(s.opt.SolverOpts, solver.OverrideHint(bsb22ID, s.bsb22Hint)) +} + +// Computing and verifying Bsb22 multi-commits explained in https://hackmd.io/x8KsadW3RRyX7YTCFJIkHg +func (s *instance) bsb22Hint(_ *big.Int, ins, outs []*big.Int) error { + var err error + commDepth := int(ins[0].Int64()) + ins = ins[1:] + + res := &s.commitmentVal[commDepth] + + commitmentInfo := s.spr.CommitmentInfo.(constraint.PlonkCommitments)[commDepth] + committedValues := make([]fr.Element, s.domain0.Cardinality) + offset := s.spr.GetNbPublicVariables() + for i := range ins { + committedValues[offset+commitmentInfo.Committed[i]].SetBigInt(ins[i]) + } + if _, err = committedValues[offset+commitmentInfo.CommitmentIndex].SetRandom(); err != nil { // Commitment injection constraint has qcp = 0. Safe to use for blinding. + return err + } + if _, err = committedValues[offset+s.spr.GetNbConstraints()-1].SetRandom(); err != nil { // Last constraint has qcp = 0. Safe to use for blinding + return err + } + s.cCommitments[commDepth] = iop.NewPolynomial(&committedValues, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[commDepth], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[commDepth]); err != nil { + return err + } + + s.htfFunc.Write(s.proof.Bsb22Commitments[commDepth].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + res.SetBytes(hashBts[:nbBuf]) // TODO @Tabaie use CommitmentIndex for this; create a new variable CommitmentConstraintIndex for other uses + res.BigInt(outs[0]) + + return nil +} + +// solveConstraints computes the evaluation of the L, R, O polynomials in Lagrange form. +func (s *instance) solveConstraints() error { + startSolve := time.Now() + log := logger.Logger() + + var solution *cs.SparseR1CSSolution + + // Try to load raw solver values from cache (fastest path) + rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + if rawCachePath != "" { + if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { + // Reconstruct L, R, O from raw values + var sol cs.SparseR1CSSolution + if sol.L, sol.R, sol.O, err = s.spr.EvaluateLROSmallDomainFromValues(rawValues); err != nil { + log.Warn().Err(err).Str("file", rawCachePath).Msg("ignoring raw solver cache") + } else { + log.Debug().Dur("took", time.Since(startSolve)).Int("wires", len(rawValues)).Msg("loaded raw solver values from cache") + solution = &sol + } + + // Load cached BSB22 cCommitments polynomials + cacheDir := filepath.Dir(rawCachePath) + for i := 0; solution != nil && i < len(s.commitmentInfo); i++ { + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + coeffs, err := cs.LoadRawSolverValues(bsb22Path) + if err != nil { + log.Warn().Err(err).Int("i", i).Msg("ignoring raw solver cache: missing BSB22 commitment sidecar") + solution = nil + break + } + coeffSlice := []fr.Element(coeffs) + s.cCommitments[i] = iop.NewPolynomial(&coeffSlice, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[i], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[i]); err != nil { + return err + } + s.htfFunc.Write(s.proof.Bsb22Commitments[i].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + s.commitmentVal[i].SetBytes(hashBts[:nbBuf]) + } + } + } + + if solution == nil { + _solution, err := s.spr.SolveAndSaveRawValues(s.fullWitness, rawCachePath, s.opt.SolverOpts...) + if err != nil { + log.Debug().Dur("took", time.Since(startSolve)).Err(err).Msg("solveConstraints: spr.Solve") + return err + } + log.Debug().Dur("took", time.Since(startSolve)).Msg("solveConstraints: spr.Solve") + if isProfileMode { + s.timings.record("solveConstraints: spr.Solve", time.Since(startSolve)) + } + solution = _solution.(*cs.SparseR1CSSolution) + + // Save cCommitments polynomial coefficients for BSB22 reconstruction + if rawCachePath != "" && len(s.commitmentInfo) > 0 { + cacheDir := filepath.Dir(rawCachePath) + for i := range s.commitmentInfo { + if s.cCommitments[i] != nil { + coeffs := s.cCommitments[i].Coefficients() + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + if err := cs.SaveRawSolverValues(bsb22Path, coeffs); err != nil { + log.Warn().Err(err).Int("i", i).Msg("failed to save BSB22 commitment polynomial") + } + } + } + } + } + + evaluationLDomainSmall := []fr.Element(solution.L) + evaluationRDomainSmall := []fr.Element(solution.R) + evaluationODomainSmall := []fr.Element(solution.O) + var wg sync.WaitGroup + wg.Add(2) + go func() { + s.polyL = iop.NewPolynomial(&evaluationLDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + go func() { + s.polyR = iop.NewPolynomial(&evaluationRDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + + s.polyO = iop.NewPolynomial(&evaluationODomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + + wg.Wait() + if _, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}); err != nil { + return err + } + + // commit to l, r, o and add blinding factors + if err := s.commitToLRO(); err != nil { + return err + } + close(s.chLRO) + return nil +} + +func (s *instance) completeQk() error { + qk := s.trace.Qk.Clone() + qkCoeffs := qk.Coefficients() + + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + copy(qkCoeffs, wWitness[:len(s.spr.Public)]) + + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + for i := range s.commitmentInfo { + qkCoeffs[s.spr.GetNbPublicVariables()+s.commitmentInfo[i].CommitmentIndex] = s.commitmentVal[i] + } + + s.polyQk = qk + close(s.chQk) + + return nil +} + +func (s *instance) commitToLRO() error { + var startCommitLRO time.Time + if isProfileMode { + startCommitLRO = time.Now() + } + sequentialLRO := s.domain0 != nil && s.domain0.Cardinality >= (1<<22) + if _, ok := os.LookupEnv("ICICLE_LRO_COMMIT_SEQUENTIAL"); ok { + sequentialLRO = envEnabled("ICICLE_LRO_COMMIT_SEQUENTIAL", true) + } + + if !useBlinding { + // When blinding is disabled, commit directly without waiting for blinding polynomials + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + if err != nil { + return err + } + } else { + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (no blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (no blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil + } + + // wait for blinding polynomials to be initialized or context to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chbp: + } + + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + if err != nil { + return err + } + } else { + // Run the three commitments concurrently. + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (with blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil +} + +// deriveGammaAndBeta (copy constraint) +func (s *instance) deriveGammaAndBeta() error { + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + if err := bindPublicData(s.fs, "gamma", s.pk.VerifyingKey().(*plonk_bls12381.VerifyingKey), wWitness[:len(s.spr.Public)]); err != nil { + return err + } + + // wait for LRO to be committed + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + gamma, err := deriveRandomness(s.fs, "gamma", &s.proof.LRO[0], &s.proof.LRO[1], &s.proof.LRO[2]) + if err != nil { + return err + } + + bbeta, err := s.fs.ComputeChallenge("beta") + if err != nil { + return err + } + s.gamma = gamma + s.beta.SetBytes(bbeta) + + close(s.chGammaBeta) + + return nil +} + +// commitToPolyAndBlinding computes the KZG commitment of a polynomial p +// in Lagrange form (large degree) +// and add the contribution of a blinding polynomial b (small degree) +// /!\ The polynomial p is supposed to be in Lagrange form. +// Only used when blinding is enabled (the default). +func (s *instance) commitToPolyAndBlinding(p, b *iop.Polynomial) (commit curve.G1Affine, err error) { + // Commit over the Lagrange SRS using shared device-resident polynomial data. + gpuCommit, err := s.commitLagrangePolynomialOnGPU(p) + if err != nil { + return curve.G1Affine{}, err + } + + // add CPU blinding contribution (two MSMs on canonical SRS) + n := int(s.domain0.Cardinality) + cb := commitBlindingFactor(n, b, s.pk.Kzg) + gpuCommit.Add(&gpuCommit, &cb) + return gpuCommit, nil +} + +func (s *instance) commitLagrangePolynomialOnGPU(p *iop.Polynomial) (curve.G1Affine, error) { + if p == nil { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: nil polynomial") + } + if p.Basis != iop.Lagrange { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial must be in Lagrange basis, got %v", p.Basis) + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{p}) + if err != nil { + return curve.G1Affine{}, err + } + idx, ok := gpuState.polyToIdx[p] + if !ok || idx < 0 || idx >= len(gpuState.deviceSlices) { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial is missing from shared GPU state") + } + if gpuState.deviceSlices[idx].IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: empty device slice for polynomial") + } + + // Keep the polynomial in its native Lagrange basis; large MSMs are split + // into device-side chunks inside commitOnGPULagrangeDevice. + return commitOnGPULagrangeDevice(gpuState.deviceSlices[idx], &s.device, s.pk) +} + +func (s *instance) deriveAlpha() (err error) { + alphaDeps := make([]*curve.G1Affine, len(s.proof.Bsb22Commitments)+1) + for i := range s.proof.Bsb22Commitments { + alphaDeps[i] = &s.proof.Bsb22Commitments[i] + } + alphaDeps[len(alphaDeps)-1] = &s.proof.Z + s.alpha, err = deriveRandomness(s.fs, "alpha", alphaDeps...) + return err +} + +func (s *instance) deriveZeta() (err error) { + s.zeta, err = deriveRandomness(s.fs, "zeta", &s.proof.H[0], &s.proof.H[1], &s.proof.H[2]) + return +} + +func (s *instance) computeQuotient() (err error) { + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + if isProfileMode { + var startComputeQuotient time.Time + startComputeQuotient = time.Now() + defer func() { + l := logger.Logger() + if err != nil { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Err(err).Msg("computeQuotient: total (with error)") + } else { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Msg("computeQuotient: total") + } + s.timings.record("computeQuotient (total)", time.Since(startComputeQuotient)) + }() + } + + // wait for Z to be committed or context done + doneWaitZ := profileStep("computeQuotient: wait Z commit") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZ: + } + doneWaitZ() + + // derive alpha + if err = s.deriveAlpha(); err != nil { + return err + } + + if err := s.waitForComputeNumeratorQk(); err != nil { + return err + } + if s.polyQk == nil { + return fmt.Errorf("computeQuotient: missing completed Qk polynomial") + } + + doneEnsureGPUState := profileStep("computeQuotient: ensure shared GPU state") + gpuState, err := s.ensurePolysOnSharedGPU(s.buildComputeNumeratorGPUBatch()) + if err != nil { + return err + } + doneEnsureGPUState() + + // compute Z shifted by one for copy-constraint terms. + if s.polyZ == nil { + return fmt.Errorf("computeQuotient: missing Z polynomial") + } + s.polyZS = s.polyZ.ShallowClone().Shift(1) + + var numeratorGPU *gpuNumeratorPolynomial + var quotientGPU *gpuQuotientPolynomial + var e error + doneComputeNumerator := profileStep("computeQuotient: computeNumerator") + numeratorGPU, e = s.computeNumerator(gpuState) + if e != nil { + return e + } + doneComputeNumerator() + + doneDivideByZH := profileStep("computeQuotient: divideByZHOnGPU") + quotientGPU, e = s.divideByZHOnGPU(numeratorGPU, [2]*fft.Domain{s.domain0, s.domain1}) + if e != nil { + return e + } + doneDivideByZH() + s.hGPU = quotientGPU + + // Shared state slices were mutated during numerator coset iterations and are no + // longer needed now; computeLinearizedPolynomial uses the immutable snapshot. + s.releaseSharedGPUState() + close(s.chRestoreLRO) + + doneCommitH := profileStep("computeQuotient: commit H from device") + if err := s.commitToQuotientGPUFromDevice(s.hGPU); err != nil { + return err + } + doneCommitH() + + if err := s.deriveZeta(); err != nil { + return err + } + + donePrepareLinearizedEval := profileStep("computeQuotient: prepare linearized eval GPU state") + if err := s.prepareLinearizedEvalGPUStateFromHost(); err != nil { + return fmt.Errorf("computeQuotient: prepare linearized eval GPU state failed: %w", err) + } + donePrepareLinearizedEval() + + close(s.chH) + + return nil +} + +func (s *instance) buildRatioCopyConstraint() (err error) { + // wait for gamma and beta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chGammaBeta: + } + + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("buildRatioCopyConstraint: missing L/R/O polynomials") + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}) + if err != nil { + return err + } + dL, err := getStateDeviceSlice(gpuState, s.polyL, "L") + if err != nil { + return err + } + dR, err := getStateDeviceSlice(gpuState, s.polyR, "R") + if err != nil { + return err + } + dO, err := getStateDeviceSlice(gpuState, s.polyO, "O") + if err != nil { + return err + } + + var startBuildRatioCopyConstraintIcicle time.Time + if isProfileMode { + startBuildRatioCopyConstraintIcicle = time.Now() + } + s.polyZ, err = s.BuildRatioCopyConstraintIcicle( + []icicle_core.DeviceSlice{dL, dR, dO}, + s.trace.S, + s.beta, + s.gamma, + s.domain0, + gpuState, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startBuildRatioCopyConstraintIcicle)).Msg("buildRatioCopyConstraint: BuildRatioCopyConstraintIcicle") + s.timings.record("buildRatioCopyConstraint (perm Z)", time.Since(startBuildRatioCopyConstraintIcicle)) + } + + dZ, err := getStateDeviceSlice(gpuState, s.polyZ, "Z") + if err != nil { + return err + } + copyDone := make(chan error, 1) + var dPersist icicle_core.DeviceSlice + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + var allocErr error + dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) + if allocErr != nil { + finish(fmt.Errorf("buildRatioCopyConstraint: alloc persist Z failed: %w", allocErr)) + return + } + if e := copyDeviceSliceIntoOnCurrentDevice(dPersist, dZ, cfg); e != icicle_runtime.Success { + _ = dPersist.Free() + dPersist = icicle_core.DeviceSlice{} + finish(fmt.Errorf("buildRatioCopyConstraint: persist Z copy failed: %s", e.AsString())) + return + } + finish(nil) + }) + if err := <-copyDone; err != nil { + return err + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + s.polyZLagrangeGPU = dPersist + + // commit to Z (with or without blinding) + var startCommitZ time.Time + if isProfileMode { + startCommitZ = time.Now() + } + if useBlinding { + s.proof.Z, err = s.commitToPolyAndBlinding(s.polyZ, s.bp[id_Bz]) + } else { + s.proof.Z, err = s.commitLagrangePolynomialOnGPU(s.polyZ) + } + if isProfileMode { + l := logger.Logger() + if useBlinding { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (no blinding)") + } + } + s.freeIdleTempGPUMemoryOnDevice() + + close(s.chZ) + + return +} + +// open Z (blinded) at ωζ +func (s *instance) openZ() (err error) { + // wait for H to be committed and zeta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + + var zetaShifted fr.Element + zetaShifted.Mul(&s.zeta, &s.pk.Vk.Generator) + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("openZ: missing GPU Z polynomial") + } + + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + + dZLagrange := s.polyZLagrangeGPU + if dZLagrange.Len() <= 1 { + return fmt.Errorf("openZ: invalid Z size %d", dZLagrange.Len()) + } + + blindSize := order_blinding_Z + 1 + if useBlinding { + if len(s.bp) <= id_Bz || s.bp[id_Bz] == nil { + return fmt.Errorf("openZ: missing Z blinding polynomial") + } + blindSize = len(s.bp[id_Bz].Coefficients()) + if blindSize == 0 { + return fmt.Errorf("openZ: empty Z blinding polynomial") + } + } + + var dBlindedCanonical icicle_core.DeviceSlice + buildDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + buildDone <- cfgErr + return + } + finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + // Async boundary for canonicalization/blinding before exposing output. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + buildDone <- runErr + } + + n := dZLagrange.Len() + dZCanonical := s.getTempDeviceSlice(n) + if err := copyDeviceSliceIntoOnCurrentDevice(dZCanonical, dZLagrange, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy Z to canonical buffer failed: %s", err.AsString()), dZCanonical, false) + return + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + cfgNtt.Ordering = icicle_core.KNN // regular lagrange -> regular canonical + if err := icicle_ntt.Ntt(dZCanonical, icicle_core.KInverse, &cfgNtt, dZCanonical); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: inverse NTT on Z failed: %s", err.AsString()), dZCanonical, false) + return + } + + dBlindedCanonical = s.getTempDeviceSlice(n + blindSize) + dBlindedPrefix := (&dBlindedCanonical).Range(0, n, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedPrefix, dZCanonical, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy canonical Z into blinded buffer failed: %s", err.AsString()), dZCanonical, true) + return + } + + if useBlinding { + dBp := uploadVector(s.bp[id_Bz].Coefficients()) + dBlindedHead := (&dBlindedPrefix).Range(0, blindSize, false) + if err := icicle_vecops.VecOp(dBlindedHead, dBp, dBlindedHead, cfgVec, icicle_core.Sub); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: subtract Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedTail, dBp, cfgVec); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: append Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + _ = dBp.FreeAsync(stream) + } else { + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := zeroDeviceSliceOnCurrentDevice(dBlindedTail, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: zero-pad non-blinded Z failed: %s", err.AsString()), dZCanonical, true) + return + } + } + + finalize(nil, dZCanonical, false) + }) + if err := <-buildDone; err != nil { + return err + } + s.blindedZCanonicalGPU = dBlindedCanonical + + // open z at zeta*w. + var startKzgOpen time.Time + if isProfileMode { + startKzgOpen = time.Now() + } + s.proof.ZShiftedOpening, err = s.openPolynomialOnGPUCanonicalDevice(s.blindedZCanonicalGPU, zetaShifted) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startKzgOpen)).Msg("openZ: open polynomial on GPU") + s.timings.record("openZ (KZG open on GPU)", time.Since(startKzgOpen)) + } + if err != nil { + return err + } + close(s.chZOpening) + return nil +} + +func (s *instance) h1() []fr.Element { + var h1 []fr.Element + if !s.opt.StatisticalZK { + h1 = s.h.Coefficients()[:s.domain0.Cardinality+2] + } else { + h1 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h1, s.h.Coefficients()[:s.domain0.Cardinality+2]) + h1[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[0]) + } + return h1 +} + +func (s *instance) h2() []fr.Element { + var h2 []fr.Element + if !s.opt.StatisticalZK { + h2 = s.h.Coefficients()[s.domain0.Cardinality+2 : 2*(s.domain0.Cardinality+2)] + } else { + h2 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h2, s.h.Coefficients()[s.domain0.Cardinality+2:2*(s.domain0.Cardinality+2)]) + h2[0].Sub(&h2[0], &s.quotientShardsRandomizers[0]) + h2[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[1]) + } + return h2 +} + +func (s *instance) h3() []fr.Element { + var h3 []fr.Element + if !s.opt.StatisticalZK { + h3 = s.h.Coefficients()[2*(s.domain0.Cardinality+2) : 3*(s.domain0.Cardinality+2)] + } else { + h3 = make([]fr.Element, s.domain0.Cardinality+2) + copy(h3, s.h.Coefficients()[2*(s.domain0.Cardinality+2):3*(s.domain0.Cardinality+2)]) + h3[0].Sub(&h3[0], &s.quotientShardsRandomizers[1]) + } + return h3 +} + +// witnessEvalAtZeta holds the scalar evaluations of witness and constraint +// polynomials at the challenge point zeta, as needed by the linearized +// polynomial computation. +type witnessEvalAtZeta struct { + blzeta, brzeta, bozeta fr.Element + s1zeta, s2zeta fr.Element + qcpzeta []fr.Element +} + +type linearizedSelectorScales struct { + s3, ql, qr, qm, qo, qk fr.Element + qcp []fr.Element +} + +// evaluateWitnessPolynomialsAtZeta evaluates L, R, O (with optional blinding), +// S1, S2, and all Qcp polynomials at the point zeta using the GPU-resident +// polynomial state. +func (s *instance) evaluateWitnessPolynomialsAtZeta( + evalGPUState *gpuPolysState, + zeta fr.Element, +) (witnessEvalAtZeta, error) { + doneTotal := profileStep("evaluateWitnessPolynomialsAtZeta: total") + defer doneTotal() + + var result witnessEvalAtZeta + var err error + + result.qcpzeta = make([]fr.Element, len(s.commitmentInfo)) + var startQcp time.Time + if isProfileMode { + startQcp = time.Now() + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i >= len(s.trace.Qcp) || s.trace.Qcp[i] == nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: missing Qcp polynomial at index %d", i) + } + var startQcpItem time.Time + if isProfileMode { + startQcpItem = time.Now() + } + result.qcpzeta[i], err = s.evalPolynomialInCurrentFormOnGPU(s.trace.Qcp[i], evalGPUState, zeta) + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: qcp[%d] GPU evaluation failed: %w", i, err) + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("idx", i).Dur("took", time.Since(startQcpItem)).Msg("evaluateWitnessPolynomialsAtZeta: qcp eval item") + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startQcp)).Msg("evaluateWitnessPolynomialsAtZeta: qcpZeta evaluate on GPU") + } + + if useBlinding { + result.blzeta, err = s.evaluateBlindedOnGPU(s.polyL, s.bp[id_Bl], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: L evaluate on GPU") + result.blzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyL, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: blzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.brzeta, err = s.evaluateBlindedOnGPU(s.polyR, s.bp[id_Br], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: R evaluate on GPU") + result.brzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyR, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: brzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.bozeta, err = s.evaluateBlindedOnGPU(s.polyO, s.bp[id_Bo], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: O evaluate on GPU") + result.bozeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyO, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: bozeta GPU evaluation failed: %w", err) + } + + doneS1 := profileStep("evaluateWitnessPolynomialsAtZeta: S1 evaluate on GPU") + result.s1zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S1, evalGPUState, zeta) + doneS1() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s1(zeta) GPU evaluation failed: %w", err) + } + doneS2 := profileStep("evaluateWitnessPolynomialsAtZeta: S2 evaluate on GPU") + result.s2zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S2, evalGPUState, zeta) + doneS2() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s2(zeta) GPU evaluation failed: %w", err) + } + + return result, nil +} + +func (s *instance) computeLinearizedPolynomial() error { + + // wait for H to be committed and zeta to be derived (or ctx.Done()) + var startWaitH time.Time + if isProfileMode { + startWaitH = time.Now() + } + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startWaitH)).Msg("computeLinearizedPolynomial: wait H and zeta") + s.timings.record("computeLinearizedPoly (wait H+zeta, overlaps)", time.Since(startWaitH)) + } + if s.opt.StatisticalZK { + return fmt.Errorf("computeLinearizedPolynomial: GPU-only opening path does not support StatisticalZK=true") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil || s.polyZ == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing required polynomials") + } + + // Reuse the immutable snapshot prepared in computeQuotient before numerator + // coset iterations mutate shared state slices. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("computeLinearizedPolynomial: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized eval prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: required polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing blinding polynomials for GPU evaluation") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized blinding prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: blinding polynomial is not on GPU: %w", e) + } + } + } + + doneEvaluate := profileStep("computeLinearizedPolynomial: evaluate witness polynomials") + evals, err := s.evaluateWitnessPolynomialsAtZeta(evalGPUState, s.zeta) + doneEvaluate() + if err != nil { + return err + } + + // wait for Z to be opened at zeta (or ctx.Done()) + doneWaitZOpening := profileStep("computeLinearizedPolynomial: wait Z opening") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZOpening: + } + doneWaitZOpening() + if s.blindedZCanonicalGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing canonical blinded Z on GPU") + } + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing lagrange Z on GPU") + } + defer func() { + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + }() + bzuzeta := s.proof.ZShiftedOpening.ClaimedValue + + if s.hGPU == nil || s.hGPU.coeffs.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing GPU quotient polynomial") + } + + doneBuild := profileStep("computeLinearizedPolynomial: build selector terms on GPU") + dLin, err := s.buildLinearizedSelectorTermsOnGPU(evals, bzuzeta, s.blindedZCanonicalGPU.Len()) + doneBuild() + if err != nil { + return err + } + + doneAddZ := profileStep("computeLinearizedPolynomial: add Z contribution on GPU") + err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) + doneAddZ() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + s.linearizedPolynomialGPU = dLin + + doneEvalClaim := profileStep("computeLinearizedPolynomial: evaluate linearized claim") + claim, err := s.evalDevicePolynomialAtPoint(dLin, s.zeta) + doneEvalClaim() + if err != nil { + return err + } + s.linearizedPolynomialClaim = claim + + // Commit the linearized polynomial over the canonical SRS. + var startMSM time.Time + if isProfileMode { + startMSM = time.Now() + } + s.linearizedPolynomialDigest, err = commitOnGPUCanonicalDevice(dLin, &s.device, s.pk) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startMSM)).Msg("computeLinearizedPolynomial: KZG commit") + s.timings.record("computeLinearizedPoly (KZG commit)", time.Since(startMSM)) + } + close(s.chLinearizedPolynomial) + return nil +} + +func (s *instance) batchOpening() error { + + // wait for linearizedPolynomial to be computed (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLinearizedPolynomial: + } + + defer func() { + freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } + }() + + if s.linearizedPolynomialGPU.IsEmpty() { + return fmt.Errorf("batchOpening: missing GPU linearized polynomial") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("batchOpening: missing L/R/O polynomials") + } + + // Reuse immutable GPU snapshot prepared before quotient iterations. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("batchOpening: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening eval prereq"); e != nil { + return fmt.Errorf("batchOpening: required polynomial is not on GPU: %w", e) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] == nil { + return fmt.Errorf("batchOpening: missing Qcp polynomial at index %d", i) + } + if _, e := getStateDeviceSlice(evalGPUState, s.trace.Qcp[i], "batchOpening qcp prereq"); e != nil { + return fmt.Errorf("batchOpening: Qcp polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("batchOpening: missing blinding polynomials") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening blinding prereq"); e != nil { + return fmt.Errorf("batchOpening: blinding polynomial is not on GPU: %w", e) + } + } + } + + devicePolys, ownedPolys, claimed, err := s.prepareBatchOpeningPolynomialsOnGPU(evalGPUState, s.zeta) + if err != nil { + return err + } + defer func() { + for i := 0; i < len(devicePolys); i++ { + if i < len(ownedPolys) && ownedPolys[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + } + } + s.releaseLinearizedEvalGPUState() + }() + + digestsToOpen := make([]curve.G1Affine, len(s.pk.Vk.Qcp)+6) + copy(digestsToOpen[6:], s.pk.Vk.Qcp) + digestsToOpen[0] = s.linearizedPolynomialDigest + digestsToOpen[1] = s.proof.LRO[0] + digestsToOpen[2] = s.proof.LRO[1] + digestsToOpen[3] = s.proof.LRO[2] + digestsToOpen[4] = s.pk.Vk.S[0] + digestsToOpen[5] = s.pk.Vk.S[1] + if len(claimed) != len(digestsToOpen) { + return fmt.Errorf("batchOpening: claimed size mismatch (%d != %d)", len(claimed), len(digestsToOpen)) + } + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + var dFold icicle_core.DeviceSlice + foldDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening fold", foldDone) + + dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: copy linearized polynomial failed: %s", e.AsString())) + return + } + + gammaPow := gamma + for i := 1; i < len(devicePolys); i++ { + dPoly := devicePolys[i] + if dPoly.IsEmpty() { + gammaPow.Mul(&gammaPow, &gamma) + continue + } + dScaled := s.getTempDeviceSlice(dPoly.Len()) + dGammaStd := uploadScalarStdOnCurrentDevice(gammaPow, cfg) + eMul := icicle_vecops.ScalarMulVec(dGammaStd, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dGammaStd.FreeAsync(cfg.StreamHandle) + } else { + _ = dGammaStd.Free() + } + if eMul != icicle_runtime.Success { + if cfg.IsAsync { + _ = icicle_runtime.SynchronizeStream(cfg.StreamHandle) + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: scale polynomial %d failed: %s", i, eMul.AsString())) + return + } + + dPrefix := (&dFold).Range(0, dPoly.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add) + if cfg.IsAsync { + // dScaled is recycled each iteration; wait before returning to pool. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: synchronize stream failed: %s", eSync.AsString())) + return + } + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + if eAdd != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: fold add polynomial %d failed: %s", i, eAdd.AsString())) + return + } + gammaPow.Mul(&gammaPow, &gamma) + } + finish(nil) + }) + if err := <-foldDone; err != nil { + if !dFold.IsEmpty() { + s.putTempDeviceSlice(dFold, dFold.Len()) + } + return err + } + var dWitness icicle_core.DeviceSlice + divDone := make(chan error, 1) + witnessSize := dFold.Len() - 1 + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dFold, dPoint, dWitness, cfg) + if cfg.IsAsync { + _ = dPoint.FreeAsync(cfg.StreamHandle) + } else { + _ = dPoint.Free() + } + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + finish(fmt.Errorf("batchOpening: divide by (x-zeta) failed: %s", eDiv.AsString())) + return + } + finish(nil) + }) + if err := <-divDone; err != nil { + s.putTempDeviceSlice(dFold, dFold.Len()) + return err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dFold, dFold.Len()) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return err + } + + s.proof.BatchedProof = kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + } + if err := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); err != nil { + l := logger.Logger() + l.Warn().Err(err).Msg("batchOpening: GPU folded opening failed raw-G2 validation; falling back to host fold with CPU KZG commitment") + fallbackProof, fallbackErr := s.batchOpeningHostFoldGPUCommitFromDevicePolys(devicePolys, digestsToOpen) + if fallbackErr != nil { + return fallbackErr + } + s.proof.BatchedProof = fallbackProof + if verifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); verifyErr != nil { + nativeProof, nativeErr := s.batchOpeningNativeCPUFromDevicePolys(devicePolys, digestsToOpen) + if nativeErr != nil { + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed: %v", verifyErr, nativeErr) + } + s.proof.BatchedProof = nativeProof + if nativeVerifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); nativeVerifyErr != nil { + diagnostic, diagnosticErr := s.diagnoseBatchOpeningDevicePolynomials(devicePolys, digestsToOpen, claimed) + if diagnosticErr != nil { + diagnostic = fmt.Sprintf("batch opening diagnostic failed: %v", diagnosticErr) + } + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed raw-G2 validation: %v; %s", verifyErr, nativeVerifyErr, diagnostic) + } + l.Warn().Msg("batchOpening: native CPU KZG fallback produced a valid proof after host-fold fallback failed") + } + } + _ = foldedEval // kept for parity with kzg.BatchOpenSinglePoint flow. + return nil +} + +func (s *instance) verifyBatchOpeningWithRecomputedLines(digestsToOpen []curve.G1Affine) error { + vk := s.pk.Vk.Kzg + vk.Lines[0] = curve.PrecomputeLines(vk.G2[0]) + vk.Lines[1] = curve.PrecomputeLines(vk.G2[1]) + return kzg.BatchVerifySinglePoint( + digestsToOpen, + &s.proof.BatchedProof, + s.zeta, + s.kzgFoldingHash, + vk, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningHostFoldGPUCommitFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningNativeCPUFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningNativeCPUFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return kzg.BatchOpenSinglePoint( + polysToOpen, + digestsToOpen, + s.zeta, + s.kzgFoldingHash, + s.pk.Kzg, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) diagnoseBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, + gpuClaimed []fr.Element, +) (string, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "diagnoseBatchOpeningDevicePolynomials", + ) + if err != nil { + return "", err + } + if len(gpuClaimed) != len(polysToOpen) { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: claimed/polynomial mismatch (%d != %d)", len(gpuClaimed), len(polysToOpen)) + } + + l := logger.Logger() + claimMismatches := make([]string, 0) + commitMismatches := make([]string, 0) + for i := range polysToOpen { + label := batchOpeningPolynomialLabel(i) + cpuClaim := evalCanonicalAtPoint(polysToOpen[i], s.zeta) + if !cpuClaim.Equal(&gpuClaimed[i]) { + claimMismatches = append(claimMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("gpuClaim", frFingerprint(gpuClaimed[i])). + Str("cpuClaim", frFingerprint(cpuClaim)). + Msg("batchOpening diagnostic: GPU claim differs from CPU evaluation") + } + + cpuDigest, err := kzg.Commit(polysToOpen[i], s.pk.Kzg) + if err != nil { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: commit %s: %w", label, err) + } + if !cpuDigest.Equal(&digestsToOpen[i]) { + commitMismatches = append(commitMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("expectedDigest", g1Fingerprint(digestsToOpen[i])). + Str("cpuDigest", g1Fingerprint(cpuDigest)). + Msg("batchOpening diagnostic: CPU commitment differs from proof digest") + } + } + + if len(claimMismatches) == 0 && len(commitMismatches) == 0 { + return "batch opening diagnostic found no per-polynomial claim or commitment mismatch", nil + } + return fmt.Sprintf( + "batch opening diagnostic claim mismatches=[%s] commitment mismatches=[%s]", + strings.Join(claimMismatches, ","), + strings.Join(commitMismatches, ","), + ), nil +} + +func batchOpeningPolynomialLabel(index int) string { + switch index { + case 0: + return "linearized" + case 1: + return "L" + case 2: + return "R" + case 3: + return "O" + case 4: + return "S1" + case 5: + return "S2" + default: + return fmt.Sprintf("Qcp[%d]", index-6) + } +} + +func frFingerprint(v fr.Element) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func g1Fingerprint(v curve.G1Affine) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func (s *instance) downloadBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + expectedDigests int, + label string, +) ([][]fr.Element, error) { + if len(devicePolys) != expectedDigests { + return nil, fmt.Errorf("%s: polynomial/digest mismatch (%d != %d)", label, len(devicePolys), expectedDigests) + } + + polysToOpen := make([][]fr.Element, len(devicePolys)) + for i := range devicePolys { + var err error + polysToOpen[i], err = s.downloadCanonicalDeviceCoefficients( + devicePolys[i], + fmt.Sprintf("%s[%d]", label, i), + ) + if err != nil { + return nil, err + } + } + return polysToOpen, nil +} + +func (s *instance) batchOpeningHostFoldGPUCommit(digestsToOpen []curve.G1Affine) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.batchOpeningHostPolynomials() + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( + polysToOpen [][]fr.Element, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + if len(polysToOpen) != len(digestsToOpen) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: polynomial/digest mismatch (%d != %d)", len(polysToOpen), len(digestsToOpen)) + } + + largestPoly := 0 + for i := range polysToOpen { + if len(polysToOpen[i]) == 0 || len(polysToOpen[i]) > len(s.pk.Kzg.G1) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: invalid polynomial %d size %d", i, len(polysToOpen[i])) + } + if len(polysToOpen[i]) > largestPoly { + largestPoly = len(polysToOpen[i]) + } + } + + claimed := make([]fr.Element, len(polysToOpen)) + utils.Parallelize(len(polysToOpen), func(start, end int) { + for i := start; i < end; i++ { + claimed[i] = evalCanonicalAtPoint(polysToOpen[i], s.zeta) + } + }) + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + foldedPolynomials := make([]fr.Element, largestPoly) + copy(foldedPolynomials, polysToOpen[0]) + gammaPow := gamma + for i := 1; i < len(polysToOpen); i++ { + poly := polysToOpen[i] + scale := gammaPow + utils.Parallelize(len(poly), func(start, end int) { + var term fr.Element + for j := start; j < end; j++ { + term.Mul(&poly[j], &scale) + foldedPolynomials[j].Add(&foldedPolynomials[j], &term) + } + }) + gammaPow.Mul(&gammaPow, &gamma) + } + + hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) + var dWitness icicle_core.DeviceSlice + uploadDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + dWitness = uploadVector(hCoeffs) + close(uploadDone) + }) + <-uploadDone + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + freeSliceOnDevice(&dWitness, &s.device) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + return kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + }, nil +} + +func (s *instance) batchOpeningHostPolynomials() ([][]fr.Element, error) { + total := 6 + len(s.trace.Qcp) + polysToOpen := make([][]fr.Element, total) + + var err error + polysToOpen[0], err = s.downloadCanonicalDeviceCoefficients( + s.linearizedPolynomialGPU, + "batchOpeningHostPolynomials linearized", + ) + if err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare linearized: %w", err) + } + + prepareLRO := func(p, bp *iop.Polynomial, blindingOrder int) ([]fr.Element, error) { + base, err := canonicalRegularCoefficientsCopy(p, s.domain0) + if err != nil { + return nil, err + } + if useBlinding { + if bp == nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomial") + } + blind, err := canonicalRegularCoefficientsCopy(bp, s.domain0) + if err != nil { + return nil, err + } + out := make([]fr.Element, len(base)+len(blind)) + copy(out, base) + copy(out[len(base):], blind) + for i := range blind { + out[i].Sub(&out[i], &blind[i]) + } + return out, nil + } + out := make([]fr.Element, len(base)+blindingOrder+1) + copy(out, base) + return out, nil + } + + var bpL, bpR, bpO *iop.Polynomial + if useBlinding { + if len(s.bp) <= id_Bo { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomials") + } + bpL, bpR, bpO = s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo] + } + + if polysToOpen[1], err = prepareLRO(s.polyL, bpL, order_blinding_L); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare L: %w", err) + } + if polysToOpen[2], err = prepareLRO(s.polyR, bpR, order_blinding_R); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare R: %w", err) + } + if polysToOpen[3], err = prepareLRO(s.polyO, bpO, order_blinding_O); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare O: %w", err) + } + if polysToOpen[4], err = canonicalRegularCoefficientsCopy(s.trace.S1, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S1: %w", err) + } + if polysToOpen[5], err = canonicalRegularCoefficientsCopy(s.trace.S2, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S2: %w", err) + } + for i := range s.trace.Qcp { + if polysToOpen[6+i], err = canonicalRegularCoefficientsCopy(s.trace.Qcp[i], s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare Qcp[%d]: %w", i, err) + } + } + return polysToOpen, nil +} + +func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceSlice, label string) ([]fr.Element, error) { + if dPoly.IsEmpty() { + return nil, fmt.Errorf("%s: empty device polynomial", label) + } + + coeffs := make([]fr.Element, dPoly.Len()) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + }) + if err := <-done; err != nil { + return nil, err + } + return coeffs, nil +} + +func canonicalRegularCoefficientsCopy(p *iop.Polynomial, domain *fft.Domain) ([]fr.Element, error) { + if p == nil { + return nil, fmt.Errorf("nil polynomial") + } + cp := p.Clone() + cp.ToCanonical(domain).ToRegular() + coeffs := cp.Coefficients() + out := make([]fr.Element, len(coeffs)) + copy(out, coeffs) + return out, nil +} + +func dividePolyByXMinusAHost(f []fr.Element, fa, a fr.Element) []fr.Element { + f[0].Sub(&f[0], &fa) + var t fr.Element + for i := len(f) - 2; i >= 0; i-- { + t.Mul(&f[i+1], &a) + f[i].Add(&f[i], &t) + } + return f[1:] +} + +// evaluate the full set of constraints on the GPU-resident polynomial state. +type computeNumeratorLoopContext struct { + n int + rho int + mm uint64 + bn *big.Int + shifters []fr.Element + twiddles0 []fr.Element + dTwiddles0 icicle_core.DeviceSlice + dPrecomputedDenominators *icicle_core.DeviceSlice + scalingVector []fr.Element + scalingVectorRev []fr.Element + gpuState *gpuPolysState + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + one fr.Element + cs fr.Element + css fr.Element + nbBsbGates int + numeratorShards []icicle_core.DeviceSlice +} + +type gpuNumeratorPolynomial struct { + shards []icicle_core.DeviceSlice + n int + rho int + mm uint64 +} + +type gpuQuotientPolynomial struct { + coeffs icicle_core.DeviceSlice + size int +} + +func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolynomial, error) { + twiddles0, err := s.buildComputeNumeratorTwiddles() + if err != nil { + return nil, err + } + if err := s.validateComputeNumeratorGPUState(gpuState); err != nil { + return nil, err + } + + var startComputeNumerator time.Time + if isProfileMode { + startComputeNumerator = time.Now() + } + + n := s.domain0.Cardinality + nbBsbGates := len(s.proof.Bsb22Commitments) + + var cs, css fr.Element + cs.Set(&s.domain1.FrMultiplicativeGen) + css.Square(&cs) + + bn := big.NewInt(int64(n)) + + rho := int(s.domain1.Cardinality / n) + shifters := make([]fr.Element, rho) + shifters[0].Set(&s.domain1.FrMultiplicativeGen) + for i := 1; i < rho; i++ { + shifters[i].Set(&s.domain1.Generator) + } + + cosetTable, err := s.domain0.CosetTable() + if err != nil { + return nil, err + } + + // for the first iteration, the scalingVector is the coset table + scalingVector := cosetTable + scalingVectorRev := make([]fr.Element, len(cosetTable)) + copy(scalingVectorRev, cosetTable) + fft.BitReverse(scalingVectorRev) + + // pre-computed to compute the bit reverse index + // of the result polynomial + m := uint64(s.domain1.Cardinality) + mm := uint64(64 - bits.TrailingZeros64(m)) + + var dPrecomputedDenominators icicle_core.DeviceSlice + defer func() { + if !dPrecomputedDenominators.IsEmpty() { + freeDone := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + freeDone <- dPrecomputedDenominators.Free() + }) + if err := <-freeDone; err != icicle_runtime.Success { + panic(fmt.Sprintf("computeNumerator: failed to free dPrecomputedDenominators: %s", err.AsString())) + } + } + }() + + var coset, cosetExponentiatedToNMinusOne, one fr.Element + coset.SetOne() + one.SetOne() + + dTwiddles0, err := s.uploadComputeNumeratorTwiddles(twiddles0) + if err != nil { + return nil, err + } + + loopCtx := &computeNumeratorLoopContext{ + n: int(n), + rho: rho, + mm: mm, + bn: bn, + shifters: shifters, + twiddles0: twiddles0, + dTwiddles0: dTwiddles0, + dPrecomputedDenominators: &dPrecomputedDenominators, + scalingVector: scalingVector, + scalingVectorRev: scalingVectorRev, + gpuState: gpuState, + coset: coset, + cosetExponentiatedToNMinusOne: cosetExponentiatedToNMinusOne, + one: one, + cs: cs, + css: css, + nbBsbGates: nbBsbGates, + numeratorShards: make([]icicle_core.DeviceSlice, rho), + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: setup before iteration loop") + } + if err := s.executeComputeNumeratorCosetIterations(loopCtx); err != nil { + return nil, err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: main body (post-wait)") + } + + return &gpuNumeratorPolynomial{ + shards: loopCtx.numeratorShards, + n: loopCtx.n, + rho: loopCtx.rho, + mm: loopCtx.mm, + }, nil + +} + +func (s *instance) buildComputeNumeratorTwiddles() ([]fr.Element, error) { + n := s.domain0.Cardinality + var startTwiddles time.Time + if isProfileMode { + startTwiddles = time.Now() + } + twiddles0 := make([]fr.Element, n) + if n == 1 { + // edge case + twiddles0[0].SetOne() + } else { + twiddles, err := s.domain0.Twiddles() + if err != nil { + return nil, err + } + copy(twiddles0, twiddles[0]) + w := twiddles0[1] + for i := len(twiddles[0]); i < len(twiddles0); i++ { + twiddles0[i].Mul(&twiddles0[i-1], &w) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startTwiddles)).Msg("computeNumerator: build twiddles") + } + return twiddles0, nil +} + +func (s *instance) waitForComputeNumeratorQk() error { + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chQk: + } + return nil +} + +func (s *instance) buildLinearizedEvalGPUBatch() []*iop.Polynomial { + baseCap := 5 + len(s.trace.Qcp) + if useBlinding { + baseCap += 3 + } + polys := make([]*iop.Polynomial, 0, baseCap) + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + } + if useBlinding && len(s.bp) > id_Bo { + for _, bpPoly := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if bpPoly != nil { + polys = append(polys, bpPoly) + } + } + } + return polys +} + +func (s *instance) buildComputeNumeratorGPUBatch() []*iop.Polynomial { + polys := make([]*iop.Polynomial, 0, 13+2*len(s.commitmentInfo)) + for _, p := range []*iop.Polynomial{ + s.polyL, s.polyR, s.polyO, s.polyZ, + s.trace.Ql, s.trace.Qr, s.trace.Qm, s.trace.Qo, s.polyQk, + s.trace.S1, s.trace.S2, s.trace.S3, + } { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i < len(s.trace.Qcp) && s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + if i < len(s.cCommitments) && s.cCommitments[i] != nil { + polys = append(polys, s.cCommitments[i]) + } + } + return polys +} + +func (s *instance) validateComputeNumeratorGPUState(state *gpuPolysState) error { + if state == nil { + return fmt.Errorf("computeNumerator: shared GPU state is nil") + } + required := s.buildComputeNumeratorGPUBatch() + if len(required) == 0 { + return fmt.Errorf("computeNumerator: no polynomials prepared for GPU batch") + } + for _, p := range required { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; !ok { + return fmt.Errorf("computeNumerator: polynomial ptr=%p is missing from shared GPU state", p) + } + } + return nil +} + +func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicle_core.DeviceSlice, error) { + var dTwiddles0 icicle_core.DeviceSlice + uploadTwiddlesDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + if s.tempGPUMemPool != nil { + s.tempGPUMemPool.FreeAll() + } + host := icicle_core.HostSliceFromElements(twiddles0) + var allocErr error + dTwiddles0, allocErr = allocDeviceUninitialized(len(twiddles0)) + if allocErr != nil { + uploadTwiddlesDone <- fmt.Errorf("uploadComputeNumeratorTwiddles: %w", allocErr) + return + } + host.CopyToDevice(&dTwiddles0, false) + uploadTwiddlesDone <- nil + }) + if err := <-uploadTwiddlesDone; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dTwiddles0, nil +} + +// executeComputeNumeratorCosetIterations runs the rho coset iterations. +func (s *instance) executeComputeNumeratorCosetIterations(loopCtx *computeNumeratorLoopContext) error { + var startIterLoop time.Time + if isProfileMode { + startIterLoop = time.Now() + } + + for i := 0; i < loopCtx.rho; i++ { + if err := s.computeNumeratorIteration(i, loopCtx); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + freeSliceOnDevice(&loopCtx.dTwiddles0, &s.device) + return err + } + } + + // Free twiddles0 device slice (uploaded once before the loop). + freeTwiddlesDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + loopCtx.dTwiddles0.Free() + close(freeTwiddlesDone) + }) + <-freeTwiddlesDone + + if useBlinding { + var startRestoreBlindingPolys time.Time + if isProfileMode { + startRestoreBlindingPolys = time.Now() + } + csInv := inverseShifterProduct(loopCtx.shifters) + if err := s.restoreBlindingPolynomials(csInv); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startRestoreBlindingPolys)).Msg("computeNumerator: restore blinding polys") + } + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startIterLoop)).Msg("computeNumerator: full iteration loop") + } + return nil +} + +func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoopContext) error { + loopCtx.coset.Mul(&loopCtx.coset, &loopCtx.shifters[i]) + loopCtx.cosetExponentiatedToNMinusOne.Exp(loopCtx.coset, loopCtx.bn). + Sub(&loopCtx.cosetExponentiatedToNMinusOne, &loopCtx.one) + + batchInvertDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) + }) + if err := <-batchInvertDone; err != nil { + return err + } + + s.applyNumeratorBlindingScale(i, loopCtx) + if i == 1 { + // We have to update the scalingVector; instead of scaling by + // cosets we scale by the twiddles of the large domain. + w := s.domain1.Generator + loopCtx.scalingVector = make([]fr.Element, loopCtx.n) + fft.BuildExpTable(w, loopCtx.scalingVector) + + // Reuse memory. + copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) + fft.BitReverse(loopCtx.scalingVectorRev) + } + + // We do **a lot** of FFT here, but on the small domain. + // Note that for all the polynomials in the proving key + // (Ql, Qr, Qm, Qo, S1, S2, S3, Qcp, Qc) and ID, LOne + // we could pre-compute these rho*2 FFTs and store them + // at the cost of a huge memory footprint. + var startGpuInverseScaleForward time.Time + if isProfileMode { + startGpuInverseScaleForward = time.Now() + } + + // Inverse NTT -> Scale -> Forward NTT all on GPU using persistent GPU memory. + if err := s.gpuNTTInverseScaleForwardOnDevice(loopCtx.gpuState, loopCtx.scalingVector, loopCtx.scalingVectorRev, s.pk); err != nil { + return err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startGpuInverseScaleForward)).Msg("computeNumerator: gpuNTTInverseScaleForwardOnDevice") + } + + // Evaluate constraints on GPU. + constraintParams := gpuConstraintEvalParams{ + beta: s.beta, + gamma: s.gamma, + alpha: s.alpha, + coset: loopCtx.coset, + cosetExponentiatedToNMinusOne: loopCtx.cosetExponentiatedToNMinusOne, + cs: loopCtx.cs, + css: loopCtx.css, + cardinalityInv: s.domain0.CardinalityInv, + n: loopCtx.n, + nbBsbGates: loopCtx.nbBsbGates, + } + var startEvalConstraints time.Time + if isProfileMode { + startEvalConstraints = time.Now() + } + dNumeratorShard, err := s.gpuEvaluateConstraints( + loopCtx.gpuState, + constraintParams, + loopCtx.twiddles0, // CPU version for computeBlindingPolynomials + loopCtx.dTwiddles0, // GPU version for computeOrderingConstraint + *loopCtx.dPrecomputedDenominators, + s.bp, + nil, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startEvalConstraints)).Msg("computeNumerator: gpuEvaluateConstraints") + } + loopCtx.numeratorShards[i] = dNumeratorShard + + loopCtx.cosetExponentiatedToNMinusOne. + Inverse(&loopCtx.cosetExponentiatedToNMinusOne) + s.applyNumeratorBlindingUnscale(i, loopCtx) + return nil +} + +func (s *instance) buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx *computeNumeratorLoopContext) error { + if loopCtx == nil || loopCtx.dPrecomputedDenominators == nil { + return fmt.Errorf("computeNumerator: nil denominator device slice") + } + if loopCtx.dTwiddles0.IsEmpty() || loopCtx.dTwiddles0.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid twiddles device slice size %d, expected %d", loopCtx.dTwiddles0.Len(), loopCtx.n) + } + + if loopCtx.dPrecomputedDenominators.IsEmpty() { + dDenominators, err := allocDeviceUninitialized(loopCtx.n) + if err != nil { + return err + } + *loopCtx.dPrecomputedDenominators = dDenominators + } else if loopCtx.dPrecomputedDenominators.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid denominator device slice size %d, expected %d", loopCtx.dPrecomputedDenominators.Len(), loopCtx.n) + } + + cfg := icicle_core.DefaultVecOpsConfig() + + // dTwiddles0 is a Montgomery scalar vector in domain0 regular order. + // ScalarMulVec expects the scalar in standard form and preserves the + // Montgomery representation of the vector result. + dCosetStd := uploadScalarStdOnCurrentDevice(loopCtx.coset, cfg) + defer dCosetStd.Free() + if err := icicle_vecops.ScalarMulVec( + dCosetStd, + loopCtx.dTwiddles0, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators coset*twiddles failed: %s", err.AsString()) + } + + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneMont := uploadScalarMontOnCurrentDevice(minusOne, cfg) + defer dMinusOneMont.Free() + if err := icicle_vecops.ScalarAddVec( + dMinusOneMont, + *loopCtx.dPrecomputedDenominators, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators subtract one failed: %s", err.AsString()) + } + + if err := s.batchInvertOnCurrentDevice(*loopCtx.dPrecomputedDenominators); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: batchInvert failed: %s", err.AsString()) + } + return nil +} + +func (s *instance) applyNumeratorBlindingScale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindScale time.Time + if isProfileMode { + startBlindScale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + acc := loopCtx.cosetExponentiatedToNMinusOne + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &acc) + acc.Mul(&acc, &loopCtx.shifters[i]) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindScale)).Msg("computeNumerator: scale blinding polys") + } +} + +func (s *instance) applyNumeratorBlindingUnscale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindUnscale time.Time + if isProfileMode { + startBlindUnscale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &loopCtx.cosetExponentiatedToNMinusOne) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindUnscale)).Msg("computeNumerator: unscale blinding polys") + } +} + +func (s *instance) restoreBlindingPolynomials(csInv fr.Element) error { + for _, q := range s.bp { + if q == nil { + continue + } + cp := q.Coefficients() + if len(cp) == 0 { + continue + } + var acc fr.Element + acc.SetOne() + for i := 0; i < len(cp); i++ { + cp[i].Mul(&cp[i], &acc) + acc.Mul(&acc, &csInv) + } + } + return nil +} + +func inverseShifterProduct(shifters []fr.Element) fr.Element { + var acc fr.Element + acc.SetOne() + for i := 0; i < len(shifters); i++ { + acc.Mul(&acc, &shifters[i]) + } + acc.Inverse(&acc) + return acc +} + +func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial) (_ *iop.Polynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("downloadNumeratorFromGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("downloadNumeratorFromGPU: invalid dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard count mismatch: got %d, expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + + defer func() { + if err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + } + }() + + for i := 0; i < gpuNumerator.rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard %d is empty", i) + } + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + var dMerged icicle_core.DeviceSlice + + mergeDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + return + } + // Async boundary before merged slice is consumed by host copy. + mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + }) + if mergeErr := <-mergeDone; mergeErr != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, mergeErr + } + + cres := make([]fr.Element, totalSize) + cresHost := icicle_core.HostSliceFromElements(cres) + downloadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + }) + if err := <-downloadDone; err != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, err + } + s.putTempDeviceSlice(dMerged, totalSize) + + s.freeNumeratorShards(gpuNumerator.shards) + return iop.NewPolynomial(&cres, iop.Form{Basis: iop.LagrangeCoset, Layout: iop.BitReverse}), nil +} + +func (s *instance) freeNumeratorShards(shards []icicle_core.DeviceSlice) { + if len(shards) == 0 { + return + } + for i := 0; i < len(shards); i++ { + if !shards[i].IsEmpty() { + s.putTempDeviceSlice(shards[i], shards[i].Len()) + shards[i] = icicle_core.DeviceSlice{} + } + } +} + +func (s *instance) batchInvert(dVec icicle_core.DeviceSlice) { + if dVec.Len() == 0 { + return + } + + done := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + done <- s.batchInvertOnCurrentDevice(dVec) + }) + if err := <-done; err != icicle_runtime.Success { + panic(fmt.Sprintf("batchInvert: BatchInverseVec failed: %s", err.AsString())) + } +} + +// batchInvertOnCurrentDevice assumes caller already runs on the active device thread. +func (s *instance) batchInvertOnCurrentDevice(dVec icicle_core.DeviceSlice) icicle_runtime.EIcicleError { + if dVec.Len() == 0 { + return icicle_runtime.Success + } + err := icicle_bls12381.FromMontgomery(dVec) + if err == icicle_runtime.Success { + cfg := icicle_core.DefaultVecOpsConfig() + err = icicle_vecops.BatchInverseVec(dVec, dVec, cfg) + } + if err == icicle_runtime.Success { + err = icicle_bls12381.ToMontgomery(dVec) + } + return err +} + +// gpuPolysState holds GPU-resident polynomial data to avoid repeated CPU-GPU transfers. +// Use ensurePolysOnSharedGPU to populate/reuse and freeGPUPolys to release GPU memory. +type gpuPolysState struct { + deviceSlices []icicle_core.DeviceSlice + hostSlices []icicle_core.HostSlice[fr.Element] + polys []*iop.Polynomial + originalForm []iop.Form + polyToIdx map[*iop.Polynomial]int +} + +func (s *instance) sharedGPUStateInitialCap(extra int) int { + base := 16 + len(s.bp) + 2*len(s.commitmentInfo) + if extra > 0 { + base += extra + } + return base +} + +func (s *instance) initSharedGPUState() { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + if s.sharedGPUState != nil { + return + } + initialCap := s.sharedGPUStateInitialCap(0) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } +} + +func (s *instance) releaseSharedGPUState() { + s.gpuStateMu.Lock() + state := s.sharedGPUState + s.sharedGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) releaseLinearizedEvalGPUState() { + s.gpuStateMu.Lock() + state := s.linearizedEvalGPUState + s.linearizedEvalGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) freeIdleTempGPUMemoryOnDevice() { + if s == nil || s.tempGPUMemPool == nil { + return + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +func (s *instance) prepareLinearizedEvalGPUState(source *gpuPolysState) error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.clonePolysOnGPUFromState(source, s.buildLinearizedEvalGPUBatch()) + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) prepareLinearizedEvalGPUStateFromHost() error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.uploadPolysToGPUState(s.buildLinearizedEvalGPUBatch(), "prepareLinearizedEvalGPUStateFromHost") + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop.Polynomial) (*gpuPolysState, error) { + if source == nil { + return nil, fmt.Errorf("clonePolysOnGPUFromState: nil source state") + } + + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("clonePolysOnGPUFromState: empty polynomial batch") + } + + srcSlices := make([]icicle_core.DeviceSlice, len(unique)) + useSource := make([]bool, len(unique)) + for i, p := range unique { + idx, ok := source.polyToIdx[p] + if ok && idx >= 0 && idx < len(source.deviceSlices) && !source.deviceSlices[idx].IsEmpty() { + srcSlices[i] = source.deviceSlices[idx] + useSource[i] = true + } + } + + snapshot := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + snapshot.polys[i] = p + snapshot.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + snapshot.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("clonePolysOnGPUFromState") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for snapshot cloning before handing state to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i := range srcSlices { + if useSource[i] { + dst, allocErr := allocDeviceUninitialized(srcSlices[i].Len()) + if allocErr != nil { + runErr = fmt.Errorf("clonePolysOnGPUFromState: alloc failed at index %d: %w", i, allocErr) + return + } + if err := copyDeviceSliceIntoOnCurrentDevice(dst, srcSlices[i], cfg); err != icicle_runtime.Success { + _ = dst.Free() + runErr = fmt.Errorf("clonePolysOnGPUFromState: device copy failed at index %d: %s", i, err.AsString()) + return + } + snapshot.deviceSlices[i] = dst + continue + } + + coeffs := unique[i].Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("clonePolysOnGPUFromState: empty host coefficients at index %d", i) + return + } + host := icicle_core.HostSliceFromElements(coeffs) + var dst icicle_core.DeviceSlice + host.CopyToDeviceAsync(&dst, cfg.StreamHandle, true) + if dst.IsEmpty() { + runErr = fmt.Errorf("clonePolysOnGPUFromState: host upload failed at index %d", i) + return + } + snapshot.deviceSlices[i] = dst + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(snapshot) + return nil, err + } + return snapshot, nil +} + +func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) (*gpuPolysState, error) { + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("%s: empty polynomial batch", label) + } + + state := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + state.polys[i] = p + state.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + state.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label) + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i, p := range unique { + coeffs := p.Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("%s: empty host coefficients at index %d", label, i) + return + } + state.hostSlices[i] = icicle_core.HostSliceFromElements(coeffs) + state.hostSlices[i].CopyToDeviceAsync(&state.deviceSlices[i], cfg.StreamHandle, true) + if state.deviceSlices[i].IsEmpty() { + runErr = fmt.Errorf("%s: host upload failed at index %d", label, i) + return + } + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(state) + return nil, err + } + return state, nil +} + +func (s *instance) getTempDeviceSlice(n int) icicle_core.DeviceSlice { + if s == nil { + panic("getTempDeviceSlice: nil instance") + } + if s.tempGPUMemPool == nil { + panic("getTempDeviceSlice: temp GPU memory pool is not initialized") + } + return s.tempGPUMemPool.Get(n) +} + +func (s *instance) putTempDeviceSlice(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + if s != nil && s.tempGPUMemPool != nil { + s.tempGPUMemPool.Put(ds, n) + return + } + _ = ds.Free() +} + +func (s *instance) releaseTempGPUMemoryPool() { + if s == nil || s.tempGPUMemPool == nil { + return + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +// ensurePolysOnSharedGPU uploads missing polynomials once and reuses already-uploaded slices. +func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysState, error) { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if s.sharedGPUState == nil { + initialCap := s.sharedGPUStateInitialCap(len(polys)) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } + } + state := s.sharedGPUState + if state == nil { + return nil, fmt.Errorf("ensurePolysOnSharedGPU: shared GPU state is nil") + } + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int, len(polys)) + } + + newIndices := make([]int, 0, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; ok { + continue + } + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, icicle_core.DeviceSlice{}) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + newIndices = append(newIndices, idx) + } + if len(newIndices) == 0 { + return state, nil + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("ensurePolysOnSharedGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for GPU uploads in ensurePolysOnSharedGPU. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + for _, idx := range newIndices { + p := state.polys[idx] + if p == nil { + continue + } + cp := p.Coefficients() + state.hostSlices[idx] = icicle_core.HostSliceFromElements(cp) + state.hostSlices[idx].CopyToDeviceAsync(&state.deviceSlices[idx], cfg.StreamHandle, true) + } + }) + if err := <-done; err != nil { + return nil, err + } + return state, nil +} + +func getStateDeviceSlice(state *gpuPolysState, p *iop.Polynomial, label string) (icicle_core.DeviceSlice, error) { + if state == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil GPU state", label) + } + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil polynomial", label) + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): polynomial is not registered on GPU", label) + } + ds := state.deviceSlices[idx] + if ds.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): empty device slice", label) + } + return ds, nil +} + +// gpuNTTInverseScaleForwardOnDevice performs inverse NTT → scale → forward NTT +// on GPU-resident polynomial data without CPU-GPU transfers for polynomial data. +// The scaling vectors are uploaded each call (they may change between iterations). +func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scalingVector, scalingVectorRev []fr.Element, pk *ProvingKey) error { + if state == nil || len(state.polys) == 0 { + return nil + } + + device := &s.device + var scalingVectorDevice, scalingVectorRevDevice icicle_core.DeviceSlice + + // Upload scaling vectors to GPU + uploadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) + scalingHost := icicle_core.HostSliceFromElements(scalingVector) + scalingHost.CopyToDevice(&scalingVectorDevice, false) + scalingVectorRevDevice = s.getTempDeviceSlice(len(scalingVectorRev)) + scalingRevHost := icicle_core.HostSliceFromElements(scalingVectorRev) + scalingRevHost.CopyToDevice(&scalingVectorRevDevice, false) + + // Convert scaling vectors from Montgomery form to standard form + if err := icicle_bls12381.FromMontgomery(scalingVectorDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVector failed: %s", err.AsString()) + return + } + if err := icicle_bls12381.FromMontgomery(scalingVectorRevDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVectorRev failed: %s", err.AsString()) + return + } + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return err + } + + doneChans := make([]chan error, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + + done := make(chan error, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Step 1: Inverse NTT + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + return + } + + // Step 2: Scale by vector using GPU vecOps + vecCfg := icicle_core.DefaultVecOpsConfig() + vecCfg.StreamHandle = stream + vecCfg.IsAsync = true + + var scaleDevice icicle_core.DeviceSlice + if layout == iop.Regular { + // After KNR inverse, output is BitReverse → use scalingVectorRev + scaleDevice = scalingVectorRevDevice + } else { + // After KRN inverse, output is Regular → use scalingVector + scaleDevice = scalingVectorDevice + } + + if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + return + } + + // Step 3: Forward NTT to Lagrange + one := icicle_ntt.GetDefaultNttConfig().CosetGen + cfg.CosetGen = one + if layout == iop.Regular { + cfg.Ordering = icicle_core.KRN // BitReverse → Regular + } else { + cfg.Ordering = icicle_core.KNR // Regular → BitReverse + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + return + } + done <- nil + }) + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + if err := <-doneChans[i]; err != nil { + return err + } + } + } + + // Update polynomial metadata: final result is in Lagrange, same layout as original + for _, p := range state.polys { + if p != nil { + p.Basis = iop.Lagrange + } + } + + // Free scaling vectors from device + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + return nil +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +// +// NOTE: The prover hot path should use the state-based API and keep data on device. +// This wrapper exists for compatibility/testing where callers expect host coefficients +// to be materialized after the transform. +func (s *instance) gpuNTTInverseBatch(polys []*iop.Polynomial, pk *ProvingKey) { + if len(polys) == 0 { + return + } + state, err := s.ensurePolysOnSharedGPU(polys) + if err != nil { + panic(fmt.Sprintf("gpuNTTInverseBatch: ensurePolysOnSharedGPU failed: %v", err)) + } + + s.gpuNTTInverseBatchOnState(state, pk) + + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + for _, p := range polys { + if p == nil { + continue + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) || state.deviceSlices[idx].IsEmpty() { + continue + } + cp := p.Coefficients() + host := icicle_core.HostSliceFromElements(cp) + host.CopyFromDevice(&state.deviceSlices[idx]) + copy(cp, ([]fr.Element)(host)) + } + close(done) + }) + <-done +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKey) { + if state == nil || len(state.polys) == 0 { + return + } + + device := &s.device + doneChans := make([]chan struct{}, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + switch p.Basis { + case iop.Canonical: + continue // already in canonical form + case iop.Lagrange, iop.LagrangeCoset: + // Schedule GPU work + done := make(chan struct{}, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Select ordering and coset generator depending on basis and input layout + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + + // Base-domain inverse: + // - Regular input → KNR (output BitReverse) + // - BitReverse input → KRN (output Regular) + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + + // Run NTT inverse directly on the existing device slice (in-place) + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) + } + icicle_runtime.SynchronizeStream(stream) + + // Update metadata inside closure to avoid race + p.Basis = iop.Canonical + if layout == iop.Regular { + p.Layout = iop.BitReverse + } else { + p.Layout = iop.Regular + } + close(done) + }) + default: + panic("unsupported polynomial basis") + } + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + <-doneChans[i] + } + } +} + +// freeGPUPolys releases GPU memory for polynomial data. +func (s *instance) freeGPUPolys(state *gpuPolysState) { + if state == nil { + return + } + + device := &s.device + freeDone := make(chan struct{}) + + icicle_runtime.RunOnDevice(device, func(args ...any) { + for i := range state.polys { + if state.deviceSlices[i].IsEmpty() { + continue + } + _ = state.deviceSlices[i].Free() + } + close(freeDone) + }) + <-freeDone +} + +// gpuMemoryPool manages a pool of reusable device slices to avoid repeated allocations. +// Must be used within RunOnDevice context to ensure thread safety per device. +type gpuMemoryPool struct { + freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice + mu sync.Mutex +} + +// newGPUMemoryPool creates a new GPU memory pool. +func newGPUMemoryPool() *gpuMemoryPool { + return &gpuMemoryPool{ + freeSlices: make(map[int][]icicle_core.DeviceSlice), + } +} + +// Get returns a device slice of the specified size, either from the pool or newly allocated. +func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { + p.mu.Lock() + defer p.mu.Unlock() + + // Check if we have a free slice of this size + if slices, ok := p.freeSlices[n]; ok && len(slices) > 0 { + // Reuse the last slice + slice := slices[len(slices)-1] + p.freeSlices[n] = slices[:len(slices)-1] + return slice + } + + // No free slice available, allocate a new one. + // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + // Free all currently cached (idle) slices to reduce memory pressure. + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + panic(fmt.Sprintf("gpuMemoryPool.Get: allocation failed for size %d after clearing idle cache", n)) +} + +// Put returns a device slice to the pool for reuse instead of freeing it. +func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + + p.mu.Lock() + defer p.mu.Unlock() + + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// FreeAll releases all pooled device slices. +func (p *gpuMemoryPool) FreeAll() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) +} + +// allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. +// Use when the destination is fully overwritten by a kernel. +func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { + var ds icicle_core.DeviceSlice + if _, err := ds.Malloc(int(unsafe.Sizeof(fr.Element{})), n); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("allocDeviceUninitialized: malloc failed for size %d: %s", n, err.AsString()) + } + return ds, nil +} + +// mustAllocDeviceUninitialized is like allocDeviceUninitialized but panics on failure. +// Use only in contexts where error propagation is impractical (e.g. upload helpers). +func mustAllocDeviceUninitialized(n int) icicle_core.DeviceSlice { + ds, err := allocDeviceUninitialized(n) + if err != nil { + panic(err) + } + return ds +} + +// freeDeviceSlice frees a device slice if non-empty and zeroes the pointer. +// Use for directly-allocated slices, NOT pool-allocated ones (use putTempDeviceSlice for those). +func freeDeviceSlice(ds *icicle_core.DeviceSlice) { + if ds != nil && !ds.IsEmpty() { + _ = ds.Free() + *ds = icicle_core.DeviceSlice{} + } +} + +// freeSliceOnDevice frees a device slice on the specified device and blocks +// until complete. Use outside RunOnDevice closures. Zeroes the slice after freeing. +func freeSliceOnDevice(ds *icicle_core.DeviceSlice, device *icicle_runtime.Device) { + if ds == nil || ds.IsEmpty() { + return + } + d := *ds + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + _ = d.Free() + close(done) + }) + <-done + *ds = icicle_core.DeviceSlice{} +} + +// copyDeviceSliceIntoOnCurrentDevice copies src into dst entirely on GPU. +func copyDeviceSliceIntoOnCurrentDevice( + dst, src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if src.IsEmpty() || src.Len() <= 0 || dst.IsEmpty() || dst.Len() < src.Len() { + return icicle_runtime.InvalidArgument + } + src.CheckDevice() + dst.CheckDevice() + + srcElemSize := src.SizeOfElement() + dstElemSize := dst.SizeOfElement() + if srcElemSize <= 0 || dstElemSize <= 0 || srcElemSize != dstElemSize { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(src.Len() * srcElemSize) + if cfg.IsAsync { + return icicle_runtime.CopyAsync(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen, cfg.StreamHandle) + } + _, err := icicle_runtime.Copy(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen) + return err +} + +// zeroDeviceSliceOnCurrentDevice zero-fills dst entirely on GPU. +func zeroDeviceSliceOnCurrentDevice( + dst icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if dst.IsEmpty() || dst.Len() <= 0 { + return icicle_runtime.InvalidArgument + } + dst.CheckDevice() + + elemSize := dst.SizeOfElement() + if elemSize <= 0 { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(dst.Len() * elemSize) + if cfg.IsAsync { + return icicle_runtime.MemSetAsync(dst.AsUnsafePointer(), 0, byteLen, cfg.StreamHandle) + } + return icicle_runtime.MemSet(dst.AsUnsafePointer(), 0, byteLen) +} + +func createAsyncVecOpsConfigOnCurrentDevice(label string) (icicle_core.VecOpsConfig, icicle_runtime.Stream, error) { + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + return icicle_core.VecOpsConfig{}, nil, fmt.Errorf("%s: create stream failed: %s", label, eStream.AsString()) + } + cfg := icicle_core.DefaultVecOpsConfig() + cfg.StreamHandle = stream + cfg.IsAsync = true + return cfg, stream, nil +} + +func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label string) error { + if stream == nil { + return nil + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + _ = icicle_runtime.DestroyStream(stream) + return fmt.Errorf("%s: synchronize stream failed: %s", label, eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success { + return fmt.Errorf("%s: destroy stream failed: %s", label, eDestroy.AsString()) + } + return nil +} + +// makeFinisher returns a closure that synchronizes and destroys the stream, +// then sends the (possibly merged) error to done. Use inside RunOnDevice closures. +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { + return func(runErr error) { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + } +} + +// uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. +// Use this for additions where the vector is already in Montgomery form. +func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + return ds +} + +func uploadScalarMontOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + if cfg.IsAsync { + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&ds, false) + } + return ds +} + +// uploadScalarStd uploads a scalar in STANDARD form (not Montgomery) as a single-element device slice. +// For use with ScalarMulVec: (a*R) * b_std = (a*b)*R +func uploadScalarStd(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + // Convert from Montgomery form to standard form + if err := icicle_bls12381.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +func uploadScalarStdOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + ds := uploadScalarMontOnCurrentDevice(scalar, cfg) + // Fallback to sync conversion for compatibility with ICICLE wrappers + // that do not expose *_WithConfig APIs. + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bls12381.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +// uploadVectorStd uploads a vector and converts to standard form. +func uploadVectorStd(vec []fr.Element) icicle_core.DeviceSlice { + ds := mustAllocDeviceUninitialized(len(vec)) + uploadVectorStdInto(&ds, vec) + return ds +} + +// uploadVectorStdInto uploads vec into an existing device slice and converts it to standard form. +// The destination must already be allocated with enough capacity for len(vec) elements. +func uploadVectorStdInto(dst *icicle_core.DeviceSlice, vec []fr.Element) { + cfg := icicle_core.DefaultVecOpsConfig() + uploadVectorStdIntoOnCurrentDevice(dst, vec, cfg) +} + +// uploadVectorStdIntoOnCurrentDevice uploads vec into an existing device slice and converts it +// to standard form while honoring the provided vector-op config/stream. +func uploadVectorStdIntoOnCurrentDevice( + dst *icicle_core.DeviceSlice, + vec []fr.Element, + cfg icicle_core.VecOpsConfig, +) { + host := icicle_core.HostSliceFromElements(vec) + if cfg.IsAsync { + host.CopyToDeviceAsync(dst, cfg.StreamHandle, false) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } else { + host.CopyToDevice(dst, false) + } + // Convert from Montgomery form to standard form + if err := icicle_bls12381.FromMontgomery(*dst); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +// uploadVector uploads a vector (keeps Montgomery form for additions). +func uploadVector(vec []fr.Element) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(len(vec)) + host.CopyToDevice(&ds, false) + return ds +} + +// uploadInt64Vector uploads int64 indices to a device slice. +func uploadInt64Vector(vec []int64) icicle_core.DeviceSlice { + cfg := icicle_core.DefaultVecOpsConfig() + return uploadInt64VectorOnCurrentDevice(vec, cfg) +} + +func uploadInt64VectorOnCurrentDevice(vec []int64, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + var ds icicle_core.DeviceSlice + if cfg.IsAsync { + if _, err := ds.MallocAsync(int(unsafe.Sizeof(int64(0))), len(vec), cfg.StreamHandle); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc async failed: %s", err.AsString())) + } + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + return ds + } + if _, err := ds.Malloc(int(unsafe.Sizeof(int64(0))), len(vec)); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc failed: %s", err.AsString())) + } + host.CopyToDevice(&ds, false) + return ds +} + +// toStandardFormInPlace converts a device slice to standard form in-place (modifies the source). +// Use this for temporary vectors that won't be needed in Montgomery form. +func toStandardFormInPlace(src icicle_core.DeviceSlice) { + cfg := icicle_core.DefaultVecOpsConfig() + toStandardFormInPlaceWithCfg(src, cfg) +} + +func toStandardFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bls12381.FromMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +func toMontgomeryFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bls12381.ToMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("ToMontgomery failed: %s", err.AsString())) + } +} + +// multiplyMontgomerySlices multiplies two device slices that are both in Montgomery form. +// It creates a copy of dSlice1Mont, converts the copy to standard form, and then multiplies +// it with dSlice2Mont (which remains in Montgomery form). The result is stored in dResult +// and will be in Montgomery form. +// +// Parameters: +// - dSlice1Mont: first device slice in Montgomery form (not modified) +// - dSlice2Mont: second device slice in Montgomery form (not modified) +// - dResult: destination device slice for the result (must be pre-allocated) +// - state: GPU state with memory pool and vector configuration +// - n: size of the slices +func multiplyMontgomerySlices( + dSlice1Mont, dSlice2Mont icicle_core.DeviceSlice, + dResult icicle_core.DeviceSlice, + state *gpuConstraintEvalState, + vecCfg icicle_core.VecOpsConfig, + n int, +) error { + // Copy dSlice1Mont to standard form + dSlice1Std := state.getTempDeviceSlice(n) + defer state.putTempDeviceSlice(dSlice1Std, n) + + if err := copyDeviceSliceIntoOnCurrentDevice(dSlice1Std, dSlice1Mont, vecCfg); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: device copy failed: %s", err.AsString()) + } + toStandardFormInPlace(dSlice1Std) + + // Multiply: dSlice1Std (standard) * dSlice2Mont (Montgomery) = dResult (Montgomery) + if err := icicle_vecops.VecOp(dSlice1Std, dSlice2Mont, dResult, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: VecOp multiplication failed: %s", err.AsString()) + } + return nil +} + +// gpuConstraintEvalParams holds parameters for GPU constraint evaluation +type gpuConstraintEvalParams struct { + beta fr.Element + gamma fr.Element + alpha fr.Element + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + cs fr.Element // domain1.FrMultiplicativeGen + css fr.Element // cs^2 + cardinalityInv fr.Element + n int + nbBsbGates int +} + +// gpuConstraintEvalState holds intermediate state during constraint evaluation +type gpuConstraintEvalState struct { + // Polynomial device slices (may point to gpuState or allocated buffers) + dL, dR, dO, dZ, dZS icicle_core.DeviceSlice + dQl, dQr, dQm, dQo, dQk icicle_core.DeviceSlice + dS1, dS2, dS3 icicle_core.DeviceSlice + // Intermediate results + dGate, dOrdering, dLocal, dResult icicle_core.DeviceSlice + // Scalar device slices + dGammaScalar icicle_core.DeviceSlice + // Configuration + vecCfg icicle_core.VecOpsConfig + // Helper function to get device slices + getDeviceSlice func(int) icicle_core.DeviceSlice + // Shared prover-level temporary GPU memory pool accessors + getTempDeviceSlice func(int) icicle_core.DeviceSlice + putTempDeviceSlice func(icicle_core.DeviceSlice, int) + // Track allocated polynomial buffers for automatic cleanup + allocatedPolyBuffers []struct { + slice icicle_core.DeviceSlice + size int + } +} + +// allocate allocates a new device slice from the memory pool and tracks it for automatic cleanup. +// Returns the allocated device slice. +func (s *gpuConstraintEvalState) allocate(size int) icicle_core.DeviceSlice { + slice := s.getTempDeviceSlice(size) + s.allocatedPolyBuffers = append(s.allocatedPolyBuffers, struct { + slice icicle_core.DeviceSlice + size int + }{slice, size}) + return slice +} + +// freeAllocatedPolyBuffers returns all allocated polynomial buffers to the memory pool. +// This should be called during cleanup to free all buffers allocated via allocate(). +func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { + for _, buf := range s.allocatedPolyBuffers { + s.putTempDeviceSlice(buf.slice, buf.size) + } + s.allocatedPolyBuffers = nil +} + +// computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. +// Returns device slices for the blinding polynomials. +func computeBlindingPolynomials( + n int, + twiddles0 []fr.Element, + bp []*iop.Polynomial, +) (dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice) { + blindL := make([]fr.Element, n) + blindR := make([]fr.Element, n) + blindO := make([]fr.Element, n) + blindZ := make([]fr.Element, n) + blindZS := make([]fr.Element, n) // ZS uses shifted index + + // TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove this. Since we don't normally use blindings, we will + // not make this optimization. + utils.Parallelize(n, func(start, end int) { + for i := start; i < end; i++ { + blindL[i] = bp[id_Bl].Evaluate(twiddles0[i]) + blindR[i] = bp[id_Br].Evaluate(twiddles0[i]) + blindO[i] = bp[id_Bo].Evaluate(twiddles0[i]) + blindZ[i] = bp[id_Bz].Evaluate(twiddles0[i]) + blindZS[i] = bp[id_Bz].Evaluate(twiddles0[(i+1)%n]) + } + }) + + dBlindL = uploadVector(blindL) + dBlindR = uploadVector(blindR) + dBlindO = uploadVector(blindO) + dBlindZ = uploadVector(blindZ) + dBlindZS = uploadVector(blindZS) + + return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS +} + +// applyBlindingToPolynomials applies blinding to polynomials L, R, O, Z, ZS. +// Allocates new buffers for L, R, O, Z (tracked for cleanup) and modifies ZS in-place. +// The original slices in gpuState remain unchanged. +func applyBlindingToPolynomials( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice, +) error { + // L' = L + blindL (allocate new buffer, tracked for cleanup) + dLBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dL, dBlindL, dLBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add L failed: %s", err.AsString()) + } + state.dL = dLBlinded + + // R' = R + blindR (allocate new buffer, tracked for cleanup) + dRBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dR, dBlindR, dRBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add R failed: %s", err.AsString()) + } + state.dR = dRBlinded + + // O' = O + blindO (allocate new buffer, tracked for cleanup) + dOBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dO, dBlindO, dOBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add O failed: %s", err.AsString()) + } + state.dO = dOBlinded + + // Z' = Z + blindZ (allocate new buffer, tracked for cleanup) + dZBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dZ, dBlindZ, dZBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add Z failed: %s", err.AsString()) + } + state.dZ = dZBlinded + + // ZS' = ZS + blindZS + // Note: dZS is a temporary buffer created inside gpuEvaluateConstraints, + // so it's safe to modify it in-place. + if err := icicle_vecops.VecOp(state.dZS, dBlindZS, state.dZS, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) + } + + // Free blinding vectors - no longer needed after creating blinded polynomials + dBlindL.Free() + dBlindR.Free() + dBlindO.Free() + dBlindZ.Free() + dBlindZS.Free() + return nil +} + +// scaleSVectorsByBeta scales S1, S2, S3 by beta. +// Allocates new buffers for S1, S2, S3 (tracked for cleanup). +func scaleSVectorsByBeta( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, +) error { + // S1' = S1 * beta (need to scale S1, S2, S3 by beta for ordering constraint) + // Use standard form for beta so: S1_mont * beta_std = (S1*beta)_mont + dBetaStd := uploadScalarStd(params.beta) + + // S1' = S1 * beta (allocate new buffer, tracked for cleanup) + dS1Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dS1Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S1 failed: %s", err.AsString()) + } + state.dS1 = dS1Scaled + + // S2' = S2 * beta (allocate new buffer, tracked for cleanup) + dS2Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dS2Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S2 failed: %s", err.AsString()) + } + state.dS2 = dS2Scaled + + // S3' = S3 * beta (allocate new buffer, tracked for cleanup) + dS3Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dS3Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S3 failed: %s", err.AsString()) + } + state.dS3 = dS3Scaled + + // Free dBetaStd - no longer needed after scaling S vectors + dBetaStd.Free() + return nil +} + +// computeGateConstraint computes the gate constraint. +// Returns dGate device slice. +func computeGateConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // gate = Ql*L' + Qr*R' + Qm*L'*R' + Qo*O' + Qk + sum(Qci*Pi) + // We use multiplyMontgomerySlices for all poly×poly multiplications. + // Note: dL, dR, dO are used later in ordering constraint, so we preserve them. + + dGate := state.getTempDeviceSlice(params.n) + dTmp := state.getTempDeviceSlice(params.n) + + // Ql * L' + if err := multiplyMontgomerySlices(state.dQl, state.dL, dGate, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Ql*L: %w", err) + } + + // + Qr * R' + if err := multiplyMontgomerySlices(state.dQr, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qr*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qr*R failed: %s", err.AsString()) + } + + // + Qm * L' * R' (need two multiplications) + // First: Qm * L' = dTmp (Montgomery) + if err := multiplyMontgomerySlices(state.dQm, state.dL, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qm*L: %w", err) + } + // Second: dTmp (Montgomery) * R' (Montgomery) + if err := multiplyMontgomerySlices(dTmp, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: (Qm*L)*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qm*L*R failed: %s", err.AsString()) + } + + // + Qo * O' + if err := multiplyMontgomerySlices(state.dQo, state.dO, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qo*O: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qo*O failed: %s", err.AsString()) + } + + // + Qk + if err := icicle_vecops.VecOp(dGate, state.dQk, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qk failed: %s", err.AsString()) + } + + // + BSB gates: sum(Qci[2*i] * Qci[2*i+1]) + for i := 0; i < params.nbBsbGates; i++ { + origQci0 := state.getDeviceSlice(id_Qci + 2*i) + origQci1 := state.getDeviceSlice(id_Qci + 2*i + 1) + if !origQci0.IsEmpty() && !origQci1.IsEmpty() { + // Use helper to multiply Qci0 * Qci1 without modifying original values + if err := multiplyMontgomerySlices(origQci0, origQci1, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qci[%d]: %w", i, err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qci[%d] failed: %s", i, err.AsString()) + } + } + } + + // Return temporary buffer to pool - no longer needed after Step 3 + state.putTempDeviceSlice(dTmp, params.n) + + return dGate, nil +} + +// computeOrderingConstraint computes the ordering constraint. +// Returns dOrdering device slice. +func computeOrderingConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dTwiddles0 icicle_core.DeviceSlice, // twiddles0 already on GPU + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // This is complex: involves ID computation, gamma, beta, Z, ZS, S1, S2, S3 + // id = twiddles[i] * coset * beta + // a = gamma + L' + id + // b = gamma + R' + id*cs + // c = gamma + O' + id*css + // r = a * b * c * Z' + // + // a2 = gamma + L' + S1*beta + // b2 = gamma + R' + S2*beta + // c2 = gamma + O' + S3*beta + // l = a2 * b2 * c2 * ZS' + // + // ordering = l - r + + // Compute ID vector: twiddles * coset * beta (computed on GPU) + // dTwiddles0 is already on GPU (passed as parameter, don't free it here) + + // Compute coset * beta on CPU, then upload as scalar in standard form + var cosetTimesBeta fr.Element + cosetTimesBeta.Mul(¶ms.coset, ¶ms.beta) + dCosetTimesBetaStd := uploadScalarStd(cosetTimesBeta) + dBetaStd := uploadScalarStd(params.beta) + + // Multiply twiddles0 by cosetTimesBeta on GPU: dID = (cosetTimesBeta * twiddles0) * R (Montgomery form) + dID := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dCosetTimesBetaStd, dTwiddles0, dID, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec coset*beta*twiddles failed: %s", err.AsString()) + } + + // Free temporary device slice (dTwiddles0 is owned by caller, don't free it) + dCosetTimesBetaStd.Free() + + // id * cs - use standard form for cs + dIDcs := state.getTempDeviceSlice(params.n) + dCsStd := uploadScalarStd(params.cs) + if err := icicle_vecops.ScalarMulVec(dCsStd, dID, dIDcs, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*cs failed: %s", err.AsString()) + } + + // id * css - use standard form for css + dIDcss := state.getTempDeviceSlice(params.n) + dCssStd := uploadScalarStd(params.css) + if err := icicle_vecops.ScalarMulVec(dCssStd, dID, dIDcss, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*css failed: %s", err.AsString()) + } + + // a = gamma + L' + id (dL now contains L' after in-place blinding) + dA := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dID, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a+id failed: %s", err.AsString()) + } + + // b = gamma + R' + id*cs (dR now contains R' after in-place blinding) + dB := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dIDcs, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b+id*cs failed: %s", err.AsString()) + } + + // c = gamma + O' + id*css (dO now contains O' after in-place blinding) + dC := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dIDcss, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c+id*css failed: %s", err.AsString()) + } + + // Return to pool: dID, dIDcs, dIDcss - no longer needed after computing a, b, c + state.putTempDeviceSlice(dID, params.n) + state.putTempDeviceSlice(dIDcs, params.n) + state.putTempDeviceSlice(dIDcss, params.n) + dCsStd.Free() + dCssStd.Free() + + // r = a * b * c * Z' (dZ now contains Z' after in-place blinding) + // For chain multiplication, convert operands to std form in-place when possible + dR_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a*b failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dR_ord, dC, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*c failed: %s", err.AsString()) + } + // Convert dR_ord to standard form in-place (temporary result, dZ needs to be preserved) + toStandardFormInPlace(dR_ord) + if err := icicle_vecops.VecOp(dR_ord, state.dZ, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*Z failed: %s", err.AsString()) + } + + // Reuse dA, dB, dC for a2, b2, c2 instead of freeing and reallocating. + // To reduce peak memory, we scale S vectors by beta on-demand through a single temp buffer. + dScaledS := state.getTempDeviceSlice(params.n) + + // a2 = gamma + L' + S1*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L (a2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S1 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dScaledS, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2+S1*beta failed: %s", err.AsString()) + } + + // b2 = gamma + R' + S2*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R (b2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S2 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dScaledS, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b2+S2*beta failed: %s", err.AsString()) + } + + // c2 = gamma + O' + S3*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O (c2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S3 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dScaledS, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c2+S3*beta failed: %s", err.AsString()) + } + + // Free dGammaScalar - no longer needed after computing a2, b2, c2 + // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. + state.dGammaScalar.Free() + state.putTempDeviceSlice(dScaledS, params.n) + dBetaStd.Free() + + // l = a2 * b2 * c2 * ZS' (dZS now contains ZS' after in-place blinding) + dL_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2*b2 failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dL_ord, dC, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*c2 failed: %s", err.AsString()) + } + // Convert dZS to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(state.dZS) + if err := icicle_vecops.VecOp(dL_ord, state.dZS, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*ZS failed: %s", err.AsString()) + } + + // Return temporary buffers to pool - no longer needed after computing l + state.putTempDeviceSlice(dA, params.n) + state.putTempDeviceSlice(dB, params.n) + state.putTempDeviceSlice(dC, params.n) + state.putTempDeviceSlice(state.dZS, params.n) + state.dZS = icicle_core.DeviceSlice{} + + // ordering = l - r, reuse dL_ord as the final ordering vector + if err := icicle_vecops.VecOp(dL_ord, dR_ord, dL_ord, vecCfg, icicle_core.Sub); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l-r failed: %s", err.AsString()) + } + + // Return dR_ord to pool - no longer needed after computing ordering + state.putTempDeviceSlice(dR_ord, params.n) + + // Return dL_ord as ordering (caller is responsible for freeing) + return dL_ord, nil +} + +// computeLocalConstraint computes the local constraint. +// Returns dLocal device slice. +func computeLocalConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dPrecomputedDenominators icicle_core.DeviceSlice, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // local = (Z' - 1) * LagrangeOne + // where LagrangeOne[i] = cosetExpMinusOne * cardinalityInv / (coset*twiddles0[i] - 1) + + if dPrecomputedDenominators.IsEmpty() || dPrecomputedDenominators.Len() < params.n { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: invalid denominator device slice size %d, expected at least %d", dPrecomputedDenominators.Len(), params.n) + } + + // Compute LagrangeOne on device. dPrecomputedDenominators is already in + // Montgomery form after batch inversion; ScalarMulVec expects the scalar in + // standard form and preserves a Montgomery vector result. + var lagrangeCoeff fr.Element + lagrangeCoeff.Mul(¶ms.cosetExponentiatedToNMinusOne, ¶ms.cardinalityInv) + dLagrangeCoeffStd := uploadScalarStdOnCurrentDevice(lagrangeCoeff, vecCfg) + dLagrangeOneStd := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dLagrangeCoeffStd, dPrecomputedDenominators, dLagrangeOneStd, vecCfg); err != icicle_runtime.Success { + dLagrangeCoeffStd.Free() + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarMulVec lagrangeOne failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dLagrangeOneStd, vecCfg) + dLagrangeCoeffStd.Free() + + // Z' - 1 using ScalarAddVec with minus one + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneScalar := uploadScalarMont(minusOne) + + dZMinusOne := state.getTempDeviceSlice(params.n) + // dZ now contains Z' after in-place blinding + if err := icicle_vecops.ScalarAddVec(dMinusOneScalar, state.dZ, dZMinusOne, vecCfg); err != icicle_runtime.Success { + dMinusOneScalar.Free() + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarAddVec Z-1 failed: %s", err.AsString()) + } + + // Free dMinusOneScalar - no longer needed after computing Z' - 1 + // Note: dZ is part of gpuState and will be freed later + dMinusOneScalar.Free() + + // local = (Z' - 1) * LagrangeOne_std + dLocal := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.VecOp(dZMinusOne, dLagrangeOneStd, dLocal, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: VecOp (Z-1)*LagrangeOne failed: %s", err.AsString()) + } + + // Return temporary buffers to pool + state.putTempDeviceSlice(dZMinusOne, params.n) + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + + return dLocal, nil +} + +// createGetDeviceSliceFunc creates a function to get device slices for polynomials. +// It returns a function that maps polynomial indices to their device slices. +func (s *instance) polyByID(polyIdx int) *iop.Polynomial { + switch polyIdx { + case id_L: + return s.polyL + case id_R: + return s.polyR + case id_O: + return s.polyO + case id_Z: + return s.polyZ + case id_ZS: + return s.polyZS + case id_Ql: + return s.trace.Ql + case id_Qr: + return s.trace.Qr + case id_Qm: + return s.trace.Qm + case id_Qo: + return s.trace.Qo + case id_Qk: + return s.polyQk + case id_S1: + return s.trace.S1 + case id_S2: + return s.trace.S2 + case id_S3: + return s.trace.S3 + default: + if polyIdx < id_Qci { + return nil + } + offset := polyIdx - id_Qci + i := offset / 2 + if i < 0 { + return nil + } + if offset%2 == 0 { + if i < len(s.trace.Qcp) { + return s.trace.Qcp[i] + } + return nil + } + if i < len(s.cCommitments) { + return s.cCommitments[i] + } + return nil + } +} + +func createGetDeviceSliceFunc( + gpuState *gpuPolysState, + polyToIdx map[*iop.Polynomial]int, + resolvePoly func(int) *iop.Polynomial, +) func(int) icicle_core.DeviceSlice { + return func(polyIdx int) icicle_core.DeviceSlice { + p := resolvePoly(polyIdx) + if p == nil { + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d is nil", polyIdx)) + } + if idx, ok := polyToIdx[p]; ok { + return gpuState.deviceSlices[idx] + } + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d (ptr=%p) not found in polyToIdx (map has %d entries)", polyIdx, p, len(polyToIdx))) + } +} + +// initializeConstraintEvalState initializes the GPU constraint evaluation state. +// It sets up all device slices. +// The slices in gpuState are treated as read-only; helper functions will allocate +// separate working buffers whenever they need to modify data. +func initializeConstraintEvalState( + getDeviceSlice func(int) icicle_core.DeviceSlice, + getTempDeviceSlice func(int) icicle_core.DeviceSlice, + putTempDeviceSlice func(icicle_core.DeviceSlice, int), +) *gpuConstraintEvalState { + vecCfg := icicle_core.DefaultVecOpsConfig() + + state := &gpuConstraintEvalState{ + dL: getDeviceSlice(id_L), + dR: getDeviceSlice(id_R), + dO: getDeviceSlice(id_O), + dZ: getDeviceSlice(id_Z), + dQl: getDeviceSlice(id_Ql), + dQr: getDeviceSlice(id_Qr), + dQm: getDeviceSlice(id_Qm), + dQo: getDeviceSlice(id_Qo), + dQk: getDeviceSlice(id_Qk), + dS1: getDeviceSlice(id_S1), + dS2: getDeviceSlice(id_S2), + dS3: getDeviceSlice(id_S3), + vecCfg: vecCfg, + getDeviceSlice: getDeviceSlice, + getTempDeviceSlice: getTempDeviceSlice, + putTempDeviceSlice: putTempDeviceSlice, + } + + return state +} + +// gpuEvaluateConstraints evaluates all PLONK constraints on GPU. +// It takes polynomials already on GPU (via gpuState), computes blinding polynomial evaluations, +// and evaluates gate, ordering, and local constraints entirely on GPU. +// If result is non-nil, it downloads into result and returns an empty device slice. +// If result is nil, it returns a persistent device slice with the result. +// TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove 'twiddles0'. Since we don't +// normally use blindings, we will not make this optimization. +func (s *instance) gpuEvaluateConstraints( + gpuState *gpuPolysState, + params gpuConstraintEvalParams, + twiddles0 []fr.Element, // CPU vector for computeBlindingPolynomials + dTwiddles0 icicle_core.DeviceSlice, // GPU vector for computeOrderingConstraint + dPrecomputedDenominators icicle_core.DeviceSlice, + bp []*iop.Polynomial, // blinding polynomials (already scaled for this iteration) + result []fr.Element, +) (icicle_core.DeviceSlice, error) { + if gpuState == nil || len(gpuState.polys) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("gpuState is nil or empty") + } + + n := params.n + device := &s.device + + // Create a map from polynomial to its index in gpuState + polyToIdx := make(map[*iop.Polynomial]int) + for i, p := range gpuState.polys { + if p != nil { + polyToIdx[p] = i + } + } + + // Get device slices for the polynomials we need. + getDeviceSlice := createGetDeviceSliceFunc(gpuState, polyToIdx, s.polyByID) + + done := make(chan error, 1) + var resultOnDevice icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(device, func(args ...any) { + state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). + // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). + state.dGammaScalar = uploadScalarMont(params.gamma) + + state.dZS = state.getTempDeviceSlice(n) + if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + return + } + + // Step 1: Compute and apply blinding polynomial evaluations (if enabled) + if useBlinding { + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + done <- err + return + } + } + + // Step 2-4: Compute gate, ordering, and local constraints sequentially on a + // single synchronous stream, folding them into dResult as + // gate + alpha*ordering + alpha^2*local. Computing one family at a time + // keeps peak device allocation low, and sequential is not a compromise: + // the family kernels are memory-bandwidth-bound and each already saturates + // the device, so the parallel three-stream variant this replaces measured + // identical timings (111ms/iteration at n=2^23) — while racing on the + // shared temp-slice pool and lazily materialized inputs (it corrupted the + // numerator at every circuit size). + seqVecCfg := state.vecCfg + seqVecCfg.IsAsync = false + + // Compute ordering first to minimize peak memory before gate/local allocations. + var err error + state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + return + } + + // dResult = alpha * ordering + state.dResult = state.getTempDeviceSlice(params.n) + dAlphaStd := uploadScalarStd(params.alpha) + if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { + dAlphaStd.Free() + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + return + } + dAlphaStd.Free() + state.putTempDeviceSlice(state.dOrdering, params.n) + + // dResult += alpha^2 * local + state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + return + } + var alphaSquared fr.Element + alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) + dAlphaSquaredStd := uploadScalarStd(alphaSquared) + dTmp := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + return + } + dAlphaSquaredStd.Free() + if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dLocal, params.n) + state.putTempDeviceSlice(dTmp, params.n) + + // dResult += gate + state.dGate, err = computeGateConstraint(state, params, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + return + } + if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dGate, params.n) + + // Step 5: materialize result either on host or as a persistent device slice. + if result != nil { + resultHost := icicle_core.HostSliceFromElements(result) + resultHost.CopyFromDevice(&state.dResult) + } else { + resultOnDevice = s.getTempDeviceSlice(params.n) + if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // Return dResult pool slice after materialization. + state.putTempDeviceSlice(state.dResult, params.n) + + // Return all allocated polynomial buffers to the pool. + // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). + // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are + // returned to pool inside computeGateConstraint. dQk and the original gpuState slices + // are owned by gpuState and will be freed separately. + state.freeAllocatedPolyBuffers() + + done <- nil + }) + + err := <-done + + if err != nil { + if !resultOnDevice.IsEmpty() { + s.putTempDeviceSlice(resultOnDevice, resultOnDevice.Len()) + } + return icicle_core.DeviceSlice{}, err + } + return resultOnDevice, nil +} + +func evaluateBlinded(p, bp *iop.Polynomial, zeta fr.Element) fr.Element { + // Get the size of the polynomial + n := big.NewInt(int64(p.Size())) + + var pEvaluatedAtZeta fr.Element + + // Evaluate the polynomial and blinded polynomial at zeta + chP := make(chan struct{}, 1) + go func() { + pEvaluatedAtZeta = p.Evaluate(zeta) + close(chP) + }() + + bpEvaluatedAtZeta := bp.Evaluate(zeta) + + // Multiply the evaluated blinded polynomial by tempElement + var t fr.Element + one := fr.One() + t.Exp(zeta, n).Sub(&t, &one) + bpEvaluatedAtZeta.Mul(&bpEvaluatedAtZeta, &t) + + // Add the evaluated polynomial and the evaluated blinded polynomial + <-chP + pEvaluatedAtZeta.Add(&pEvaluatedAtZeta, &bpEvaluatedAtZeta) + + // Return the result + return pEvaluatedAtZeta +} + +// /!\ modifies the size +func getBlindedCoefficients(p, bp *iop.Polynomial) []fr.Element { + cp := p.Coefficients() + cbp := bp.Coefficients() + cp = append(cp, cbp...) + for i := 0; i < len(cbp); i++ { + cp[i].Sub(&cp[i], &cbp[i]) + } + return cp +} + +// getNonBlindedCoefficients returns a padded copy of polynomial coefficients +// to match the size they would have with blinding enabled. +// The padding size is blindingOrder+1 (e.g., order 2 → 3 coefficients). +func getNonBlindedCoefficients(p *iop.Polynomial, blindingOrder int) []fr.Element { + cp := p.Coefficients() + padded := make([]fr.Element, len(cp)+blindingOrder+1) + copy(padded, cp) + return padded +} + +// commits to a polynomial of the form b*(Xⁿ-1) where b is of small degree +func commitBlindingFactor(n int, b *iop.Polynomial, key kzg.ProvingKey) curve.G1Affine { + cp := b.Coefficients() + np := b.Size() + + // lo + var tmp curve.G1Affine + tmp.MultiExp(key.G1[:np], cp, ecc.MultiExpConfig{}) + + // hi + var res curve.G1Affine + res.MultiExp(key.G1[n:n+np], cp, ecc.MultiExpConfig{}) + res.Sub(&res, &tmp) + return res +} + +// return a random polynomial of degree n, if n==-1 cancel the blinding +func getRandomPolynomial(n int) *iop.Polynomial { + var a []fr.Element + if n == -1 { + a := make([]fr.Element, 1) + a[0].SetZero() + } else { + a = make([]fr.Element, n+1) + for i := 0; i <= n; i++ { + a[i].SetRandom() + } + } + res := iop.NewPolynomial(&a, iop.Form{ + Basis: iop.Canonical, Layout: iop.Regular}) + return res +} + +func coefficients(p []*iop.Polynomial) [][]fr.Element { + res := make([][]fr.Element, len(p)) + for i, pI := range p { + res[i] = pI.Coefficients() + } + return res +} + +func (s *instance) freeGPUQuotient(quotient *gpuQuotientPolynomial) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return + } + s.putTempDeviceSlice(quotient.coeffs, quotient.coeffs.Len()) + quotient.coeffs = icicle_core.DeviceSlice{} + quotient.size = 0 +} + +// commitToQuotientGPUFromDevice commits H1/H2/H3 directly from device memory. +// For StatisticalZK=true we materialize adjusted device vectors for h1/h2/h3 +// and commit those without downloading quotient coefficients to host. +// prepareStatisticalZKQuotientShards constructs blinded quotient polynomial +// shards h1, h2, h3 on the GPU for the Statistical ZK path. Each shard is +// randomized so that the quotient split h = h1 + X^(n+2)*h2 + X^(2(n+2))*h3 +// hides the original polynomial. +// +// Caller is responsible for returning dH1, dH2, dH3 to the temp pool: +// - dH1 and dH2 have size nPlus2+1 +// - dH3 has size nPlus2 +func (s *instance) prepareStatisticalZKQuotientShards( + h1Device, h2Device, h3Device icicle_core.DeviceSlice, + nPlus2 int, +) (dH1, dH2, dH3 icicle_core.DeviceSlice, err error) { + nPlus3 := nPlus2 + 1 + + prepareDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + dH1 = s.getTempDeviceSlice(nPlus3) + dH2 = s.getTempDeviceSlice(nPlus3) + dH3 = s.getTempDeviceSlice(nPlus2) + + // h1 = base h1 with extra randomizer coefficient at degree n+2. + dH1Prefix := (&dH1).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH1Prefix, h1Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h1 failed: %s", e.AsString()) + return + } + dH1Tail := (&dH1).Range(nPlus2, nPlus3, false) + r0Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[0]}) + r0Host.CopyToDevice(&dH1Tail, false) + + // h2 = base h2 with first coefficient adjusted by -r0 and tail = r1. + dH2Prefix := (&dH2).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH2Prefix, h2Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h2 failed: %s", e.AsString()) + return + } + dH2First := (&dH2).Range(0, 1, false) + var negR0 fr.Element + negR0.Neg(&s.quotientShardsRandomizers[0]) + dNegR0 := uploadScalarMont(negR0) + if e := icicle_vecops.ScalarAddVec(dNegR0, dH2First, dH2First, cfg); e != icicle_runtime.Success { + _ = dNegR0.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h2[0] failed: %s", e.AsString()) + return + } + _ = dNegR0.Free() + dH2Tail := (&dH2).Range(nPlus2, nPlus3, false) + r1Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[1]}) + r1Host.CopyToDevice(&dH2Tail, false) + + // h3 = base h3 with first coefficient adjusted by -r1. + if e := copyDeviceSliceIntoOnCurrentDevice(dH3, h3Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h3 failed: %s", e.AsString()) + return + } + dH3First := (&dH3).Range(0, 1, false) + var negR1 fr.Element + negR1.Neg(&s.quotientShardsRandomizers[1]) + dNegR1 := uploadScalarMont(negR1) + if e := icicle_vecops.ScalarAddVec(dNegR1, dH3First, dH3First, cfg); e != icicle_runtime.Success { + _ = dNegR1.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h3[0] failed: %s", e.AsString()) + return + } + _ = dNegR1.Free() + prepareDone <- nil + }) + if err := <-prepareDone; err != nil { + if !dH1.IsEmpty() { + s.putTempDeviceSlice(dH1, nPlus3) + } + if !dH2.IsEmpty() { + s.putTempDeviceSlice(dH2, nPlus3) + } + if !dH3.IsEmpty() { + s.putTempDeviceSlice(dH3, nPlus2) + } + return icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, err + } + return dH1, dH2, dH3, nil +} + +func (s *instance) commitToQuotientGPUFromDevice(quotient *gpuQuotientPolynomial) error { + if quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("commitToQuotientGPUFromDevice: empty quotient") + } + + nPlus2 := int(s.domain0.Cardinality) + 2 + required := 3 * nPlus2 + if quotient.coeffs.Len() < required { + return fmt.Errorf("commitToQuotientGPUFromDevice: quotient too small: got %d need >= %d", quotient.coeffs.Len(), required) + } + + h1Device := ("ient.coeffs).Range(0, nPlus2, false) + h2Device := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3Device := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + if s.opt.StatisticalZK { + nPlus3 := nPlus2 + 1 + dH1, dH2, dH3, err := s.prepareStatisticalZKQuotientShards(h1Device, h2Device, h3Device, nPlus2) + if err != nil { + return err + } + defer s.putTempDeviceSlice(dH1, nPlus3) + defer s.putTempDeviceSlice(dH2, nPlus3) + defer s.putTempDeviceSlice(dH3, nPlus2) + + c0, err := commitOnGPUCanonicalDevice(dH1, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(dH2, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(dH3, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + return nil + } + + // Commit sequentially to avoid 3-way concurrent MSM memory spikes. + c0, err := commitOnGPUCanonicalDevice(h1Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(h2Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(h3Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + + return nil +} + +func (s *instance) inverseAndMergeShards( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (icicle_core.DeviceSlice, error) { + if gpuNumerator == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: nil numerator") + } + n := gpuNumerator.n + rho := gpuNumerator.rho + if n <= 0 || rho <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: invalid n=%d rho=%d", n, rho) + } + if len(gpuNumerator.shards) != rho { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), rho) + } + for i := 0; i < rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard %d is empty", i) + } + } + + expo := big.NewInt(int64(n)) + + // Per-shard cosets: c_i = c * g^i where c=FrMultiplicativeGen, g=Generator. + cosets := make([]fr.Element, rho) + cosets[0].Set(&domains[1].FrMultiplicativeGen) + for i := 1; i < rho; i++ { + cosets[i].Mul(&cosets[i-1], &domains[1].Generator) + } + invCosets := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + invCosets[i].Inverse(&cosets[i]) + } + + // ν = g^n is a rho-th root, used for the rho-point inverse DFT in combine. + var nu, nuInv fr.Element + nu.Exp(domains[1].Generator, expo) + nuInv.Inverse(&nu) + nuInvPowers := make([]fr.Element, rho) + nuInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + nuInvPowers[i].Mul(&nuInvPowers[i-1], &nuInv) + } + + // cN = c^n. Recover original coefficient blocks by scaling with cN^{-t}. + var cN, cNInv fr.Element + cN.Exp(domains[1].FrMultiplicativeGen, expo) + cNInv.Inverse(&cN) + cNInvPowers := make([]fr.Element, rho) + cNInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + cNInvPowers[i].Mul(&cNInvPowers[i-1], &cNInv) + } + + // Each shard inverse contributes a 1/n factor; apply extra 1/rho. + var rhoFr, invRho fr.Element + rhoFr.SetUint64(uint64(rho)) + invRho.Inverse(&rhoFr) + combineScales := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + combineScales[i].Mul(&invRho, &cNInvPowers[i]) + } + + totalSize := rho * n + done := make(chan error, 1) + var dMerged icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish := makeFinisher(stream, "inverseAndMergeShards", done) + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + // KNR is often faster than KNN; we restore regular output explicitly + // by bit-reversing each shard after the inverse NTT. + cfgNtt.Ordering = icicle_core.KNR + + ext := config_extension.Create() + defer config_extension.Delete(ext) + alg := nttAlgorithmFromEnv("ICICLE_DIVIDE_BY_ZH_NTT_ALGO", icicle_core.MixedRadix) + ext.SetInt(icicle_core.CUDA_NTT_ALGORITHM, int(alg)) + cfgNtt.Ext = ext.AsUnsafePointer() + + // Step 1: inverse NTT each shard without coset, reorder to regular, + // then unscale by (c*g^i)^t to recover the coset-inverse equivalent. + nn := uint64(64 - bits.TrailingZeros64(uint64(n))) + invPowers := make([]fr.Element, n) + for i := 0; i < rho; i++ { + if nttErr := icicle_ntt.Ntt(gpuNumerator.shards[i], icicle_core.KInverse, &cfgNtt, gpuNumerator.shards[i]); nttErr != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: inverse NTT failed at shard %d: %s", i, nttErr.AsString())) + return + } + + // KNR outputs bit-reversed coefficients. Reorder back to regular. + dRegular := s.getTempDeviceSlice(n) + mergeErr := icicle_vecops.MergeShardsBitReverse( + []icicle_core.DeviceSlice{gpuNumerator.shards[i]}, + n, + nn, + dRegular, + cfgVec, + ) + if mergeErr != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: reorder failed at shard %d: %s", i, mergeErr.AsString())) + return + } + // gpuNumerator.shards[i] is returned to pool and replaced; wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(gpuNumerator.shards[i], n) + gpuNumerator.shards[i] = dRegular + + fft.BuildExpTable(invCosets[i], invPowers) + dInvPowers := s.getTempDeviceSlice(n) + uploadVectorStdIntoOnCurrentDevice(&dInvPowers, invPowers, cfgVec) + if e := icicle_vecops.VecOp(gpuNumerator.shards[i], dInvPowers, gpuNumerator.shards[i], cfgVec, icicle_core.Mul); e != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: normalize shard %d failed: %s", i, e.AsString())) + return + } + // dInvPowers is temporary and returned to pool each iteration. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(dInvPowers, n) + } + + // Step 2: combine shard results via a size-rho inverse DFT per coefficient index. + dMerged = s.getTempDeviceSlice(totalSize) + keepMerged := false + defer func() { + if !keepMerged && !dMerged.IsEmpty() { + s.putTempDeviceSlice(dMerged, totalSize) + } + }() + + dTmp := s.getTempDeviceSlice(n) + defer func() { + if !dTmp.IsEmpty() { + s.putTempDeviceSlice(dTmp, n) + } + }() + + dNuWeights := make([]icicle_core.DeviceSlice, rho) + dCombineScales := make([]icicle_core.DeviceSlice, rho) + for i := 0; i < rho; i++ { + dNuWeights[i] = uploadScalarStdOnCurrentDevice(nuInvPowers[i], cfgVec) + dCombineScales[i] = uploadScalarStdOnCurrentDevice(combineScales[i], cfgVec) + } + defer func() { + for i := 0; i < rho; i++ { + if !dNuWeights[i].IsEmpty() { + _ = dNuWeights[i].Free() + } + if !dCombineScales[i].IsEmpty() { + _ = dCombineScales[i].Free() + } + } + }() + + for t := 0; t < rho; t++ { + outT := (&dMerged).Range(t*n, (t+1)*n, false) + if e := copyDeviceSliceIntoOnCurrentDevice(outT, gpuNumerator.shards[0], cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: init out[%d] failed: %s", t, e.AsString())) + return + } + for i := 1; i < rho; i++ { + weightIdx := (i * t) % rho + if weightIdx == 0 { + if e := icicle_vecops.VecOp(outT, gpuNumerator.shards[i], outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: add shard %d to out[%d] failed: %s", i, t, e.AsString())) + return + } + continue + } + if e := icicle_vecops.ScalarMulVec(dNuWeights[weightIdx], gpuNumerator.shards[i], dTmp, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: weight shard %d for out[%d] failed: %s", i, t, e.AsString())) + return + } + if e := icicle_vecops.VecOp(outT, dTmp, outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: accumulate shard %d into out[%d] failed: %s", i, t, e.AsString())) + return + } + } + if e := icicle_vecops.ScalarMulVec(dCombineScales[t], outT, outT, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: scale out[%d] failed: %s", t, e.AsString())) + return + } + } + keepMerged = true + finish(nil) + }) + + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dMerged, nil +} + +func (s *instance) divideByZHOnGPU( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (_ *gpuQuotientPolynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("divideByZHOnGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("divideByZHOnGPU: invalid numerator dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + for i := range gpuNumerator.shards { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("divideByZHOnGPU: shard %d is empty", i) + } + } + + rho := int(domains[1].Cardinality / domains[0].Cardinality) + if rho != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: rho mismatch domains=%d numerator=%d", rho, gpuNumerator.rho) + } + + // Evaluate 1/(X^n-1) over the large-domain coset values used by this quotient. + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + + // In bit-reversed merged layout, each shard maps to a fixed (iRev % rho) bucket. + // So we can divide by Z_H by scaling each shard with its corresponding inverse. + scaleDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + for i := 0; i < gpuNumerator.rho; i++ { + dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) + vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if vecErr != icicle_runtime.Success { + finish(fmt.Errorf("divideByZHOnGPU: shard scaling failed at %d: %s", i, vecErr.AsString())) + return + } + } + finish(nil) + }) + if err := <-scaleDone; err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, err + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + dMerged, splitErr := s.inverseAndMergeShards(gpuNumerator, domains) + if splitErr != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, splitErr + } + // Shards are not needed after split inverse+merge. + s.freeNumeratorShards(gpuNumerator.shards) + return &gpuQuotientPolynomial{coeffs: dMerged, size: totalSize}, nil +} + +func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*iop.Polynomial, error) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return nil, fmt.Errorf("downloadQuotientFromGPU: empty quotient") + } + if quotient.size <= 0 { + return nil, fmt.Errorf("downloadQuotientFromGPU: invalid quotient size %d", quotient.size) + } + + coeffs := make([]fr.Element, quotient.size) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) + // Async boundary for host materialization of quotient coefficients. + done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + }) + if err := <-done; err != nil { + return nil, err + } + + return iop.NewPolynomial(&coeffs, iop.Form{Basis: iop.Canonical, Layout: iop.Regular}), nil +} + +func commitOnGPUWithDeviceBases( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, +) (curve.G1Affine, error) { + return commitOnGPUWithDeviceBasesChunked(scalarsDevice, basesDevice, device, icicleMSMChunkSize()) +} + +func commitOnGPUWithDeviceBasesChunked( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, + chunkSize int, +) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty scalar slice") + } + if basesDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty basis slice") + } + if scalarsDevice.Len() > basesDevice.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: invalid scalar size %d", scalarsDevice.Len()) + } + if chunkSize <= 0 || chunkSize > scalarsDevice.Len() { + chunkSize = scalarsDevice.Len() + } + + var commit curve.G1Affine + var msmErr error + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) + close(done) + }) + <-done + if msmErr != nil { + return curve.G1Affine{}, fmt.Errorf("icicle: MSM commit from device bases failed (%d scalars): %w", scalarsDevice.Len(), msmErr) + } + return commit, nil +} + +func commitOnGPUWithDeviceBasesChunkedOnCurrentDevice( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + chunkSize int, +) (curve.G1Affine, error) { + var commit curve.G1Affine + for start := 0; start < scalarsDevice.Len(); start += chunkSize { + end := start + chunkSize + if end > scalarsDevice.Len() { + end = scalarsDevice.Len() + } + + // Each chunk must pair with exactly bases[start:end]: ICICLE treats a + // bases slice longer than the scalars as a batched MSM (and requires + // divisibility), so the full bases buffer cannot be passed as-is when + // it is longer than the scalar vector. + scalarsChunk := scalarsDevice + if start != 0 || end != scalarsDevice.Len() { + scalarsChunk = (&scalarsDevice).Range(start, end, false) + } + basesChunk := basesDevice + if start != 0 || end != basesDevice.Len() { + basesChunk = (&basesDevice).Range(start, end, false) + } + + res := make(icicle_core.HostSlice[icicle_bls12381.Projective], 1) + cfg := icicle_msm.GetDefaultMSMConfig() + cfg.AreBasesMontgomeryForm = true + cfg.AreScalarsMontgomeryForm = true + e := icicle_msm.Msm(scalarsChunk, basesChunk, &cfg, res) + if e != icicle_runtime.Success { + return curve.G1Affine{}, fmt.Errorf("icicle MSM failed for chunk [%d:%d]: %s", start, end, e.AsString()) + } + + chunkCommit, err := projectiveToGnarkAffine(res[0]) + if err != nil { + return curve.G1Affine{}, fmt.Errorf("convert chunk [%d:%d]: %w", start, end, err) + } + commit.Add(&commit, &chunkCommit) + } + return commit, nil +} + +func icicleMSMChunkSize() int { + // Production-sized MSMs still need chunking, but tiny chunks add thousands of + // ICICLE calls. 4M-point chunks passed the gnark replay profile; 8M did not. + const defaultChunkSize = 1 << 22 + v := strings.TrimSpace(os.Getenv("ICICLE_MSM_CHUNK_SIZE")) + if v == "" { + return defaultChunkSize + } + n, err := strconv.Atoi(v) + if err != nil || n <= 0 { + return defaultChunkSize + } + return n +} + +func commitOnGPUCanonicalDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: canonical SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgDevice.G1, device) +} + +func evalCanonicalAtPoint(coeffs []fr.Element, point fr.Element) fr.Element { + var acc fr.Element + if len(coeffs) == 0 { + return acc + } + acc.Set(&coeffs[len(coeffs)-1]) + for i := len(coeffs) - 2; i >= 0; i-- { + acc.Mul(&acc, &point).Add(&acc, &coeffs[i]) + } + return acc +} + +func deriveBatchOpeningGamma( + point fr.Element, + digests []curve.G1Affine, + claimedValues []fr.Element, + hf hash.Hash, + dataTranscript ...[]byte, +) (fr.Element, error) { + fs := fiatshamir.NewTranscript(hf, "gamma") + if err := fs.Bind("gamma", point.Marshal()); err != nil { + return fr.Element{}, err + } + for i := range digests { + if err := fs.Bind("gamma", digests[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := range claimedValues { + if err := fs.Bind("gamma", claimedValues[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := 0; i < len(dataTranscript); i++ { + if err := fs.Bind("gamma", dataTranscript[i]); err != nil { + return fr.Element{}, err + } + } + gammaByte, err := fs.ComputeChallenge("gamma") + if err != nil { + return fr.Element{}, err + } + var gamma fr.Element + gamma.SetBytes(gammaByte) + return gamma, nil +} + +func (s *instance) evalDevicePolynomialAtPointOnCurrentDevice( + coeffsDevice icicle_core.DeviceSlice, + point fr.Element, + useBitReverse bool, + cfg icicle_core.VecOpsConfig, +) (fr.Element, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: empty coefficients") + } + + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(point, cfg) + defer dPoint.Free() + + dOut := s.getTempDeviceSlice(1) + defer s.putTempDeviceSlice(dOut, 1) + + opName := "PolyEvalAt" + var eEval icicle_runtime.EIcicleError + if useBitReverse { + opName = "PolyEvalAtBitReverse" + mm := uint64(64 - bits.TrailingZeros64(uint64(coeffsDevice.Len()))) + eEval = icicle_vecops.PolyEvalAtBitReverse(coeffsDevice, dPoint, mm, dOut, cfg) + } else { + eEval = icicle_vecops.PolyEvalAt(coeffsDevice, dPoint, dOut, cfg) + } + if eEval != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: %s failed: %s", opName, eEval.AsString()) + } + + var out fr.Element + hostOut := icicle_core.HostSliceFromElements([]fr.Element{out}) + if cfg.IsAsync { + hostOut.CopyFromDeviceAsync(&dOut, cfg.StreamHandle) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } else { + hostOut.CopyFromDevice(&dOut) + } + return ([]fr.Element)(hostOut)[0], nil +} + +func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (fr.Element, error) { + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + var runErr error + out, runErr = s.evalDevicePolynomialAtPointOnCurrentDevice(coeffsDevice, point, false, cfg) + done <- runErr + }) + return out, <-done +} + +func (s *instance) copyDeviceSliceOnCurrentDevice( + src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, + label string, +) (icicle_core.DeviceSlice, error) { + if src.IsEmpty() || src.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: empty source slice", label) + } + dst := s.getTempDeviceSlice(src.Len()) + eCopy := copyDeviceSliceIntoOnCurrentDevice(dst, src, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dst, src.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: copy failed: %s", label, eCopy.AsString()) + } + return dst, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromStateOnCurrentDevice( + p *iop.Polynomial, + state *gpuPolysState, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "batchOpening poly") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + dCanon, err := s.copyDeviceSliceOnCurrentDevice(dSrc, cfg, "materializePolynomialCanonicalRegularFromStateOnCurrentDevice") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + // Canonical bit-reverse -> canonical regular. + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + dRegular := s.getTempDeviceSlice(dCanon.Len()) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNN // regular -> regular canonical + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular canonical + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) buildBlindedCanonicalPolynomialOnCurrentDevice( + dBaseCanon, dBlindCanon icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBlindCanon.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: empty input") + } + n := dBaseCanon.Len() + blindLen := dBlindCanon.Len() + dOut := s.getTempDeviceSlice(n + blindLen) + + dPrefix := (&dOut).Range(0, n, false) + eCopyBase := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopyBase != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopyBase.AsString()) + } + + dTail := (&dOut).Range(n, n+blindLen, false) + eCopyBlind := copyDeviceSliceIntoOnCurrentDevice(dTail, dBlindCanon, cfg) + if eCopyBlind != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy tail failed: %s", eCopyBlind.AsString()) + } + + dHead := (&dOut).Range(0, blindLen, false) + if eSub := icicle_vecops.VecOp(dHead, dBlindCanon, dHead, cfg, icicle_core.Sub); eSub != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: subtract blind from head failed: %s", eSub.AsString()) + } + return dOut, nil +} + +func (s *instance) buildPaddedCanonicalPolynomialOnCurrentDevice( + dBaseCanon icicle_core.DeviceSlice, + padLen int, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBaseCanon.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: empty base") + } + if padLen < 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: negative pad length %d", padLen) + } + n := dBaseCanon.Len() + dOut := s.getTempDeviceSlice(n + padLen) + dPrefix := (&dOut).Range(0, n, false) + + eCopy := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopy.AsString()) + } + if padLen == 0 { + return dOut, nil + } + + dTail := (&dOut).Range(n, n+padLen, false) + eZero := zeroDeviceSliceOnCurrentDevice(dTail, cfg) + if eZero != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: zero tail failed: %s", eZero.AsString()) + } + return dOut, nil +} + +func (s *instance) prepareBatchOpeningPolynomialsOnGPU( + state *gpuPolysState, + point fr.Element, +) (devicePolys []icicle_core.DeviceSlice, owned []bool, claimed []fr.Element, err error) { + if state == nil { + return nil, nil, nil, fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: nil GPU state") + } + + total := 6 + len(s.trace.Qcp) + devicePolys = make([]icicle_core.DeviceSlice, total) + owned = make([]bool, total) + claimed = make([]fr.Element, total) + devicePolys[0] = s.linearizedPolynomialGPU + claimed[0] = s.linearizedPolynomialClaim + + prepDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + + cleanupOwned := func(from int) { + for i := from; i < len(devicePolys); i++ { + if owned[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + owned[i] = false + } + } + } + + prepareLRORow := func(dstIdx int, p, bp *iop.Polynomial, padLen int) error { + dBase, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + defer s.putTempDeviceSlice(dBase, dBase.Len()) + + var dFinal icicle_core.DeviceSlice + if useBlinding { + if bp == nil { + return fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: missing blinding polynomial") + } + dBlind, eBlind := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(bp, state, cfg) + if eBlind != nil { + return eBlind + } + defer s.putTempDeviceSlice(dBlind, dBlind.Len()) + dFinal, e = s.buildBlindedCanonicalPolynomialOnCurrentDevice(dBase, dBlind, cfg) + } else { + dFinal, e = s.buildPaddedCanonicalPolynomialOnCurrentDevice(dBase, padLen, cfg) + } + if e != nil { + return e + } + devicePolys[dstIdx] = dFinal + owned[dstIdx] = true + + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dFinal, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareLRORow(1, s.polyL, s.bp[id_Bl], order_blinding_L+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare L failed: %w", e)) + return + } + if e := prepareLRORow(2, s.polyR, s.bp[id_Br], order_blinding_R+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare R failed: %w", e)) + return + } + if e := prepareLRORow(3, s.polyO, s.bp[id_Bo], order_blinding_O+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare O failed: %w", e)) + return + } + + prepareDirect := func(dstIdx int, p *iop.Polynomial, label string) error { + dPoly, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + devicePolys[dstIdx] = dPoly + owned[dstIdx] = true + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dPoly, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareDirect(4, s.trace.S1, "S1"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S1 failed: %w", e)) + return + } + if e := prepareDirect(5, s.trace.S2, "S2"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S2 failed: %w", e)) + return + } + + for i := 0; i < len(s.trace.Qcp); i++ { + idx := 6 + i + if e := prepareDirect(idx, s.trace.Qcp[i], "Qcp"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare Qcp[%d] failed: %w", i, e)) + return + } + } + + finish(nil) + }) + if err := <-prepDone; err != nil { + return nil, nil, nil, err + } + return devicePolys, owned, claimed, nil +} + +type evalPolynomialInputPreparationResult struct { + dEval icicle_core.DeviceSlice + ownedLen int + useBitReverseEval bool +} + +func (s *instance) prepareEvalPolynomialInputOnCurrentDevice( + p *iop.Polynomial, + dSrc icicle_core.DeviceSlice, + cfgVec icicle_core.VecOpsConfig, +) (evalPolynomialInputPreparationResult, error) { + if p == nil { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: nil polynomial") + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: empty source polynomial") + } + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] prepareEvalInput begin n=%d basis=%v layout=%v step_profile=%q ntt_trace=%q ntt_profile_full=%q ntt_profile_arbitrary=%q\n", + dSrc.Len(), + p.Basis, + p.Layout, + os.Getenv("ICICLE_STEP_PROFILE"), + os.Getenv("ICICLE_NTT_TRACE"), + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + + result := evalPolynomialInputPreparationResult{ + dEval: dSrc, + } + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return result, nil + } + + releaseOwned := func() { + if result.ownedLen > 0 && !result.dEval.IsEmpty() { + s.putTempDeviceSlice(result.dEval, result.ownedLen) + result.dEval = icicle_core.DeviceSlice{} + result.ownedLen = 0 + } + } + + dWork := s.getTempDeviceSlice(dSrc.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dWork, dSrc, cfgVec); e != icicle_runtime.Success { + s.putTempDeviceSlice(dWork, dSrc.Len()) + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: copy source polynomial failed: %s", e.AsString()) + } + + result.dEval = dWork + result.ownedLen = dSrc.Len() + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfgVec.IsAsync + cfgNtt.StreamHandle = cfgVec.StreamHandle + + var ownsNttStream bool + destroyOwnedNttStream := func() error { + if !ownsNttStream { + return nil + } + return syncAndDestroyStreamOnCurrentDevice(cfgNtt.StreamHandle, "prepareEvalPolynomialInputOnCurrentDevice") + } + + switch p.Basis { + case iop.Canonical: + // No transform required. + result.useBitReverseEval = p.Layout == iop.BitReverse + case iop.Lagrange, iop.LagrangeCoset: + if cfgNtt.StreamHandle == nil { + nttStream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: create NTT stream failed: %s", eStream.AsString()) + } + cfgNtt.StreamHandle = nttStream + cfgNtt.IsAsync = true + ownsNttStream = true + } + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNR // regular -> bitreverse on inverse + result.useBitReverseEval = true + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular on inverse + result.useBitReverseEval = false + } + startNtt := time.Now() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT launch n=%d ordering=%v has_coset=%t basis=%v layout=%v\n", + dSrc.Len(), + cfgNtt.Ordering, + p.Basis == iop.LagrangeCoset, + p.Basis, + p.Layout, + ) + } + + // Martun: This call to Ntt takes about 3 seconds, because Ntt is reusing NTT domain data that + // gets prepared during InitDomain (once per device), not re-derived every call. + // Inside ICICLE, InitDomain precomputes: domain.twiddles (main roots-of-unity table, N+1) + // internal_twiddles and basic_twiddles for mixed-radix kernels + // if fast mode is on (it is by default here), extra forward+inverse fast twiddle tables (fast_external/internal/basic and _inv) — comment says this costs ~4N extra memory + // CPU-side coset_index map (root -> index), then reused by later Ntt calls + eNtt := icicle_ntt.Ntt(result.dEval, icicle_core.KInverse, &cfgNtt, result.dEval) + nttElapsed := time.Since(startNtt) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT done status=%s took=%s\n", + eNtt.AsString(), + nttElapsed, + ) + } + if eNtt != icicle_runtime.Success { + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + l := logger.Logger() + l.Warn().Err(errDestroy).Msg("prepareEvalPolynomialInputOnCurrentDevice") + } + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + // Async boundary for this helper when it owns the stream. + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + releaseOwned() + return evalPolynomialInputPreparationResult{}, errDestroy + } + default: + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: unsupported basis %v", p.Basis) + } + + return result, nil +} + +// evalPolynomialInCurrentFormOnGPU evaluates a polynomial at a point directly +// from the shared GPU state regardless of its current basis/layout by converting +// a temporary device copy to canonical/regular when needed. +func (s *instance) evalPolynomialInCurrentFormOnGPU( + p *iop.Polynomial, + state *gpuPolysState, + point fr.Element, +) (fr.Element, error) { + if p == nil { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "eval") + if err != nil { + return fr.Element{}, err + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: empty device polynomial") + } + + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + dEval := dSrc + ownedLen := 0 + releaseEval := func() { + if ownedLen > 0 && !dEval.IsEmpty() { + s.putTempDeviceSlice(dEval, ownedLen) + dEval = icicle_core.DeviceSlice{} + ownedLen = 0 + } + } + finish := func(runErr error) { + // Async boundary for eval path before handing result back to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + releaseEval() + done <- runErr + } + + prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) + dEval = prepareResult.dEval + ownedLen = prepareResult.ownedLen + if prepErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", prepErr)) + return + } + + evalOut, evalErr := s.evalDevicePolynomialAtPointOnCurrentDevice(dEval, point, prepareResult.useBitReverseEval, cfgVec) + if evalErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", evalErr)) + return + } + out = evalOut + finish(nil) + }) + return out, <-done +} + +func (s *instance) evaluateBlindedOnGPU( + p, bp *iop.Polynomial, + state *gpuPolysState, + zeta fr.Element, +) (fr.Element, error) { + if p == nil || bp == nil { + return fr.Element{}, fmt.Errorf("evaluateBlindedOnGPU: nil polynomial") + } + pAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(p, state, zeta) + if err != nil { + return fr.Element{}, err + } + bpAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(bp, state, zeta) + if err != nil { + return fr.Element{}, err + } + + var t, one fr.Element + one.SetOne() + t.Exp(zeta, big.NewInt(int64(p.Size()))).Sub(&t, &one) + bpAtZeta.Mul(&bpAtZeta, &t) + pAtZeta.Add(&pAtZeta, &bpAtZeta) + return pAtZeta, nil +} + +func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.Element) (kzg.OpeningProof, error) { + if len(coeffs) < 2 || len(coeffs) > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonical: invalid polynomial size %d", len(coeffs)) + } + claimed := evalCanonicalAtPoint(coeffs, point) + + var dWitness icicle_core.DeviceSlice + witnessSize := len(coeffs) - 1 + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dCoeffs := uploadVector(coeffs) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dCoeffs, dPoint, dWitness, cfg) + _ = dCoeffs.Free() + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonical: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + if err := <-divDone; err != nil { + return kzg.OpeningProof{}, err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (kzg.OpeningProof, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() < 2 || coeffsDevice.Len() > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonicalDevice: invalid polynomial size %d", coeffsDevice.Len()) + } + n := coeffsDevice.Len() + + var startEval time.Time + if isProfileMode { + startEval = time.Now() + } + claimed, err := s.evalDevicePolynomialAtPoint(coeffsDevice, point) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startEval)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint") + } + } + if err != nil { + return kzg.OpeningProof{}, err + } + + var dWitness icicle_core.DeviceSlice + witnessSize := coeffsDevice.Len() - 1 + var startDivideByXMinusA time.Time + if isProfileMode { + startDivideByXMinusA = time.Now() + } + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(coeffsDevice, dPoint, dWitness, cfg) + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonicalDevice: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + divideErr := <-divDone + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startDivideByXMinusA)) + if divideErr != nil { + ev.Err(divideErr).Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA") + } + } + if divideErr != nil { + return kzg.OpeningProof{}, divideErr + } + + var startCommit time.Time + if isProfileMode { + startCommit = time.Now() + } + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startCommit)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice") + } + } + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) linearizedZContributionScale(lZeta, rZeta, oZeta fr.Element) fr.Element { + var s2, tmp fr.Element + var uzeta, uuzeta fr.Element + uzeta.Mul(&s.zeta, &s.pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &s.pk.Vk.CosetShift) + + s2.Mul(&s.beta, &s.zeta).Add(&s2, &lZeta).Add(&s2, &s.gamma) + tmp.Mul(&s.beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp) + tmp.Mul(&s.beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp).Neg(&s2).Mul(&s2, &s.alpha) + + var one, alphaSquareLagrangeZero, den fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&s.zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) + den.Sub(&s.zeta, &one).Inverse(&den) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) + + s2.Add(&s2, &alphaSquareLagrangeZero) + return s2 +} + +func (s *instance) linearizedSelectorScales(evals witnessEvalAtZeta, zu fr.Element) linearizedSelectorScales { + var scales linearizedSelectorScales + + // S3 scale: + // alpha * beta * Z(mu*zeta) * + // (L(zeta) + beta*S1(zeta) + gamma) * + // (R(zeta) + beta*S2(zeta) + gamma) + var tmp fr.Element + scales.s3.Mul(&evals.s1zeta, &s.beta).Add(&scales.s3, &evals.blzeta).Add(&scales.s3, &s.gamma) + tmp.Mul(&evals.s2zeta, &s.beta).Add(&tmp, &evals.brzeta).Add(&tmp, &s.gamma) + scales.s3.Mul(&scales.s3, &tmp).Mul(&scales.s3, &zu).Mul(&scales.s3, &s.beta).Mul(&scales.s3, &s.alpha) + + scales.ql.Set(&evals.blzeta) + scales.qr.Set(&evals.brzeta) + scales.qm.Mul(&evals.brzeta, &evals.blzeta) + scales.qo.Set(&evals.bozeta) + scales.qk.SetOne() + scales.qcp = append(scales.qcp, evals.qcpzeta...) + + return scales +} + +func (s *instance) buildLinearizedSelectorTermsOnGPU( + evals witnessEvalAtZeta, + zu fr.Element, + linearizedLen int, +) (icicle_core.DeviceSlice, error) { + if linearizedLen <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: invalid length %d", linearizedLen) + } + if len(evals.qcpzeta) > len(s.cCommitments) { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: qcp/cCommitments mismatch (%d > %d)", len(evals.qcpzeta), len(s.cCommitments)) + } + + scales := s.linearizedSelectorScales(evals, zu) + + var dLinearized icicle_core.DeviceSlice + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildLinearizedSelectorTermsOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if runErr != nil && !dLinearized.IsEmpty() { + s.putTempDeviceSlice(dLinearized, dLinearized.Len()) + dLinearized = icicle_core.DeviceSlice{} + } + done <- runErr + }() + + dLinearized = s.getTempDeviceSlice(linearizedLen) + if eZero := zeroDeviceSliceOnCurrentDevice(dLinearized, cfg); eZero != icicle_runtime.Success { + runErr = fmt.Errorf("buildLinearizedSelectorTermsOnGPU: zero output failed: %s", eZero.AsString()) + return + } + + addTerm := func(p *iop.Polynomial, scale fr.Element, label string) error { + if p == nil { + return fmt.Errorf("missing polynomial %s", label) + } + if scale.IsZero() { + return nil + } + + start := time.Now() + dPoly, err := s.materializePolynomialCanonicalRegularFromHostOnCurrentDevice(p, cfg) + if err != nil { + return fmt.Errorf("%s: %w", label, err) + } + defer s.putTempDeviceSlice(dPoly, dPoly.Len()) + if dPoly.Len() > dLinearized.Len() { + return fmt.Errorf("%s: polynomial too large (%d > %d)", label, dPoly.Len(), dLinearized.Len()) + } + + dScale := uploadScalarStdOnCurrentDevice(scale, cfg) + dScaled := s.getTempDeviceSlice(dPoly.Len()) + defer s.putTempDeviceSlice(dScaled, dScaled.Len()) + + eScale := icicle_vecops.ScalarMulVec(dScale, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if eScale != icicle_runtime.Success { + return fmt.Errorf("%s: scale failed: %s", label, eScale.AsString()) + } + + dPrefix := (&dLinearized).Range(0, dPoly.Len(), false) + if eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add); eAdd != icicle_runtime.Success { + return fmt.Errorf("%s: add failed: %s", label, eAdd.AsString()) + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Str("term", label).Int("n", dPoly.Len()).Dur("took", time.Since(start)).Msg("computeLinearizedPolynomial: add selector term on GPU") + } + return nil + } + + if runErr = addTerm(s.trace.S3, scales.s3, "S3"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Ql, scales.ql, "Ql"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qm, scales.qm, "Qm"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qr, scales.qr, "Qr"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qo, scales.qo, "Qo"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qk, scales.qk, "Qk"); runErr != nil { + return + } + for i := range scales.qcp { + if runErr = addTerm(s.cCommitments[i], scales.qcp[i], fmt.Sprintf("Qcp[%d]", i)); runErr != nil { + return + } + } + }) + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dLinearized, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromHostOnCurrentDevice( + p *iop.Polynomial, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: nil polynomial") + } + coeffs := p.Coefficients() + if len(coeffs) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: empty polynomial") + } + + dCanon := s.getTempDeviceSlice(len(coeffs)) + host := icicle_core.HostSliceFromElements(coeffs) + if cfg.IsAsync { + host.CopyToDeviceAsync(&dCanon, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&dCanon, false) + } + if dCanon.IsEmpty() { + s.putTempDeviceSlice(dCanon, len(coeffs)) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: host upload failed") + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + dRegular := s.getTempDeviceSlice(dCanon.Len()) + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + switch p.Layout { + case iop.Regular: + cfgNtt.Ordering = icicle_core.KNN + case iop.BitReverse: + cfgNtt.Ordering = icicle_core.KRN + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported layout %v", p.Layout) + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) addZContributionToLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + dBlindedZCanonical icicle_core.DeviceSlice, + lZeta, rZeta, oZeta fr.Element, +) error { + if dLinearized.IsEmpty() || dBlindedZCanonical.IsEmpty() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: empty input") + } + if dLinearized.Len() < dBlindedZCanonical.Len() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: linearized too small (%d < %d)", dLinearized.Len(), dBlindedZCanonical.Len()) + } + + zScale := s.linearizedZContributionScale(lZeta, rZeta, oZeta) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("addZContributionToLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dScale icicle_core.DeviceSlice + var dScaledZ icicle_core.DeviceSlice + defer func() { + // Async boundary before returning temporary buffers to the pool. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dScale) + if !dScaledZ.IsEmpty() { + s.putTempDeviceSlice(dScaledZ, dScaledZ.Len()) + } + done <- runErr + }() + + dScale = uploadScalarStdOnCurrentDevice(zScale, cfg) + dScaledZ = s.getTempDeviceSlice(dBlindedZCanonical.Len()) + eMul := icicle_vecops.ScalarMulVec(dScale, dBlindedZCanonical, dScaledZ, cfg) + if eMul != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: scale Z failed: %s", eMul.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, dBlindedZCanonical.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaledZ, dPrefix, cfg, icicle_core.Add) + if eAdd != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: add scaled Z failed: %s", eAdd.AsString()) + return + } + }) + return <-done +} + +func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + quotient *gpuQuotientPolynomial, +) error { + if dLinearized.IsEmpty() || quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: empty input") + } + nPlus2 := int(s.domain0.Cardinality) + 2 + if dLinearized.Len() < nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: linearized too small") + } + if quotient.coeffs.Len() < 3*nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: quotient too small") + } + + var one fr.Element + one.SetOne() + var zetaN, zetaNPlusTwo, zhZeta fr.Element + zetaN.Exp(s.zeta, big.NewInt(int64(s.domain0.Cardinality))) + zhZeta.Sub(&zetaN, &one) + zetaNPlusTwo.Mul(&zetaN, &s.zeta).Mul(&zetaNPlusTwo, &s.zeta) + + h1 := ("ient.coeffs).Range(0, nPlus2, false) + h2 := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3 := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("subtractQuotientContributionFromLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dAcc icicle_core.DeviceSlice + var dZetaStd icicle_core.DeviceSlice + var dZhStd icicle_core.DeviceSlice + defer func() { + // Async boundary before reusing temporary quotient vectors. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dZhStd) + freeDeviceSlice(&dZetaStd) + if !dAcc.IsEmpty() { + s.putTempDeviceSlice(dAcc, dAcc.Len()) + } + done <- runErr + }() + + dAcc = s.getTempDeviceSlice(nPlus2) + dZetaStd = uploadScalarStdOnCurrentDevice(zetaNPlusTwo, cfg) + eMulH3 := icicle_vecops.ScalarMulVec(dZetaStd, h3, dAcc, cfg) + if eMulH3 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale h3 failed: %s", eMulH3.AsString()) + return + } + if eAddH2 := icicle_vecops.VecOp(dAcc, h2, dAcc, cfg, icicle_core.Add); eAddH2 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h2 failed: %s", eAddH2.AsString()) + return + } + if eMulPow := icicle_vecops.ScalarMulVec(dZetaStd, dAcc, dAcc, cfg); eMulPow != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by zeta^(n+2) failed: %s", eMulPow.AsString()) + return + } + if eAddH1 := icicle_vecops.VecOp(dAcc, h1, dAcc, cfg, icicle_core.Add); eAddH1 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h1 failed: %s", eAddH1.AsString()) + return + } + + dZhStd = uploadScalarStdOnCurrentDevice(zhZeta, cfg) + eScaleZh := icicle_vecops.ScalarMulVec(dZhStd, dAcc, dAcc, cfg) + if eScaleZh != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by Z_H(zeta) failed: %s", eScaleZh.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, nPlus2, false) + eSub := icicle_vecops.VecOp(dPrefix, dAcc, dPrefix, cfg, icicle_core.Sub) + if eSub != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: subtract term failed: %s", eSub.AsString()) + return + } + }) + return <-done +} + +// divideByZH +// The input must be in LagrangeCoset. +// The result is in Canonical Regular. (in place using a) +func (s *instance) divideByZH(a *iop.Polynomial, domains [2]*fft.Domain) (*iop.Polynomial, error) { + + // check that the basis is LagrangeCoset + if a.Basis != iop.LagrangeCoset || a.Layout != iop.BitReverse { + return nil, errors.New("invalid form") + } + + // prepare the evaluations of x^n-1 on the big domain's coset + var startEvaluateXnMinusOne time.Time + if isProfileMode { + startEvaluateXnMinusOne = time.Now() + } + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startEvaluateXnMinusOne)).Msg("divideByZH: evaluateXnMinusOneDomainBigCoset") + } + rho := int(domains[1].Cardinality / domains[0].Cardinality) + + r := a.Coefficients() + n := uint64(len(r)) + nn := uint64(64 - bits.TrailingZeros64(n)) + + var startParallelizeMul time.Time + if isProfileMode { + startParallelizeMul = time.Now() + } + utils.Parallelize(len(r), func(start, end int) { + for i := start; i < end; i++ { + iRev := bits.Reverse64(uint64(i)) >> nn + r[i].Mul(&r[i], &xnMinusOneInverseLagrangeCoset[int(iRev)%rho]) + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallelizeMul)).Msg("divideByZH: parallelize multiply coefficients") + } + + // Replace CPU FFT inverse by ICICLE NTT inverse. + var startGpuNTTInverse time.Time + if isProfileMode { + startGpuNTTInverse = time.Now() + } + // It's faster on CPU. + // s.gpuNTTInverse(a) + a.ToCanonical(domains[1]).ToRegular() + if isProfileMode { + l := logger.Logger() + l.Debug(). + Int("size", a.Size()). + Dur("took", time.Since(startGpuNTTInverse)). + Msg("divideByZH: gpuNTTInverse on input of size n") + } + + return a, nil +} + +// evaluateXnMinusOneDomainBigCoset evaluates Xᵐ-1 on DomainBig coset +func evaluateXnMinusOneDomainBigCoset(domains [2]*fft.Domain) []fr.Element { + + rho := domains[1].Cardinality / domains[0].Cardinality + + res := make([]fr.Element, rho) + + expo := big.NewInt(int64(domains[0].Cardinality)) + res[0].Exp(domains[1].FrMultiplicativeGen, expo) + + var t fr.Element + t.Exp(domains[1].Generator, expo) + + one := fr.One() + + for i := 1; i < int(rho); i++ { + res[i].Mul(&res[i-1], &t) + res[i-1].Sub(&res[i-1], &one) + } + res[len(res)-1].Sub(&res[len(res)-1], &one) + + res = fr.BatchInvert(res) + + return res +} + +// innerComputeLinearizedPoly computes the linearized polynomial in canonical basis. +// The purpose is to commit and open all in one ql, qr, qm, qo, qk. +// * lZeta, rZeta, oZeta are the evaluation of l, r, o at zeta +// * z is the permutation polynomial, zu is Z(μX), the shifted version of Z +// * pk is the proving key: the linearized polynomial is a linear combination of ql, qr, qm, qo, qk. +// +// The Linearized polynomial is: +// +// α²*L₁(ζ)*Z(X) +// + α*( (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*(β*s3(X))*Z(μζ) - Z(X)*(l(ζ)+β*id1(ζ)+γ)*(r(ζ)+β*id2(ζ)+γ)*(o(ζ)+β*id3(ζ)+γ)) +// + l(ζ)*Ql(X) + l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) +// - Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) +// +// /!\ blindedZCanonical is modified +func (s *instance) innerComputeLinearizedPoly( + lZeta, rZeta, oZeta, s1Zeta, s2Zeta, + alpha, beta, gamma, zeta, zu fr.Element, + qcpZeta, blindedZCanonical []fr.Element, + pi2Canonical [][]fr.Element, + pk *ProvingKey, +) []fr.Element { + + // l(ζ)r(ζ) + var rl fr.Element + rl.Mul(&rZeta, &lZeta) + + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + // the linearised polynomial is + // α²*L₁(ζ)*Z(X) + + // s1*s3(X)+s2*Z(X) + l(ζ)*Ql(X) + + // l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) - + // Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + var s1, s2, tmp fr.Element + s1.Mul(&s1Zeta, &beta).Add(&s1, &lZeta).Add(&s1, &gamma) // (l(ζ)+β*s1(ζ)+γ) + tmp.Mul(&s2Zeta, &beta).Add(&tmp, &rZeta).Add(&tmp, &gamma) + s1.Mul(&s1, &tmp).Mul(&s1, &zu).Mul(&s1, &beta).Mul(&s1, &alpha) // (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*β*Z(μζ)*α + + var uzeta, uuzeta fr.Element + uzeta.Mul(&zeta, &pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &pk.Vk.CosetShift) + + s2.Mul(&beta, &zeta).Add(&s2, &lZeta).Add(&s2, &gamma) // (l(ζ)+β*ζ+γ) + tmp.Mul(&beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &gamma) // (r(ζ)+β*u*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ) + tmp.Mul(&beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &gamma) // (o(ζ)+β*u²*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + s2.Neg(&s2).Mul(&s2, &alpha) + + // Z_h(ζ), ζⁿ⁺², L₁(ζ)*α²*Z + var zhZeta, zetaNPlusTwo, alphaSquareLagrangeZero, one, den, frNbElmt fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) // ζⁿ + zetaNPlusTwo.Mul(&alphaSquareLagrangeZero, &zeta).Mul(&zetaNPlusTwo, &zeta) // ζⁿ⁺² + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) // ζⁿ - 1 + zhZeta.Set(&alphaSquareLagrangeZero) // Z_h(ζ) = ζⁿ - 1 + frNbElmt.SetUint64(uint64(nbElmt)) + den.Sub(&zeta, &one).Inverse(&den) // 1/(ζ-1) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). // L₁ = (ζⁿ - 1)/(ζ-1) + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) // α²*L₁(ζ) + + s3canonical := s.trace.S3.Coefficients() + // Qk is prepared in canonical/regular form by computeLinearizedPolynomial. + + // len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 when Statistical ZK is activated + // len(h1)=len(h2)=len(h3)=len(blindedZCanonical)-1 when Statistical ZK is deactivated + h1 := s.h1() + h2 := s.h2() + h3 := s.h3() + + // at this stage we have + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + startParallel := time.Now() + utils.Parallelize(len(blindedZCanonical), func(start, end int) { + + cql := s.trace.Ql.Coefficients() + cqr := s.trace.Qr.Coefficients() + cqm := s.trace.Qm.Coefficients() + cqo := s.trace.Qo.Coefficients() + cqk := s.trace.Qk.Coefficients() + + var t, t0, t1 fr.Element + + for i := start; i < end; i++ { + t.Mul(&blindedZCanonical[i], &s2) // -Z(X)*α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + if i < len(s3canonical) { + t0.Mul(&s3canonical[i], &s1) // α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ)*β*s3(X) + t.Add(&t, &t0) + } + if i < len(cqm) { + t1.Mul(&cqm[i], &rl) // l(ζ)r(ζ)*Qm(X) + t.Add(&t, &t1) // linPol += l(ζ)r(ζ)*Qm(X) + t0.Mul(&cql[i], &lZeta) // l(ζ)Q_l(X) + t.Add(&t, &t0) // linPol += l(ζ)*Ql(X) + t0.Mul(&cqr[i], &rZeta) //r(ζ)*Qr(X) + t.Add(&t, &t0) // linPol += r(ζ)*Qr(X) + t0.Mul(&cqo[i], &oZeta) // o(ζ)*Qo(X) + t.Add(&t, &t0) // linPol += o(ζ)*Qo(X) + t.Add(&t, &cqk[i]) // linPol += Qk(X) + for j := range qcpZeta { // linPol += ∑ᵢQcp_(ζ)Pi_(X) + t0.Mul(&pi2Canonical[j][i], &qcpZeta[j]) + t.Add(&t, &t0) + } + } + + t0.Mul(&blindedZCanonical[i], &alphaSquareLagrangeZero) // α²L₁(ζ)Z(X) + blindedZCanonical[i].Add(&t, &t0) // linPol += α²L₁(ζ)Z(X) + + // if statistical zeroknowledge is deactivated, len(h1)=len(h2)=len(h3)=len(blindedZ)-1. + // Else len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 + if i < len(h3) { + t.Mul(&h3[i], &zetaNPlusTwo). + Add(&t, &h2[i]). + Mul(&t, &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } else { + if s.opt.StatisticalZK { + t.Mul(&h2[i], &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } + } + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallel)).Msg("computeLinearizedPolynomial: inner parallel loop") + } + + return blindedZCanonical +} + +var errContextDone = errors.New("context done") + +// local copies of verification-time helpers used by prover transcript +func bindPublicData(fs *fiatshamir.Transcript, challenge string, vk *plonk_bls12381.VerifyingKey, publicInputs []fr.Element) error { + if err := fs.Bind(challenge, vk.S[0].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[1].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[2].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Ql.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qr.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qm.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qo.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qk.Marshal()); err != nil { + return err + } + for i := range vk.Qcp { + if err := fs.Bind(challenge, vk.Qcp[i].Marshal()); err != nil { + return err + } + } + for i := 0; i < len(publicInputs); i++ { + if err := fs.Bind(challenge, publicInputs[i].Marshal()); err != nil { + return err + } + } + return nil +} + +func deriveRandomness(fs *fiatshamir.Transcript, challenge string, points ...*curve.G1Affine) (fr.Element, error) { + var buf [curve.SizeOfG1AffineUncompressed]byte + var r fr.Element + for _, p := range points { + buf = p.RawBytes() + if err := fs.Bind(challenge, buf[:]); err != nil { + return r, err + } + } + b, err := fs.ComputeChallenge(challenge) + if err != nil { + return r, err + } + r.SetBytes(b) + return r, nil +} + +// -------------------- GPU helpers and device setup -------------------- + +func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { + if pk.deviceInfo != nil { + return nil + } + pk.deviceInfo = &deviceInfo{} + + // Initialize ICICLE NTT domain (root of unity) and store coset generator for coset NTTs. + // ICICLE InitDomain expects a primitive root of unity; for coset transforms we use 𝔽ᵣ* generator. + + var gen fr.Element + var err error + if pk.Vk.Size < 6 { + gen, err = fft.Generator(8 * pk.Vk.Size) + if err != nil { + return err + } + } else { + gen, err = fft.Generator(4 * pk.Vk.Size) + if err != nil { + return err + } + } + genBits := gen.Bits() + limbs := icicle_core.ConvertUint64ArrToUint32Arr(genBits[:]) + // Initialize ICICLE NTT domain with root of unity + var rouIcicle icicle_bls12381.ScalarField + rouIcicle.FromLimbs(limbs) + + // Store coset generator = generator of 𝔽ᵣ* (matches CPU ToLagrangeCoset) + { + cosetGen := fft.GeneratorFullMultiplicativeGroup() + cosetBits := cosetGen.Bits() + cosetLimbs := icicle_core.ConvertUint64ArrToUint32Arr(cosetBits[:]) + copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) + } + + chInitDomain := make(chan struct{}) + initDomainQueuedAt := time.Now() + icicle_runtime.RunOnDevice(device, func(args ...any) { + initDomainStartedAt := time.Now() + initCfg := icicle_core.GetDefaultNTTInitDomainConfig() + ext := config_extension.Create() + defer config_extension.Delete(ext) + fastTwiddles := envEnabled("ICICLE_NTT_FAST_TWIDDLES", true) + ext.SetBool(icicle_core.CUDA_NTT_FAST_TWIDDLES_MODE, fastTwiddles) + initCfg.Ext = ext.AsUnsafePointer() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain start vk_size=%d fast_twiddles=%t profile_full=%q profile_arbitrary=%q\n", + pk.Vk.Size, + fastTwiddles, + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + e := icicle_ntt.InitDomain(rouIcicle, initCfg) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain end status=%s call_took=%s\n", + e.AsString(), + time.Since(initDomainStartedAt), + ) + } + if e != icicle_runtime.Success { + panic("icicle: InitDomain failed") + } + close(chInitDomain) + }) + + <-chInitDomain + if isNttTrace { + fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) + } + + chLag := make(chan struct{}) + chCan := make(chan struct{}) + + if len(pk.KzgLagrange.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) + g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) + close(chLag) + }) + } else { + close(chLag) + } + + if len(pk.Kzg.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) + g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) + close(chCan) + }) + } else { + close(chCan) + } + + <-chLag + <-chCan + return nil +} + +func projectiveToGnarkAffine(p icicle_bls12381.Projective) (curve.G1Affine, error) { + x, err := icicleBaseFieldToGnarkFp(p.X) + if err != nil { + return curve.G1Affine{}, err + } + y, err := icicleBaseFieldToGnarkFp(p.Y) + if err != nil { + return curve.G1Affine{}, err + } + z, err := icicleBaseFieldToGnarkFp(p.Z) + if err != nil { + return curve.G1Affine{}, err + } + if z.IsZero() { + return curve.G1Affine{}, nil + } + + var zInv fp.Element + zInv.Inverse(&z) + x.Mul(&x, &zInv) + y.Mul(&y, &zInv) + return curve.G1Affine{X: x, Y: y}, nil +} + +func icicleBaseFieldToGnarkFp(v icicle_bls12381.BaseField) (fp.Element, error) { + bytes := v.ToBytesLittleEndian() + if len(bytes) != fp.Bytes { + return fp.Element{}, fmt.Errorf("invalid ICICLE base field byte length %d", len(bytes)) + } + var buf [fp.Bytes]byte + copy(buf[:], bytes) + return fp.LittleEndian.Element(&buf) +} + +func commitOnGPULagrangeDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgLagrangeDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: lagrange SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgLagrangeDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgLagrangeDevice.G1, device) +} + +func (s *instance) registerDevicePolynomialInSharedState(state *gpuPolysState, p *iop.Polynomial, dSlice icicle_core.DeviceSlice) error { + if state == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil shared state") + } + if p == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil polynomial") + } + if dSlice.IsEmpty() { + return fmt.Errorf("registerDevicePolynomialInSharedState: empty device slice") + } + + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int) + } + if idx, ok := state.polyToIdx[p]; ok { + if idx < 0 || idx >= len(state.deviceSlices) { + return fmt.Errorf("registerDevicePolynomialInSharedState: invalid index %d", idx) + } + state.deviceSlices[idx] = dSlice + state.hostSlices[idx] = nil + state.originalForm[idx] = iop.Form{Basis: p.Basis, Layout: p.Layout} + return nil + } + + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, dSlice) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + return nil +} + +func (s *instance) gpuInclusivePrefixProductOnCurrentDevice( + dVec icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + if dVec.IsEmpty() || dVec.Len() <= 1 { + return nil + } + + n := dVec.Len() + for step := 1; step < n; step <<= 1 { + src := (&dVec).Range(0, n-step, false) + dst := (&dVec).Range(step, n, false) + + tmpStd := s.getTempDeviceSlice(n - step) + if err := copyDeviceSliceIntoOnCurrentDevice(tmpStd, src, cfg); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: copy stage failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(tmpStd, cfg) + if err := icicle_vecops.VecOp(tmpStd, dst, dst, cfg, icicle_core.Mul); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: multiply stage failed: %s", err.AsString()) + } + if cfg.IsAsync { + // tmpStd is returned to pool each stage, so we must wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } + s.putTempDeviceSlice(tmpStd, n-step) + } + return nil +} + +// buildPermutationGatherIndices prepares the subset of permutation indices that +// are consumed by the copy-constraint ratio loop (only rows [0, n-1) per copy). +func buildPermutationGatherIndices(permutation []int64, nbPolynomials, n, supportLen int) ([]int64, error) { + if n <= 1 { + return nil, nil + } + total := nbPolynomials * (n - 1) + indices := make([]int64, total) + + var permBuildErr error + var permBuildErrOnce sync.Once + utils.Parallelize(total, func(start, end int) { + for k := start; k < end; k++ { + j := k / (n - 1) + i := k % (n - 1) + base := j * n + permIdx := permutation[base+i] + if permIdx < 0 || int(permIdx) >= supportLen { + jj, ii, bad := j, i, permIdx + permBuildErrOnce.Do(func() { + permBuildErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: permutation index out of range at (%d,%d): %d", jj, ii, bad) + }) + continue + } + indices[k] = permIdx + } + }) + if permBuildErr != nil { + return nil, permBuildErr + } + return indices, nil +} + +func (s *instance) prepareCopyConstraintSupportsOnCurrentDevice( + n, nbPolynomials int, + domain *fft.Domain, + permGatherIndices []int64, + cfg icicle_core.VecOpsConfig, +) (dSupportFlat, dPermFlat icicle_core.DeviceSlice, err error) { + defer func() { + if err == nil { + return + } + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + }() + + if len(permGatherIndices) == 0 { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: empty permutation gather indices") + return + } + + dOmegaStd := uploadScalarStdOnCurrentDevice(domain.Generator, cfg) + defer dOmegaStd.Free() + dShiftStd := uploadScalarStdOnCurrentDevice(domain.FrMultiplicativeGen, cfg) + defer dShiftStd.Free() + + dSupportFlat, err = allocDeviceUninitialized(nbPolynomials * n) + if err != nil { + return + } + if e := icicle_vecops.SupportIdentity(dOmegaStd, dShiftStd, n, nbPolynomials, dSupportFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: generate identity support on GPU failed: %s", e.AsString()) + return + } + toMontgomeryFormInPlaceWithCfg(dSupportFlat, cfg) + + dPermIndicesDevice := uploadInt64VectorOnCurrentDevice(permGatherIndices, cfg) + defer dPermIndicesDevice.Free() + + dPermFlat, err = allocDeviceUninitialized(len(permGatherIndices)) + if err != nil { + return + } + if e := icicle_vecops.GatherByIndices(dSupportFlat, dPermIndicesDevice, dPermFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: gather permutation support on GPU failed: %s", e.AsString()) + return + } + if cfg.IsAsync { + // Ensure temporary support/index slices are safe to free on return. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + return + } + } + + return +} + +func (s *instance) accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail icicle_core.DeviceSlice, + dBetaStd, dGammaMont icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + nMinusOne := dEntryTail.Len() + dScaled := s.getTempDeviceSlice(nMinusOne) + dTerm := s.getTempDeviceSlice(nMinusOne) + defer func() { + s.putTempDeviceSlice(dScaled, nMinusOne) + s.putTempDeviceSlice(dTerm, nMinusOne) + }() + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dID, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale identity support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add beta*id failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dNumTail, dNumTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply numerator term failed: %s", err.AsString()) + } + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dPerm, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale permutation support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add beta*sigma failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dDenTail, dDenTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply denominator term failed: %s", err.AsString()) + } + if cfg.IsAsync { + // Temp vectors are released at function exit, so ensure queued work is complete. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + } + } + + return nil +} + +// validateDeviceEntries checks that all entries are non-empty and have consistent length. +// Returns the common length n. +func validateDeviceEntries(entries []icicle_core.DeviceSlice, label string) (int, error) { + if len(entries) == 0 { + return 0, fmt.Errorf("%s: no entries", label) + } + n := entries[0].Len() + if n == 0 { + return 0, fmt.Errorf("%s: empty device entry 0", label) + } + for i := range entries { + if entries[i].IsEmpty() { + return 0, fmt.Errorf("%s: empty device entry %d", label, i) + } + if entries[i].Len() != n { + return 0, fmt.Errorf("%s: inconsistent device entry size at %d (%d != %d)", label, i, entries[i].Len(), n) + } + } + return n, nil +} + +// BuildRatioCopyConstraintIcicle builds the accumulating ratio polynomial to prove that +// [P₁ ∥ .. ∥ P_{n—1}] is invariant by the permutation \sigma. +// Namely it returns the polynomial Z whose evaluation on the j-th root of unity is +// Z(ω^j) = Π_{i 1 { + dNumTail := (&dNum).Range(1, n, false) + dDenTail := (&dDen).Range(1, n, false) + var supportErr error + dSupportFlat, dPermFlat, supportErr = s.prepareCopyConstraintSupportsOnCurrentDevice(n, nbPolynomials, domain, permGatherIndices, cfg) + if supportErr != nil { + runErr = supportErr + return + } + + dBetaStd := uploadScalarStdOnCurrentDevice(beta, cfg) + dGammaMont := uploadScalarMontOnCurrentDevice(gamma, cfg) + + for j := 0; j < nbPolynomials; j++ { + dEntryTail := (&entriesDevice[j]).Range(0, n-1, false) + baseID := j * n + dID := (&dSupportFlat).Range(baseID, baseID+n-1, false) + basePerm := j * (n - 1) + dPerm := (&dPermFlat).Range(basePerm, basePerm+(n-1), false) + if err := s.accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail, dBetaStd, dGammaMont, cfg, + ); err != nil { + runErr = err + return + } + } + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize before releasing copy-constraint workspace failed: %s", eSync.AsString()) + return + } + } + _ = dBetaStd.Free() + _ = dGammaMont.Free() + + // Support vectors and loop temps are only needed for term accumulation. + // Free them before prefix products and batch inversion, whose ICICLE + // kernels allocate additional full-domain workspace internally. + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + s.tempGPUMemPool.FreeAll() + } + + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dNum, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dDen, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + + if invErr := s.batchInvertOnCurrentDevice(dDen); invErr != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: GPU batch inversion failed: %s", invErr.AsString()) + return + } + + toStandardFormInPlace(dDen) + if err := icicle_vecops.VecOp(dDen, dNum, dNum, cfg, icicle_core.Mul); err != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: final numerator*denominatorInv multiplication failed: %s", err.AsString()) + return + } + + dResult = dNum + dNum = icicle_core.DeviceSlice{} // transfer ownership to dResult + }) + if err := <-buildDone; err != nil { + return nil, err + } + + hostMirror := make([]fr.Element, n) + if len(hostMirror) > 0 { + hostMirror[0].SetOne() + } + res := iop.NewPolynomial(&hostMirror, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if err := s.registerDevicePolynomialInSharedState(gpuState, res, dResult); err != nil { + freeSliceOnDevice(&dResult, &s.device) + return nil, err + } + + return res, nil +} diff --git a/backend/accelerated/icicle/plonk/bls12-381/provingkey.go b/backend/accelerated/icicle/plonk/bls12-381/provingkey.go new file mode 100644 index 0000000000..6eb8948d0c --- /dev/null +++ b/backend/accelerated/icicle/plonk/bls12-381/provingkey.go @@ -0,0 +1,100 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +//go:build icicle + +package bls12381 + +import ( + "sync" + "time" + + "github.com/consensys/gnark-crypto/ecc/bls12-381/fr" + "github.com/consensys/gnark-crypto/ecc/bls12-381/fr/fft" + plonk_bls12381 "github.com/consensys/gnark/backend/plonk/bls12-381" + cs "github.com/consensys/gnark/constraint/bls12-381" + "github.com/consensys/gnark/logger" + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" +) + +// deviceInfo holds device-resident buffers for GPU acceleration. +type deviceInfo struct { + CosetGenerator [fr.Limbs * 2]uint32 + KzgDevice struct { + G1 icicle_core.DeviceSlice + } + KzgLagrangeDevice struct { + G1 icicle_core.DeviceSlice + } +} + +// hostSetup holds host-side, witness-independent prover state derived from the +// constraint system: the FFT domains and the PLONK trace (selector + +// permutation polynomials). Building the trace walks every constraint (~3s at +// 23M constraints), so it is computed once per proving key and shared across +// proofs. Everything here is read-only during proving: the prover clones Qk +// before patching public inputs into it, and every basis conversion of a trace +// polynomial copies first (see canonicalRegularCoefficientsCopy). +type hostSetup struct { + sizeSystem uint64 + domain0 *fft.Domain + domain1 *fft.Domain + trace *plonk_bls12381.Trace +} + +// ProvingKey wraps the native PLONK proving key with device-resident state +// (KZG bases, NTT domains, cached trace) that is uploaded once and reused +// across Prove calls. +// +// Concurrency: Prove calls sharing the same ProvingKey must be serialized by +// the caller. The device state hangs off the key and proofs share a single +// GPU; concurrent proves against the same key are not safe. +type ProvingKey struct { + plonk_bls12381.ProvingKey + *deviceInfo + hostSetupOnce sync.Once + hostSetup *hostSetup +} + +func buildHostSetup(spr *cs.SparseR1CS, sizeSystem uint64) *hostSetup { + domain0 := fft.NewDomain(sizeSystem) + + // h, the quotient polynomial is of degree 3(n+1)+2, so it's in a 3(n+2) dim + // vector space, the domain is the next power of 2 superior to 3(n+2). + // 4*domainNum is enough in all cases except when n<6. + var domain1 *fft.Domain + if sizeSystem < 6 { + domain1 = fft.NewDomain(8*sizeSystem, fft.WithoutPrecompute()) + } else { + domain1 = fft.NewDomain(4*sizeSystem, fft.WithoutPrecompute()) + } + + return &hostSetup{ + sizeSystem: sizeSystem, + domain0: domain0, + domain1: domain1, + trace: plonk_bls12381.NewTrace(spr, domain0), + } +} + +// hostSetupFor returns the FFT domains and trace for spr, building them on +// first use and caching them on the proving key. A PLONK proving key is bound +// to exactly one constraint system, so per-key caching is sound; as a +// defensive measure a system-size mismatch falls back to an uncached build +// rather than ever serving another circuit's trace. +func (pk *ProvingKey) hostSetupFor(spr *cs.SparseR1CS) *hostSetup { + nbConstraints := spr.GetNbConstraints() + sizeSystem := uint64(nbConstraints + len(spr.Public)) // len(spr.Public) is for the placeholder constraints + pk.hostSetupOnce.Do(func() { + start := time.Now() + pk.hostSetup = buildHostSetup(spr, sizeSystem) + log := logger.Logger() + log.Debug().Dur("took", time.Since(start)).Msg("built prover host setup (fft domains + trace)") + }) + if pk.hostSetup.sizeSystem != sizeSystem { + return buildHostSetup(spr, sizeSystem) + } + return pk.hostSetup +} diff --git a/backend/accelerated/icicle/plonk/bn254/doc.go b/backend/accelerated/icicle/plonk/bn254/doc.go new file mode 100644 index 0000000000..d461dc2957 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bn254/doc.go @@ -0,0 +1,7 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +// Package bn254 implements ICICLE acceleration for BN254 PLONK backend. +package bn254 diff --git a/backend/accelerated/icicle/plonk/bn254/icicle.go b/backend/accelerated/icicle/plonk/bn254/icicle.go new file mode 100644 index 0000000000..718ad8d37d --- /dev/null +++ b/backend/accelerated/icicle/plonk/bn254/icicle.go @@ -0,0 +1,7093 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +//go:build icicle + +package bn254 + +import ( + "context" + "errors" + "fmt" + "hash" + "io" + "math/big" + "math/bits" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + "unsafe" + + "golang.org/x/sync/errgroup" + + "github.com/consensys/gnark/backend" + plonk_bn254 "github.com/consensys/gnark/backend/plonk/bn254" + "github.com/consensys/gnark/backend/witness" + constraint "github.com/consensys/gnark/constraint" + cs "github.com/consensys/gnark/constraint/bn254" + "github.com/consensys/gnark/constraint/solver" + fcs "github.com/consensys/gnark/frontend/cs" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + + "github.com/consensys/gnark-crypto/ecc" + curve "github.com/consensys/gnark-crypto/ecc/bn254" + "github.com/consensys/gnark-crypto/ecc/bn254/fp" + "github.com/consensys/gnark-crypto/ecc/bn254/fr" + "github.com/consensys/gnark-crypto/ecc/bn254/fr/fft" + "github.com/consensys/gnark-crypto/ecc/bn254/fr/hash_to_field" + "github.com/consensys/gnark-crypto/ecc/bn254/fr/iop" + "github.com/consensys/gnark-crypto/ecc/bn254/kzg" + fiatshamir "github.com/consensys/gnark-crypto/fiat-shamir" + + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" + icicle_bn254 "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bn254" + icicle_msm "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bn254/msm" + icicle_ntt "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bn254/ntt" + icicle_vecops "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bn254/vecOps" + icicle_runtime "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime" + "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime/config_extension" +) + +const HasIcicle = true + +var isProfileMode bool + +var useBlinding bool + +var isNttTrace bool + +func init() { + _, isProfileMode = os.LookupEnv("ICICLE_STEP_PROFILE") + // Blinding polynomials (zero-knowledge) are enabled by default, matching the + // native prover. Set GNARK_DISABLE_BLINDING to trade zero-knowledge for a + // faster, deterministic prover (e.g. when the witness is not secret). + _, disableBlinding := os.LookupEnv("GNARK_DISABLE_BLINDING") + useBlinding = !disableBlinding + isNttTrace = envEnabled("ICICLE_NTT_TRACE", false) +} + +// profileStep returns a function that, when called, logs the elapsed time since +// profileStep was invoked. If profiling is disabled, it returns a no-op. +// Usage: done := profileStep("label"); defer done() +func profileStep(msg string) func() { + if !isProfileMode { + return func() {} + } + start := time.Now() + return func() { + l := logger.Logger() + l.Debug().Dur("took", time.Since(start)).Msg(msg) + } +} + +// stageTiming is a single recorded prover stage and its wall-clock duration. +type stageTiming struct { + name string + dur time.Duration +} + +// stageTimings is a concurrency-safe, ordered recorder of prover stage +// durations. The PLONK prover runs its stages as concurrent goroutines, so the +// recorded durations OVERLAP and do not sum to the total — the printed table +// flags this. +type stageTimings struct { + mu sync.Mutex + entries []stageTiming +} + +// record appends a (stage, duration) entry. Safe to call from any goroutine and +// safe on a nil receiver (records nothing). +func (t *stageTimings) record(name string, d time.Duration) { + if t == nil { + return + } + t.mu.Lock() + t.entries = append(t.entries, stageTiming{name: name, dur: d}) + t.mu.Unlock() +} + +// printTable writes an aligned breakdown of the recorded stages to w, sorted by +// duration (largest first), followed by the overall prover total. Stages run +// concurrently, so the rows overlap and intentionally do not sum to the total. +func (t *stageTimings) printTable(w io.Writer, total time.Duration) { + if t == nil { + return + } + t.mu.Lock() + rows := make([]stageTiming, len(t.entries)) + copy(rows, t.entries) + t.mu.Unlock() + + sort.SliceStable(rows, func(i, j int) bool { return rows[i].dur > rows[j].dur }) + + nameW := len("TOTAL (prover done)") + for _, r := range rows { + if len(r.name) > nameW { + nameW = len(r.name) + } + } + + fmt.Fprintln(w, "") + fmt.Fprintln(w, "================ gnark PLONK prove breakdown (GPU) ================") + fmt.Fprintln(w, "(stages run concurrently — durations overlap and do not sum to TOTAL)") + fmt.Fprintf(w, " %-*s %12s %6s\n", nameW, "STAGE", "TIME", "%TOTAL") + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + for _, r := range rows { + pct := 0.0 + if total > 0 { + pct = 100 * float64(r.dur) / float64(total) + } + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, r.name, r.dur.Round(time.Millisecond), pct) + } + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, "TOTAL (prover done)", total.Round(time.Millisecond), 100.0) + fmt.Fprintln(w, "===================================================================") + fmt.Fprintln(w, "") +} + +func envEnabled(key string, defaultVal bool) bool { + v, ok := os.LookupEnv(key) + if !ok { + return defaultVal + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "0", "false", "off", "no": + return false + default: + return true + } +} + +func nttAlgorithmFromEnv(key string, fallback icicle_core.NttAlgorithm) icicle_core.NttAlgorithm { + v, ok := os.LookupEnv(key) + if !ok { + return fallback + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "auto", "0": + return icicle_core.Auto + case "radix2", "radix-2", "r2", "1": + return icicle_core.Radix2 + case "mixed", "mixedradix", "mixed-radix", "2": + return icicle_core.MixedRadix + default: + return fallback + } +} + +const ( + id_L int = iota + id_R + id_O + id_Z + id_ZS + id_Ql + id_Qr + id_Qm + id_Qo + id_Qk + id_S1 + id_S2 + id_S3 + id_Qci // [ .. , Qc_i, Pi_i, ...] +) + +// blinding factors +const ( + id_Bl int = iota + id_Br + id_Bo + id_Bz + nb_blinding_polynomials +) + +// blinding orders (-1 to deactivate) +const ( + order_blinding_L = 1 + order_blinding_R = 1 + order_blinding_O = 1 + order_blinding_Z = 2 +) + +// Prove generates a PLONK proof. When the accelerator option is not set to +// "icicle", we delegate to the native prover. Otherwise, we run a local copy +// of the CPU prover logic to enable incremental GPU adaptation. +func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts ...backend.ProverOption) (*plonk_bn254.Proof, error) { + opt, err := backend.NewProverConfig(opts...) + if err != nil { + return nil, err + } + + log := logger.Logger().With(). + Str("curve", spr.CurveID().String()). + Int("nbConstraints", spr.GetNbConstraints()). + Str("backend", "plonk").Logger() + + // parse the options + if opt.HashToFieldFn == nil { + opt.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + + // When blinding is disabled (GNARK_DISABLE_BLINDING), also disable StatisticalZK, it makes no sense + // to use statistical zero knowledge when we don't use blinding. + if !useBlinding { + opt.StatisticalZK = false + } + + start := time.Now() + + // Initialize device and preload KZG bases once per proving key + device := icicle_runtime.CreateDevice("CUDA", 0) + if pk.deviceInfo == nil { + if err := pk.setupDevicePointers(&device); err != nil { + return nil, err + } + } + + // init instance + g, ctx := errgroup.WithContext(context.Background()) + instance, err := newInstance(ctx, spr, pk, fullWitness, &opt) + if err != nil { + return nil, fmt.Errorf("new instance: %w", err) + } + // attach device to instance for GPU calls + instance.device = device + instance.initSharedGPUState() + defer instance.releaseTempGPUMemoryPool() + defer instance.releaseSharedGPUState() + defer instance.releaseLinearizedEvalGPUState() + + // solve constraints + g.Go(instance.solveConstraints) + + // complete qk + g.Go(instance.completeQk) + + // init blinding polynomials + g.Go(instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + g.Go(instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + g.Go(instance.buildRatioCopyConstraint) + + // compute h + g.Go(instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + g.Go(instance.openZ) + + // linearized polynomial + g.Go(instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + g.Go(func() error { + startBatchOpening := time.Now() + err := instance.batchOpening() + if isProfileMode { + instance.timings.record("batchOpening (folded KZG)", time.Since(startBatchOpening)) + } + return err + }) + + if err := g.Wait(); err != nil { + return nil, err + } + + total := time.Since(start) + log.Debug().Dur("took", total).Msg("prover done") + if isProfileMode { + instance.timings.printTable(os.Stderr, total) + } + return instance.proof, nil +} + +// represents a Prover instance +type instance struct { + ctx context.Context + + pk *ProvingKey + proof *plonk_bn254.Proof + spr *cs.SparseR1CS + opt *backend.ProverConfig + + fs *fiatshamir.Transcript + kzgFoldingHash hash.Hash // for KZG folding + htfFunc hash.Hash // hash to field function + + // polynomials + polyL, polyR, polyO *iop.Polynomial + polyZ, polyZS, polyQk *iop.Polynomial + bp []*iop.Polynomial // blinding polynomials + h *iop.Polynomial // h is the quotient polynomial + hGPU *gpuQuotientPolynomial + polyZLagrangeGPU icicle_core.DeviceSlice + blindedZCanonicalGPU icicle_core.DeviceSlice + quotientShardsRandomizers [2]fr.Element // random elements for blinding the shards of the quotient + + linearizedPolynomial []fr.Element + linearizedPolynomialGPU icicle_core.DeviceSlice + linearizedPolynomialClaim fr.Element + linearizedPolynomialDigest kzg.Digest + + fullWitness witness.Witness + + // bsb22 commitment stuff + commitmentInfo constraint.PlonkCommitments + commitmentVal []fr.Element + cCommitments []*iop.Polynomial + + // challenges + gamma, beta, alpha, zeta fr.Element + + // channel to wait for the steps + chLRO, + chQk, + chbp, + chZ, + chH, + chRestoreLRO, + chZOpening, + chLinearizedPolynomial, + chGammaBeta chan struct{} + + domain0, domain1 *fft.Domain + + trace *plonk_bn254.Trace + + // GPU device handle + device icicle_runtime.Device + + // Shared GPU polynomial context reused across buildRatioCopyConstraint + // and computeQuotient to avoid repeated host<->device uploads. + gpuStateMu sync.Mutex + sharedGPUState *gpuPolysState + // Snapshot of immutable polynomial slices used by computeLinearizedPolynomial + // for zeta evaluations after computeQuotient mutates/frees shared state. + linearizedEvalGPUState *gpuPolysState + + // Reusable temporary GPU memory pool for non-state buffers. + tempGPUMemPool *gpuMemoryPool + + // Per-prove stage-timing recorder (used to print the breakdown table when + // ICICLE_STEP_PROFILE is set). + timings *stageTimings +} + +func newInstance(ctx context.Context, spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts *backend.ProverConfig) (*instance, error) { + if opts.HashToFieldFn == nil { + opts.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + s := instance{ + ctx: ctx, + pk: pk, + proof: &plonk_bn254.Proof{}, + spr: spr, + opt: opts, + fullWitness: fullWitness, + bp: make([]*iop.Polynomial, nb_blinding_polynomials), + fs: fiatshamir.NewTranscript(opts.ChallengeHash, "gamma", "beta", "alpha", "zeta"), + kzgFoldingHash: opts.KZGFoldingHash, + htfFunc: opts.HashToFieldFn, + chLRO: make(chan struct{}, 1), + chQk: make(chan struct{}, 1), + chbp: make(chan struct{}, 1), + chGammaBeta: make(chan struct{}, 1), + chZ: make(chan struct{}, 1), + chH: make(chan struct{}, 1), + chZOpening: make(chan struct{}, 1), + chLinearizedPolynomial: make(chan struct{}, 1), + chRestoreLRO: make(chan struct{}, 1), + tempGPUMemPool: newGPUMemoryPool(), + timings: &stageTimings{}, + } + s.initBSB22Commitments() + + // FFT domains and the PLONK trace are witness-independent and expensive to + // build at large n (NewTrace walks every constraint), so they are cached + // on the proving key and shared read-only across proofs. + setup := pk.hostSetupFor(spr) + s.domain0 = setup.domain0 + s.domain1 = setup.domain1 + s.trace = setup.trace + + // sampling random numbers for blinding the quotient + if opts.StatisticalZK { + s.quotientShardsRandomizers[0].SetRandom() + s.quotientShardsRandomizers[1].SetRandom() + } + + return &s, nil +} + +func (s *instance) initBlindingPolynomials() error { + if !useBlinding { + // When blinding is disabled (GNARK_DISABLE_BLINDING), skip creating blinding polynomials entirely + // Just close the channel to unblock any goroutines waiting on it + close(s.chbp) + return nil + } + + s.bp[id_Bl] = getRandomPolynomial(order_blinding_L) + s.bp[id_Br] = getRandomPolynomial(order_blinding_R) + s.bp[id_Bo] = getRandomPolynomial(order_blinding_O) + s.bp[id_Bz] = getRandomPolynomial(order_blinding_Z) + close(s.chbp) + return nil +} + +func (s *instance) initBSB22Commitments() { + s.commitmentInfo = s.spr.CommitmentInfo.(constraint.PlonkCommitments) + s.commitmentVal = make([]fr.Element, len(s.commitmentInfo)) // TODO @Tabaie get rid of this + s.cCommitments = make([]*iop.Polynomial, len(s.commitmentInfo)) + s.proof.Bsb22Commitments = make([]kzg.Digest, len(s.commitmentInfo)) + + // override the hint for the commitment constraints + bsb22ID := solver.GetHintID(fcs.Bsb22CommitmentComputePlaceholder) + s.opt.SolverOpts = append(s.opt.SolverOpts, solver.OverrideHint(bsb22ID, s.bsb22Hint)) +} + +// Computing and verifying Bsb22 multi-commits explained in https://hackmd.io/x8KsadW3RRyX7YTCFJIkHg +func (s *instance) bsb22Hint(_ *big.Int, ins, outs []*big.Int) error { + var err error + commDepth := int(ins[0].Int64()) + ins = ins[1:] + + res := &s.commitmentVal[commDepth] + + commitmentInfo := s.spr.CommitmentInfo.(constraint.PlonkCommitments)[commDepth] + committedValues := make([]fr.Element, s.domain0.Cardinality) + offset := s.spr.GetNbPublicVariables() + for i := range ins { + committedValues[offset+commitmentInfo.Committed[i]].SetBigInt(ins[i]) + } + if _, err = committedValues[offset+commitmentInfo.CommitmentIndex].SetRandom(); err != nil { // Commitment injection constraint has qcp = 0. Safe to use for blinding. + return err + } + if _, err = committedValues[offset+s.spr.GetNbConstraints()-1].SetRandom(); err != nil { // Last constraint has qcp = 0. Safe to use for blinding + return err + } + s.cCommitments[commDepth] = iop.NewPolynomial(&committedValues, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[commDepth], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[commDepth]); err != nil { + return err + } + + s.htfFunc.Write(s.proof.Bsb22Commitments[commDepth].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + res.SetBytes(hashBts[:nbBuf]) // TODO @Tabaie use CommitmentIndex for this; create a new variable CommitmentConstraintIndex for other uses + res.BigInt(outs[0]) + + return nil +} + +// solveConstraints computes the evaluation of the L, R, O polynomials in Lagrange form. +func (s *instance) solveConstraints() error { + startSolve := time.Now() + log := logger.Logger() + + var solution *cs.SparseR1CSSolution + + // Try to load raw solver values from cache (fastest path) + rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + if rawCachePath != "" { + if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { + // Reconstruct L, R, O from raw values + var sol cs.SparseR1CSSolution + if sol.L, sol.R, sol.O, err = s.spr.EvaluateLROSmallDomainFromValues(rawValues); err != nil { + log.Warn().Err(err).Str("file", rawCachePath).Msg("ignoring raw solver cache") + } else { + log.Debug().Dur("took", time.Since(startSolve)).Int("wires", len(rawValues)).Msg("loaded raw solver values from cache") + solution = &sol + } + + // Load cached BSB22 cCommitments polynomials + cacheDir := filepath.Dir(rawCachePath) + for i := 0; solution != nil && i < len(s.commitmentInfo); i++ { + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + coeffs, err := cs.LoadRawSolverValues(bsb22Path) + if err != nil { + log.Warn().Err(err).Int("i", i).Msg("ignoring raw solver cache: missing BSB22 commitment sidecar") + solution = nil + break + } + coeffSlice := []fr.Element(coeffs) + s.cCommitments[i] = iop.NewPolynomial(&coeffSlice, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[i], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[i]); err != nil { + return err + } + s.htfFunc.Write(s.proof.Bsb22Commitments[i].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + s.commitmentVal[i].SetBytes(hashBts[:nbBuf]) + } + } + } + + if solution == nil { + _solution, err := s.spr.SolveAndSaveRawValues(s.fullWitness, rawCachePath, s.opt.SolverOpts...) + if err != nil { + log.Debug().Dur("took", time.Since(startSolve)).Err(err).Msg("solveConstraints: spr.Solve") + return err + } + log.Debug().Dur("took", time.Since(startSolve)).Msg("solveConstraints: spr.Solve") + if isProfileMode { + s.timings.record("solveConstraints: spr.Solve", time.Since(startSolve)) + } + solution = _solution.(*cs.SparseR1CSSolution) + + // Save cCommitments polynomial coefficients for BSB22 reconstruction + if rawCachePath != "" && len(s.commitmentInfo) > 0 { + cacheDir := filepath.Dir(rawCachePath) + for i := range s.commitmentInfo { + if s.cCommitments[i] != nil { + coeffs := s.cCommitments[i].Coefficients() + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + if err := cs.SaveRawSolverValues(bsb22Path, coeffs); err != nil { + log.Warn().Err(err).Int("i", i).Msg("failed to save BSB22 commitment polynomial") + } + } + } + } + } + + evaluationLDomainSmall := []fr.Element(solution.L) + evaluationRDomainSmall := []fr.Element(solution.R) + evaluationODomainSmall := []fr.Element(solution.O) + var wg sync.WaitGroup + wg.Add(2) + go func() { + s.polyL = iop.NewPolynomial(&evaluationLDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + go func() { + s.polyR = iop.NewPolynomial(&evaluationRDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + + s.polyO = iop.NewPolynomial(&evaluationODomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + + wg.Wait() + if _, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}); err != nil { + return err + } + + // commit to l, r, o and add blinding factors + if err := s.commitToLRO(); err != nil { + return err + } + close(s.chLRO) + return nil +} + +func (s *instance) completeQk() error { + qk := s.trace.Qk.Clone() + qkCoeffs := qk.Coefficients() + + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + copy(qkCoeffs, wWitness[:len(s.spr.Public)]) + + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + for i := range s.commitmentInfo { + qkCoeffs[s.spr.GetNbPublicVariables()+s.commitmentInfo[i].CommitmentIndex] = s.commitmentVal[i] + } + + s.polyQk = qk + close(s.chQk) + + return nil +} + +func (s *instance) commitToLRO() error { + var startCommitLRO time.Time + if isProfileMode { + startCommitLRO = time.Now() + } + sequentialLRO := s.domain0 != nil && s.domain0.Cardinality >= (1<<22) + if _, ok := os.LookupEnv("ICICLE_LRO_COMMIT_SEQUENTIAL"); ok { + sequentialLRO = envEnabled("ICICLE_LRO_COMMIT_SEQUENTIAL", true) + } + + if !useBlinding { + // When blinding is disabled, commit directly without waiting for blinding polynomials + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + if err != nil { + return err + } + } else { + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (no blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (no blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil + } + + // wait for blinding polynomials to be initialized or context to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chbp: + } + + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + if err != nil { + return err + } + } else { + // Run the three commitments concurrently. + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (with blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil +} + +// deriveGammaAndBeta (copy constraint) +func (s *instance) deriveGammaAndBeta() error { + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + if err := bindPublicData(s.fs, "gamma", s.pk.VerifyingKey().(*plonk_bn254.VerifyingKey), wWitness[:len(s.spr.Public)]); err != nil { + return err + } + + // wait for LRO to be committed + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + gamma, err := deriveRandomness(s.fs, "gamma", &s.proof.LRO[0], &s.proof.LRO[1], &s.proof.LRO[2]) + if err != nil { + return err + } + + bbeta, err := s.fs.ComputeChallenge("beta") + if err != nil { + return err + } + s.gamma = gamma + s.beta.SetBytes(bbeta) + + close(s.chGammaBeta) + + return nil +} + +// commitToPolyAndBlinding computes the KZG commitment of a polynomial p +// in Lagrange form (large degree) +// and add the contribution of a blinding polynomial b (small degree) +// /!\ The polynomial p is supposed to be in Lagrange form. +// Only used when blinding is enabled (the default). +func (s *instance) commitToPolyAndBlinding(p, b *iop.Polynomial) (commit curve.G1Affine, err error) { + // Commit over the Lagrange SRS using shared device-resident polynomial data. + gpuCommit, err := s.commitLagrangePolynomialOnGPU(p) + if err != nil { + return curve.G1Affine{}, err + } + + // add CPU blinding contribution (two MSMs on canonical SRS) + n := int(s.domain0.Cardinality) + cb := commitBlindingFactor(n, b, s.pk.Kzg) + gpuCommit.Add(&gpuCommit, &cb) + return gpuCommit, nil +} + +func (s *instance) commitLagrangePolynomialOnGPU(p *iop.Polynomial) (curve.G1Affine, error) { + if p == nil { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: nil polynomial") + } + if p.Basis != iop.Lagrange { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial must be in Lagrange basis, got %v", p.Basis) + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{p}) + if err != nil { + return curve.G1Affine{}, err + } + idx, ok := gpuState.polyToIdx[p] + if !ok || idx < 0 || idx >= len(gpuState.deviceSlices) { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial is missing from shared GPU state") + } + if gpuState.deviceSlices[idx].IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: empty device slice for polynomial") + } + + // Keep the polynomial in its native Lagrange basis; large MSMs are split + // into device-side chunks inside commitOnGPULagrangeDevice. + return commitOnGPULagrangeDevice(gpuState.deviceSlices[idx], &s.device, s.pk) +} + +func (s *instance) deriveAlpha() (err error) { + alphaDeps := make([]*curve.G1Affine, len(s.proof.Bsb22Commitments)+1) + for i := range s.proof.Bsb22Commitments { + alphaDeps[i] = &s.proof.Bsb22Commitments[i] + } + alphaDeps[len(alphaDeps)-1] = &s.proof.Z + s.alpha, err = deriveRandomness(s.fs, "alpha", alphaDeps...) + return err +} + +func (s *instance) deriveZeta() (err error) { + s.zeta, err = deriveRandomness(s.fs, "zeta", &s.proof.H[0], &s.proof.H[1], &s.proof.H[2]) + return +} + +func (s *instance) computeQuotient() (err error) { + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + if isProfileMode { + var startComputeQuotient time.Time + startComputeQuotient = time.Now() + defer func() { + l := logger.Logger() + if err != nil { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Err(err).Msg("computeQuotient: total (with error)") + } else { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Msg("computeQuotient: total") + } + s.timings.record("computeQuotient (total)", time.Since(startComputeQuotient)) + }() + } + + // wait for Z to be committed or context done + doneWaitZ := profileStep("computeQuotient: wait Z commit") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZ: + } + doneWaitZ() + + // derive alpha + if err = s.deriveAlpha(); err != nil { + return err + } + + if err := s.waitForComputeNumeratorQk(); err != nil { + return err + } + if s.polyQk == nil { + return fmt.Errorf("computeQuotient: missing completed Qk polynomial") + } + + doneEnsureGPUState := profileStep("computeQuotient: ensure shared GPU state") + gpuState, err := s.ensurePolysOnSharedGPU(s.buildComputeNumeratorGPUBatch()) + if err != nil { + return err + } + doneEnsureGPUState() + + // compute Z shifted by one for copy-constraint terms. + if s.polyZ == nil { + return fmt.Errorf("computeQuotient: missing Z polynomial") + } + s.polyZS = s.polyZ.ShallowClone().Shift(1) + + var numeratorGPU *gpuNumeratorPolynomial + var quotientGPU *gpuQuotientPolynomial + var e error + doneComputeNumerator := profileStep("computeQuotient: computeNumerator") + numeratorGPU, e = s.computeNumerator(gpuState) + if e != nil { + return e + } + doneComputeNumerator() + + doneDivideByZH := profileStep("computeQuotient: divideByZHOnGPU") + quotientGPU, e = s.divideByZHOnGPU(numeratorGPU, [2]*fft.Domain{s.domain0, s.domain1}) + if e != nil { + return e + } + doneDivideByZH() + s.hGPU = quotientGPU + + // Shared state slices were mutated during numerator coset iterations and are no + // longer needed now; computeLinearizedPolynomial uses the immutable snapshot. + s.releaseSharedGPUState() + close(s.chRestoreLRO) + + doneCommitH := profileStep("computeQuotient: commit H from device") + if err := s.commitToQuotientGPUFromDevice(s.hGPU); err != nil { + return err + } + doneCommitH() + + if err := s.deriveZeta(); err != nil { + return err + } + + donePrepareLinearizedEval := profileStep("computeQuotient: prepare linearized eval GPU state") + if err := s.prepareLinearizedEvalGPUStateFromHost(); err != nil { + return fmt.Errorf("computeQuotient: prepare linearized eval GPU state failed: %w", err) + } + donePrepareLinearizedEval() + + close(s.chH) + + return nil +} + +func (s *instance) buildRatioCopyConstraint() (err error) { + // wait for gamma and beta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chGammaBeta: + } + + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("buildRatioCopyConstraint: missing L/R/O polynomials") + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}) + if err != nil { + return err + } + dL, err := getStateDeviceSlice(gpuState, s.polyL, "L") + if err != nil { + return err + } + dR, err := getStateDeviceSlice(gpuState, s.polyR, "R") + if err != nil { + return err + } + dO, err := getStateDeviceSlice(gpuState, s.polyO, "O") + if err != nil { + return err + } + + var startBuildRatioCopyConstraintIcicle time.Time + if isProfileMode { + startBuildRatioCopyConstraintIcicle = time.Now() + } + s.polyZ, err = s.BuildRatioCopyConstraintIcicle( + []icicle_core.DeviceSlice{dL, dR, dO}, + s.trace.S, + s.beta, + s.gamma, + s.domain0, + gpuState, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startBuildRatioCopyConstraintIcicle)).Msg("buildRatioCopyConstraint: BuildRatioCopyConstraintIcicle") + s.timings.record("buildRatioCopyConstraint (perm Z)", time.Since(startBuildRatioCopyConstraintIcicle)) + } + + dZ, err := getStateDeviceSlice(gpuState, s.polyZ, "Z") + if err != nil { + return err + } + copyDone := make(chan error, 1) + var dPersist icicle_core.DeviceSlice + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + var allocErr error + dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) + if allocErr != nil { + finish(fmt.Errorf("buildRatioCopyConstraint: alloc persist Z failed: %w", allocErr)) + return + } + if e := copyDeviceSliceIntoOnCurrentDevice(dPersist, dZ, cfg); e != icicle_runtime.Success { + _ = dPersist.Free() + dPersist = icicle_core.DeviceSlice{} + finish(fmt.Errorf("buildRatioCopyConstraint: persist Z copy failed: %s", e.AsString())) + return + } + finish(nil) + }) + if err := <-copyDone; err != nil { + return err + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + s.polyZLagrangeGPU = dPersist + + // commit to Z (with or without blinding) + var startCommitZ time.Time + if isProfileMode { + startCommitZ = time.Now() + } + if useBlinding { + s.proof.Z, err = s.commitToPolyAndBlinding(s.polyZ, s.bp[id_Bz]) + } else { + s.proof.Z, err = s.commitLagrangePolynomialOnGPU(s.polyZ) + } + if isProfileMode { + l := logger.Logger() + if useBlinding { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (no blinding)") + } + } + s.freeIdleTempGPUMemoryOnDevice() + + close(s.chZ) + + return +} + +// open Z (blinded) at ωζ +func (s *instance) openZ() (err error) { + // wait for H to be committed and zeta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + + var zetaShifted fr.Element + zetaShifted.Mul(&s.zeta, &s.pk.Vk.Generator) + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("openZ: missing GPU Z polynomial") + } + + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + + dZLagrange := s.polyZLagrangeGPU + if dZLagrange.Len() <= 1 { + return fmt.Errorf("openZ: invalid Z size %d", dZLagrange.Len()) + } + + blindSize := order_blinding_Z + 1 + if useBlinding { + if len(s.bp) <= id_Bz || s.bp[id_Bz] == nil { + return fmt.Errorf("openZ: missing Z blinding polynomial") + } + blindSize = len(s.bp[id_Bz].Coefficients()) + if blindSize == 0 { + return fmt.Errorf("openZ: empty Z blinding polynomial") + } + } + + var dBlindedCanonical icicle_core.DeviceSlice + buildDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + buildDone <- cfgErr + return + } + finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + // Async boundary for canonicalization/blinding before exposing output. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + buildDone <- runErr + } + + n := dZLagrange.Len() + dZCanonical := s.getTempDeviceSlice(n) + if err := copyDeviceSliceIntoOnCurrentDevice(dZCanonical, dZLagrange, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy Z to canonical buffer failed: %s", err.AsString()), dZCanonical, false) + return + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + cfgNtt.Ordering = icicle_core.KNN // regular lagrange -> regular canonical + if err := icicle_ntt.Ntt(dZCanonical, icicle_core.KInverse, &cfgNtt, dZCanonical); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: inverse NTT on Z failed: %s", err.AsString()), dZCanonical, false) + return + } + + dBlindedCanonical = s.getTempDeviceSlice(n + blindSize) + dBlindedPrefix := (&dBlindedCanonical).Range(0, n, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedPrefix, dZCanonical, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy canonical Z into blinded buffer failed: %s", err.AsString()), dZCanonical, true) + return + } + + if useBlinding { + dBp := uploadVector(s.bp[id_Bz].Coefficients()) + dBlindedHead := (&dBlindedPrefix).Range(0, blindSize, false) + if err := icicle_vecops.VecOp(dBlindedHead, dBp, dBlindedHead, cfgVec, icicle_core.Sub); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: subtract Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedTail, dBp, cfgVec); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: append Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + _ = dBp.FreeAsync(stream) + } else { + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := zeroDeviceSliceOnCurrentDevice(dBlindedTail, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: zero-pad non-blinded Z failed: %s", err.AsString()), dZCanonical, true) + return + } + } + + finalize(nil, dZCanonical, false) + }) + if err := <-buildDone; err != nil { + return err + } + s.blindedZCanonicalGPU = dBlindedCanonical + + // open z at zeta*w. + var startKzgOpen time.Time + if isProfileMode { + startKzgOpen = time.Now() + } + s.proof.ZShiftedOpening, err = s.openPolynomialOnGPUCanonicalDevice(s.blindedZCanonicalGPU, zetaShifted) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startKzgOpen)).Msg("openZ: open polynomial on GPU") + s.timings.record("openZ (KZG open on GPU)", time.Since(startKzgOpen)) + } + if err != nil { + return err + } + close(s.chZOpening) + return nil +} + +func (s *instance) h1() []fr.Element { + var h1 []fr.Element + if !s.opt.StatisticalZK { + h1 = s.h.Coefficients()[:s.domain0.Cardinality+2] + } else { + h1 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h1, s.h.Coefficients()[:s.domain0.Cardinality+2]) + h1[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[0]) + } + return h1 +} + +func (s *instance) h2() []fr.Element { + var h2 []fr.Element + if !s.opt.StatisticalZK { + h2 = s.h.Coefficients()[s.domain0.Cardinality+2 : 2*(s.domain0.Cardinality+2)] + } else { + h2 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h2, s.h.Coefficients()[s.domain0.Cardinality+2:2*(s.domain0.Cardinality+2)]) + h2[0].Sub(&h2[0], &s.quotientShardsRandomizers[0]) + h2[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[1]) + } + return h2 +} + +func (s *instance) h3() []fr.Element { + var h3 []fr.Element + if !s.opt.StatisticalZK { + h3 = s.h.Coefficients()[2*(s.domain0.Cardinality+2) : 3*(s.domain0.Cardinality+2)] + } else { + h3 = make([]fr.Element, s.domain0.Cardinality+2) + copy(h3, s.h.Coefficients()[2*(s.domain0.Cardinality+2):3*(s.domain0.Cardinality+2)]) + h3[0].Sub(&h3[0], &s.quotientShardsRandomizers[1]) + } + return h3 +} + +// witnessEvalAtZeta holds the scalar evaluations of witness and constraint +// polynomials at the challenge point zeta, as needed by the linearized +// polynomial computation. +type witnessEvalAtZeta struct { + blzeta, brzeta, bozeta fr.Element + s1zeta, s2zeta fr.Element + qcpzeta []fr.Element +} + +type linearizedSelectorScales struct { + s3, ql, qr, qm, qo, qk fr.Element + qcp []fr.Element +} + +// evaluateWitnessPolynomialsAtZeta evaluates L, R, O (with optional blinding), +// S1, S2, and all Qcp polynomials at the point zeta using the GPU-resident +// polynomial state. +func (s *instance) evaluateWitnessPolynomialsAtZeta( + evalGPUState *gpuPolysState, + zeta fr.Element, +) (witnessEvalAtZeta, error) { + doneTotal := profileStep("evaluateWitnessPolynomialsAtZeta: total") + defer doneTotal() + + var result witnessEvalAtZeta + var err error + + result.qcpzeta = make([]fr.Element, len(s.commitmentInfo)) + var startQcp time.Time + if isProfileMode { + startQcp = time.Now() + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i >= len(s.trace.Qcp) || s.trace.Qcp[i] == nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: missing Qcp polynomial at index %d", i) + } + var startQcpItem time.Time + if isProfileMode { + startQcpItem = time.Now() + } + result.qcpzeta[i], err = s.evalPolynomialInCurrentFormOnGPU(s.trace.Qcp[i], evalGPUState, zeta) + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: qcp[%d] GPU evaluation failed: %w", i, err) + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("idx", i).Dur("took", time.Since(startQcpItem)).Msg("evaluateWitnessPolynomialsAtZeta: qcp eval item") + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startQcp)).Msg("evaluateWitnessPolynomialsAtZeta: qcpZeta evaluate on GPU") + } + + if useBlinding { + result.blzeta, err = s.evaluateBlindedOnGPU(s.polyL, s.bp[id_Bl], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: L evaluate on GPU") + result.blzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyL, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: blzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.brzeta, err = s.evaluateBlindedOnGPU(s.polyR, s.bp[id_Br], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: R evaluate on GPU") + result.brzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyR, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: brzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.bozeta, err = s.evaluateBlindedOnGPU(s.polyO, s.bp[id_Bo], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: O evaluate on GPU") + result.bozeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyO, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: bozeta GPU evaluation failed: %w", err) + } + + doneS1 := profileStep("evaluateWitnessPolynomialsAtZeta: S1 evaluate on GPU") + result.s1zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S1, evalGPUState, zeta) + doneS1() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s1(zeta) GPU evaluation failed: %w", err) + } + doneS2 := profileStep("evaluateWitnessPolynomialsAtZeta: S2 evaluate on GPU") + result.s2zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S2, evalGPUState, zeta) + doneS2() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s2(zeta) GPU evaluation failed: %w", err) + } + + return result, nil +} + +func (s *instance) computeLinearizedPolynomial() error { + + // wait for H to be committed and zeta to be derived (or ctx.Done()) + var startWaitH time.Time + if isProfileMode { + startWaitH = time.Now() + } + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startWaitH)).Msg("computeLinearizedPolynomial: wait H and zeta") + s.timings.record("computeLinearizedPoly (wait H+zeta, overlaps)", time.Since(startWaitH)) + } + if s.opt.StatisticalZK { + return fmt.Errorf("computeLinearizedPolynomial: GPU-only opening path does not support StatisticalZK=true") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil || s.polyZ == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing required polynomials") + } + + // Reuse the immutable snapshot prepared in computeQuotient before numerator + // coset iterations mutate shared state slices. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("computeLinearizedPolynomial: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized eval prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: required polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing blinding polynomials for GPU evaluation") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized blinding prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: blinding polynomial is not on GPU: %w", e) + } + } + } + + doneEvaluate := profileStep("computeLinearizedPolynomial: evaluate witness polynomials") + evals, err := s.evaluateWitnessPolynomialsAtZeta(evalGPUState, s.zeta) + doneEvaluate() + if err != nil { + return err + } + + // wait for Z to be opened at zeta (or ctx.Done()) + doneWaitZOpening := profileStep("computeLinearizedPolynomial: wait Z opening") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZOpening: + } + doneWaitZOpening() + if s.blindedZCanonicalGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing canonical blinded Z on GPU") + } + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing lagrange Z on GPU") + } + defer func() { + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + }() + bzuzeta := s.proof.ZShiftedOpening.ClaimedValue + + if s.hGPU == nil || s.hGPU.coeffs.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing GPU quotient polynomial") + } + + doneBuild := profileStep("computeLinearizedPolynomial: build selector terms on GPU") + dLin, err := s.buildLinearizedSelectorTermsOnGPU(evals, bzuzeta, s.blindedZCanonicalGPU.Len()) + doneBuild() + if err != nil { + return err + } + + doneAddZ := profileStep("computeLinearizedPolynomial: add Z contribution on GPU") + err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) + doneAddZ() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + s.linearizedPolynomialGPU = dLin + + doneEvalClaim := profileStep("computeLinearizedPolynomial: evaluate linearized claim") + claim, err := s.evalDevicePolynomialAtPoint(dLin, s.zeta) + doneEvalClaim() + if err != nil { + return err + } + s.linearizedPolynomialClaim = claim + + // Commit the linearized polynomial over the canonical SRS. + var startMSM time.Time + if isProfileMode { + startMSM = time.Now() + } + s.linearizedPolynomialDigest, err = commitOnGPUCanonicalDevice(dLin, &s.device, s.pk) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startMSM)).Msg("computeLinearizedPolynomial: KZG commit") + s.timings.record("computeLinearizedPoly (KZG commit)", time.Since(startMSM)) + } + close(s.chLinearizedPolynomial) + return nil +} + +func (s *instance) batchOpening() error { + + // wait for linearizedPolynomial to be computed (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLinearizedPolynomial: + } + + defer func() { + freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } + }() + + if s.linearizedPolynomialGPU.IsEmpty() { + return fmt.Errorf("batchOpening: missing GPU linearized polynomial") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("batchOpening: missing L/R/O polynomials") + } + + // Reuse immutable GPU snapshot prepared before quotient iterations. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("batchOpening: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening eval prereq"); e != nil { + return fmt.Errorf("batchOpening: required polynomial is not on GPU: %w", e) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] == nil { + return fmt.Errorf("batchOpening: missing Qcp polynomial at index %d", i) + } + if _, e := getStateDeviceSlice(evalGPUState, s.trace.Qcp[i], "batchOpening qcp prereq"); e != nil { + return fmt.Errorf("batchOpening: Qcp polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("batchOpening: missing blinding polynomials") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening blinding prereq"); e != nil { + return fmt.Errorf("batchOpening: blinding polynomial is not on GPU: %w", e) + } + } + } + + devicePolys, ownedPolys, claimed, err := s.prepareBatchOpeningPolynomialsOnGPU(evalGPUState, s.zeta) + if err != nil { + return err + } + defer func() { + for i := 0; i < len(devicePolys); i++ { + if i < len(ownedPolys) && ownedPolys[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + } + } + s.releaseLinearizedEvalGPUState() + }() + + digestsToOpen := make([]curve.G1Affine, len(s.pk.Vk.Qcp)+6) + copy(digestsToOpen[6:], s.pk.Vk.Qcp) + digestsToOpen[0] = s.linearizedPolynomialDigest + digestsToOpen[1] = s.proof.LRO[0] + digestsToOpen[2] = s.proof.LRO[1] + digestsToOpen[3] = s.proof.LRO[2] + digestsToOpen[4] = s.pk.Vk.S[0] + digestsToOpen[5] = s.pk.Vk.S[1] + if len(claimed) != len(digestsToOpen) { + return fmt.Errorf("batchOpening: claimed size mismatch (%d != %d)", len(claimed), len(digestsToOpen)) + } + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + var dFold icicle_core.DeviceSlice + foldDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening fold", foldDone) + + dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: copy linearized polynomial failed: %s", e.AsString())) + return + } + + gammaPow := gamma + for i := 1; i < len(devicePolys); i++ { + dPoly := devicePolys[i] + if dPoly.IsEmpty() { + gammaPow.Mul(&gammaPow, &gamma) + continue + } + dScaled := s.getTempDeviceSlice(dPoly.Len()) + dGammaStd := uploadScalarStdOnCurrentDevice(gammaPow, cfg) + eMul := icicle_vecops.ScalarMulVec(dGammaStd, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dGammaStd.FreeAsync(cfg.StreamHandle) + } else { + _ = dGammaStd.Free() + } + if eMul != icicle_runtime.Success { + if cfg.IsAsync { + _ = icicle_runtime.SynchronizeStream(cfg.StreamHandle) + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: scale polynomial %d failed: %s", i, eMul.AsString())) + return + } + + dPrefix := (&dFold).Range(0, dPoly.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add) + if cfg.IsAsync { + // dScaled is recycled each iteration; wait before returning to pool. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: synchronize stream failed: %s", eSync.AsString())) + return + } + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + if eAdd != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: fold add polynomial %d failed: %s", i, eAdd.AsString())) + return + } + gammaPow.Mul(&gammaPow, &gamma) + } + finish(nil) + }) + if err := <-foldDone; err != nil { + if !dFold.IsEmpty() { + s.putTempDeviceSlice(dFold, dFold.Len()) + } + return err + } + var dWitness icicle_core.DeviceSlice + divDone := make(chan error, 1) + witnessSize := dFold.Len() - 1 + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dFold, dPoint, dWitness, cfg) + if cfg.IsAsync { + _ = dPoint.FreeAsync(cfg.StreamHandle) + } else { + _ = dPoint.Free() + } + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + finish(fmt.Errorf("batchOpening: divide by (x-zeta) failed: %s", eDiv.AsString())) + return + } + finish(nil) + }) + if err := <-divDone; err != nil { + s.putTempDeviceSlice(dFold, dFold.Len()) + return err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dFold, dFold.Len()) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return err + } + + s.proof.BatchedProof = kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + } + if err := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); err != nil { + l := logger.Logger() + l.Warn().Err(err).Msg("batchOpening: GPU folded opening failed raw-G2 validation; falling back to host fold with CPU KZG commitment") + fallbackProof, fallbackErr := s.batchOpeningHostFoldGPUCommitFromDevicePolys(devicePolys, digestsToOpen) + if fallbackErr != nil { + return fallbackErr + } + s.proof.BatchedProof = fallbackProof + if verifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); verifyErr != nil { + nativeProof, nativeErr := s.batchOpeningNativeCPUFromDevicePolys(devicePolys, digestsToOpen) + if nativeErr != nil { + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed: %v", verifyErr, nativeErr) + } + s.proof.BatchedProof = nativeProof + if nativeVerifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); nativeVerifyErr != nil { + diagnostic, diagnosticErr := s.diagnoseBatchOpeningDevicePolynomials(devicePolys, digestsToOpen, claimed) + if diagnosticErr != nil { + diagnostic = fmt.Sprintf("batch opening diagnostic failed: %v", diagnosticErr) + } + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed raw-G2 validation: %v; %s", verifyErr, nativeVerifyErr, diagnostic) + } + l.Warn().Msg("batchOpening: native CPU KZG fallback produced a valid proof after host-fold fallback failed") + } + } + _ = foldedEval // kept for parity with kzg.BatchOpenSinglePoint flow. + return nil +} + +func (s *instance) verifyBatchOpeningWithRecomputedLines(digestsToOpen []curve.G1Affine) error { + vk := s.pk.Vk.Kzg + vk.Lines[0] = curve.PrecomputeLines(vk.G2[0]) + vk.Lines[1] = curve.PrecomputeLines(vk.G2[1]) + return kzg.BatchVerifySinglePoint( + digestsToOpen, + &s.proof.BatchedProof, + s.zeta, + s.kzgFoldingHash, + vk, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningHostFoldGPUCommitFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningNativeCPUFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningNativeCPUFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return kzg.BatchOpenSinglePoint( + polysToOpen, + digestsToOpen, + s.zeta, + s.kzgFoldingHash, + s.pk.Kzg, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) diagnoseBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, + gpuClaimed []fr.Element, +) (string, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "diagnoseBatchOpeningDevicePolynomials", + ) + if err != nil { + return "", err + } + if len(gpuClaimed) != len(polysToOpen) { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: claimed/polynomial mismatch (%d != %d)", len(gpuClaimed), len(polysToOpen)) + } + + l := logger.Logger() + claimMismatches := make([]string, 0) + commitMismatches := make([]string, 0) + for i := range polysToOpen { + label := batchOpeningPolynomialLabel(i) + cpuClaim := evalCanonicalAtPoint(polysToOpen[i], s.zeta) + if !cpuClaim.Equal(&gpuClaimed[i]) { + claimMismatches = append(claimMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("gpuClaim", frFingerprint(gpuClaimed[i])). + Str("cpuClaim", frFingerprint(cpuClaim)). + Msg("batchOpening diagnostic: GPU claim differs from CPU evaluation") + } + + cpuDigest, err := kzg.Commit(polysToOpen[i], s.pk.Kzg) + if err != nil { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: commit %s: %w", label, err) + } + if !cpuDigest.Equal(&digestsToOpen[i]) { + commitMismatches = append(commitMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("expectedDigest", g1Fingerprint(digestsToOpen[i])). + Str("cpuDigest", g1Fingerprint(cpuDigest)). + Msg("batchOpening diagnostic: CPU commitment differs from proof digest") + } + } + + if len(claimMismatches) == 0 && len(commitMismatches) == 0 { + return "batch opening diagnostic found no per-polynomial claim or commitment mismatch", nil + } + return fmt.Sprintf( + "batch opening diagnostic claim mismatches=[%s] commitment mismatches=[%s]", + strings.Join(claimMismatches, ","), + strings.Join(commitMismatches, ","), + ), nil +} + +func batchOpeningPolynomialLabel(index int) string { + switch index { + case 0: + return "linearized" + case 1: + return "L" + case 2: + return "R" + case 3: + return "O" + case 4: + return "S1" + case 5: + return "S2" + default: + return fmt.Sprintf("Qcp[%d]", index-6) + } +} + +func frFingerprint(v fr.Element) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func g1Fingerprint(v curve.G1Affine) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func (s *instance) downloadBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + expectedDigests int, + label string, +) ([][]fr.Element, error) { + if len(devicePolys) != expectedDigests { + return nil, fmt.Errorf("%s: polynomial/digest mismatch (%d != %d)", label, len(devicePolys), expectedDigests) + } + + polysToOpen := make([][]fr.Element, len(devicePolys)) + for i := range devicePolys { + var err error + polysToOpen[i], err = s.downloadCanonicalDeviceCoefficients( + devicePolys[i], + fmt.Sprintf("%s[%d]", label, i), + ) + if err != nil { + return nil, err + } + } + return polysToOpen, nil +} + +func (s *instance) batchOpeningHostFoldGPUCommit(digestsToOpen []curve.G1Affine) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.batchOpeningHostPolynomials() + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( + polysToOpen [][]fr.Element, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + if len(polysToOpen) != len(digestsToOpen) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: polynomial/digest mismatch (%d != %d)", len(polysToOpen), len(digestsToOpen)) + } + + largestPoly := 0 + for i := range polysToOpen { + if len(polysToOpen[i]) == 0 || len(polysToOpen[i]) > len(s.pk.Kzg.G1) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: invalid polynomial %d size %d", i, len(polysToOpen[i])) + } + if len(polysToOpen[i]) > largestPoly { + largestPoly = len(polysToOpen[i]) + } + } + + claimed := make([]fr.Element, len(polysToOpen)) + utils.Parallelize(len(polysToOpen), func(start, end int) { + for i := start; i < end; i++ { + claimed[i] = evalCanonicalAtPoint(polysToOpen[i], s.zeta) + } + }) + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + foldedPolynomials := make([]fr.Element, largestPoly) + copy(foldedPolynomials, polysToOpen[0]) + gammaPow := gamma + for i := 1; i < len(polysToOpen); i++ { + poly := polysToOpen[i] + scale := gammaPow + utils.Parallelize(len(poly), func(start, end int) { + var term fr.Element + for j := start; j < end; j++ { + term.Mul(&poly[j], &scale) + foldedPolynomials[j].Add(&foldedPolynomials[j], &term) + } + }) + gammaPow.Mul(&gammaPow, &gamma) + } + + hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) + var dWitness icicle_core.DeviceSlice + uploadDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + dWitness = uploadVector(hCoeffs) + close(uploadDone) + }) + <-uploadDone + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + freeSliceOnDevice(&dWitness, &s.device) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + return kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + }, nil +} + +func (s *instance) batchOpeningHostPolynomials() ([][]fr.Element, error) { + total := 6 + len(s.trace.Qcp) + polysToOpen := make([][]fr.Element, total) + + var err error + polysToOpen[0], err = s.downloadCanonicalDeviceCoefficients( + s.linearizedPolynomialGPU, + "batchOpeningHostPolynomials linearized", + ) + if err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare linearized: %w", err) + } + + prepareLRO := func(p, bp *iop.Polynomial, blindingOrder int) ([]fr.Element, error) { + base, err := canonicalRegularCoefficientsCopy(p, s.domain0) + if err != nil { + return nil, err + } + if useBlinding { + if bp == nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomial") + } + blind, err := canonicalRegularCoefficientsCopy(bp, s.domain0) + if err != nil { + return nil, err + } + out := make([]fr.Element, len(base)+len(blind)) + copy(out, base) + copy(out[len(base):], blind) + for i := range blind { + out[i].Sub(&out[i], &blind[i]) + } + return out, nil + } + out := make([]fr.Element, len(base)+blindingOrder+1) + copy(out, base) + return out, nil + } + + var bpL, bpR, bpO *iop.Polynomial + if useBlinding { + if len(s.bp) <= id_Bo { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomials") + } + bpL, bpR, bpO = s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo] + } + + if polysToOpen[1], err = prepareLRO(s.polyL, bpL, order_blinding_L); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare L: %w", err) + } + if polysToOpen[2], err = prepareLRO(s.polyR, bpR, order_blinding_R); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare R: %w", err) + } + if polysToOpen[3], err = prepareLRO(s.polyO, bpO, order_blinding_O); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare O: %w", err) + } + if polysToOpen[4], err = canonicalRegularCoefficientsCopy(s.trace.S1, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S1: %w", err) + } + if polysToOpen[5], err = canonicalRegularCoefficientsCopy(s.trace.S2, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S2: %w", err) + } + for i := range s.trace.Qcp { + if polysToOpen[6+i], err = canonicalRegularCoefficientsCopy(s.trace.Qcp[i], s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare Qcp[%d]: %w", i, err) + } + } + return polysToOpen, nil +} + +func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceSlice, label string) ([]fr.Element, error) { + if dPoly.IsEmpty() { + return nil, fmt.Errorf("%s: empty device polynomial", label) + } + + coeffs := make([]fr.Element, dPoly.Len()) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + }) + if err := <-done; err != nil { + return nil, err + } + return coeffs, nil +} + +func canonicalRegularCoefficientsCopy(p *iop.Polynomial, domain *fft.Domain) ([]fr.Element, error) { + if p == nil { + return nil, fmt.Errorf("nil polynomial") + } + cp := p.Clone() + cp.ToCanonical(domain).ToRegular() + coeffs := cp.Coefficients() + out := make([]fr.Element, len(coeffs)) + copy(out, coeffs) + return out, nil +} + +func dividePolyByXMinusAHost(f []fr.Element, fa, a fr.Element) []fr.Element { + f[0].Sub(&f[0], &fa) + var t fr.Element + for i := len(f) - 2; i >= 0; i-- { + t.Mul(&f[i+1], &a) + f[i].Add(&f[i], &t) + } + return f[1:] +} + +// evaluate the full set of constraints on the GPU-resident polynomial state. +type computeNumeratorLoopContext struct { + n int + rho int + mm uint64 + bn *big.Int + shifters []fr.Element + twiddles0 []fr.Element + dTwiddles0 icicle_core.DeviceSlice + dPrecomputedDenominators *icicle_core.DeviceSlice + scalingVector []fr.Element + scalingVectorRev []fr.Element + gpuState *gpuPolysState + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + one fr.Element + cs fr.Element + css fr.Element + nbBsbGates int + numeratorShards []icicle_core.DeviceSlice +} + +type gpuNumeratorPolynomial struct { + shards []icicle_core.DeviceSlice + n int + rho int + mm uint64 +} + +type gpuQuotientPolynomial struct { + coeffs icicle_core.DeviceSlice + size int +} + +func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolynomial, error) { + twiddles0, err := s.buildComputeNumeratorTwiddles() + if err != nil { + return nil, err + } + if err := s.validateComputeNumeratorGPUState(gpuState); err != nil { + return nil, err + } + + var startComputeNumerator time.Time + if isProfileMode { + startComputeNumerator = time.Now() + } + + n := s.domain0.Cardinality + nbBsbGates := len(s.proof.Bsb22Commitments) + + var cs, css fr.Element + cs.Set(&s.domain1.FrMultiplicativeGen) + css.Square(&cs) + + bn := big.NewInt(int64(n)) + + rho := int(s.domain1.Cardinality / n) + shifters := make([]fr.Element, rho) + shifters[0].Set(&s.domain1.FrMultiplicativeGen) + for i := 1; i < rho; i++ { + shifters[i].Set(&s.domain1.Generator) + } + + cosetTable, err := s.domain0.CosetTable() + if err != nil { + return nil, err + } + + // for the first iteration, the scalingVector is the coset table + scalingVector := cosetTable + scalingVectorRev := make([]fr.Element, len(cosetTable)) + copy(scalingVectorRev, cosetTable) + fft.BitReverse(scalingVectorRev) + + // pre-computed to compute the bit reverse index + // of the result polynomial + m := uint64(s.domain1.Cardinality) + mm := uint64(64 - bits.TrailingZeros64(m)) + + var dPrecomputedDenominators icicle_core.DeviceSlice + defer func() { + if !dPrecomputedDenominators.IsEmpty() { + freeDone := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + freeDone <- dPrecomputedDenominators.Free() + }) + if err := <-freeDone; err != icicle_runtime.Success { + panic(fmt.Sprintf("computeNumerator: failed to free dPrecomputedDenominators: %s", err.AsString())) + } + } + }() + + var coset, cosetExponentiatedToNMinusOne, one fr.Element + coset.SetOne() + one.SetOne() + + dTwiddles0, err := s.uploadComputeNumeratorTwiddles(twiddles0) + if err != nil { + return nil, err + } + + loopCtx := &computeNumeratorLoopContext{ + n: int(n), + rho: rho, + mm: mm, + bn: bn, + shifters: shifters, + twiddles0: twiddles0, + dTwiddles0: dTwiddles0, + dPrecomputedDenominators: &dPrecomputedDenominators, + scalingVector: scalingVector, + scalingVectorRev: scalingVectorRev, + gpuState: gpuState, + coset: coset, + cosetExponentiatedToNMinusOne: cosetExponentiatedToNMinusOne, + one: one, + cs: cs, + css: css, + nbBsbGates: nbBsbGates, + numeratorShards: make([]icicle_core.DeviceSlice, rho), + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: setup before iteration loop") + } + if err := s.executeComputeNumeratorCosetIterations(loopCtx); err != nil { + return nil, err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: main body (post-wait)") + } + + return &gpuNumeratorPolynomial{ + shards: loopCtx.numeratorShards, + n: loopCtx.n, + rho: loopCtx.rho, + mm: loopCtx.mm, + }, nil + +} + +func (s *instance) buildComputeNumeratorTwiddles() ([]fr.Element, error) { + n := s.domain0.Cardinality + var startTwiddles time.Time + if isProfileMode { + startTwiddles = time.Now() + } + twiddles0 := make([]fr.Element, n) + if n == 1 { + // edge case + twiddles0[0].SetOne() + } else { + twiddles, err := s.domain0.Twiddles() + if err != nil { + return nil, err + } + copy(twiddles0, twiddles[0]) + w := twiddles0[1] + for i := len(twiddles[0]); i < len(twiddles0); i++ { + twiddles0[i].Mul(&twiddles0[i-1], &w) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startTwiddles)).Msg("computeNumerator: build twiddles") + } + return twiddles0, nil +} + +func (s *instance) waitForComputeNumeratorQk() error { + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chQk: + } + return nil +} + +func (s *instance) buildLinearizedEvalGPUBatch() []*iop.Polynomial { + baseCap := 5 + len(s.trace.Qcp) + if useBlinding { + baseCap += 3 + } + polys := make([]*iop.Polynomial, 0, baseCap) + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + } + if useBlinding && len(s.bp) > id_Bo { + for _, bpPoly := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if bpPoly != nil { + polys = append(polys, bpPoly) + } + } + } + return polys +} + +func (s *instance) buildComputeNumeratorGPUBatch() []*iop.Polynomial { + polys := make([]*iop.Polynomial, 0, 13+2*len(s.commitmentInfo)) + for _, p := range []*iop.Polynomial{ + s.polyL, s.polyR, s.polyO, s.polyZ, + s.trace.Ql, s.trace.Qr, s.trace.Qm, s.trace.Qo, s.polyQk, + s.trace.S1, s.trace.S2, s.trace.S3, + } { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i < len(s.trace.Qcp) && s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + if i < len(s.cCommitments) && s.cCommitments[i] != nil { + polys = append(polys, s.cCommitments[i]) + } + } + return polys +} + +func (s *instance) validateComputeNumeratorGPUState(state *gpuPolysState) error { + if state == nil { + return fmt.Errorf("computeNumerator: shared GPU state is nil") + } + required := s.buildComputeNumeratorGPUBatch() + if len(required) == 0 { + return fmt.Errorf("computeNumerator: no polynomials prepared for GPU batch") + } + for _, p := range required { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; !ok { + return fmt.Errorf("computeNumerator: polynomial ptr=%p is missing from shared GPU state", p) + } + } + return nil +} + +func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicle_core.DeviceSlice, error) { + var dTwiddles0 icicle_core.DeviceSlice + uploadTwiddlesDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + if s.tempGPUMemPool != nil { + s.tempGPUMemPool.FreeAll() + } + host := icicle_core.HostSliceFromElements(twiddles0) + var allocErr error + dTwiddles0, allocErr = allocDeviceUninitialized(len(twiddles0)) + if allocErr != nil { + uploadTwiddlesDone <- fmt.Errorf("uploadComputeNumeratorTwiddles: %w", allocErr) + return + } + host.CopyToDevice(&dTwiddles0, false) + uploadTwiddlesDone <- nil + }) + if err := <-uploadTwiddlesDone; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dTwiddles0, nil +} + +// executeComputeNumeratorCosetIterations runs the rho coset iterations. +func (s *instance) executeComputeNumeratorCosetIterations(loopCtx *computeNumeratorLoopContext) error { + var startIterLoop time.Time + if isProfileMode { + startIterLoop = time.Now() + } + + for i := 0; i < loopCtx.rho; i++ { + if err := s.computeNumeratorIteration(i, loopCtx); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + freeSliceOnDevice(&loopCtx.dTwiddles0, &s.device) + return err + } + } + + // Free twiddles0 device slice (uploaded once before the loop). + freeTwiddlesDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + loopCtx.dTwiddles0.Free() + close(freeTwiddlesDone) + }) + <-freeTwiddlesDone + + if useBlinding { + var startRestoreBlindingPolys time.Time + if isProfileMode { + startRestoreBlindingPolys = time.Now() + } + csInv := inverseShifterProduct(loopCtx.shifters) + if err := s.restoreBlindingPolynomials(csInv); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startRestoreBlindingPolys)).Msg("computeNumerator: restore blinding polys") + } + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startIterLoop)).Msg("computeNumerator: full iteration loop") + } + return nil +} + +func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoopContext) error { + loopCtx.coset.Mul(&loopCtx.coset, &loopCtx.shifters[i]) + loopCtx.cosetExponentiatedToNMinusOne.Exp(loopCtx.coset, loopCtx.bn). + Sub(&loopCtx.cosetExponentiatedToNMinusOne, &loopCtx.one) + + batchInvertDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) + }) + if err := <-batchInvertDone; err != nil { + return err + } + + s.applyNumeratorBlindingScale(i, loopCtx) + if i == 1 { + // We have to update the scalingVector; instead of scaling by + // cosets we scale by the twiddles of the large domain. + w := s.domain1.Generator + loopCtx.scalingVector = make([]fr.Element, loopCtx.n) + fft.BuildExpTable(w, loopCtx.scalingVector) + + // Reuse memory. + copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) + fft.BitReverse(loopCtx.scalingVectorRev) + } + + // We do **a lot** of FFT here, but on the small domain. + // Note that for all the polynomials in the proving key + // (Ql, Qr, Qm, Qo, S1, S2, S3, Qcp, Qc) and ID, LOne + // we could pre-compute these rho*2 FFTs and store them + // at the cost of a huge memory footprint. + var startGpuInverseScaleForward time.Time + if isProfileMode { + startGpuInverseScaleForward = time.Now() + } + + // Inverse NTT -> Scale -> Forward NTT all on GPU using persistent GPU memory. + if err := s.gpuNTTInverseScaleForwardOnDevice(loopCtx.gpuState, loopCtx.scalingVector, loopCtx.scalingVectorRev, s.pk); err != nil { + return err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startGpuInverseScaleForward)).Msg("computeNumerator: gpuNTTInverseScaleForwardOnDevice") + } + + // Evaluate constraints on GPU. + constraintParams := gpuConstraintEvalParams{ + beta: s.beta, + gamma: s.gamma, + alpha: s.alpha, + coset: loopCtx.coset, + cosetExponentiatedToNMinusOne: loopCtx.cosetExponentiatedToNMinusOne, + cs: loopCtx.cs, + css: loopCtx.css, + cardinalityInv: s.domain0.CardinalityInv, + n: loopCtx.n, + nbBsbGates: loopCtx.nbBsbGates, + } + var startEvalConstraints time.Time + if isProfileMode { + startEvalConstraints = time.Now() + } + dNumeratorShard, err := s.gpuEvaluateConstraints( + loopCtx.gpuState, + constraintParams, + loopCtx.twiddles0, // CPU version for computeBlindingPolynomials + loopCtx.dTwiddles0, // GPU version for computeOrderingConstraint + *loopCtx.dPrecomputedDenominators, + s.bp, + nil, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startEvalConstraints)).Msg("computeNumerator: gpuEvaluateConstraints") + } + loopCtx.numeratorShards[i] = dNumeratorShard + + loopCtx.cosetExponentiatedToNMinusOne. + Inverse(&loopCtx.cosetExponentiatedToNMinusOne) + s.applyNumeratorBlindingUnscale(i, loopCtx) + return nil +} + +func (s *instance) buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx *computeNumeratorLoopContext) error { + if loopCtx == nil || loopCtx.dPrecomputedDenominators == nil { + return fmt.Errorf("computeNumerator: nil denominator device slice") + } + if loopCtx.dTwiddles0.IsEmpty() || loopCtx.dTwiddles0.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid twiddles device slice size %d, expected %d", loopCtx.dTwiddles0.Len(), loopCtx.n) + } + + if loopCtx.dPrecomputedDenominators.IsEmpty() { + dDenominators, err := allocDeviceUninitialized(loopCtx.n) + if err != nil { + return err + } + *loopCtx.dPrecomputedDenominators = dDenominators + } else if loopCtx.dPrecomputedDenominators.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid denominator device slice size %d, expected %d", loopCtx.dPrecomputedDenominators.Len(), loopCtx.n) + } + + cfg := icicle_core.DefaultVecOpsConfig() + + // dTwiddles0 is a Montgomery scalar vector in domain0 regular order. + // ScalarMulVec expects the scalar in standard form and preserves the + // Montgomery representation of the vector result. + dCosetStd := uploadScalarStdOnCurrentDevice(loopCtx.coset, cfg) + defer dCosetStd.Free() + if err := icicle_vecops.ScalarMulVec( + dCosetStd, + loopCtx.dTwiddles0, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators coset*twiddles failed: %s", err.AsString()) + } + + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneMont := uploadScalarMontOnCurrentDevice(minusOne, cfg) + defer dMinusOneMont.Free() + if err := icicle_vecops.ScalarAddVec( + dMinusOneMont, + *loopCtx.dPrecomputedDenominators, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators subtract one failed: %s", err.AsString()) + } + + if err := s.batchInvertOnCurrentDevice(*loopCtx.dPrecomputedDenominators); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: batchInvert failed: %s", err.AsString()) + } + return nil +} + +func (s *instance) applyNumeratorBlindingScale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindScale time.Time + if isProfileMode { + startBlindScale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + acc := loopCtx.cosetExponentiatedToNMinusOne + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &acc) + acc.Mul(&acc, &loopCtx.shifters[i]) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindScale)).Msg("computeNumerator: scale blinding polys") + } +} + +func (s *instance) applyNumeratorBlindingUnscale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindUnscale time.Time + if isProfileMode { + startBlindUnscale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &loopCtx.cosetExponentiatedToNMinusOne) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindUnscale)).Msg("computeNumerator: unscale blinding polys") + } +} + +func (s *instance) restoreBlindingPolynomials(csInv fr.Element) error { + for _, q := range s.bp { + if q == nil { + continue + } + cp := q.Coefficients() + if len(cp) == 0 { + continue + } + var acc fr.Element + acc.SetOne() + for i := 0; i < len(cp); i++ { + cp[i].Mul(&cp[i], &acc) + acc.Mul(&acc, &csInv) + } + } + return nil +} + +func inverseShifterProduct(shifters []fr.Element) fr.Element { + var acc fr.Element + acc.SetOne() + for i := 0; i < len(shifters); i++ { + acc.Mul(&acc, &shifters[i]) + } + acc.Inverse(&acc) + return acc +} + +func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial) (_ *iop.Polynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("downloadNumeratorFromGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("downloadNumeratorFromGPU: invalid dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard count mismatch: got %d, expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + + defer func() { + if err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + } + }() + + for i := 0; i < gpuNumerator.rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard %d is empty", i) + } + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + var dMerged icicle_core.DeviceSlice + + mergeDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + return + } + // Async boundary before merged slice is consumed by host copy. + mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + }) + if mergeErr := <-mergeDone; mergeErr != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, mergeErr + } + + cres := make([]fr.Element, totalSize) + cresHost := icicle_core.HostSliceFromElements(cres) + downloadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + }) + if err := <-downloadDone; err != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, err + } + s.putTempDeviceSlice(dMerged, totalSize) + + s.freeNumeratorShards(gpuNumerator.shards) + return iop.NewPolynomial(&cres, iop.Form{Basis: iop.LagrangeCoset, Layout: iop.BitReverse}), nil +} + +func (s *instance) freeNumeratorShards(shards []icicle_core.DeviceSlice) { + if len(shards) == 0 { + return + } + for i := 0; i < len(shards); i++ { + if !shards[i].IsEmpty() { + s.putTempDeviceSlice(shards[i], shards[i].Len()) + shards[i] = icicle_core.DeviceSlice{} + } + } +} + +func (s *instance) batchInvert(dVec icicle_core.DeviceSlice) { + if dVec.Len() == 0 { + return + } + + done := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + done <- s.batchInvertOnCurrentDevice(dVec) + }) + if err := <-done; err != icicle_runtime.Success { + panic(fmt.Sprintf("batchInvert: BatchInverseVec failed: %s", err.AsString())) + } +} + +// batchInvertOnCurrentDevice assumes caller already runs on the active device thread. +func (s *instance) batchInvertOnCurrentDevice(dVec icicle_core.DeviceSlice) icicle_runtime.EIcicleError { + if dVec.Len() == 0 { + return icicle_runtime.Success + } + err := icicle_bn254.FromMontgomery(dVec) + if err == icicle_runtime.Success { + cfg := icicle_core.DefaultVecOpsConfig() + err = icicle_vecops.BatchInverseVec(dVec, dVec, cfg) + } + if err == icicle_runtime.Success { + err = icicle_bn254.ToMontgomery(dVec) + } + return err +} + +// gpuPolysState holds GPU-resident polynomial data to avoid repeated CPU-GPU transfers. +// Use ensurePolysOnSharedGPU to populate/reuse and freeGPUPolys to release GPU memory. +type gpuPolysState struct { + deviceSlices []icicle_core.DeviceSlice + hostSlices []icicle_core.HostSlice[fr.Element] + polys []*iop.Polynomial + originalForm []iop.Form + polyToIdx map[*iop.Polynomial]int +} + +func (s *instance) sharedGPUStateInitialCap(extra int) int { + base := 16 + len(s.bp) + 2*len(s.commitmentInfo) + if extra > 0 { + base += extra + } + return base +} + +func (s *instance) initSharedGPUState() { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + if s.sharedGPUState != nil { + return + } + initialCap := s.sharedGPUStateInitialCap(0) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } +} + +func (s *instance) releaseSharedGPUState() { + s.gpuStateMu.Lock() + state := s.sharedGPUState + s.sharedGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) releaseLinearizedEvalGPUState() { + s.gpuStateMu.Lock() + state := s.linearizedEvalGPUState + s.linearizedEvalGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) freeIdleTempGPUMemoryOnDevice() { + if s == nil || s.tempGPUMemPool == nil { + return + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +func (s *instance) prepareLinearizedEvalGPUState(source *gpuPolysState) error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.clonePolysOnGPUFromState(source, s.buildLinearizedEvalGPUBatch()) + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) prepareLinearizedEvalGPUStateFromHost() error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.uploadPolysToGPUState(s.buildLinearizedEvalGPUBatch(), "prepareLinearizedEvalGPUStateFromHost") + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop.Polynomial) (*gpuPolysState, error) { + if source == nil { + return nil, fmt.Errorf("clonePolysOnGPUFromState: nil source state") + } + + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("clonePolysOnGPUFromState: empty polynomial batch") + } + + srcSlices := make([]icicle_core.DeviceSlice, len(unique)) + useSource := make([]bool, len(unique)) + for i, p := range unique { + idx, ok := source.polyToIdx[p] + if ok && idx >= 0 && idx < len(source.deviceSlices) && !source.deviceSlices[idx].IsEmpty() { + srcSlices[i] = source.deviceSlices[idx] + useSource[i] = true + } + } + + snapshot := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + snapshot.polys[i] = p + snapshot.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + snapshot.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("clonePolysOnGPUFromState") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for snapshot cloning before handing state to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i := range srcSlices { + if useSource[i] { + dst, allocErr := allocDeviceUninitialized(srcSlices[i].Len()) + if allocErr != nil { + runErr = fmt.Errorf("clonePolysOnGPUFromState: alloc failed at index %d: %w", i, allocErr) + return + } + if err := copyDeviceSliceIntoOnCurrentDevice(dst, srcSlices[i], cfg); err != icicle_runtime.Success { + _ = dst.Free() + runErr = fmt.Errorf("clonePolysOnGPUFromState: device copy failed at index %d: %s", i, err.AsString()) + return + } + snapshot.deviceSlices[i] = dst + continue + } + + coeffs := unique[i].Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("clonePolysOnGPUFromState: empty host coefficients at index %d", i) + return + } + host := icicle_core.HostSliceFromElements(coeffs) + var dst icicle_core.DeviceSlice + host.CopyToDeviceAsync(&dst, cfg.StreamHandle, true) + if dst.IsEmpty() { + runErr = fmt.Errorf("clonePolysOnGPUFromState: host upload failed at index %d", i) + return + } + snapshot.deviceSlices[i] = dst + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(snapshot) + return nil, err + } + return snapshot, nil +} + +func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) (*gpuPolysState, error) { + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("%s: empty polynomial batch", label) + } + + state := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + state.polys[i] = p + state.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + state.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label) + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i, p := range unique { + coeffs := p.Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("%s: empty host coefficients at index %d", label, i) + return + } + state.hostSlices[i] = icicle_core.HostSliceFromElements(coeffs) + state.hostSlices[i].CopyToDeviceAsync(&state.deviceSlices[i], cfg.StreamHandle, true) + if state.deviceSlices[i].IsEmpty() { + runErr = fmt.Errorf("%s: host upload failed at index %d", label, i) + return + } + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(state) + return nil, err + } + return state, nil +} + +func (s *instance) getTempDeviceSlice(n int) icicle_core.DeviceSlice { + if s == nil { + panic("getTempDeviceSlice: nil instance") + } + if s.tempGPUMemPool == nil { + panic("getTempDeviceSlice: temp GPU memory pool is not initialized") + } + return s.tempGPUMemPool.Get(n) +} + +func (s *instance) putTempDeviceSlice(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + if s != nil && s.tempGPUMemPool != nil { + s.tempGPUMemPool.Put(ds, n) + return + } + _ = ds.Free() +} + +func (s *instance) releaseTempGPUMemoryPool() { + if s == nil || s.tempGPUMemPool == nil { + return + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +// ensurePolysOnSharedGPU uploads missing polynomials once and reuses already-uploaded slices. +func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysState, error) { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if s.sharedGPUState == nil { + initialCap := s.sharedGPUStateInitialCap(len(polys)) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } + } + state := s.sharedGPUState + if state == nil { + return nil, fmt.Errorf("ensurePolysOnSharedGPU: shared GPU state is nil") + } + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int, len(polys)) + } + + newIndices := make([]int, 0, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; ok { + continue + } + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, icicle_core.DeviceSlice{}) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + newIndices = append(newIndices, idx) + } + if len(newIndices) == 0 { + return state, nil + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("ensurePolysOnSharedGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for GPU uploads in ensurePolysOnSharedGPU. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + for _, idx := range newIndices { + p := state.polys[idx] + if p == nil { + continue + } + cp := p.Coefficients() + state.hostSlices[idx] = icicle_core.HostSliceFromElements(cp) + state.hostSlices[idx].CopyToDeviceAsync(&state.deviceSlices[idx], cfg.StreamHandle, true) + } + }) + if err := <-done; err != nil { + return nil, err + } + return state, nil +} + +func getStateDeviceSlice(state *gpuPolysState, p *iop.Polynomial, label string) (icicle_core.DeviceSlice, error) { + if state == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil GPU state", label) + } + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil polynomial", label) + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): polynomial is not registered on GPU", label) + } + ds := state.deviceSlices[idx] + if ds.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): empty device slice", label) + } + return ds, nil +} + +// gpuNTTInverseScaleForwardOnDevice performs inverse NTT → scale → forward NTT +// on GPU-resident polynomial data without CPU-GPU transfers for polynomial data. +// The scaling vectors are uploaded each call (they may change between iterations). +func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scalingVector, scalingVectorRev []fr.Element, pk *ProvingKey) error { + if state == nil || len(state.polys) == 0 { + return nil + } + + device := &s.device + var scalingVectorDevice, scalingVectorRevDevice icicle_core.DeviceSlice + + // Upload scaling vectors to GPU + uploadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) + scalingHost := icicle_core.HostSliceFromElements(scalingVector) + scalingHost.CopyToDevice(&scalingVectorDevice, false) + scalingVectorRevDevice = s.getTempDeviceSlice(len(scalingVectorRev)) + scalingRevHost := icicle_core.HostSliceFromElements(scalingVectorRev) + scalingRevHost.CopyToDevice(&scalingVectorRevDevice, false) + + // Convert scaling vectors from Montgomery form to standard form + if err := icicle_bn254.FromMontgomery(scalingVectorDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVector failed: %s", err.AsString()) + return + } + if err := icicle_bn254.FromMontgomery(scalingVectorRevDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVectorRev failed: %s", err.AsString()) + return + } + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return err + } + + doneChans := make([]chan error, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + + done := make(chan error, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Step 1: Inverse NTT + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + return + } + + // Step 2: Scale by vector using GPU vecOps + vecCfg := icicle_core.DefaultVecOpsConfig() + vecCfg.StreamHandle = stream + vecCfg.IsAsync = true + + var scaleDevice icicle_core.DeviceSlice + if layout == iop.Regular { + // After KNR inverse, output is BitReverse → use scalingVectorRev + scaleDevice = scalingVectorRevDevice + } else { + // After KRN inverse, output is Regular → use scalingVector + scaleDevice = scalingVectorDevice + } + + if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + return + } + + // Step 3: Forward NTT to Lagrange + one := icicle_ntt.GetDefaultNttConfig().CosetGen + cfg.CosetGen = one + if layout == iop.Regular { + cfg.Ordering = icicle_core.KRN // BitReverse → Regular + } else { + cfg.Ordering = icicle_core.KNR // Regular → BitReverse + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + return + } + done <- nil + }) + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + if err := <-doneChans[i]; err != nil { + return err + } + } + } + + // Update polynomial metadata: final result is in Lagrange, same layout as original + for _, p := range state.polys { + if p != nil { + p.Basis = iop.Lagrange + } + } + + // Free scaling vectors from device + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + return nil +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +// +// NOTE: The prover hot path should use the state-based API and keep data on device. +// This wrapper exists for compatibility/testing where callers expect host coefficients +// to be materialized after the transform. +func (s *instance) gpuNTTInverseBatch(polys []*iop.Polynomial, pk *ProvingKey) { + if len(polys) == 0 { + return + } + state, err := s.ensurePolysOnSharedGPU(polys) + if err != nil { + panic(fmt.Sprintf("gpuNTTInverseBatch: ensurePolysOnSharedGPU failed: %v", err)) + } + + s.gpuNTTInverseBatchOnState(state, pk) + + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + for _, p := range polys { + if p == nil { + continue + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) || state.deviceSlices[idx].IsEmpty() { + continue + } + cp := p.Coefficients() + host := icicle_core.HostSliceFromElements(cp) + host.CopyFromDevice(&state.deviceSlices[idx]) + copy(cp, ([]fr.Element)(host)) + } + close(done) + }) + <-done +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKey) { + if state == nil || len(state.polys) == 0 { + return + } + + device := &s.device + doneChans := make([]chan struct{}, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + switch p.Basis { + case iop.Canonical: + continue // already in canonical form + case iop.Lagrange, iop.LagrangeCoset: + // Schedule GPU work + done := make(chan struct{}, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Select ordering and coset generator depending on basis and input layout + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + + // Base-domain inverse: + // - Regular input → KNR (output BitReverse) + // - BitReverse input → KRN (output Regular) + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + + // Run NTT inverse directly on the existing device slice (in-place) + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) + } + icicle_runtime.SynchronizeStream(stream) + + // Update metadata inside closure to avoid race + p.Basis = iop.Canonical + if layout == iop.Regular { + p.Layout = iop.BitReverse + } else { + p.Layout = iop.Regular + } + close(done) + }) + default: + panic("unsupported polynomial basis") + } + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + <-doneChans[i] + } + } +} + +// freeGPUPolys releases GPU memory for polynomial data. +func (s *instance) freeGPUPolys(state *gpuPolysState) { + if state == nil { + return + } + + device := &s.device + freeDone := make(chan struct{}) + + icicle_runtime.RunOnDevice(device, func(args ...any) { + for i := range state.polys { + if state.deviceSlices[i].IsEmpty() { + continue + } + _ = state.deviceSlices[i].Free() + } + close(freeDone) + }) + <-freeDone +} + +// gpuMemoryPool manages a pool of reusable device slices to avoid repeated allocations. +// Must be used within RunOnDevice context to ensure thread safety per device. +type gpuMemoryPool struct { + freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice + mu sync.Mutex +} + +// newGPUMemoryPool creates a new GPU memory pool. +func newGPUMemoryPool() *gpuMemoryPool { + return &gpuMemoryPool{ + freeSlices: make(map[int][]icicle_core.DeviceSlice), + } +} + +// Get returns a device slice of the specified size, either from the pool or newly allocated. +func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { + p.mu.Lock() + defer p.mu.Unlock() + + // Check if we have a free slice of this size + if slices, ok := p.freeSlices[n]; ok && len(slices) > 0 { + // Reuse the last slice + slice := slices[len(slices)-1] + p.freeSlices[n] = slices[:len(slices)-1] + return slice + } + + // No free slice available, allocate a new one. + // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + // Free all currently cached (idle) slices to reduce memory pressure. + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + panic(fmt.Sprintf("gpuMemoryPool.Get: allocation failed for size %d after clearing idle cache", n)) +} + +// Put returns a device slice to the pool for reuse instead of freeing it. +func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + + p.mu.Lock() + defer p.mu.Unlock() + + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// FreeAll releases all pooled device slices. +func (p *gpuMemoryPool) FreeAll() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) +} + +// allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. +// Use when the destination is fully overwritten by a kernel. +func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { + var ds icicle_core.DeviceSlice + if _, err := ds.Malloc(int(unsafe.Sizeof(fr.Element{})), n); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("allocDeviceUninitialized: malloc failed for size %d: %s", n, err.AsString()) + } + return ds, nil +} + +// mustAllocDeviceUninitialized is like allocDeviceUninitialized but panics on failure. +// Use only in contexts where error propagation is impractical (e.g. upload helpers). +func mustAllocDeviceUninitialized(n int) icicle_core.DeviceSlice { + ds, err := allocDeviceUninitialized(n) + if err != nil { + panic(err) + } + return ds +} + +// freeDeviceSlice frees a device slice if non-empty and zeroes the pointer. +// Use for directly-allocated slices, NOT pool-allocated ones (use putTempDeviceSlice for those). +func freeDeviceSlice(ds *icicle_core.DeviceSlice) { + if ds != nil && !ds.IsEmpty() { + _ = ds.Free() + *ds = icicle_core.DeviceSlice{} + } +} + +// freeSliceOnDevice frees a device slice on the specified device and blocks +// until complete. Use outside RunOnDevice closures. Zeroes the slice after freeing. +func freeSliceOnDevice(ds *icicle_core.DeviceSlice, device *icicle_runtime.Device) { + if ds == nil || ds.IsEmpty() { + return + } + d := *ds + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + _ = d.Free() + close(done) + }) + <-done + *ds = icicle_core.DeviceSlice{} +} + +// copyDeviceSliceIntoOnCurrentDevice copies src into dst entirely on GPU. +func copyDeviceSliceIntoOnCurrentDevice( + dst, src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if src.IsEmpty() || src.Len() <= 0 || dst.IsEmpty() || dst.Len() < src.Len() { + return icicle_runtime.InvalidArgument + } + src.CheckDevice() + dst.CheckDevice() + + srcElemSize := src.SizeOfElement() + dstElemSize := dst.SizeOfElement() + if srcElemSize <= 0 || dstElemSize <= 0 || srcElemSize != dstElemSize { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(src.Len() * srcElemSize) + if cfg.IsAsync { + return icicle_runtime.CopyAsync(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen, cfg.StreamHandle) + } + _, err := icicle_runtime.Copy(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen) + return err +} + +// zeroDeviceSliceOnCurrentDevice zero-fills dst entirely on GPU. +func zeroDeviceSliceOnCurrentDevice( + dst icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if dst.IsEmpty() || dst.Len() <= 0 { + return icicle_runtime.InvalidArgument + } + dst.CheckDevice() + + elemSize := dst.SizeOfElement() + if elemSize <= 0 { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(dst.Len() * elemSize) + if cfg.IsAsync { + return icicle_runtime.MemSetAsync(dst.AsUnsafePointer(), 0, byteLen, cfg.StreamHandle) + } + return icicle_runtime.MemSet(dst.AsUnsafePointer(), 0, byteLen) +} + +func createAsyncVecOpsConfigOnCurrentDevice(label string) (icicle_core.VecOpsConfig, icicle_runtime.Stream, error) { + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + return icicle_core.VecOpsConfig{}, nil, fmt.Errorf("%s: create stream failed: %s", label, eStream.AsString()) + } + cfg := icicle_core.DefaultVecOpsConfig() + cfg.StreamHandle = stream + cfg.IsAsync = true + return cfg, stream, nil +} + +func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label string) error { + if stream == nil { + return nil + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + _ = icicle_runtime.DestroyStream(stream) + return fmt.Errorf("%s: synchronize stream failed: %s", label, eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success { + return fmt.Errorf("%s: destroy stream failed: %s", label, eDestroy.AsString()) + } + return nil +} + +// makeFinisher returns a closure that synchronizes and destroys the stream, +// then sends the (possibly merged) error to done. Use inside RunOnDevice closures. +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { + return func(runErr error) { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + } +} + +// uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. +// Use this for additions where the vector is already in Montgomery form. +func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + return ds +} + +func uploadScalarMontOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + if cfg.IsAsync { + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&ds, false) + } + return ds +} + +// uploadScalarStd uploads a scalar in STANDARD form (not Montgomery) as a single-element device slice. +// For use with ScalarMulVec: (a*R) * b_std = (a*b)*R +func uploadScalarStd(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + // Convert from Montgomery form to standard form + if err := icicle_bn254.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +func uploadScalarStdOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + ds := uploadScalarMontOnCurrentDevice(scalar, cfg) + // Fallback to sync conversion for compatibility with ICICLE wrappers + // that do not expose *_WithConfig APIs. + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bn254.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +// uploadVectorStd uploads a vector and converts to standard form. +func uploadVectorStd(vec []fr.Element) icicle_core.DeviceSlice { + ds := mustAllocDeviceUninitialized(len(vec)) + uploadVectorStdInto(&ds, vec) + return ds +} + +// uploadVectorStdInto uploads vec into an existing device slice and converts it to standard form. +// The destination must already be allocated with enough capacity for len(vec) elements. +func uploadVectorStdInto(dst *icicle_core.DeviceSlice, vec []fr.Element) { + cfg := icicle_core.DefaultVecOpsConfig() + uploadVectorStdIntoOnCurrentDevice(dst, vec, cfg) +} + +// uploadVectorStdIntoOnCurrentDevice uploads vec into an existing device slice and converts it +// to standard form while honoring the provided vector-op config/stream. +func uploadVectorStdIntoOnCurrentDevice( + dst *icicle_core.DeviceSlice, + vec []fr.Element, + cfg icicle_core.VecOpsConfig, +) { + host := icicle_core.HostSliceFromElements(vec) + if cfg.IsAsync { + host.CopyToDeviceAsync(dst, cfg.StreamHandle, false) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } else { + host.CopyToDevice(dst, false) + } + // Convert from Montgomery form to standard form + if err := icicle_bn254.FromMontgomery(*dst); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +// uploadVector uploads a vector (keeps Montgomery form for additions). +func uploadVector(vec []fr.Element) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(len(vec)) + host.CopyToDevice(&ds, false) + return ds +} + +// uploadInt64Vector uploads int64 indices to a device slice. +func uploadInt64Vector(vec []int64) icicle_core.DeviceSlice { + cfg := icicle_core.DefaultVecOpsConfig() + return uploadInt64VectorOnCurrentDevice(vec, cfg) +} + +func uploadInt64VectorOnCurrentDevice(vec []int64, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + var ds icicle_core.DeviceSlice + if cfg.IsAsync { + if _, err := ds.MallocAsync(int(unsafe.Sizeof(int64(0))), len(vec), cfg.StreamHandle); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc async failed: %s", err.AsString())) + } + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + return ds + } + if _, err := ds.Malloc(int(unsafe.Sizeof(int64(0))), len(vec)); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc failed: %s", err.AsString())) + } + host.CopyToDevice(&ds, false) + return ds +} + +// toStandardFormInPlace converts a device slice to standard form in-place (modifies the source). +// Use this for temporary vectors that won't be needed in Montgomery form. +func toStandardFormInPlace(src icicle_core.DeviceSlice) { + cfg := icicle_core.DefaultVecOpsConfig() + toStandardFormInPlaceWithCfg(src, cfg) +} + +func toStandardFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bn254.FromMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +func toMontgomeryFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bn254.ToMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("ToMontgomery failed: %s", err.AsString())) + } +} + +// multiplyMontgomerySlices multiplies two device slices that are both in Montgomery form. +// It creates a copy of dSlice1Mont, converts the copy to standard form, and then multiplies +// it with dSlice2Mont (which remains in Montgomery form). The result is stored in dResult +// and will be in Montgomery form. +// +// Parameters: +// - dSlice1Mont: first device slice in Montgomery form (not modified) +// - dSlice2Mont: second device slice in Montgomery form (not modified) +// - dResult: destination device slice for the result (must be pre-allocated) +// - state: GPU state with memory pool and vector configuration +// - n: size of the slices +func multiplyMontgomerySlices( + dSlice1Mont, dSlice2Mont icicle_core.DeviceSlice, + dResult icicle_core.DeviceSlice, + state *gpuConstraintEvalState, + vecCfg icicle_core.VecOpsConfig, + n int, +) error { + // Copy dSlice1Mont to standard form + dSlice1Std := state.getTempDeviceSlice(n) + defer state.putTempDeviceSlice(dSlice1Std, n) + + if err := copyDeviceSliceIntoOnCurrentDevice(dSlice1Std, dSlice1Mont, vecCfg); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: device copy failed: %s", err.AsString()) + } + toStandardFormInPlace(dSlice1Std) + + // Multiply: dSlice1Std (standard) * dSlice2Mont (Montgomery) = dResult (Montgomery) + if err := icicle_vecops.VecOp(dSlice1Std, dSlice2Mont, dResult, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: VecOp multiplication failed: %s", err.AsString()) + } + return nil +} + +// gpuConstraintEvalParams holds parameters for GPU constraint evaluation +type gpuConstraintEvalParams struct { + beta fr.Element + gamma fr.Element + alpha fr.Element + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + cs fr.Element // domain1.FrMultiplicativeGen + css fr.Element // cs^2 + cardinalityInv fr.Element + n int + nbBsbGates int +} + +// gpuConstraintEvalState holds intermediate state during constraint evaluation +type gpuConstraintEvalState struct { + // Polynomial device slices (may point to gpuState or allocated buffers) + dL, dR, dO, dZ, dZS icicle_core.DeviceSlice + dQl, dQr, dQm, dQo, dQk icicle_core.DeviceSlice + dS1, dS2, dS3 icicle_core.DeviceSlice + // Intermediate results + dGate, dOrdering, dLocal, dResult icicle_core.DeviceSlice + // Scalar device slices + dGammaScalar icicle_core.DeviceSlice + // Configuration + vecCfg icicle_core.VecOpsConfig + // Helper function to get device slices + getDeviceSlice func(int) icicle_core.DeviceSlice + // Shared prover-level temporary GPU memory pool accessors + getTempDeviceSlice func(int) icicle_core.DeviceSlice + putTempDeviceSlice func(icicle_core.DeviceSlice, int) + // Track allocated polynomial buffers for automatic cleanup + allocatedPolyBuffers []struct { + slice icicle_core.DeviceSlice + size int + } +} + +// allocate allocates a new device slice from the memory pool and tracks it for automatic cleanup. +// Returns the allocated device slice. +func (s *gpuConstraintEvalState) allocate(size int) icicle_core.DeviceSlice { + slice := s.getTempDeviceSlice(size) + s.allocatedPolyBuffers = append(s.allocatedPolyBuffers, struct { + slice icicle_core.DeviceSlice + size int + }{slice, size}) + return slice +} + +// freeAllocatedPolyBuffers returns all allocated polynomial buffers to the memory pool. +// This should be called during cleanup to free all buffers allocated via allocate(). +func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { + for _, buf := range s.allocatedPolyBuffers { + s.putTempDeviceSlice(buf.slice, buf.size) + } + s.allocatedPolyBuffers = nil +} + +// computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. +// Returns device slices for the blinding polynomials. +func computeBlindingPolynomials( + n int, + twiddles0 []fr.Element, + bp []*iop.Polynomial, +) (dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice) { + blindL := make([]fr.Element, n) + blindR := make([]fr.Element, n) + blindO := make([]fr.Element, n) + blindZ := make([]fr.Element, n) + blindZS := make([]fr.Element, n) // ZS uses shifted index + + // TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove this. Since we don't normally use blindings, we will + // not make this optimization. + utils.Parallelize(n, func(start, end int) { + for i := start; i < end; i++ { + blindL[i] = bp[id_Bl].Evaluate(twiddles0[i]) + blindR[i] = bp[id_Br].Evaluate(twiddles0[i]) + blindO[i] = bp[id_Bo].Evaluate(twiddles0[i]) + blindZ[i] = bp[id_Bz].Evaluate(twiddles0[i]) + blindZS[i] = bp[id_Bz].Evaluate(twiddles0[(i+1)%n]) + } + }) + + dBlindL = uploadVector(blindL) + dBlindR = uploadVector(blindR) + dBlindO = uploadVector(blindO) + dBlindZ = uploadVector(blindZ) + dBlindZS = uploadVector(blindZS) + + return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS +} + +// applyBlindingToPolynomials applies blinding to polynomials L, R, O, Z, ZS. +// Allocates new buffers for L, R, O, Z (tracked for cleanup) and modifies ZS in-place. +// The original slices in gpuState remain unchanged. +func applyBlindingToPolynomials( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice, +) error { + // L' = L + blindL (allocate new buffer, tracked for cleanup) + dLBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dL, dBlindL, dLBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add L failed: %s", err.AsString()) + } + state.dL = dLBlinded + + // R' = R + blindR (allocate new buffer, tracked for cleanup) + dRBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dR, dBlindR, dRBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add R failed: %s", err.AsString()) + } + state.dR = dRBlinded + + // O' = O + blindO (allocate new buffer, tracked for cleanup) + dOBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dO, dBlindO, dOBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add O failed: %s", err.AsString()) + } + state.dO = dOBlinded + + // Z' = Z + blindZ (allocate new buffer, tracked for cleanup) + dZBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dZ, dBlindZ, dZBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add Z failed: %s", err.AsString()) + } + state.dZ = dZBlinded + + // ZS' = ZS + blindZS + // Note: dZS is a temporary buffer created inside gpuEvaluateConstraints, + // so it's safe to modify it in-place. + if err := icicle_vecops.VecOp(state.dZS, dBlindZS, state.dZS, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) + } + + // Free blinding vectors - no longer needed after creating blinded polynomials + dBlindL.Free() + dBlindR.Free() + dBlindO.Free() + dBlindZ.Free() + dBlindZS.Free() + return nil +} + +// scaleSVectorsByBeta scales S1, S2, S3 by beta. +// Allocates new buffers for S1, S2, S3 (tracked for cleanup). +func scaleSVectorsByBeta( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, +) error { + // S1' = S1 * beta (need to scale S1, S2, S3 by beta for ordering constraint) + // Use standard form for beta so: S1_mont * beta_std = (S1*beta)_mont + dBetaStd := uploadScalarStd(params.beta) + + // S1' = S1 * beta (allocate new buffer, tracked for cleanup) + dS1Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dS1Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S1 failed: %s", err.AsString()) + } + state.dS1 = dS1Scaled + + // S2' = S2 * beta (allocate new buffer, tracked for cleanup) + dS2Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dS2Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S2 failed: %s", err.AsString()) + } + state.dS2 = dS2Scaled + + // S3' = S3 * beta (allocate new buffer, tracked for cleanup) + dS3Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dS3Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S3 failed: %s", err.AsString()) + } + state.dS3 = dS3Scaled + + // Free dBetaStd - no longer needed after scaling S vectors + dBetaStd.Free() + return nil +} + +// computeGateConstraint computes the gate constraint. +// Returns dGate device slice. +func computeGateConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // gate = Ql*L' + Qr*R' + Qm*L'*R' + Qo*O' + Qk + sum(Qci*Pi) + // We use multiplyMontgomerySlices for all poly×poly multiplications. + // Note: dL, dR, dO are used later in ordering constraint, so we preserve them. + + dGate := state.getTempDeviceSlice(params.n) + dTmp := state.getTempDeviceSlice(params.n) + + // Ql * L' + if err := multiplyMontgomerySlices(state.dQl, state.dL, dGate, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Ql*L: %w", err) + } + + // + Qr * R' + if err := multiplyMontgomerySlices(state.dQr, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qr*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qr*R failed: %s", err.AsString()) + } + + // + Qm * L' * R' (need two multiplications) + // First: Qm * L' = dTmp (Montgomery) + if err := multiplyMontgomerySlices(state.dQm, state.dL, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qm*L: %w", err) + } + // Second: dTmp (Montgomery) * R' (Montgomery) + if err := multiplyMontgomerySlices(dTmp, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: (Qm*L)*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qm*L*R failed: %s", err.AsString()) + } + + // + Qo * O' + if err := multiplyMontgomerySlices(state.dQo, state.dO, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qo*O: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qo*O failed: %s", err.AsString()) + } + + // + Qk + if err := icicle_vecops.VecOp(dGate, state.dQk, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qk failed: %s", err.AsString()) + } + + // + BSB gates: sum(Qci[2*i] * Qci[2*i+1]) + for i := 0; i < params.nbBsbGates; i++ { + origQci0 := state.getDeviceSlice(id_Qci + 2*i) + origQci1 := state.getDeviceSlice(id_Qci + 2*i + 1) + if !origQci0.IsEmpty() && !origQci1.IsEmpty() { + // Use helper to multiply Qci0 * Qci1 without modifying original values + if err := multiplyMontgomerySlices(origQci0, origQci1, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qci[%d]: %w", i, err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qci[%d] failed: %s", i, err.AsString()) + } + } + } + + // Return temporary buffer to pool - no longer needed after Step 3 + state.putTempDeviceSlice(dTmp, params.n) + + return dGate, nil +} + +// computeOrderingConstraint computes the ordering constraint. +// Returns dOrdering device slice. +func computeOrderingConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dTwiddles0 icicle_core.DeviceSlice, // twiddles0 already on GPU + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // This is complex: involves ID computation, gamma, beta, Z, ZS, S1, S2, S3 + // id = twiddles[i] * coset * beta + // a = gamma + L' + id + // b = gamma + R' + id*cs + // c = gamma + O' + id*css + // r = a * b * c * Z' + // + // a2 = gamma + L' + S1*beta + // b2 = gamma + R' + S2*beta + // c2 = gamma + O' + S3*beta + // l = a2 * b2 * c2 * ZS' + // + // ordering = l - r + + // Compute ID vector: twiddles * coset * beta (computed on GPU) + // dTwiddles0 is already on GPU (passed as parameter, don't free it here) + + // Compute coset * beta on CPU, then upload as scalar in standard form + var cosetTimesBeta fr.Element + cosetTimesBeta.Mul(¶ms.coset, ¶ms.beta) + dCosetTimesBetaStd := uploadScalarStd(cosetTimesBeta) + dBetaStd := uploadScalarStd(params.beta) + + // Multiply twiddles0 by cosetTimesBeta on GPU: dID = (cosetTimesBeta * twiddles0) * R (Montgomery form) + dID := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dCosetTimesBetaStd, dTwiddles0, dID, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec coset*beta*twiddles failed: %s", err.AsString()) + } + + // Free temporary device slice (dTwiddles0 is owned by caller, don't free it) + dCosetTimesBetaStd.Free() + + // id * cs - use standard form for cs + dIDcs := state.getTempDeviceSlice(params.n) + dCsStd := uploadScalarStd(params.cs) + if err := icicle_vecops.ScalarMulVec(dCsStd, dID, dIDcs, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*cs failed: %s", err.AsString()) + } + + // id * css - use standard form for css + dIDcss := state.getTempDeviceSlice(params.n) + dCssStd := uploadScalarStd(params.css) + if err := icicle_vecops.ScalarMulVec(dCssStd, dID, dIDcss, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*css failed: %s", err.AsString()) + } + + // a = gamma + L' + id (dL now contains L' after in-place blinding) + dA := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dID, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a+id failed: %s", err.AsString()) + } + + // b = gamma + R' + id*cs (dR now contains R' after in-place blinding) + dB := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dIDcs, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b+id*cs failed: %s", err.AsString()) + } + + // c = gamma + O' + id*css (dO now contains O' after in-place blinding) + dC := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dIDcss, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c+id*css failed: %s", err.AsString()) + } + + // Return to pool: dID, dIDcs, dIDcss - no longer needed after computing a, b, c + state.putTempDeviceSlice(dID, params.n) + state.putTempDeviceSlice(dIDcs, params.n) + state.putTempDeviceSlice(dIDcss, params.n) + dCsStd.Free() + dCssStd.Free() + + // r = a * b * c * Z' (dZ now contains Z' after in-place blinding) + // For chain multiplication, convert operands to std form in-place when possible + dR_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a*b failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dR_ord, dC, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*c failed: %s", err.AsString()) + } + // Convert dR_ord to standard form in-place (temporary result, dZ needs to be preserved) + toStandardFormInPlace(dR_ord) + if err := icicle_vecops.VecOp(dR_ord, state.dZ, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*Z failed: %s", err.AsString()) + } + + // Reuse dA, dB, dC for a2, b2, c2 instead of freeing and reallocating. + // To reduce peak memory, we scale S vectors by beta on-demand through a single temp buffer. + dScaledS := state.getTempDeviceSlice(params.n) + + // a2 = gamma + L' + S1*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L (a2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S1 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dScaledS, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2+S1*beta failed: %s", err.AsString()) + } + + // b2 = gamma + R' + S2*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R (b2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S2 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dScaledS, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b2+S2*beta failed: %s", err.AsString()) + } + + // c2 = gamma + O' + S3*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O (c2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S3 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dScaledS, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c2+S3*beta failed: %s", err.AsString()) + } + + // Free dGammaScalar - no longer needed after computing a2, b2, c2 + // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. + state.dGammaScalar.Free() + state.putTempDeviceSlice(dScaledS, params.n) + dBetaStd.Free() + + // l = a2 * b2 * c2 * ZS' (dZS now contains ZS' after in-place blinding) + dL_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2*b2 failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dL_ord, dC, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*c2 failed: %s", err.AsString()) + } + // Convert dZS to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(state.dZS) + if err := icicle_vecops.VecOp(dL_ord, state.dZS, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*ZS failed: %s", err.AsString()) + } + + // Return temporary buffers to pool - no longer needed after computing l + state.putTempDeviceSlice(dA, params.n) + state.putTempDeviceSlice(dB, params.n) + state.putTempDeviceSlice(dC, params.n) + state.putTempDeviceSlice(state.dZS, params.n) + state.dZS = icicle_core.DeviceSlice{} + + // ordering = l - r, reuse dL_ord as the final ordering vector + if err := icicle_vecops.VecOp(dL_ord, dR_ord, dL_ord, vecCfg, icicle_core.Sub); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l-r failed: %s", err.AsString()) + } + + // Return dR_ord to pool - no longer needed after computing ordering + state.putTempDeviceSlice(dR_ord, params.n) + + // Return dL_ord as ordering (caller is responsible for freeing) + return dL_ord, nil +} + +// computeLocalConstraint computes the local constraint. +// Returns dLocal device slice. +func computeLocalConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dPrecomputedDenominators icicle_core.DeviceSlice, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // local = (Z' - 1) * LagrangeOne + // where LagrangeOne[i] = cosetExpMinusOne * cardinalityInv / (coset*twiddles0[i] - 1) + + if dPrecomputedDenominators.IsEmpty() || dPrecomputedDenominators.Len() < params.n { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: invalid denominator device slice size %d, expected at least %d", dPrecomputedDenominators.Len(), params.n) + } + + // Compute LagrangeOne on device. dPrecomputedDenominators is already in + // Montgomery form after batch inversion; ScalarMulVec expects the scalar in + // standard form and preserves a Montgomery vector result. + var lagrangeCoeff fr.Element + lagrangeCoeff.Mul(¶ms.cosetExponentiatedToNMinusOne, ¶ms.cardinalityInv) + dLagrangeCoeffStd := uploadScalarStdOnCurrentDevice(lagrangeCoeff, vecCfg) + dLagrangeOneStd := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dLagrangeCoeffStd, dPrecomputedDenominators, dLagrangeOneStd, vecCfg); err != icicle_runtime.Success { + dLagrangeCoeffStd.Free() + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarMulVec lagrangeOne failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dLagrangeOneStd, vecCfg) + dLagrangeCoeffStd.Free() + + // Z' - 1 using ScalarAddVec with minus one + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneScalar := uploadScalarMont(minusOne) + + dZMinusOne := state.getTempDeviceSlice(params.n) + // dZ now contains Z' after in-place blinding + if err := icicle_vecops.ScalarAddVec(dMinusOneScalar, state.dZ, dZMinusOne, vecCfg); err != icicle_runtime.Success { + dMinusOneScalar.Free() + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarAddVec Z-1 failed: %s", err.AsString()) + } + + // Free dMinusOneScalar - no longer needed after computing Z' - 1 + // Note: dZ is part of gpuState and will be freed later + dMinusOneScalar.Free() + + // local = (Z' - 1) * LagrangeOne_std + dLocal := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.VecOp(dZMinusOne, dLagrangeOneStd, dLocal, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: VecOp (Z-1)*LagrangeOne failed: %s", err.AsString()) + } + + // Return temporary buffers to pool + state.putTempDeviceSlice(dZMinusOne, params.n) + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + + return dLocal, nil +} + +// createGetDeviceSliceFunc creates a function to get device slices for polynomials. +// It returns a function that maps polynomial indices to their device slices. +func (s *instance) polyByID(polyIdx int) *iop.Polynomial { + switch polyIdx { + case id_L: + return s.polyL + case id_R: + return s.polyR + case id_O: + return s.polyO + case id_Z: + return s.polyZ + case id_ZS: + return s.polyZS + case id_Ql: + return s.trace.Ql + case id_Qr: + return s.trace.Qr + case id_Qm: + return s.trace.Qm + case id_Qo: + return s.trace.Qo + case id_Qk: + return s.polyQk + case id_S1: + return s.trace.S1 + case id_S2: + return s.trace.S2 + case id_S3: + return s.trace.S3 + default: + if polyIdx < id_Qci { + return nil + } + offset := polyIdx - id_Qci + i := offset / 2 + if i < 0 { + return nil + } + if offset%2 == 0 { + if i < len(s.trace.Qcp) { + return s.trace.Qcp[i] + } + return nil + } + if i < len(s.cCommitments) { + return s.cCommitments[i] + } + return nil + } +} + +func createGetDeviceSliceFunc( + gpuState *gpuPolysState, + polyToIdx map[*iop.Polynomial]int, + resolvePoly func(int) *iop.Polynomial, +) func(int) icicle_core.DeviceSlice { + return func(polyIdx int) icicle_core.DeviceSlice { + p := resolvePoly(polyIdx) + if p == nil { + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d is nil", polyIdx)) + } + if idx, ok := polyToIdx[p]; ok { + return gpuState.deviceSlices[idx] + } + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d (ptr=%p) not found in polyToIdx (map has %d entries)", polyIdx, p, len(polyToIdx))) + } +} + +// initializeConstraintEvalState initializes the GPU constraint evaluation state. +// It sets up all device slices. +// The slices in gpuState are treated as read-only; helper functions will allocate +// separate working buffers whenever they need to modify data. +func initializeConstraintEvalState( + getDeviceSlice func(int) icicle_core.DeviceSlice, + getTempDeviceSlice func(int) icicle_core.DeviceSlice, + putTempDeviceSlice func(icicle_core.DeviceSlice, int), +) *gpuConstraintEvalState { + vecCfg := icicle_core.DefaultVecOpsConfig() + + state := &gpuConstraintEvalState{ + dL: getDeviceSlice(id_L), + dR: getDeviceSlice(id_R), + dO: getDeviceSlice(id_O), + dZ: getDeviceSlice(id_Z), + dQl: getDeviceSlice(id_Ql), + dQr: getDeviceSlice(id_Qr), + dQm: getDeviceSlice(id_Qm), + dQo: getDeviceSlice(id_Qo), + dQk: getDeviceSlice(id_Qk), + dS1: getDeviceSlice(id_S1), + dS2: getDeviceSlice(id_S2), + dS3: getDeviceSlice(id_S3), + vecCfg: vecCfg, + getDeviceSlice: getDeviceSlice, + getTempDeviceSlice: getTempDeviceSlice, + putTempDeviceSlice: putTempDeviceSlice, + } + + return state +} + +// gpuEvaluateConstraints evaluates all PLONK constraints on GPU. +// It takes polynomials already on GPU (via gpuState), computes blinding polynomial evaluations, +// and evaluates gate, ordering, and local constraints entirely on GPU. +// If result is non-nil, it downloads into result and returns an empty device slice. +// If result is nil, it returns a persistent device slice with the result. +// TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove 'twiddles0'. Since we don't +// normally use blindings, we will not make this optimization. +func (s *instance) gpuEvaluateConstraints( + gpuState *gpuPolysState, + params gpuConstraintEvalParams, + twiddles0 []fr.Element, // CPU vector for computeBlindingPolynomials + dTwiddles0 icicle_core.DeviceSlice, // GPU vector for computeOrderingConstraint + dPrecomputedDenominators icicle_core.DeviceSlice, + bp []*iop.Polynomial, // blinding polynomials (already scaled for this iteration) + result []fr.Element, +) (icicle_core.DeviceSlice, error) { + if gpuState == nil || len(gpuState.polys) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("gpuState is nil or empty") + } + + n := params.n + device := &s.device + + // Create a map from polynomial to its index in gpuState + polyToIdx := make(map[*iop.Polynomial]int) + for i, p := range gpuState.polys { + if p != nil { + polyToIdx[p] = i + } + } + + // Get device slices for the polynomials we need. + getDeviceSlice := createGetDeviceSliceFunc(gpuState, polyToIdx, s.polyByID) + + done := make(chan error, 1) + var resultOnDevice icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(device, func(args ...any) { + state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). + // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). + state.dGammaScalar = uploadScalarMont(params.gamma) + + state.dZS = state.getTempDeviceSlice(n) + if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + return + } + + // Step 1: Compute and apply blinding polynomial evaluations (if enabled) + if useBlinding { + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + done <- err + return + } + } + + // Step 2-4: Compute gate, ordering, and local constraints sequentially on a + // single synchronous stream, folding them into dResult as + // gate + alpha*ordering + alpha^2*local. Computing one family at a time + // keeps peak device allocation low, and sequential is not a compromise: + // the family kernels are memory-bandwidth-bound and each already saturates + // the device, so the parallel three-stream variant this replaces measured + // identical timings (111ms/iteration at n=2^23) — while racing on the + // shared temp-slice pool and lazily materialized inputs (it corrupted the + // numerator at every circuit size). + seqVecCfg := state.vecCfg + seqVecCfg.IsAsync = false + + // Compute ordering first to minimize peak memory before gate/local allocations. + var err error + state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + return + } + + // dResult = alpha * ordering + state.dResult = state.getTempDeviceSlice(params.n) + dAlphaStd := uploadScalarStd(params.alpha) + if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { + dAlphaStd.Free() + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + return + } + dAlphaStd.Free() + state.putTempDeviceSlice(state.dOrdering, params.n) + + // dResult += alpha^2 * local + state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + return + } + var alphaSquared fr.Element + alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) + dAlphaSquaredStd := uploadScalarStd(alphaSquared) + dTmp := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + return + } + dAlphaSquaredStd.Free() + if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dLocal, params.n) + state.putTempDeviceSlice(dTmp, params.n) + + // dResult += gate + state.dGate, err = computeGateConstraint(state, params, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + return + } + if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dGate, params.n) + + // Step 5: materialize result either on host or as a persistent device slice. + if result != nil { + resultHost := icicle_core.HostSliceFromElements(result) + resultHost.CopyFromDevice(&state.dResult) + } else { + resultOnDevice = s.getTempDeviceSlice(params.n) + if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // Return dResult pool slice after materialization. + state.putTempDeviceSlice(state.dResult, params.n) + + // Return all allocated polynomial buffers to the pool. + // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). + // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are + // returned to pool inside computeGateConstraint. dQk and the original gpuState slices + // are owned by gpuState and will be freed separately. + state.freeAllocatedPolyBuffers() + + done <- nil + }) + + err := <-done + + if err != nil { + if !resultOnDevice.IsEmpty() { + s.putTempDeviceSlice(resultOnDevice, resultOnDevice.Len()) + } + return icicle_core.DeviceSlice{}, err + } + return resultOnDevice, nil +} + +func evaluateBlinded(p, bp *iop.Polynomial, zeta fr.Element) fr.Element { + // Get the size of the polynomial + n := big.NewInt(int64(p.Size())) + + var pEvaluatedAtZeta fr.Element + + // Evaluate the polynomial and blinded polynomial at zeta + chP := make(chan struct{}, 1) + go func() { + pEvaluatedAtZeta = p.Evaluate(zeta) + close(chP) + }() + + bpEvaluatedAtZeta := bp.Evaluate(zeta) + + // Multiply the evaluated blinded polynomial by tempElement + var t fr.Element + one := fr.One() + t.Exp(zeta, n).Sub(&t, &one) + bpEvaluatedAtZeta.Mul(&bpEvaluatedAtZeta, &t) + + // Add the evaluated polynomial and the evaluated blinded polynomial + <-chP + pEvaluatedAtZeta.Add(&pEvaluatedAtZeta, &bpEvaluatedAtZeta) + + // Return the result + return pEvaluatedAtZeta +} + +// /!\ modifies the size +func getBlindedCoefficients(p, bp *iop.Polynomial) []fr.Element { + cp := p.Coefficients() + cbp := bp.Coefficients() + cp = append(cp, cbp...) + for i := 0; i < len(cbp); i++ { + cp[i].Sub(&cp[i], &cbp[i]) + } + return cp +} + +// getNonBlindedCoefficients returns a padded copy of polynomial coefficients +// to match the size they would have with blinding enabled. +// The padding size is blindingOrder+1 (e.g., order 2 → 3 coefficients). +func getNonBlindedCoefficients(p *iop.Polynomial, blindingOrder int) []fr.Element { + cp := p.Coefficients() + padded := make([]fr.Element, len(cp)+blindingOrder+1) + copy(padded, cp) + return padded +} + +// commits to a polynomial of the form b*(Xⁿ-1) where b is of small degree +func commitBlindingFactor(n int, b *iop.Polynomial, key kzg.ProvingKey) curve.G1Affine { + cp := b.Coefficients() + np := b.Size() + + // lo + var tmp curve.G1Affine + tmp.MultiExp(key.G1[:np], cp, ecc.MultiExpConfig{}) + + // hi + var res curve.G1Affine + res.MultiExp(key.G1[n:n+np], cp, ecc.MultiExpConfig{}) + res.Sub(&res, &tmp) + return res +} + +// return a random polynomial of degree n, if n==-1 cancel the blinding +func getRandomPolynomial(n int) *iop.Polynomial { + var a []fr.Element + if n == -1 { + a := make([]fr.Element, 1) + a[0].SetZero() + } else { + a = make([]fr.Element, n+1) + for i := 0; i <= n; i++ { + a[i].SetRandom() + } + } + res := iop.NewPolynomial(&a, iop.Form{ + Basis: iop.Canonical, Layout: iop.Regular}) + return res +} + +func coefficients(p []*iop.Polynomial) [][]fr.Element { + res := make([][]fr.Element, len(p)) + for i, pI := range p { + res[i] = pI.Coefficients() + } + return res +} + +func (s *instance) freeGPUQuotient(quotient *gpuQuotientPolynomial) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return + } + s.putTempDeviceSlice(quotient.coeffs, quotient.coeffs.Len()) + quotient.coeffs = icicle_core.DeviceSlice{} + quotient.size = 0 +} + +// commitToQuotientGPUFromDevice commits H1/H2/H3 directly from device memory. +// For StatisticalZK=true we materialize adjusted device vectors for h1/h2/h3 +// and commit those without downloading quotient coefficients to host. +// prepareStatisticalZKQuotientShards constructs blinded quotient polynomial +// shards h1, h2, h3 on the GPU for the Statistical ZK path. Each shard is +// randomized so that the quotient split h = h1 + X^(n+2)*h2 + X^(2(n+2))*h3 +// hides the original polynomial. +// +// Caller is responsible for returning dH1, dH2, dH3 to the temp pool: +// - dH1 and dH2 have size nPlus2+1 +// - dH3 has size nPlus2 +func (s *instance) prepareStatisticalZKQuotientShards( + h1Device, h2Device, h3Device icicle_core.DeviceSlice, + nPlus2 int, +) (dH1, dH2, dH3 icicle_core.DeviceSlice, err error) { + nPlus3 := nPlus2 + 1 + + prepareDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + dH1 = s.getTempDeviceSlice(nPlus3) + dH2 = s.getTempDeviceSlice(nPlus3) + dH3 = s.getTempDeviceSlice(nPlus2) + + // h1 = base h1 with extra randomizer coefficient at degree n+2. + dH1Prefix := (&dH1).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH1Prefix, h1Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h1 failed: %s", e.AsString()) + return + } + dH1Tail := (&dH1).Range(nPlus2, nPlus3, false) + r0Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[0]}) + r0Host.CopyToDevice(&dH1Tail, false) + + // h2 = base h2 with first coefficient adjusted by -r0 and tail = r1. + dH2Prefix := (&dH2).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH2Prefix, h2Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h2 failed: %s", e.AsString()) + return + } + dH2First := (&dH2).Range(0, 1, false) + var negR0 fr.Element + negR0.Neg(&s.quotientShardsRandomizers[0]) + dNegR0 := uploadScalarMont(negR0) + if e := icicle_vecops.ScalarAddVec(dNegR0, dH2First, dH2First, cfg); e != icicle_runtime.Success { + _ = dNegR0.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h2[0] failed: %s", e.AsString()) + return + } + _ = dNegR0.Free() + dH2Tail := (&dH2).Range(nPlus2, nPlus3, false) + r1Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[1]}) + r1Host.CopyToDevice(&dH2Tail, false) + + // h3 = base h3 with first coefficient adjusted by -r1. + if e := copyDeviceSliceIntoOnCurrentDevice(dH3, h3Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h3 failed: %s", e.AsString()) + return + } + dH3First := (&dH3).Range(0, 1, false) + var negR1 fr.Element + negR1.Neg(&s.quotientShardsRandomizers[1]) + dNegR1 := uploadScalarMont(negR1) + if e := icicle_vecops.ScalarAddVec(dNegR1, dH3First, dH3First, cfg); e != icicle_runtime.Success { + _ = dNegR1.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h3[0] failed: %s", e.AsString()) + return + } + _ = dNegR1.Free() + prepareDone <- nil + }) + if err := <-prepareDone; err != nil { + if !dH1.IsEmpty() { + s.putTempDeviceSlice(dH1, nPlus3) + } + if !dH2.IsEmpty() { + s.putTempDeviceSlice(dH2, nPlus3) + } + if !dH3.IsEmpty() { + s.putTempDeviceSlice(dH3, nPlus2) + } + return icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, err + } + return dH1, dH2, dH3, nil +} + +func (s *instance) commitToQuotientGPUFromDevice(quotient *gpuQuotientPolynomial) error { + if quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("commitToQuotientGPUFromDevice: empty quotient") + } + + nPlus2 := int(s.domain0.Cardinality) + 2 + required := 3 * nPlus2 + if quotient.coeffs.Len() < required { + return fmt.Errorf("commitToQuotientGPUFromDevice: quotient too small: got %d need >= %d", quotient.coeffs.Len(), required) + } + + h1Device := ("ient.coeffs).Range(0, nPlus2, false) + h2Device := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3Device := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + if s.opt.StatisticalZK { + nPlus3 := nPlus2 + 1 + dH1, dH2, dH3, err := s.prepareStatisticalZKQuotientShards(h1Device, h2Device, h3Device, nPlus2) + if err != nil { + return err + } + defer s.putTempDeviceSlice(dH1, nPlus3) + defer s.putTempDeviceSlice(dH2, nPlus3) + defer s.putTempDeviceSlice(dH3, nPlus2) + + c0, err := commitOnGPUCanonicalDevice(dH1, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(dH2, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(dH3, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + return nil + } + + // Commit sequentially to avoid 3-way concurrent MSM memory spikes. + c0, err := commitOnGPUCanonicalDevice(h1Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(h2Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(h3Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + + return nil +} + +func (s *instance) inverseAndMergeShards( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (icicle_core.DeviceSlice, error) { + if gpuNumerator == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: nil numerator") + } + n := gpuNumerator.n + rho := gpuNumerator.rho + if n <= 0 || rho <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: invalid n=%d rho=%d", n, rho) + } + if len(gpuNumerator.shards) != rho { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), rho) + } + for i := 0; i < rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard %d is empty", i) + } + } + + expo := big.NewInt(int64(n)) + + // Per-shard cosets: c_i = c * g^i where c=FrMultiplicativeGen, g=Generator. + cosets := make([]fr.Element, rho) + cosets[0].Set(&domains[1].FrMultiplicativeGen) + for i := 1; i < rho; i++ { + cosets[i].Mul(&cosets[i-1], &domains[1].Generator) + } + invCosets := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + invCosets[i].Inverse(&cosets[i]) + } + + // ν = g^n is a rho-th root, used for the rho-point inverse DFT in combine. + var nu, nuInv fr.Element + nu.Exp(domains[1].Generator, expo) + nuInv.Inverse(&nu) + nuInvPowers := make([]fr.Element, rho) + nuInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + nuInvPowers[i].Mul(&nuInvPowers[i-1], &nuInv) + } + + // cN = c^n. Recover original coefficient blocks by scaling with cN^{-t}. + var cN, cNInv fr.Element + cN.Exp(domains[1].FrMultiplicativeGen, expo) + cNInv.Inverse(&cN) + cNInvPowers := make([]fr.Element, rho) + cNInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + cNInvPowers[i].Mul(&cNInvPowers[i-1], &cNInv) + } + + // Each shard inverse contributes a 1/n factor; apply extra 1/rho. + var rhoFr, invRho fr.Element + rhoFr.SetUint64(uint64(rho)) + invRho.Inverse(&rhoFr) + combineScales := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + combineScales[i].Mul(&invRho, &cNInvPowers[i]) + } + + totalSize := rho * n + done := make(chan error, 1) + var dMerged icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish := makeFinisher(stream, "inverseAndMergeShards", done) + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + // KNR is often faster than KNN; we restore regular output explicitly + // by bit-reversing each shard after the inverse NTT. + cfgNtt.Ordering = icicle_core.KNR + + ext := config_extension.Create() + defer config_extension.Delete(ext) + alg := nttAlgorithmFromEnv("ICICLE_DIVIDE_BY_ZH_NTT_ALGO", icicle_core.MixedRadix) + ext.SetInt(icicle_core.CUDA_NTT_ALGORITHM, int(alg)) + cfgNtt.Ext = ext.AsUnsafePointer() + + // Step 1: inverse NTT each shard without coset, reorder to regular, + // then unscale by (c*g^i)^t to recover the coset-inverse equivalent. + nn := uint64(64 - bits.TrailingZeros64(uint64(n))) + invPowers := make([]fr.Element, n) + for i := 0; i < rho; i++ { + if nttErr := icicle_ntt.Ntt(gpuNumerator.shards[i], icicle_core.KInverse, &cfgNtt, gpuNumerator.shards[i]); nttErr != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: inverse NTT failed at shard %d: %s", i, nttErr.AsString())) + return + } + + // KNR outputs bit-reversed coefficients. Reorder back to regular. + dRegular := s.getTempDeviceSlice(n) + mergeErr := icicle_vecops.MergeShardsBitReverse( + []icicle_core.DeviceSlice{gpuNumerator.shards[i]}, + n, + nn, + dRegular, + cfgVec, + ) + if mergeErr != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: reorder failed at shard %d: %s", i, mergeErr.AsString())) + return + } + // gpuNumerator.shards[i] is returned to pool and replaced; wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(gpuNumerator.shards[i], n) + gpuNumerator.shards[i] = dRegular + + fft.BuildExpTable(invCosets[i], invPowers) + dInvPowers := s.getTempDeviceSlice(n) + uploadVectorStdIntoOnCurrentDevice(&dInvPowers, invPowers, cfgVec) + if e := icicle_vecops.VecOp(gpuNumerator.shards[i], dInvPowers, gpuNumerator.shards[i], cfgVec, icicle_core.Mul); e != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: normalize shard %d failed: %s", i, e.AsString())) + return + } + // dInvPowers is temporary and returned to pool each iteration. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(dInvPowers, n) + } + + // Step 2: combine shard results via a size-rho inverse DFT per coefficient index. + dMerged = s.getTempDeviceSlice(totalSize) + keepMerged := false + defer func() { + if !keepMerged && !dMerged.IsEmpty() { + s.putTempDeviceSlice(dMerged, totalSize) + } + }() + + dTmp := s.getTempDeviceSlice(n) + defer func() { + if !dTmp.IsEmpty() { + s.putTempDeviceSlice(dTmp, n) + } + }() + + dNuWeights := make([]icicle_core.DeviceSlice, rho) + dCombineScales := make([]icicle_core.DeviceSlice, rho) + for i := 0; i < rho; i++ { + dNuWeights[i] = uploadScalarStdOnCurrentDevice(nuInvPowers[i], cfgVec) + dCombineScales[i] = uploadScalarStdOnCurrentDevice(combineScales[i], cfgVec) + } + defer func() { + for i := 0; i < rho; i++ { + if !dNuWeights[i].IsEmpty() { + _ = dNuWeights[i].Free() + } + if !dCombineScales[i].IsEmpty() { + _ = dCombineScales[i].Free() + } + } + }() + + for t := 0; t < rho; t++ { + outT := (&dMerged).Range(t*n, (t+1)*n, false) + if e := copyDeviceSliceIntoOnCurrentDevice(outT, gpuNumerator.shards[0], cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: init out[%d] failed: %s", t, e.AsString())) + return + } + for i := 1; i < rho; i++ { + weightIdx := (i * t) % rho + if weightIdx == 0 { + if e := icicle_vecops.VecOp(outT, gpuNumerator.shards[i], outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: add shard %d to out[%d] failed: %s", i, t, e.AsString())) + return + } + continue + } + if e := icicle_vecops.ScalarMulVec(dNuWeights[weightIdx], gpuNumerator.shards[i], dTmp, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: weight shard %d for out[%d] failed: %s", i, t, e.AsString())) + return + } + if e := icicle_vecops.VecOp(outT, dTmp, outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: accumulate shard %d into out[%d] failed: %s", i, t, e.AsString())) + return + } + } + if e := icicle_vecops.ScalarMulVec(dCombineScales[t], outT, outT, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: scale out[%d] failed: %s", t, e.AsString())) + return + } + } + keepMerged = true + finish(nil) + }) + + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dMerged, nil +} + +func (s *instance) divideByZHOnGPU( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (_ *gpuQuotientPolynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("divideByZHOnGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("divideByZHOnGPU: invalid numerator dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + for i := range gpuNumerator.shards { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("divideByZHOnGPU: shard %d is empty", i) + } + } + + rho := int(domains[1].Cardinality / domains[0].Cardinality) + if rho != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: rho mismatch domains=%d numerator=%d", rho, gpuNumerator.rho) + } + + // Evaluate 1/(X^n-1) over the large-domain coset values used by this quotient. + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + + // In bit-reversed merged layout, each shard maps to a fixed (iRev % rho) bucket. + // So we can divide by Z_H by scaling each shard with its corresponding inverse. + scaleDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + for i := 0; i < gpuNumerator.rho; i++ { + dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) + vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if vecErr != icicle_runtime.Success { + finish(fmt.Errorf("divideByZHOnGPU: shard scaling failed at %d: %s", i, vecErr.AsString())) + return + } + } + finish(nil) + }) + if err := <-scaleDone; err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, err + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + dMerged, splitErr := s.inverseAndMergeShards(gpuNumerator, domains) + if splitErr != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, splitErr + } + // Shards are not needed after split inverse+merge. + s.freeNumeratorShards(gpuNumerator.shards) + return &gpuQuotientPolynomial{coeffs: dMerged, size: totalSize}, nil +} + +func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*iop.Polynomial, error) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return nil, fmt.Errorf("downloadQuotientFromGPU: empty quotient") + } + if quotient.size <= 0 { + return nil, fmt.Errorf("downloadQuotientFromGPU: invalid quotient size %d", quotient.size) + } + + coeffs := make([]fr.Element, quotient.size) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) + // Async boundary for host materialization of quotient coefficients. + done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + }) + if err := <-done; err != nil { + return nil, err + } + + return iop.NewPolynomial(&coeffs, iop.Form{Basis: iop.Canonical, Layout: iop.Regular}), nil +} + +func commitOnGPUWithDeviceBases( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, +) (curve.G1Affine, error) { + return commitOnGPUWithDeviceBasesChunked(scalarsDevice, basesDevice, device, icicleMSMChunkSize()) +} + +func commitOnGPUWithDeviceBasesChunked( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, + chunkSize int, +) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty scalar slice") + } + if basesDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty basis slice") + } + if scalarsDevice.Len() > basesDevice.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: invalid scalar size %d", scalarsDevice.Len()) + } + if chunkSize <= 0 || chunkSize > scalarsDevice.Len() { + chunkSize = scalarsDevice.Len() + } + + var commit curve.G1Affine + var msmErr error + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) + close(done) + }) + <-done + if msmErr != nil { + return curve.G1Affine{}, fmt.Errorf("icicle: MSM commit from device bases failed (%d scalars): %w", scalarsDevice.Len(), msmErr) + } + return commit, nil +} + +func commitOnGPUWithDeviceBasesChunkedOnCurrentDevice( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + chunkSize int, +) (curve.G1Affine, error) { + var commit curve.G1Affine + for start := 0; start < scalarsDevice.Len(); start += chunkSize { + end := start + chunkSize + if end > scalarsDevice.Len() { + end = scalarsDevice.Len() + } + + // Each chunk must pair with exactly bases[start:end]: ICICLE treats a + // bases slice longer than the scalars as a batched MSM (and requires + // divisibility), so the full bases buffer cannot be passed as-is when + // it is longer than the scalar vector. + scalarsChunk := scalarsDevice + if start != 0 || end != scalarsDevice.Len() { + scalarsChunk = (&scalarsDevice).Range(start, end, false) + } + basesChunk := basesDevice + if start != 0 || end != basesDevice.Len() { + basesChunk = (&basesDevice).Range(start, end, false) + } + + res := make(icicle_core.HostSlice[icicle_bn254.Projective], 1) + cfg := icicle_msm.GetDefaultMSMConfig() + cfg.AreBasesMontgomeryForm = true + cfg.AreScalarsMontgomeryForm = true + e := icicle_msm.Msm(scalarsChunk, basesChunk, &cfg, res) + if e != icicle_runtime.Success { + return curve.G1Affine{}, fmt.Errorf("icicle MSM failed for chunk [%d:%d]: %s", start, end, e.AsString()) + } + + chunkCommit, err := projectiveToGnarkAffine(res[0]) + if err != nil { + return curve.G1Affine{}, fmt.Errorf("convert chunk [%d:%d]: %w", start, end, err) + } + commit.Add(&commit, &chunkCommit) + } + return commit, nil +} + +func icicleMSMChunkSize() int { + // Production-sized MSMs still need chunking, but tiny chunks add thousands of + // ICICLE calls. 4M-point chunks passed the gnark replay profile; 8M did not. + const defaultChunkSize = 1 << 22 + v := strings.TrimSpace(os.Getenv("ICICLE_MSM_CHUNK_SIZE")) + if v == "" { + return defaultChunkSize + } + n, err := strconv.Atoi(v) + if err != nil || n <= 0 { + return defaultChunkSize + } + return n +} + +func commitOnGPUCanonicalDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: canonical SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgDevice.G1, device) +} + +func evalCanonicalAtPoint(coeffs []fr.Element, point fr.Element) fr.Element { + var acc fr.Element + if len(coeffs) == 0 { + return acc + } + acc.Set(&coeffs[len(coeffs)-1]) + for i := len(coeffs) - 2; i >= 0; i-- { + acc.Mul(&acc, &point).Add(&acc, &coeffs[i]) + } + return acc +} + +func deriveBatchOpeningGamma( + point fr.Element, + digests []curve.G1Affine, + claimedValues []fr.Element, + hf hash.Hash, + dataTranscript ...[]byte, +) (fr.Element, error) { + fs := fiatshamir.NewTranscript(hf, "gamma") + if err := fs.Bind("gamma", point.Marshal()); err != nil { + return fr.Element{}, err + } + for i := range digests { + if err := fs.Bind("gamma", digests[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := range claimedValues { + if err := fs.Bind("gamma", claimedValues[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := 0; i < len(dataTranscript); i++ { + if err := fs.Bind("gamma", dataTranscript[i]); err != nil { + return fr.Element{}, err + } + } + gammaByte, err := fs.ComputeChallenge("gamma") + if err != nil { + return fr.Element{}, err + } + var gamma fr.Element + gamma.SetBytes(gammaByte) + return gamma, nil +} + +func (s *instance) evalDevicePolynomialAtPointOnCurrentDevice( + coeffsDevice icicle_core.DeviceSlice, + point fr.Element, + useBitReverse bool, + cfg icicle_core.VecOpsConfig, +) (fr.Element, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: empty coefficients") + } + + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(point, cfg) + defer dPoint.Free() + + dOut := s.getTempDeviceSlice(1) + defer s.putTempDeviceSlice(dOut, 1) + + opName := "PolyEvalAt" + var eEval icicle_runtime.EIcicleError + if useBitReverse { + opName = "PolyEvalAtBitReverse" + mm := uint64(64 - bits.TrailingZeros64(uint64(coeffsDevice.Len()))) + eEval = icicle_vecops.PolyEvalAtBitReverse(coeffsDevice, dPoint, mm, dOut, cfg) + } else { + eEval = icicle_vecops.PolyEvalAt(coeffsDevice, dPoint, dOut, cfg) + } + if eEval != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: %s failed: %s", opName, eEval.AsString()) + } + + var out fr.Element + hostOut := icicle_core.HostSliceFromElements([]fr.Element{out}) + if cfg.IsAsync { + hostOut.CopyFromDeviceAsync(&dOut, cfg.StreamHandle) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } else { + hostOut.CopyFromDevice(&dOut) + } + return ([]fr.Element)(hostOut)[0], nil +} + +func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (fr.Element, error) { + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + var runErr error + out, runErr = s.evalDevicePolynomialAtPointOnCurrentDevice(coeffsDevice, point, false, cfg) + done <- runErr + }) + return out, <-done +} + +func (s *instance) copyDeviceSliceOnCurrentDevice( + src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, + label string, +) (icicle_core.DeviceSlice, error) { + if src.IsEmpty() || src.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: empty source slice", label) + } + dst := s.getTempDeviceSlice(src.Len()) + eCopy := copyDeviceSliceIntoOnCurrentDevice(dst, src, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dst, src.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: copy failed: %s", label, eCopy.AsString()) + } + return dst, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromStateOnCurrentDevice( + p *iop.Polynomial, + state *gpuPolysState, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "batchOpening poly") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + dCanon, err := s.copyDeviceSliceOnCurrentDevice(dSrc, cfg, "materializePolynomialCanonicalRegularFromStateOnCurrentDevice") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + // Canonical bit-reverse -> canonical regular. + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + dRegular := s.getTempDeviceSlice(dCanon.Len()) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNN // regular -> regular canonical + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular canonical + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) buildBlindedCanonicalPolynomialOnCurrentDevice( + dBaseCanon, dBlindCanon icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBlindCanon.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: empty input") + } + n := dBaseCanon.Len() + blindLen := dBlindCanon.Len() + dOut := s.getTempDeviceSlice(n + blindLen) + + dPrefix := (&dOut).Range(0, n, false) + eCopyBase := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopyBase != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopyBase.AsString()) + } + + dTail := (&dOut).Range(n, n+blindLen, false) + eCopyBlind := copyDeviceSliceIntoOnCurrentDevice(dTail, dBlindCanon, cfg) + if eCopyBlind != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy tail failed: %s", eCopyBlind.AsString()) + } + + dHead := (&dOut).Range(0, blindLen, false) + if eSub := icicle_vecops.VecOp(dHead, dBlindCanon, dHead, cfg, icicle_core.Sub); eSub != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: subtract blind from head failed: %s", eSub.AsString()) + } + return dOut, nil +} + +func (s *instance) buildPaddedCanonicalPolynomialOnCurrentDevice( + dBaseCanon icicle_core.DeviceSlice, + padLen int, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBaseCanon.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: empty base") + } + if padLen < 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: negative pad length %d", padLen) + } + n := dBaseCanon.Len() + dOut := s.getTempDeviceSlice(n + padLen) + dPrefix := (&dOut).Range(0, n, false) + + eCopy := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopy.AsString()) + } + if padLen == 0 { + return dOut, nil + } + + dTail := (&dOut).Range(n, n+padLen, false) + eZero := zeroDeviceSliceOnCurrentDevice(dTail, cfg) + if eZero != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: zero tail failed: %s", eZero.AsString()) + } + return dOut, nil +} + +func (s *instance) prepareBatchOpeningPolynomialsOnGPU( + state *gpuPolysState, + point fr.Element, +) (devicePolys []icicle_core.DeviceSlice, owned []bool, claimed []fr.Element, err error) { + if state == nil { + return nil, nil, nil, fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: nil GPU state") + } + + total := 6 + len(s.trace.Qcp) + devicePolys = make([]icicle_core.DeviceSlice, total) + owned = make([]bool, total) + claimed = make([]fr.Element, total) + devicePolys[0] = s.linearizedPolynomialGPU + claimed[0] = s.linearizedPolynomialClaim + + prepDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + + cleanupOwned := func(from int) { + for i := from; i < len(devicePolys); i++ { + if owned[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + owned[i] = false + } + } + } + + prepareLRORow := func(dstIdx int, p, bp *iop.Polynomial, padLen int) error { + dBase, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + defer s.putTempDeviceSlice(dBase, dBase.Len()) + + var dFinal icicle_core.DeviceSlice + if useBlinding { + if bp == nil { + return fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: missing blinding polynomial") + } + dBlind, eBlind := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(bp, state, cfg) + if eBlind != nil { + return eBlind + } + defer s.putTempDeviceSlice(dBlind, dBlind.Len()) + dFinal, e = s.buildBlindedCanonicalPolynomialOnCurrentDevice(dBase, dBlind, cfg) + } else { + dFinal, e = s.buildPaddedCanonicalPolynomialOnCurrentDevice(dBase, padLen, cfg) + } + if e != nil { + return e + } + devicePolys[dstIdx] = dFinal + owned[dstIdx] = true + + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dFinal, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareLRORow(1, s.polyL, s.bp[id_Bl], order_blinding_L+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare L failed: %w", e)) + return + } + if e := prepareLRORow(2, s.polyR, s.bp[id_Br], order_blinding_R+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare R failed: %w", e)) + return + } + if e := prepareLRORow(3, s.polyO, s.bp[id_Bo], order_blinding_O+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare O failed: %w", e)) + return + } + + prepareDirect := func(dstIdx int, p *iop.Polynomial, label string) error { + dPoly, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + devicePolys[dstIdx] = dPoly + owned[dstIdx] = true + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dPoly, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareDirect(4, s.trace.S1, "S1"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S1 failed: %w", e)) + return + } + if e := prepareDirect(5, s.trace.S2, "S2"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S2 failed: %w", e)) + return + } + + for i := 0; i < len(s.trace.Qcp); i++ { + idx := 6 + i + if e := prepareDirect(idx, s.trace.Qcp[i], "Qcp"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare Qcp[%d] failed: %w", i, e)) + return + } + } + + finish(nil) + }) + if err := <-prepDone; err != nil { + return nil, nil, nil, err + } + return devicePolys, owned, claimed, nil +} + +type evalPolynomialInputPreparationResult struct { + dEval icicle_core.DeviceSlice + ownedLen int + useBitReverseEval bool +} + +func (s *instance) prepareEvalPolynomialInputOnCurrentDevice( + p *iop.Polynomial, + dSrc icicle_core.DeviceSlice, + cfgVec icicle_core.VecOpsConfig, +) (evalPolynomialInputPreparationResult, error) { + if p == nil { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: nil polynomial") + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: empty source polynomial") + } + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] prepareEvalInput begin n=%d basis=%v layout=%v step_profile=%q ntt_trace=%q ntt_profile_full=%q ntt_profile_arbitrary=%q\n", + dSrc.Len(), + p.Basis, + p.Layout, + os.Getenv("ICICLE_STEP_PROFILE"), + os.Getenv("ICICLE_NTT_TRACE"), + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + + result := evalPolynomialInputPreparationResult{ + dEval: dSrc, + } + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return result, nil + } + + releaseOwned := func() { + if result.ownedLen > 0 && !result.dEval.IsEmpty() { + s.putTempDeviceSlice(result.dEval, result.ownedLen) + result.dEval = icicle_core.DeviceSlice{} + result.ownedLen = 0 + } + } + + dWork := s.getTempDeviceSlice(dSrc.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dWork, dSrc, cfgVec); e != icicle_runtime.Success { + s.putTempDeviceSlice(dWork, dSrc.Len()) + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: copy source polynomial failed: %s", e.AsString()) + } + + result.dEval = dWork + result.ownedLen = dSrc.Len() + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfgVec.IsAsync + cfgNtt.StreamHandle = cfgVec.StreamHandle + + var ownsNttStream bool + destroyOwnedNttStream := func() error { + if !ownsNttStream { + return nil + } + return syncAndDestroyStreamOnCurrentDevice(cfgNtt.StreamHandle, "prepareEvalPolynomialInputOnCurrentDevice") + } + + switch p.Basis { + case iop.Canonical: + // No transform required. + result.useBitReverseEval = p.Layout == iop.BitReverse + case iop.Lagrange, iop.LagrangeCoset: + if cfgNtt.StreamHandle == nil { + nttStream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: create NTT stream failed: %s", eStream.AsString()) + } + cfgNtt.StreamHandle = nttStream + cfgNtt.IsAsync = true + ownsNttStream = true + } + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNR // regular -> bitreverse on inverse + result.useBitReverseEval = true + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular on inverse + result.useBitReverseEval = false + } + startNtt := time.Now() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT launch n=%d ordering=%v has_coset=%t basis=%v layout=%v\n", + dSrc.Len(), + cfgNtt.Ordering, + p.Basis == iop.LagrangeCoset, + p.Basis, + p.Layout, + ) + } + + // Martun: This call to Ntt takes about 3 seconds, because Ntt is reusing NTT domain data that + // gets prepared during InitDomain (once per device), not re-derived every call. + // Inside ICICLE, InitDomain precomputes: domain.twiddles (main roots-of-unity table, N+1) + // internal_twiddles and basic_twiddles for mixed-radix kernels + // if fast mode is on (it is by default here), extra forward+inverse fast twiddle tables (fast_external/internal/basic and _inv) — comment says this costs ~4N extra memory + // CPU-side coset_index map (root -> index), then reused by later Ntt calls + eNtt := icicle_ntt.Ntt(result.dEval, icicle_core.KInverse, &cfgNtt, result.dEval) + nttElapsed := time.Since(startNtt) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT done status=%s took=%s\n", + eNtt.AsString(), + nttElapsed, + ) + } + if eNtt != icicle_runtime.Success { + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + l := logger.Logger() + l.Warn().Err(errDestroy).Msg("prepareEvalPolynomialInputOnCurrentDevice") + } + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + // Async boundary for this helper when it owns the stream. + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + releaseOwned() + return evalPolynomialInputPreparationResult{}, errDestroy + } + default: + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: unsupported basis %v", p.Basis) + } + + return result, nil +} + +// evalPolynomialInCurrentFormOnGPU evaluates a polynomial at a point directly +// from the shared GPU state regardless of its current basis/layout by converting +// a temporary device copy to canonical/regular when needed. +func (s *instance) evalPolynomialInCurrentFormOnGPU( + p *iop.Polynomial, + state *gpuPolysState, + point fr.Element, +) (fr.Element, error) { + if p == nil { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "eval") + if err != nil { + return fr.Element{}, err + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: empty device polynomial") + } + + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + dEval := dSrc + ownedLen := 0 + releaseEval := func() { + if ownedLen > 0 && !dEval.IsEmpty() { + s.putTempDeviceSlice(dEval, ownedLen) + dEval = icicle_core.DeviceSlice{} + ownedLen = 0 + } + } + finish := func(runErr error) { + // Async boundary for eval path before handing result back to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + releaseEval() + done <- runErr + } + + prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) + dEval = prepareResult.dEval + ownedLen = prepareResult.ownedLen + if prepErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", prepErr)) + return + } + + evalOut, evalErr := s.evalDevicePolynomialAtPointOnCurrentDevice(dEval, point, prepareResult.useBitReverseEval, cfgVec) + if evalErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", evalErr)) + return + } + out = evalOut + finish(nil) + }) + return out, <-done +} + +func (s *instance) evaluateBlindedOnGPU( + p, bp *iop.Polynomial, + state *gpuPolysState, + zeta fr.Element, +) (fr.Element, error) { + if p == nil || bp == nil { + return fr.Element{}, fmt.Errorf("evaluateBlindedOnGPU: nil polynomial") + } + pAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(p, state, zeta) + if err != nil { + return fr.Element{}, err + } + bpAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(bp, state, zeta) + if err != nil { + return fr.Element{}, err + } + + var t, one fr.Element + one.SetOne() + t.Exp(zeta, big.NewInt(int64(p.Size()))).Sub(&t, &one) + bpAtZeta.Mul(&bpAtZeta, &t) + pAtZeta.Add(&pAtZeta, &bpAtZeta) + return pAtZeta, nil +} + +func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.Element) (kzg.OpeningProof, error) { + if len(coeffs) < 2 || len(coeffs) > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonical: invalid polynomial size %d", len(coeffs)) + } + claimed := evalCanonicalAtPoint(coeffs, point) + + var dWitness icicle_core.DeviceSlice + witnessSize := len(coeffs) - 1 + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dCoeffs := uploadVector(coeffs) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dCoeffs, dPoint, dWitness, cfg) + _ = dCoeffs.Free() + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonical: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + if err := <-divDone; err != nil { + return kzg.OpeningProof{}, err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (kzg.OpeningProof, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() < 2 || coeffsDevice.Len() > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonicalDevice: invalid polynomial size %d", coeffsDevice.Len()) + } + n := coeffsDevice.Len() + + var startEval time.Time + if isProfileMode { + startEval = time.Now() + } + claimed, err := s.evalDevicePolynomialAtPoint(coeffsDevice, point) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startEval)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint") + } + } + if err != nil { + return kzg.OpeningProof{}, err + } + + var dWitness icicle_core.DeviceSlice + witnessSize := coeffsDevice.Len() - 1 + var startDivideByXMinusA time.Time + if isProfileMode { + startDivideByXMinusA = time.Now() + } + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(coeffsDevice, dPoint, dWitness, cfg) + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonicalDevice: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + divideErr := <-divDone + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startDivideByXMinusA)) + if divideErr != nil { + ev.Err(divideErr).Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA") + } + } + if divideErr != nil { + return kzg.OpeningProof{}, divideErr + } + + var startCommit time.Time + if isProfileMode { + startCommit = time.Now() + } + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startCommit)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice") + } + } + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) linearizedZContributionScale(lZeta, rZeta, oZeta fr.Element) fr.Element { + var s2, tmp fr.Element + var uzeta, uuzeta fr.Element + uzeta.Mul(&s.zeta, &s.pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &s.pk.Vk.CosetShift) + + s2.Mul(&s.beta, &s.zeta).Add(&s2, &lZeta).Add(&s2, &s.gamma) + tmp.Mul(&s.beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp) + tmp.Mul(&s.beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp).Neg(&s2).Mul(&s2, &s.alpha) + + var one, alphaSquareLagrangeZero, den fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&s.zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) + den.Sub(&s.zeta, &one).Inverse(&den) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) + + s2.Add(&s2, &alphaSquareLagrangeZero) + return s2 +} + +func (s *instance) linearizedSelectorScales(evals witnessEvalAtZeta, zu fr.Element) linearizedSelectorScales { + var scales linearizedSelectorScales + + // S3 scale: + // alpha * beta * Z(mu*zeta) * + // (L(zeta) + beta*S1(zeta) + gamma) * + // (R(zeta) + beta*S2(zeta) + gamma) + var tmp fr.Element + scales.s3.Mul(&evals.s1zeta, &s.beta).Add(&scales.s3, &evals.blzeta).Add(&scales.s3, &s.gamma) + tmp.Mul(&evals.s2zeta, &s.beta).Add(&tmp, &evals.brzeta).Add(&tmp, &s.gamma) + scales.s3.Mul(&scales.s3, &tmp).Mul(&scales.s3, &zu).Mul(&scales.s3, &s.beta).Mul(&scales.s3, &s.alpha) + + scales.ql.Set(&evals.blzeta) + scales.qr.Set(&evals.brzeta) + scales.qm.Mul(&evals.brzeta, &evals.blzeta) + scales.qo.Set(&evals.bozeta) + scales.qk.SetOne() + scales.qcp = append(scales.qcp, evals.qcpzeta...) + + return scales +} + +func (s *instance) buildLinearizedSelectorTermsOnGPU( + evals witnessEvalAtZeta, + zu fr.Element, + linearizedLen int, +) (icicle_core.DeviceSlice, error) { + if linearizedLen <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: invalid length %d", linearizedLen) + } + if len(evals.qcpzeta) > len(s.cCommitments) { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: qcp/cCommitments mismatch (%d > %d)", len(evals.qcpzeta), len(s.cCommitments)) + } + + scales := s.linearizedSelectorScales(evals, zu) + + var dLinearized icicle_core.DeviceSlice + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildLinearizedSelectorTermsOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if runErr != nil && !dLinearized.IsEmpty() { + s.putTempDeviceSlice(dLinearized, dLinearized.Len()) + dLinearized = icicle_core.DeviceSlice{} + } + done <- runErr + }() + + dLinearized = s.getTempDeviceSlice(linearizedLen) + if eZero := zeroDeviceSliceOnCurrentDevice(dLinearized, cfg); eZero != icicle_runtime.Success { + runErr = fmt.Errorf("buildLinearizedSelectorTermsOnGPU: zero output failed: %s", eZero.AsString()) + return + } + + addTerm := func(p *iop.Polynomial, scale fr.Element, label string) error { + if p == nil { + return fmt.Errorf("missing polynomial %s", label) + } + if scale.IsZero() { + return nil + } + + start := time.Now() + dPoly, err := s.materializePolynomialCanonicalRegularFromHostOnCurrentDevice(p, cfg) + if err != nil { + return fmt.Errorf("%s: %w", label, err) + } + defer s.putTempDeviceSlice(dPoly, dPoly.Len()) + if dPoly.Len() > dLinearized.Len() { + return fmt.Errorf("%s: polynomial too large (%d > %d)", label, dPoly.Len(), dLinearized.Len()) + } + + dScale := uploadScalarStdOnCurrentDevice(scale, cfg) + dScaled := s.getTempDeviceSlice(dPoly.Len()) + defer s.putTempDeviceSlice(dScaled, dScaled.Len()) + + eScale := icicle_vecops.ScalarMulVec(dScale, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if eScale != icicle_runtime.Success { + return fmt.Errorf("%s: scale failed: %s", label, eScale.AsString()) + } + + dPrefix := (&dLinearized).Range(0, dPoly.Len(), false) + if eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add); eAdd != icicle_runtime.Success { + return fmt.Errorf("%s: add failed: %s", label, eAdd.AsString()) + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Str("term", label).Int("n", dPoly.Len()).Dur("took", time.Since(start)).Msg("computeLinearizedPolynomial: add selector term on GPU") + } + return nil + } + + if runErr = addTerm(s.trace.S3, scales.s3, "S3"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Ql, scales.ql, "Ql"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qm, scales.qm, "Qm"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qr, scales.qr, "Qr"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qo, scales.qo, "Qo"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qk, scales.qk, "Qk"); runErr != nil { + return + } + for i := range scales.qcp { + if runErr = addTerm(s.cCommitments[i], scales.qcp[i], fmt.Sprintf("Qcp[%d]", i)); runErr != nil { + return + } + } + }) + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dLinearized, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromHostOnCurrentDevice( + p *iop.Polynomial, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: nil polynomial") + } + coeffs := p.Coefficients() + if len(coeffs) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: empty polynomial") + } + + dCanon := s.getTempDeviceSlice(len(coeffs)) + host := icicle_core.HostSliceFromElements(coeffs) + if cfg.IsAsync { + host.CopyToDeviceAsync(&dCanon, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&dCanon, false) + } + if dCanon.IsEmpty() { + s.putTempDeviceSlice(dCanon, len(coeffs)) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: host upload failed") + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + dRegular := s.getTempDeviceSlice(dCanon.Len()) + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + switch p.Layout { + case iop.Regular: + cfgNtt.Ordering = icicle_core.KNN + case iop.BitReverse: + cfgNtt.Ordering = icicle_core.KRN + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported layout %v", p.Layout) + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) addZContributionToLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + dBlindedZCanonical icicle_core.DeviceSlice, + lZeta, rZeta, oZeta fr.Element, +) error { + if dLinearized.IsEmpty() || dBlindedZCanonical.IsEmpty() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: empty input") + } + if dLinearized.Len() < dBlindedZCanonical.Len() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: linearized too small (%d < %d)", dLinearized.Len(), dBlindedZCanonical.Len()) + } + + zScale := s.linearizedZContributionScale(lZeta, rZeta, oZeta) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("addZContributionToLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dScale icicle_core.DeviceSlice + var dScaledZ icicle_core.DeviceSlice + defer func() { + // Async boundary before returning temporary buffers to the pool. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dScale) + if !dScaledZ.IsEmpty() { + s.putTempDeviceSlice(dScaledZ, dScaledZ.Len()) + } + done <- runErr + }() + + dScale = uploadScalarStdOnCurrentDevice(zScale, cfg) + dScaledZ = s.getTempDeviceSlice(dBlindedZCanonical.Len()) + eMul := icicle_vecops.ScalarMulVec(dScale, dBlindedZCanonical, dScaledZ, cfg) + if eMul != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: scale Z failed: %s", eMul.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, dBlindedZCanonical.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaledZ, dPrefix, cfg, icicle_core.Add) + if eAdd != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: add scaled Z failed: %s", eAdd.AsString()) + return + } + }) + return <-done +} + +func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + quotient *gpuQuotientPolynomial, +) error { + if dLinearized.IsEmpty() || quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: empty input") + } + nPlus2 := int(s.domain0.Cardinality) + 2 + if dLinearized.Len() < nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: linearized too small") + } + if quotient.coeffs.Len() < 3*nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: quotient too small") + } + + var one fr.Element + one.SetOne() + var zetaN, zetaNPlusTwo, zhZeta fr.Element + zetaN.Exp(s.zeta, big.NewInt(int64(s.domain0.Cardinality))) + zhZeta.Sub(&zetaN, &one) + zetaNPlusTwo.Mul(&zetaN, &s.zeta).Mul(&zetaNPlusTwo, &s.zeta) + + h1 := ("ient.coeffs).Range(0, nPlus2, false) + h2 := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3 := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("subtractQuotientContributionFromLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dAcc icicle_core.DeviceSlice + var dZetaStd icicle_core.DeviceSlice + var dZhStd icicle_core.DeviceSlice + defer func() { + // Async boundary before reusing temporary quotient vectors. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dZhStd) + freeDeviceSlice(&dZetaStd) + if !dAcc.IsEmpty() { + s.putTempDeviceSlice(dAcc, dAcc.Len()) + } + done <- runErr + }() + + dAcc = s.getTempDeviceSlice(nPlus2) + dZetaStd = uploadScalarStdOnCurrentDevice(zetaNPlusTwo, cfg) + eMulH3 := icicle_vecops.ScalarMulVec(dZetaStd, h3, dAcc, cfg) + if eMulH3 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale h3 failed: %s", eMulH3.AsString()) + return + } + if eAddH2 := icicle_vecops.VecOp(dAcc, h2, dAcc, cfg, icicle_core.Add); eAddH2 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h2 failed: %s", eAddH2.AsString()) + return + } + if eMulPow := icicle_vecops.ScalarMulVec(dZetaStd, dAcc, dAcc, cfg); eMulPow != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by zeta^(n+2) failed: %s", eMulPow.AsString()) + return + } + if eAddH1 := icicle_vecops.VecOp(dAcc, h1, dAcc, cfg, icicle_core.Add); eAddH1 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h1 failed: %s", eAddH1.AsString()) + return + } + + dZhStd = uploadScalarStdOnCurrentDevice(zhZeta, cfg) + eScaleZh := icicle_vecops.ScalarMulVec(dZhStd, dAcc, dAcc, cfg) + if eScaleZh != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by Z_H(zeta) failed: %s", eScaleZh.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, nPlus2, false) + eSub := icicle_vecops.VecOp(dPrefix, dAcc, dPrefix, cfg, icicle_core.Sub) + if eSub != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: subtract term failed: %s", eSub.AsString()) + return + } + }) + return <-done +} + +// divideByZH +// The input must be in LagrangeCoset. +// The result is in Canonical Regular. (in place using a) +func (s *instance) divideByZH(a *iop.Polynomial, domains [2]*fft.Domain) (*iop.Polynomial, error) { + + // check that the basis is LagrangeCoset + if a.Basis != iop.LagrangeCoset || a.Layout != iop.BitReverse { + return nil, errors.New("invalid form") + } + + // prepare the evaluations of x^n-1 on the big domain's coset + var startEvaluateXnMinusOne time.Time + if isProfileMode { + startEvaluateXnMinusOne = time.Now() + } + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startEvaluateXnMinusOne)).Msg("divideByZH: evaluateXnMinusOneDomainBigCoset") + } + rho := int(domains[1].Cardinality / domains[0].Cardinality) + + r := a.Coefficients() + n := uint64(len(r)) + nn := uint64(64 - bits.TrailingZeros64(n)) + + var startParallelizeMul time.Time + if isProfileMode { + startParallelizeMul = time.Now() + } + utils.Parallelize(len(r), func(start, end int) { + for i := start; i < end; i++ { + iRev := bits.Reverse64(uint64(i)) >> nn + r[i].Mul(&r[i], &xnMinusOneInverseLagrangeCoset[int(iRev)%rho]) + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallelizeMul)).Msg("divideByZH: parallelize multiply coefficients") + } + + // Replace CPU FFT inverse by ICICLE NTT inverse. + var startGpuNTTInverse time.Time + if isProfileMode { + startGpuNTTInverse = time.Now() + } + // It's faster on CPU. + // s.gpuNTTInverse(a) + a.ToCanonical(domains[1]).ToRegular() + if isProfileMode { + l := logger.Logger() + l.Debug(). + Int("size", a.Size()). + Dur("took", time.Since(startGpuNTTInverse)). + Msg("divideByZH: gpuNTTInverse on input of size n") + } + + return a, nil +} + +// evaluateXnMinusOneDomainBigCoset evaluates Xᵐ-1 on DomainBig coset +func evaluateXnMinusOneDomainBigCoset(domains [2]*fft.Domain) []fr.Element { + + rho := domains[1].Cardinality / domains[0].Cardinality + + res := make([]fr.Element, rho) + + expo := big.NewInt(int64(domains[0].Cardinality)) + res[0].Exp(domains[1].FrMultiplicativeGen, expo) + + var t fr.Element + t.Exp(domains[1].Generator, expo) + + one := fr.One() + + for i := 1; i < int(rho); i++ { + res[i].Mul(&res[i-1], &t) + res[i-1].Sub(&res[i-1], &one) + } + res[len(res)-1].Sub(&res[len(res)-1], &one) + + res = fr.BatchInvert(res) + + return res +} + +// innerComputeLinearizedPoly computes the linearized polynomial in canonical basis. +// The purpose is to commit and open all in one ql, qr, qm, qo, qk. +// * lZeta, rZeta, oZeta are the evaluation of l, r, o at zeta +// * z is the permutation polynomial, zu is Z(μX), the shifted version of Z +// * pk is the proving key: the linearized polynomial is a linear combination of ql, qr, qm, qo, qk. +// +// The Linearized polynomial is: +// +// α²*L₁(ζ)*Z(X) +// + α*( (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*(β*s3(X))*Z(μζ) - Z(X)*(l(ζ)+β*id1(ζ)+γ)*(r(ζ)+β*id2(ζ)+γ)*(o(ζ)+β*id3(ζ)+γ)) +// + l(ζ)*Ql(X) + l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) +// - Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) +// +// /!\ blindedZCanonical is modified +func (s *instance) innerComputeLinearizedPoly( + lZeta, rZeta, oZeta, s1Zeta, s2Zeta, + alpha, beta, gamma, zeta, zu fr.Element, + qcpZeta, blindedZCanonical []fr.Element, + pi2Canonical [][]fr.Element, + pk *ProvingKey, +) []fr.Element { + + // l(ζ)r(ζ) + var rl fr.Element + rl.Mul(&rZeta, &lZeta) + + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + // the linearised polynomial is + // α²*L₁(ζ)*Z(X) + + // s1*s3(X)+s2*Z(X) + l(ζ)*Ql(X) + + // l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) - + // Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + var s1, s2, tmp fr.Element + s1.Mul(&s1Zeta, &beta).Add(&s1, &lZeta).Add(&s1, &gamma) // (l(ζ)+β*s1(ζ)+γ) + tmp.Mul(&s2Zeta, &beta).Add(&tmp, &rZeta).Add(&tmp, &gamma) + s1.Mul(&s1, &tmp).Mul(&s1, &zu).Mul(&s1, &beta).Mul(&s1, &alpha) // (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*β*Z(μζ)*α + + var uzeta, uuzeta fr.Element + uzeta.Mul(&zeta, &pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &pk.Vk.CosetShift) + + s2.Mul(&beta, &zeta).Add(&s2, &lZeta).Add(&s2, &gamma) // (l(ζ)+β*ζ+γ) + tmp.Mul(&beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &gamma) // (r(ζ)+β*u*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ) + tmp.Mul(&beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &gamma) // (o(ζ)+β*u²*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + s2.Neg(&s2).Mul(&s2, &alpha) + + // Z_h(ζ), ζⁿ⁺², L₁(ζ)*α²*Z + var zhZeta, zetaNPlusTwo, alphaSquareLagrangeZero, one, den, frNbElmt fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) // ζⁿ + zetaNPlusTwo.Mul(&alphaSquareLagrangeZero, &zeta).Mul(&zetaNPlusTwo, &zeta) // ζⁿ⁺² + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) // ζⁿ - 1 + zhZeta.Set(&alphaSquareLagrangeZero) // Z_h(ζ) = ζⁿ - 1 + frNbElmt.SetUint64(uint64(nbElmt)) + den.Sub(&zeta, &one).Inverse(&den) // 1/(ζ-1) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). // L₁ = (ζⁿ - 1)/(ζ-1) + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) // α²*L₁(ζ) + + s3canonical := s.trace.S3.Coefficients() + // Qk is prepared in canonical/regular form by computeLinearizedPolynomial. + + // len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 when Statistical ZK is activated + // len(h1)=len(h2)=len(h3)=len(blindedZCanonical)-1 when Statistical ZK is deactivated + h1 := s.h1() + h2 := s.h2() + h3 := s.h3() + + // at this stage we have + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + startParallel := time.Now() + utils.Parallelize(len(blindedZCanonical), func(start, end int) { + + cql := s.trace.Ql.Coefficients() + cqr := s.trace.Qr.Coefficients() + cqm := s.trace.Qm.Coefficients() + cqo := s.trace.Qo.Coefficients() + cqk := s.trace.Qk.Coefficients() + + var t, t0, t1 fr.Element + + for i := start; i < end; i++ { + t.Mul(&blindedZCanonical[i], &s2) // -Z(X)*α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + if i < len(s3canonical) { + t0.Mul(&s3canonical[i], &s1) // α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ)*β*s3(X) + t.Add(&t, &t0) + } + if i < len(cqm) { + t1.Mul(&cqm[i], &rl) // l(ζ)r(ζ)*Qm(X) + t.Add(&t, &t1) // linPol += l(ζ)r(ζ)*Qm(X) + t0.Mul(&cql[i], &lZeta) // l(ζ)Q_l(X) + t.Add(&t, &t0) // linPol += l(ζ)*Ql(X) + t0.Mul(&cqr[i], &rZeta) //r(ζ)*Qr(X) + t.Add(&t, &t0) // linPol += r(ζ)*Qr(X) + t0.Mul(&cqo[i], &oZeta) // o(ζ)*Qo(X) + t.Add(&t, &t0) // linPol += o(ζ)*Qo(X) + t.Add(&t, &cqk[i]) // linPol += Qk(X) + for j := range qcpZeta { // linPol += ∑ᵢQcp_(ζ)Pi_(X) + t0.Mul(&pi2Canonical[j][i], &qcpZeta[j]) + t.Add(&t, &t0) + } + } + + t0.Mul(&blindedZCanonical[i], &alphaSquareLagrangeZero) // α²L₁(ζ)Z(X) + blindedZCanonical[i].Add(&t, &t0) // linPol += α²L₁(ζ)Z(X) + + // if statistical zeroknowledge is deactivated, len(h1)=len(h2)=len(h3)=len(blindedZ)-1. + // Else len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 + if i < len(h3) { + t.Mul(&h3[i], &zetaNPlusTwo). + Add(&t, &h2[i]). + Mul(&t, &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } else { + if s.opt.StatisticalZK { + t.Mul(&h2[i], &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } + } + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallel)).Msg("computeLinearizedPolynomial: inner parallel loop") + } + + return blindedZCanonical +} + +var errContextDone = errors.New("context done") + +// local copies of verification-time helpers used by prover transcript +func bindPublicData(fs *fiatshamir.Transcript, challenge string, vk *plonk_bn254.VerifyingKey, publicInputs []fr.Element) error { + if err := fs.Bind(challenge, vk.S[0].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[1].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[2].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Ql.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qr.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qm.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qo.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qk.Marshal()); err != nil { + return err + } + for i := range vk.Qcp { + if err := fs.Bind(challenge, vk.Qcp[i].Marshal()); err != nil { + return err + } + } + for i := 0; i < len(publicInputs); i++ { + if err := fs.Bind(challenge, publicInputs[i].Marshal()); err != nil { + return err + } + } + return nil +} + +func deriveRandomness(fs *fiatshamir.Transcript, challenge string, points ...*curve.G1Affine) (fr.Element, error) { + var buf [curve.SizeOfG1AffineUncompressed]byte + var r fr.Element + for _, p := range points { + buf = p.RawBytes() + if err := fs.Bind(challenge, buf[:]); err != nil { + return r, err + } + } + b, err := fs.ComputeChallenge(challenge) + if err != nil { + return r, err + } + r.SetBytes(b) + return r, nil +} + +// -------------------- GPU helpers and device setup -------------------- + +func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { + if pk.deviceInfo != nil { + return nil + } + pk.deviceInfo = &deviceInfo{} + + // Initialize ICICLE NTT domain (root of unity) and store coset generator for coset NTTs. + // ICICLE InitDomain expects a primitive root of unity; for coset transforms we use 𝔽ᵣ* generator. + + var gen fr.Element + var err error + if pk.Vk.Size < 6 { + gen, err = fft.Generator(8 * pk.Vk.Size) + if err != nil { + return err + } + } else { + gen, err = fft.Generator(4 * pk.Vk.Size) + if err != nil { + return err + } + } + genBits := gen.Bits() + limbs := icicle_core.ConvertUint64ArrToUint32Arr(genBits[:]) + // Initialize ICICLE NTT domain with root of unity + var rouIcicle icicle_bn254.ScalarField + rouIcicle.FromLimbs(limbs) + + // Store coset generator = generator of 𝔽ᵣ* (matches CPU ToLagrangeCoset) + { + cosetGen := fft.GeneratorFullMultiplicativeGroup() + cosetBits := cosetGen.Bits() + cosetLimbs := icicle_core.ConvertUint64ArrToUint32Arr(cosetBits[:]) + copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) + } + + chInitDomain := make(chan struct{}) + initDomainQueuedAt := time.Now() + icicle_runtime.RunOnDevice(device, func(args ...any) { + initDomainStartedAt := time.Now() + initCfg := icicle_core.GetDefaultNTTInitDomainConfig() + ext := config_extension.Create() + defer config_extension.Delete(ext) + fastTwiddles := envEnabled("ICICLE_NTT_FAST_TWIDDLES", true) + ext.SetBool(icicle_core.CUDA_NTT_FAST_TWIDDLES_MODE, fastTwiddles) + initCfg.Ext = ext.AsUnsafePointer() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain start vk_size=%d fast_twiddles=%t profile_full=%q profile_arbitrary=%q\n", + pk.Vk.Size, + fastTwiddles, + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + e := icicle_ntt.InitDomain(rouIcicle, initCfg) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain end status=%s call_took=%s\n", + e.AsString(), + time.Since(initDomainStartedAt), + ) + } + if e != icicle_runtime.Success { + panic("icicle: InitDomain failed") + } + close(chInitDomain) + }) + + <-chInitDomain + if isNttTrace { + fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) + } + + chLag := make(chan struct{}) + chCan := make(chan struct{}) + + if len(pk.KzgLagrange.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) + g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) + close(chLag) + }) + } else { + close(chLag) + } + + if len(pk.Kzg.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) + g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) + close(chCan) + }) + } else { + close(chCan) + } + + <-chLag + <-chCan + return nil +} + +func projectiveToGnarkAffine(p icicle_bn254.Projective) (curve.G1Affine, error) { + x, err := icicleBaseFieldToGnarkFp(p.X) + if err != nil { + return curve.G1Affine{}, err + } + y, err := icicleBaseFieldToGnarkFp(p.Y) + if err != nil { + return curve.G1Affine{}, err + } + z, err := icicleBaseFieldToGnarkFp(p.Z) + if err != nil { + return curve.G1Affine{}, err + } + if z.IsZero() { + return curve.G1Affine{}, nil + } + + var zInv fp.Element + zInv.Inverse(&z) + x.Mul(&x, &zInv) + y.Mul(&y, &zInv) + return curve.G1Affine{X: x, Y: y}, nil +} + +func icicleBaseFieldToGnarkFp(v icicle_bn254.BaseField) (fp.Element, error) { + bytes := v.ToBytesLittleEndian() + if len(bytes) != fp.Bytes { + return fp.Element{}, fmt.Errorf("invalid ICICLE base field byte length %d", len(bytes)) + } + var buf [fp.Bytes]byte + copy(buf[:], bytes) + return fp.LittleEndian.Element(&buf) +} + +func commitOnGPULagrangeDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgLagrangeDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: lagrange SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgLagrangeDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgLagrangeDevice.G1, device) +} + +func (s *instance) registerDevicePolynomialInSharedState(state *gpuPolysState, p *iop.Polynomial, dSlice icicle_core.DeviceSlice) error { + if state == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil shared state") + } + if p == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil polynomial") + } + if dSlice.IsEmpty() { + return fmt.Errorf("registerDevicePolynomialInSharedState: empty device slice") + } + + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int) + } + if idx, ok := state.polyToIdx[p]; ok { + if idx < 0 || idx >= len(state.deviceSlices) { + return fmt.Errorf("registerDevicePolynomialInSharedState: invalid index %d", idx) + } + state.deviceSlices[idx] = dSlice + state.hostSlices[idx] = nil + state.originalForm[idx] = iop.Form{Basis: p.Basis, Layout: p.Layout} + return nil + } + + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, dSlice) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + return nil +} + +func (s *instance) gpuInclusivePrefixProductOnCurrentDevice( + dVec icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + if dVec.IsEmpty() || dVec.Len() <= 1 { + return nil + } + + n := dVec.Len() + for step := 1; step < n; step <<= 1 { + src := (&dVec).Range(0, n-step, false) + dst := (&dVec).Range(step, n, false) + + tmpStd := s.getTempDeviceSlice(n - step) + if err := copyDeviceSliceIntoOnCurrentDevice(tmpStd, src, cfg); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: copy stage failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(tmpStd, cfg) + if err := icicle_vecops.VecOp(tmpStd, dst, dst, cfg, icicle_core.Mul); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: multiply stage failed: %s", err.AsString()) + } + if cfg.IsAsync { + // tmpStd is returned to pool each stage, so we must wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } + s.putTempDeviceSlice(tmpStd, n-step) + } + return nil +} + +// buildPermutationGatherIndices prepares the subset of permutation indices that +// are consumed by the copy-constraint ratio loop (only rows [0, n-1) per copy). +func buildPermutationGatherIndices(permutation []int64, nbPolynomials, n, supportLen int) ([]int64, error) { + if n <= 1 { + return nil, nil + } + total := nbPolynomials * (n - 1) + indices := make([]int64, total) + + var permBuildErr error + var permBuildErrOnce sync.Once + utils.Parallelize(total, func(start, end int) { + for k := start; k < end; k++ { + j := k / (n - 1) + i := k % (n - 1) + base := j * n + permIdx := permutation[base+i] + if permIdx < 0 || int(permIdx) >= supportLen { + jj, ii, bad := j, i, permIdx + permBuildErrOnce.Do(func() { + permBuildErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: permutation index out of range at (%d,%d): %d", jj, ii, bad) + }) + continue + } + indices[k] = permIdx + } + }) + if permBuildErr != nil { + return nil, permBuildErr + } + return indices, nil +} + +func (s *instance) prepareCopyConstraintSupportsOnCurrentDevice( + n, nbPolynomials int, + domain *fft.Domain, + permGatherIndices []int64, + cfg icicle_core.VecOpsConfig, +) (dSupportFlat, dPermFlat icicle_core.DeviceSlice, err error) { + defer func() { + if err == nil { + return + } + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + }() + + if len(permGatherIndices) == 0 { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: empty permutation gather indices") + return + } + + dOmegaStd := uploadScalarStdOnCurrentDevice(domain.Generator, cfg) + defer dOmegaStd.Free() + dShiftStd := uploadScalarStdOnCurrentDevice(domain.FrMultiplicativeGen, cfg) + defer dShiftStd.Free() + + dSupportFlat, err = allocDeviceUninitialized(nbPolynomials * n) + if err != nil { + return + } + if e := icicle_vecops.SupportIdentity(dOmegaStd, dShiftStd, n, nbPolynomials, dSupportFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: generate identity support on GPU failed: %s", e.AsString()) + return + } + toMontgomeryFormInPlaceWithCfg(dSupportFlat, cfg) + + dPermIndicesDevice := uploadInt64VectorOnCurrentDevice(permGatherIndices, cfg) + defer dPermIndicesDevice.Free() + + dPermFlat, err = allocDeviceUninitialized(len(permGatherIndices)) + if err != nil { + return + } + if e := icicle_vecops.GatherByIndices(dSupportFlat, dPermIndicesDevice, dPermFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: gather permutation support on GPU failed: %s", e.AsString()) + return + } + if cfg.IsAsync { + // Ensure temporary support/index slices are safe to free on return. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + return + } + } + + return +} + +func (s *instance) accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail icicle_core.DeviceSlice, + dBetaStd, dGammaMont icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + nMinusOne := dEntryTail.Len() + dScaled := s.getTempDeviceSlice(nMinusOne) + dTerm := s.getTempDeviceSlice(nMinusOne) + defer func() { + s.putTempDeviceSlice(dScaled, nMinusOne) + s.putTempDeviceSlice(dTerm, nMinusOne) + }() + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dID, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale identity support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add beta*id failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dNumTail, dNumTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply numerator term failed: %s", err.AsString()) + } + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dPerm, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale permutation support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add beta*sigma failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dDenTail, dDenTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply denominator term failed: %s", err.AsString()) + } + if cfg.IsAsync { + // Temp vectors are released at function exit, so ensure queued work is complete. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + } + } + + return nil +} + +// validateDeviceEntries checks that all entries are non-empty and have consistent length. +// Returns the common length n. +func validateDeviceEntries(entries []icicle_core.DeviceSlice, label string) (int, error) { + if len(entries) == 0 { + return 0, fmt.Errorf("%s: no entries", label) + } + n := entries[0].Len() + if n == 0 { + return 0, fmt.Errorf("%s: empty device entry 0", label) + } + for i := range entries { + if entries[i].IsEmpty() { + return 0, fmt.Errorf("%s: empty device entry %d", label, i) + } + if entries[i].Len() != n { + return 0, fmt.Errorf("%s: inconsistent device entry size at %d (%d != %d)", label, i, entries[i].Len(), n) + } + } + return n, nil +} + +// BuildRatioCopyConstraintIcicle builds the accumulating ratio polynomial to prove that +// [P₁ ∥ .. ∥ P_{n—1}] is invariant by the permutation \sigma. +// Namely it returns the polynomial Z whose evaluation on the j-th root of unity is +// Z(ω^j) = Π_{i 1 { + dNumTail := (&dNum).Range(1, n, false) + dDenTail := (&dDen).Range(1, n, false) + var supportErr error + dSupportFlat, dPermFlat, supportErr = s.prepareCopyConstraintSupportsOnCurrentDevice(n, nbPolynomials, domain, permGatherIndices, cfg) + if supportErr != nil { + runErr = supportErr + return + } + + dBetaStd := uploadScalarStdOnCurrentDevice(beta, cfg) + dGammaMont := uploadScalarMontOnCurrentDevice(gamma, cfg) + + for j := 0; j < nbPolynomials; j++ { + dEntryTail := (&entriesDevice[j]).Range(0, n-1, false) + baseID := j * n + dID := (&dSupportFlat).Range(baseID, baseID+n-1, false) + basePerm := j * (n - 1) + dPerm := (&dPermFlat).Range(basePerm, basePerm+(n-1), false) + if err := s.accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail, dBetaStd, dGammaMont, cfg, + ); err != nil { + runErr = err + return + } + } + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize before releasing copy-constraint workspace failed: %s", eSync.AsString()) + return + } + } + _ = dBetaStd.Free() + _ = dGammaMont.Free() + + // Support vectors and loop temps are only needed for term accumulation. + // Free them before prefix products and batch inversion, whose ICICLE + // kernels allocate additional full-domain workspace internally. + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + s.tempGPUMemPool.FreeAll() + } + + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dNum, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dDen, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + + if invErr := s.batchInvertOnCurrentDevice(dDen); invErr != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: GPU batch inversion failed: %s", invErr.AsString()) + return + } + + toStandardFormInPlace(dDen) + if err := icicle_vecops.VecOp(dDen, dNum, dNum, cfg, icicle_core.Mul); err != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: final numerator*denominatorInv multiplication failed: %s", err.AsString()) + return + } + + dResult = dNum + dNum = icicle_core.DeviceSlice{} // transfer ownership to dResult + }) + if err := <-buildDone; err != nil { + return nil, err + } + + hostMirror := make([]fr.Element, n) + if len(hostMirror) > 0 { + hostMirror[0].SetOne() + } + res := iop.NewPolynomial(&hostMirror, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if err := s.registerDevicePolynomialInSharedState(gpuState, res, dResult); err != nil { + freeSliceOnDevice(&dResult, &s.device) + return nil, err + } + + return res, nil +} diff --git a/backend/accelerated/icicle/plonk/bn254/icicle_msm_kzg_test.go b/backend/accelerated/icicle/plonk/bn254/icicle_msm_kzg_test.go new file mode 100644 index 0000000000..49488debad --- /dev/null +++ b/backend/accelerated/icicle/plonk/bn254/icicle_msm_kzg_test.go @@ -0,0 +1,374 @@ +//go:build icicle + +package bn254 + +import ( + "fmt" + "math/big" + "os" + "strconv" + "strings" + "testing" + + curve "github.com/consensys/gnark-crypto/ecc/bn254" + "github.com/consensys/gnark-crypto/ecc/bn254/fp" + "github.com/consensys/gnark-crypto/ecc/bn254/fr" + "github.com/consensys/gnark-crypto/ecc/bn254/kzg" + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" + icicle_bn254 "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bn254" + icicle_msm "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bn254/msm" + icicle_runtime "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime" +) + +type icicleMSMMode struct { + name string + scalarsOnDevice bool + basesOnDevice bool + scalarsMont bool + basesMont bool +} + +func TestKZGCommitmentICICLEMSMParity(t *testing.T) { + device, ok := requireICICLEDevice(t) + if !ok { + return + } + + requiredMode := os.Getenv("GNARK_ICICLE_MSM_PARITY_REQUIRED_MODE") + if requiredMode == "" { + requiredMode = "device/montgomery" + } + + for _, size := range msmParitySizes(t) { + t.Run(fmt.Sprintf("size=%d", size), func(t *testing.T) { + srs, err := kzg.NewSRS(uint64(size), big.NewInt(5)) + if err != nil { + t.Fatalf("NewSRS(%d): %v", size, err) + } + + lagrangeBases := append([]curve.G1Affine(nil), srs.Pk.G1[:size]...) + lagrangeBases, err = kzg.ToLagrangeG1(lagrangeBases) + if err != nil { + t.Fatalf("ToLagrangeG1(%d): %v", size, err) + } + + scalars := deterministicMSMScalars(size) + runKZGMSMParityCase(t, device, "canonical", scalars, srs.Pk.G1[:size], requiredMode) + runKZGMSMParityCase(t, device, "lagrange", scalars, lagrangeBases, requiredMode) + }) + } +} + +func requireICICLEDevice(t *testing.T) (*icicle_runtime.Device, bool) { + t.Helper() + if err := icicle_runtime.LoadBackendFromEnvOrDefault(); err != icicle_runtime.Success { + t.Skipf("ICICLE backend not available: %s", err.AsString()) + } + if nDev, err := icicle_runtime.GetDeviceCount(); err != icicle_runtime.Success || nDev == 0 { + t.Skip("no ICICLE devices detected") + } + device := icicle_runtime.CreateDevice("CUDA", 0) + return &device, true +} + +func msmParitySizes(t *testing.T) []int { + t.Helper() + env := strings.TrimSpace(os.Getenv("GNARK_ICICLE_MSM_PARITY_SIZES")) + if env == "" { + return []int{16, 256, 4096} + } + + parts := strings.Split(env, ",") + sizes := make([]int, 0, len(parts)) + for _, part := range parts { + part = strings.TrimSpace(part) + if part == "" { + continue + } + size, err := strconv.Atoi(part) + if err != nil || size < 2 { + t.Fatalf("invalid GNARK_ICICLE_MSM_PARITY_SIZES entry %q", part) + } + sizes = append(sizes, size) + } + if len(sizes) == 0 { + t.Fatalf("GNARK_ICICLE_MSM_PARITY_SIZES did not contain any sizes") + } + return sizes +} + +func deterministicMSMScalars(size int) []fr.Element { + scalars := make([]fr.Element, size) + for i := range scalars { + // Keep values deterministic but non-uniform to exercise bucket carries. + scalars[i].SetUint64(uint64((i+3)*(i+17)) + uint64(i%11+1)) + } + return scalars +} + +func runKZGMSMParityCase( + t *testing.T, + device *icicle_runtime.Device, + name string, + scalars []fr.Element, + bases []curve.G1Affine, + requiredMode string, +) { + t.Helper() + + want, err := kzg.Commit(scalars, kzg.ProvingKey{G1: bases}) + if err != nil { + t.Fatalf("%s: CPU KZG commit failed: %v", name, err) + } + + modes := []icicleMSMMode{ + {name: "device/montgomery", scalarsOnDevice: true, basesOnDevice: true, scalarsMont: true, basesMont: true}, + {name: "host/montgomery", scalarsOnDevice: false, basesOnDevice: false, scalarsMont: true, basesMont: true}, + {name: "device/raw", scalarsOnDevice: true, basesOnDevice: true, scalarsMont: false, basesMont: false}, + {name: "device/scalar-montgomery", scalarsOnDevice: true, basesOnDevice: true, scalarsMont: true, basesMont: false}, + {name: "device/base-montgomery", scalarsOnDevice: true, basesOnDevice: true, scalarsMont: false, basesMont: true}, + } + + matchedRequired := requiredMode == "any" + modeMatches := make([]string, 0, len(modes)) + for _, mode := range modes { + conversions, err := runICICLEMSMMode(device, scalars, bases, mode) + if err != nil { + t.Fatalf("%s/%s: ICICLE MSM failed: %v", name, mode.name, err) + } + + matches := make([]string, 0, len(conversions)) + for conversionName, got := range conversions { + if got.Equal(&want) { + matches = append(matches, conversionName) + } + } + + if len(matches) == 0 { + t.Logf("%s/%s: no conversion matched CPU KZG", name, mode.name) + continue + } + t.Logf("%s/%s: matched CPU KZG with conversions %s", name, mode.name, strings.Join(matches, ",")) + modeMatches = append(modeMatches, mode.name) + if mode.name == requiredMode { + matchedRequired = true + } + } + + if !matchedRequired { + t.Fatalf("%s: required ICICLE MSM mode %q did not match CPU KZG; matching modes: %s", name, requiredMode, strings.Join(modeMatches, ",")) + } + + chunkSize := len(scalars) / 3 + if chunkSize < 1 { + chunkSize = 1 + } + if chunkSize > 257 { + chunkSize = 257 + } + chunked, err := runICICLEChunkedKZG(device, scalars, bases, chunkSize) + if err != nil { + t.Fatalf("%s/device-chunked-%d: ICICLE MSM failed: %v", name, chunkSize, err) + } + if !chunked.Equal(&want) { + t.Fatalf("%s/device-chunked-%d: chunked ICICLE MSM did not match CPU KZG", name, chunkSize) + } + t.Logf("%s/device-chunked-%d: matched CPU KZG", name, chunkSize) +} + +func runICICLEMSMMode( + device *icicle_runtime.Device, + scalars []fr.Element, + bases []curve.G1Affine, + mode icicleMSMMode, +) (map[string]curve.G1Affine, error) { + if len(scalars) == 0 { + return nil, fmt.Errorf("empty scalar slice") + } + if len(scalars) > len(bases) { + return nil, fmt.Errorf("scalar/basis mismatch: %d > %d", len(scalars), len(bases)) + } + + out := make(chan struct { + conversions map[string]curve.G1Affine + err error + }, 1) + + icicle_runtime.RunOnDevice(device, func(args ...any) { + result, err := runICICLEMSMModeOnCurrentDevice(scalars, bases[:len(scalars)], mode) + out <- struct { + conversions map[string]curve.G1Affine + err error + }{conversions: result, err: err} + }) + + result := <-out + return result.conversions, result.err +} + +func runICICLEMSMModeOnCurrentDevice( + scalars []fr.Element, + bases []curve.G1Affine, + mode icicleMSMMode, +) (map[string]curve.G1Affine, error) { + scalarsHost := icicle_core.HostSliceFromElements(scalars) + basesHost := (icicle_core.HostSlice[curve.G1Affine])(bases) + + var scalarsInput icicle_core.HostOrDeviceSlice = scalarsHost + var basesInput icicle_core.HostOrDeviceSlice = basesHost + var scalarsDevice icicle_core.DeviceSlice + var basesDevice icicle_core.DeviceSlice + + if mode.scalarsOnDevice { + scalarsHost.CopyToDevice(&scalarsDevice, true) + defer scalarsDevice.Free() + scalarsInput = scalarsDevice + } + if mode.basesOnDevice { + basesHost.CopyToDevice(&basesDevice, true) + defer basesDevice.Free() + basesInput = basesDevice + } + + res := make(icicle_core.HostSlice[icicle_bn254.Projective], 1) + cfg := icicle_msm.GetDefaultMSMConfig() + cfg.AreScalarsMontgomeryForm = mode.scalarsMont + cfg.AreBasesMontgomeryForm = mode.basesMont + if e := icicle_msm.Msm(scalarsInput, basesInput, &cfg, res); e != icicle_runtime.Success { + return nil, fmt.Errorf("%s", e.AsString()) + } + + conversions := make(map[string]curve.G1Affine, 4) + regularHomogeneous, err := icicleProjectiveRegularHomogeneousToGnarkAffine(res[0]) + if err != nil { + return nil, fmt.Errorf("regular homogeneous conversion: %w", err) + } + conversions["regular-homogeneous"] = regularHomogeneous + + regularAffine, err := icicleProjectiveRegularViaIcicleToGnarkAffine(res[0]) + if err != nil { + return nil, fmt.Errorf("regular affine conversion: %w", err) + } + conversions["regular-icicle-affine"] = regularAffine + conversions["montgomery-homogeneous"] = icicleProjectiveMontgomeryHomogeneousToGnarkAffine(res[0]) + conversions["montgomery-jacobian"] = icicleProjectiveMontgomeryJacobianToGnarkAffine(res[0]) + return conversions, nil +} + +func runICICLEChunkedKZG( + device *icicle_runtime.Device, + scalars []fr.Element, + bases []curve.G1Affine, + chunkSize int, +) (curve.G1Affine, error) { + out := make(chan struct { + commit curve.G1Affine + err error + }, 1) + + icicle_runtime.RunOnDevice(device, func(args ...any) { + scalarsHost := icicle_core.HostSliceFromElements(scalars) + basesHost := (icicle_core.HostSlice[curve.G1Affine])(bases[:len(scalars)]) + + var scalarsDevice icicle_core.DeviceSlice + var basesDevice icicle_core.DeviceSlice + scalarsHost.CopyToDevice(&scalarsDevice, true) + basesHost.CopyToDevice(&basesDevice, true) + defer scalarsDevice.Free() + defer basesDevice.Free() + + commit, err := commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) + out <- struct { + commit curve.G1Affine + err error + }{commit: commit, err: err} + }) + + result := <-out + return result.commit, result.err +} + +func icicleProjectiveRegularViaIcicleToGnarkAffine(p icicle_bn254.Projective) (curve.G1Affine, error) { + a := p.ToAffine() + x, err := icicleBaseFieldRegularToGnarkFp(a.X) + if err != nil { + return curve.G1Affine{}, err + } + y, err := icicleBaseFieldRegularToGnarkFp(a.Y) + if err != nil { + return curve.G1Affine{}, err + } + return curve.G1Affine{X: x, Y: y}, nil +} + +func icicleProjectiveRegularHomogeneousToGnarkAffine(p icicle_bn254.Projective) (curve.G1Affine, error) { + x, err := icicleBaseFieldRegularToGnarkFp(p.X) + if err != nil { + return curve.G1Affine{}, err + } + y, err := icicleBaseFieldRegularToGnarkFp(p.Y) + if err != nil { + return curve.G1Affine{}, err + } + z, err := icicleBaseFieldRegularToGnarkFp(p.Z) + if err != nil { + return curve.G1Affine{}, err + } + if z.IsZero() { + return curve.G1Affine{}, nil + } + var zInv fp.Element + zInv.Inverse(&z) + x.Mul(&x, &zInv) + y.Mul(&y, &zInv) + return curve.G1Affine{X: x, Y: y}, nil +} + +func icicleProjectiveMontgomeryHomogeneousToGnarkAffine(p icicle_bn254.Projective) curve.G1Affine { + x := icicleBaseFieldMontgomeryToGnarkFp(p.X) + y := icicleBaseFieldMontgomeryToGnarkFp(p.Y) + z := icicleBaseFieldMontgomeryToGnarkFp(p.Z) + if z.IsZero() { + return curve.G1Affine{} + } + var zInv fp.Element + zInv.Inverse(&z) + x.Mul(&x, &zInv) + y.Mul(&y, &zInv) + return curve.G1Affine{X: x, Y: y} +} + +func icicleProjectiveMontgomeryJacobianToGnarkAffine(p icicle_bn254.Projective) curve.G1Affine { + x := icicleBaseFieldMontgomeryToGnarkFp(p.X) + y := icicleBaseFieldMontgomeryToGnarkFp(p.Y) + z := icicleBaseFieldMontgomeryToGnarkFp(p.Z) + if z.IsZero() { + return curve.G1Affine{} + } + var zInv, zInv2 fp.Element + zInv.Inverse(&z) + zInv2.Square(&zInv) + x.Mul(&x, &zInv2) + y.Mul(&y, &zInv2).Mul(&y, &zInv) + return curve.G1Affine{X: x, Y: y} +} + +func icicleBaseFieldRegularToGnarkFp(v icicle_bn254.BaseField) (fp.Element, error) { + bytes := v.ToBytesLittleEndian() + if len(bytes) != fp.Bytes { + return fp.Element{}, fmt.Errorf("invalid byte length %d", len(bytes)) + } + var buf [fp.Bytes]byte + copy(buf[:], bytes) + return fp.LittleEndian.Element(&buf) +} + +func icicleBaseFieldMontgomeryToGnarkFp(v icicle_bn254.BaseField) fp.Element { + limbs := v.GetLimbs() + return fp.Element{ + uint64(limbs[0]) | uint64(limbs[1])<<32, + uint64(limbs[2]) | uint64(limbs[3])<<32, + uint64(limbs[4]) | uint64(limbs[5])<<32, + uint64(limbs[6]) | uint64(limbs[7])<<32, + } +} diff --git a/backend/accelerated/icicle/plonk/bn254/icicle_ntt_test.go b/backend/accelerated/icicle/plonk/bn254/icicle_ntt_test.go new file mode 100644 index 0000000000..c45b04fda1 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bn254/icicle_ntt_test.go @@ -0,0 +1,125 @@ +//go:build icicle + +package bn254 + +import ( + "testing" + + "github.com/consensys/gnark-crypto/ecc/bn254/fr" + "github.com/consensys/gnark-crypto/ecc/bn254/fr/fft" + "github.com/consensys/gnark-crypto/ecc/bn254/fr/iop" + plonk_bn254 "github.com/consensys/gnark/backend/plonk/bn254" + icicle_runtime "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime" +) + +// TestGpuNTTInverseBatch_matchesCPUToCanonical validates that the GPU batch inverse NTT +// produces the exact same canonical coefficients and final layout as the CPU ToCanonical. +func TestGpuNTTInverseBatch_matchesCPUToCanonical(t *testing.T) { + // Ensure ICICLE backend is available + if err := icicle_runtime.LoadBackendFromEnvOrDefault(); err != icicle_runtime.Success { + t.Skipf("ICICLE backend not available: %s", err.AsString()) + } + if nDev, err := icicle_runtime.GetDeviceCount(); err != icicle_runtime.Success || nDev == 0 { + t.Skip("No ICICLE devices detected; skipping GPU symmetry test") + } + + device := icicle_runtime.CreateDevice("CUDA", 0) + + // Domain size for the test; small power of two is sufficient + const n uint64 = 256 + d := fft.NewDomain(n) + + // Build a minimal proving key with only Vk.Size set and initialize device pointers. + // setupDevicePointers will also initialize ICICLE NTT domain on the device. + vk := &plonk_bn254.VerifyingKey{Size: n} + pk := &ProvingKey{ProvingKey: plonk_bn254.ProvingKey{Vk: vk}} + if err := pk.setupDevicePointers(&device); err != nil { + t.Fatalf("setupDevicePointers failed: %v", err) + } + + // Create random canonical coefficients + makeRandCoeffs := func() []fr.Element { + cp := make([]fr.Element, d.Cardinality) + for i := range cp { + cp[i].SetRandom() + } + return cp + } + + // Helpers to build forms + newCanonical := func(coeffs []fr.Element) *iop.Polynomial { + return iop.NewPolynomial(&coeffs, iop.Form{Basis: iop.Canonical, Layout: iop.Regular}) + } + clonePoly := func(p *iop.Polynomial) *iop.Polynomial { + src := p.Coefficients() + dst := make([]fr.Element, len(src)) + copy(dst, src) + return iop.NewPolynomial(&dst, iop.Form{Basis: p.Basis, Layout: p.Layout}) + } + + // Prepare inputs covering: + // - Lagrange Regular + // - Lagrange BitReverse + // - LagrangeCoset Regular + // - LagrangeCoset BitReverse + base := newCanonical(makeRandCoeffs()) + + lagReg := clonePoly(base).ToLagrange(d).ToRegular() + lagBR := clonePoly(base).ToLagrange(d) // BitReverse layout by construction from canonical Regular + + cosReg := clonePoly(base).ToLagrangeCoset(d).ToRegular() + cosBR := clonePoly(base).ToLagrangeCoset(d) // BitReverse layout by construction from canonical Regular + + // CPU copies + cpuLagReg := clonePoly(lagReg) + cpuLagBR := clonePoly(lagBR) + cpuCosReg := clonePoly(cosReg) + cpuCosBR := clonePoly(cosBR) + + // GPU copies + gpuLagReg := clonePoly(lagReg) + gpuLagBR := clonePoly(lagBR) + gpuCosReg := clonePoly(cosReg) + gpuCosBR := clonePoly(cosBR) + + // CPU: convert to canonical + _ = cpuLagReg.ToCanonical(d) + _ = cpuLagBR.ToCanonical(d) + _ = cpuCosReg.ToCanonical(d) + _ = cpuCosBR.ToCanonical(d) + + // GPU: batch inverse to canonical + inst := instance{ + pk: pk, + device: device, + domain0: d, + } + inst.gpuNTTInverseBatch([]*iop.Polynomial{ + gpuLagReg, gpuLagBR, gpuCosReg, gpuCosBR, + }, pk) + + // Compare coefficients and layouts + compare := func(name string, want, got *iop.Polynomial) { + if want.Basis != iop.Canonical || got.Basis != iop.Canonical { + t.Fatalf("%s: expected Canonical basis; got want=%v got=%v", name, want.Basis, got.Basis) + } + // GPU path returns Regular layout; normalize both to Regular before comparing. + want.ToRegular() + got.ToRegular() + wc := want.Coefficients() + gc := got.Coefficients() + if len(wc) != len(gc) { + t.Fatalf("%s: coeff length mismatch; want=%d got=%d", name, len(wc), len(gc)) + } + for i := range wc { + if !wc[i].Equal(&gc[i]) { + t.Fatalf("%s: coeff[%d] mismatch", name, i) + } + } + } + + compare("Lagrange/Regular", cpuLagReg, gpuLagReg) + compare("Lagrange/BitReverse", cpuLagBR, gpuLagBR) + compare("LagrangeCoset/Regular", cpuCosReg, gpuCosReg) + compare("LagrangeCoset/BitReverse", cpuCosBR, gpuCosBR) +} diff --git a/backend/accelerated/icicle/plonk/bn254/provingkey.go b/backend/accelerated/icicle/plonk/bn254/provingkey.go new file mode 100644 index 0000000000..ceb68da958 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bn254/provingkey.go @@ -0,0 +1,100 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +//go:build icicle + +package bn254 + +import ( + "sync" + "time" + + "github.com/consensys/gnark-crypto/ecc/bn254/fr" + "github.com/consensys/gnark-crypto/ecc/bn254/fr/fft" + plonk_bn254 "github.com/consensys/gnark/backend/plonk/bn254" + cs "github.com/consensys/gnark/constraint/bn254" + "github.com/consensys/gnark/logger" + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" +) + +// deviceInfo holds device-resident buffers for GPU acceleration. +type deviceInfo struct { + CosetGenerator [fr.Limbs * 2]uint32 + KzgDevice struct { + G1 icicle_core.DeviceSlice + } + KzgLagrangeDevice struct { + G1 icicle_core.DeviceSlice + } +} + +// hostSetup holds host-side, witness-independent prover state derived from the +// constraint system: the FFT domains and the PLONK trace (selector + +// permutation polynomials). Building the trace walks every constraint (~3s at +// 23M constraints), so it is computed once per proving key and shared across +// proofs. Everything here is read-only during proving: the prover clones Qk +// before patching public inputs into it, and every basis conversion of a trace +// polynomial copies first (see canonicalRegularCoefficientsCopy). +type hostSetup struct { + sizeSystem uint64 + domain0 *fft.Domain + domain1 *fft.Domain + trace *plonk_bn254.Trace +} + +// ProvingKey wraps the native PLONK proving key with device-resident state +// (KZG bases, NTT domains, cached trace) that is uploaded once and reused +// across Prove calls. +// +// Concurrency: Prove calls sharing the same ProvingKey must be serialized by +// the caller. The device state hangs off the key and proofs share a single +// GPU; concurrent proves against the same key are not safe. +type ProvingKey struct { + plonk_bn254.ProvingKey + *deviceInfo + hostSetupOnce sync.Once + hostSetup *hostSetup +} + +func buildHostSetup(spr *cs.SparseR1CS, sizeSystem uint64) *hostSetup { + domain0 := fft.NewDomain(sizeSystem) + + // h, the quotient polynomial is of degree 3(n+1)+2, so it's in a 3(n+2) dim + // vector space, the domain is the next power of 2 superior to 3(n+2). + // 4*domainNum is enough in all cases except when n<6. + var domain1 *fft.Domain + if sizeSystem < 6 { + domain1 = fft.NewDomain(8*sizeSystem, fft.WithoutPrecompute()) + } else { + domain1 = fft.NewDomain(4*sizeSystem, fft.WithoutPrecompute()) + } + + return &hostSetup{ + sizeSystem: sizeSystem, + domain0: domain0, + domain1: domain1, + trace: plonk_bn254.NewTrace(spr, domain0), + } +} + +// hostSetupFor returns the FFT domains and trace for spr, building them on +// first use and caching them on the proving key. A PLONK proving key is bound +// to exactly one constraint system, so per-key caching is sound; as a +// defensive measure a system-size mismatch falls back to an uncached build +// rather than ever serving another circuit's trace. +func (pk *ProvingKey) hostSetupFor(spr *cs.SparseR1CS) *hostSetup { + nbConstraints := spr.GetNbConstraints() + sizeSystem := uint64(nbConstraints + len(spr.Public)) // len(spr.Public) is for the placeholder constraints + pk.hostSetupOnce.Do(func() { + start := time.Now() + pk.hostSetup = buildHostSetup(spr, sizeSystem) + log := logger.Logger() + log.Debug().Dur("took", time.Since(start)).Msg("built prover host setup (fft domains + trace)") + }) + if pk.hostSetup.sizeSystem != sizeSystem { + return buildHostSetup(spr, sizeSystem) + } + return pk.hostSetup +} diff --git a/backend/accelerated/icicle/plonk/bw6-761/doc.go b/backend/accelerated/icicle/plonk/bw6-761/doc.go new file mode 100644 index 0000000000..f2afd3da21 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bw6-761/doc.go @@ -0,0 +1,7 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +// Package bw6761 implements ICICLE acceleration for BW6-761 PLONK backend. +package bw6761 diff --git a/backend/accelerated/icicle/plonk/bw6-761/icicle.go b/backend/accelerated/icicle/plonk/bw6-761/icicle.go new file mode 100644 index 0000000000..0796f293f0 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bw6-761/icicle.go @@ -0,0 +1,7093 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +//go:build icicle + +package bw6761 + +import ( + "context" + "errors" + "fmt" + "hash" + "io" + "math/big" + "math/bits" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + "unsafe" + + "golang.org/x/sync/errgroup" + + "github.com/consensys/gnark/backend" + plonk_bw6761 "github.com/consensys/gnark/backend/plonk/bw6-761" + "github.com/consensys/gnark/backend/witness" + constraint "github.com/consensys/gnark/constraint" + cs "github.com/consensys/gnark/constraint/bw6-761" + "github.com/consensys/gnark/constraint/solver" + fcs "github.com/consensys/gnark/frontend/cs" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + + "github.com/consensys/gnark-crypto/ecc" + curve "github.com/consensys/gnark-crypto/ecc/bw6-761" + "github.com/consensys/gnark-crypto/ecc/bw6-761/fp" + "github.com/consensys/gnark-crypto/ecc/bw6-761/fr" + "github.com/consensys/gnark-crypto/ecc/bw6-761/fr/fft" + "github.com/consensys/gnark-crypto/ecc/bw6-761/fr/hash_to_field" + "github.com/consensys/gnark-crypto/ecc/bw6-761/fr/iop" + "github.com/consensys/gnark-crypto/ecc/bw6-761/kzg" + fiatshamir "github.com/consensys/gnark-crypto/fiat-shamir" + + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" + icicle_bw6761 "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bw6761" + icicle_msm "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bw6761/msm" + icicle_ntt "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bw6761/ntt" + icicle_vecops "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/curves/bw6761/vecOps" + icicle_runtime "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime" + "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime/config_extension" +) + +const HasIcicle = true + +var isProfileMode bool + +var useBlinding bool + +var isNttTrace bool + +func init() { + _, isProfileMode = os.LookupEnv("ICICLE_STEP_PROFILE") + // Blinding polynomials (zero-knowledge) are enabled by default, matching the + // native prover. Set GNARK_DISABLE_BLINDING to trade zero-knowledge for a + // faster, deterministic prover (e.g. when the witness is not secret). + _, disableBlinding := os.LookupEnv("GNARK_DISABLE_BLINDING") + useBlinding = !disableBlinding + isNttTrace = envEnabled("ICICLE_NTT_TRACE", false) +} + +// profileStep returns a function that, when called, logs the elapsed time since +// profileStep was invoked. If profiling is disabled, it returns a no-op. +// Usage: done := profileStep("label"); defer done() +func profileStep(msg string) func() { + if !isProfileMode { + return func() {} + } + start := time.Now() + return func() { + l := logger.Logger() + l.Debug().Dur("took", time.Since(start)).Msg(msg) + } +} + +// stageTiming is a single recorded prover stage and its wall-clock duration. +type stageTiming struct { + name string + dur time.Duration +} + +// stageTimings is a concurrency-safe, ordered recorder of prover stage +// durations. The PLONK prover runs its stages as concurrent goroutines, so the +// recorded durations OVERLAP and do not sum to the total — the printed table +// flags this. +type stageTimings struct { + mu sync.Mutex + entries []stageTiming +} + +// record appends a (stage, duration) entry. Safe to call from any goroutine and +// safe on a nil receiver (records nothing). +func (t *stageTimings) record(name string, d time.Duration) { + if t == nil { + return + } + t.mu.Lock() + t.entries = append(t.entries, stageTiming{name: name, dur: d}) + t.mu.Unlock() +} + +// printTable writes an aligned breakdown of the recorded stages to w, sorted by +// duration (largest first), followed by the overall prover total. Stages run +// concurrently, so the rows overlap and intentionally do not sum to the total. +func (t *stageTimings) printTable(w io.Writer, total time.Duration) { + if t == nil { + return + } + t.mu.Lock() + rows := make([]stageTiming, len(t.entries)) + copy(rows, t.entries) + t.mu.Unlock() + + sort.SliceStable(rows, func(i, j int) bool { return rows[i].dur > rows[j].dur }) + + nameW := len("TOTAL (prover done)") + for _, r := range rows { + if len(r.name) > nameW { + nameW = len(r.name) + } + } + + fmt.Fprintln(w, "") + fmt.Fprintln(w, "================ gnark PLONK prove breakdown (GPU) ================") + fmt.Fprintln(w, "(stages run concurrently — durations overlap and do not sum to TOTAL)") + fmt.Fprintf(w, " %-*s %12s %6s\n", nameW, "STAGE", "TIME", "%TOTAL") + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + for _, r := range rows { + pct := 0.0 + if total > 0 { + pct = 100 * float64(r.dur) / float64(total) + } + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, r.name, r.dur.Round(time.Millisecond), pct) + } + fmt.Fprintf(w, " %s %12s %6s\n", strings.Repeat("-", nameW), "------------", "------") + fmt.Fprintf(w, " %-*s %12s %5.1f%%\n", nameW, "TOTAL (prover done)", total.Round(time.Millisecond), 100.0) + fmt.Fprintln(w, "===================================================================") + fmt.Fprintln(w, "") +} + +func envEnabled(key string, defaultVal bool) bool { + v, ok := os.LookupEnv(key) + if !ok { + return defaultVal + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "0", "false", "off", "no": + return false + default: + return true + } +} + +func nttAlgorithmFromEnv(key string, fallback icicle_core.NttAlgorithm) icicle_core.NttAlgorithm { + v, ok := os.LookupEnv(key) + if !ok { + return fallback + } + switch strings.ToLower(strings.TrimSpace(v)) { + case "", "auto", "0": + return icicle_core.Auto + case "radix2", "radix-2", "r2", "1": + return icicle_core.Radix2 + case "mixed", "mixedradix", "mixed-radix", "2": + return icicle_core.MixedRadix + default: + return fallback + } +} + +const ( + id_L int = iota + id_R + id_O + id_Z + id_ZS + id_Ql + id_Qr + id_Qm + id_Qo + id_Qk + id_S1 + id_S2 + id_S3 + id_Qci // [ .. , Qc_i, Pi_i, ...] +) + +// blinding factors +const ( + id_Bl int = iota + id_Br + id_Bo + id_Bz + nb_blinding_polynomials +) + +// blinding orders (-1 to deactivate) +const ( + order_blinding_L = 1 + order_blinding_R = 1 + order_blinding_O = 1 + order_blinding_Z = 2 +) + +// Prove generates a PLONK proof. When the accelerator option is not set to +// "icicle", we delegate to the native prover. Otherwise, we run a local copy +// of the CPU prover logic to enable incremental GPU adaptation. +func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts ...backend.ProverOption) (*plonk_bw6761.Proof, error) { + opt, err := backend.NewProverConfig(opts...) + if err != nil { + return nil, err + } + + log := logger.Logger().With(). + Str("curve", spr.CurveID().String()). + Int("nbConstraints", spr.GetNbConstraints()). + Str("backend", "plonk").Logger() + + // parse the options + if opt.HashToFieldFn == nil { + opt.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + + // When blinding is disabled (GNARK_DISABLE_BLINDING), also disable StatisticalZK, it makes no sense + // to use statistical zero knowledge when we don't use blinding. + if !useBlinding { + opt.StatisticalZK = false + } + + start := time.Now() + + // Initialize device and preload KZG bases once per proving key + device := icicle_runtime.CreateDevice("CUDA", 0) + if pk.deviceInfo == nil { + if err := pk.setupDevicePointers(&device); err != nil { + return nil, err + } + } + + // init instance + g, ctx := errgroup.WithContext(context.Background()) + instance, err := newInstance(ctx, spr, pk, fullWitness, &opt) + if err != nil { + return nil, fmt.Errorf("new instance: %w", err) + } + // attach device to instance for GPU calls + instance.device = device + instance.initSharedGPUState() + defer instance.releaseTempGPUMemoryPool() + defer instance.releaseSharedGPUState() + defer instance.releaseLinearizedEvalGPUState() + + // solve constraints + g.Go(instance.solveConstraints) + + // complete qk + g.Go(instance.completeQk) + + // init blinding polynomials + g.Go(instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + g.Go(instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + g.Go(instance.buildRatioCopyConstraint) + + // compute h + g.Go(instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + g.Go(instance.openZ) + + // linearized polynomial + g.Go(instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + g.Go(func() error { + startBatchOpening := time.Now() + err := instance.batchOpening() + if isProfileMode { + instance.timings.record("batchOpening (folded KZG)", time.Since(startBatchOpening)) + } + return err + }) + + if err := g.Wait(); err != nil { + return nil, err + } + + total := time.Since(start) + log.Debug().Dur("took", total).Msg("prover done") + if isProfileMode { + instance.timings.printTable(os.Stderr, total) + } + return instance.proof, nil +} + +// represents a Prover instance +type instance struct { + ctx context.Context + + pk *ProvingKey + proof *plonk_bw6761.Proof + spr *cs.SparseR1CS + opt *backend.ProverConfig + + fs *fiatshamir.Transcript + kzgFoldingHash hash.Hash // for KZG folding + htfFunc hash.Hash // hash to field function + + // polynomials + polyL, polyR, polyO *iop.Polynomial + polyZ, polyZS, polyQk *iop.Polynomial + bp []*iop.Polynomial // blinding polynomials + h *iop.Polynomial // h is the quotient polynomial + hGPU *gpuQuotientPolynomial + polyZLagrangeGPU icicle_core.DeviceSlice + blindedZCanonicalGPU icicle_core.DeviceSlice + quotientShardsRandomizers [2]fr.Element // random elements for blinding the shards of the quotient + + linearizedPolynomial []fr.Element + linearizedPolynomialGPU icicle_core.DeviceSlice + linearizedPolynomialClaim fr.Element + linearizedPolynomialDigest kzg.Digest + + fullWitness witness.Witness + + // bsb22 commitment stuff + commitmentInfo constraint.PlonkCommitments + commitmentVal []fr.Element + cCommitments []*iop.Polynomial + + // challenges + gamma, beta, alpha, zeta fr.Element + + // channel to wait for the steps + chLRO, + chQk, + chbp, + chZ, + chH, + chRestoreLRO, + chZOpening, + chLinearizedPolynomial, + chGammaBeta chan struct{} + + domain0, domain1 *fft.Domain + + trace *plonk_bw6761.Trace + + // GPU device handle + device icicle_runtime.Device + + // Shared GPU polynomial context reused across buildRatioCopyConstraint + // and computeQuotient to avoid repeated host<->device uploads. + gpuStateMu sync.Mutex + sharedGPUState *gpuPolysState + // Snapshot of immutable polynomial slices used by computeLinearizedPolynomial + // for zeta evaluations after computeQuotient mutates/frees shared state. + linearizedEvalGPUState *gpuPolysState + + // Reusable temporary GPU memory pool for non-state buffers. + tempGPUMemPool *gpuMemoryPool + + // Per-prove stage-timing recorder (used to print the breakdown table when + // ICICLE_STEP_PROFILE is set). + timings *stageTimings +} + +func newInstance(ctx context.Context, spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts *backend.ProverConfig) (*instance, error) { + if opts.HashToFieldFn == nil { + opts.HashToFieldFn = hash_to_field.New([]byte("BSB22-Plonk")) + } + s := instance{ + ctx: ctx, + pk: pk, + proof: &plonk_bw6761.Proof{}, + spr: spr, + opt: opts, + fullWitness: fullWitness, + bp: make([]*iop.Polynomial, nb_blinding_polynomials), + fs: fiatshamir.NewTranscript(opts.ChallengeHash, "gamma", "beta", "alpha", "zeta"), + kzgFoldingHash: opts.KZGFoldingHash, + htfFunc: opts.HashToFieldFn, + chLRO: make(chan struct{}, 1), + chQk: make(chan struct{}, 1), + chbp: make(chan struct{}, 1), + chGammaBeta: make(chan struct{}, 1), + chZ: make(chan struct{}, 1), + chH: make(chan struct{}, 1), + chZOpening: make(chan struct{}, 1), + chLinearizedPolynomial: make(chan struct{}, 1), + chRestoreLRO: make(chan struct{}, 1), + tempGPUMemPool: newGPUMemoryPool(), + timings: &stageTimings{}, + } + s.initBSB22Commitments() + + // FFT domains and the PLONK trace are witness-independent and expensive to + // build at large n (NewTrace walks every constraint), so they are cached + // on the proving key and shared read-only across proofs. + setup := pk.hostSetupFor(spr) + s.domain0 = setup.domain0 + s.domain1 = setup.domain1 + s.trace = setup.trace + + // sampling random numbers for blinding the quotient + if opts.StatisticalZK { + s.quotientShardsRandomizers[0].SetRandom() + s.quotientShardsRandomizers[1].SetRandom() + } + + return &s, nil +} + +func (s *instance) initBlindingPolynomials() error { + if !useBlinding { + // When blinding is disabled (GNARK_DISABLE_BLINDING), skip creating blinding polynomials entirely + // Just close the channel to unblock any goroutines waiting on it + close(s.chbp) + return nil + } + + s.bp[id_Bl] = getRandomPolynomial(order_blinding_L) + s.bp[id_Br] = getRandomPolynomial(order_blinding_R) + s.bp[id_Bo] = getRandomPolynomial(order_blinding_O) + s.bp[id_Bz] = getRandomPolynomial(order_blinding_Z) + close(s.chbp) + return nil +} + +func (s *instance) initBSB22Commitments() { + s.commitmentInfo = s.spr.CommitmentInfo.(constraint.PlonkCommitments) + s.commitmentVal = make([]fr.Element, len(s.commitmentInfo)) // TODO @Tabaie get rid of this + s.cCommitments = make([]*iop.Polynomial, len(s.commitmentInfo)) + s.proof.Bsb22Commitments = make([]kzg.Digest, len(s.commitmentInfo)) + + // override the hint for the commitment constraints + bsb22ID := solver.GetHintID(fcs.Bsb22CommitmentComputePlaceholder) + s.opt.SolverOpts = append(s.opt.SolverOpts, solver.OverrideHint(bsb22ID, s.bsb22Hint)) +} + +// Computing and verifying Bsb22 multi-commits explained in https://hackmd.io/x8KsadW3RRyX7YTCFJIkHg +func (s *instance) bsb22Hint(_ *big.Int, ins, outs []*big.Int) error { + var err error + commDepth := int(ins[0].Int64()) + ins = ins[1:] + + res := &s.commitmentVal[commDepth] + + commitmentInfo := s.spr.CommitmentInfo.(constraint.PlonkCommitments)[commDepth] + committedValues := make([]fr.Element, s.domain0.Cardinality) + offset := s.spr.GetNbPublicVariables() + for i := range ins { + committedValues[offset+commitmentInfo.Committed[i]].SetBigInt(ins[i]) + } + if _, err = committedValues[offset+commitmentInfo.CommitmentIndex].SetRandom(); err != nil { // Commitment injection constraint has qcp = 0. Safe to use for blinding. + return err + } + if _, err = committedValues[offset+s.spr.GetNbConstraints()-1].SetRandom(); err != nil { // Last constraint has qcp = 0. Safe to use for blinding + return err + } + s.cCommitments[commDepth] = iop.NewPolynomial(&committedValues, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[commDepth], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[commDepth]); err != nil { + return err + } + + s.htfFunc.Write(s.proof.Bsb22Commitments[commDepth].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + res.SetBytes(hashBts[:nbBuf]) // TODO @Tabaie use CommitmentIndex for this; create a new variable CommitmentConstraintIndex for other uses + res.BigInt(outs[0]) + + return nil +} + +// solveConstraints computes the evaluation of the L, R, O polynomials in Lagrange form. +func (s *instance) solveConstraints() error { + startSolve := time.Now() + log := logger.Logger() + + var solution *cs.SparseR1CSSolution + + // Try to load raw solver values from cache (fastest path) + rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + if rawCachePath != "" { + if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { + // Reconstruct L, R, O from raw values + var sol cs.SparseR1CSSolution + if sol.L, sol.R, sol.O, err = s.spr.EvaluateLROSmallDomainFromValues(rawValues); err != nil { + log.Warn().Err(err).Str("file", rawCachePath).Msg("ignoring raw solver cache") + } else { + log.Debug().Dur("took", time.Since(startSolve)).Int("wires", len(rawValues)).Msg("loaded raw solver values from cache") + solution = &sol + } + + // Load cached BSB22 cCommitments polynomials + cacheDir := filepath.Dir(rawCachePath) + for i := 0; solution != nil && i < len(s.commitmentInfo); i++ { + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + coeffs, err := cs.LoadRawSolverValues(bsb22Path) + if err != nil { + log.Warn().Err(err).Int("i", i).Msg("ignoring raw solver cache: missing BSB22 commitment sidecar") + solution = nil + break + } + coeffSlice := []fr.Element(coeffs) + s.cCommitments[i] = iop.NewPolynomial(&coeffSlice, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if s.proof.Bsb22Commitments[i], err = s.commitLagrangePolynomialOnGPU(s.cCommitments[i]); err != nil { + return err + } + s.htfFunc.Write(s.proof.Bsb22Commitments[i].Marshal()) + hashBts := s.htfFunc.Sum(nil) + s.htfFunc.Reset() + nbBuf := fr.Bytes + if s.htfFunc.Size() < fr.Bytes { + nbBuf = s.htfFunc.Size() + } + s.commitmentVal[i].SetBytes(hashBts[:nbBuf]) + } + } + } + + if solution == nil { + _solution, err := s.spr.SolveAndSaveRawValues(s.fullWitness, rawCachePath, s.opt.SolverOpts...) + if err != nil { + log.Debug().Dur("took", time.Since(startSolve)).Err(err).Msg("solveConstraints: spr.Solve") + return err + } + log.Debug().Dur("took", time.Since(startSolve)).Msg("solveConstraints: spr.Solve") + if isProfileMode { + s.timings.record("solveConstraints: spr.Solve", time.Since(startSolve)) + } + solution = _solution.(*cs.SparseR1CSSolution) + + // Save cCommitments polynomial coefficients for BSB22 reconstruction + if rawCachePath != "" && len(s.commitmentInfo) > 0 { + cacheDir := filepath.Dir(rawCachePath) + for i := range s.commitmentInfo { + if s.cCommitments[i] != nil { + coeffs := s.cCommitments[i].Coefficients() + bsb22Path := filepath.Join(cacheDir, fmt.Sprintf("bsb22_commit_%d.bin", i)) + if err := cs.SaveRawSolverValues(bsb22Path, coeffs); err != nil { + log.Warn().Err(err).Int("i", i).Msg("failed to save BSB22 commitment polynomial") + } + } + } + } + } + + evaluationLDomainSmall := []fr.Element(solution.L) + evaluationRDomainSmall := []fr.Element(solution.R) + evaluationODomainSmall := []fr.Element(solution.O) + var wg sync.WaitGroup + wg.Add(2) + go func() { + s.polyL = iop.NewPolynomial(&evaluationLDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + go func() { + s.polyR = iop.NewPolynomial(&evaluationRDomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + wg.Done() + }() + + s.polyO = iop.NewPolynomial(&evaluationODomainSmall, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + + wg.Wait() + if _, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}); err != nil { + return err + } + + // commit to l, r, o and add blinding factors + if err := s.commitToLRO(); err != nil { + return err + } + close(s.chLRO) + return nil +} + +func (s *instance) completeQk() error { + qk := s.trace.Qk.Clone() + qkCoeffs := qk.Coefficients() + + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + copy(qkCoeffs, wWitness[:len(s.spr.Public)]) + + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + for i := range s.commitmentInfo { + qkCoeffs[s.spr.GetNbPublicVariables()+s.commitmentInfo[i].CommitmentIndex] = s.commitmentVal[i] + } + + s.polyQk = qk + close(s.chQk) + + return nil +} + +func (s *instance) commitToLRO() error { + var startCommitLRO time.Time + if isProfileMode { + startCommitLRO = time.Now() + } + sequentialLRO := s.domain0 != nil && s.domain0.Cardinality >= (1<<22) + if _, ok := os.LookupEnv("ICICLE_LRO_COMMIT_SEQUENTIAL"); ok { + sequentialLRO = envEnabled("ICICLE_LRO_COMMIT_SEQUENTIAL", true) + } + + if !useBlinding { + // When blinding is disabled, commit directly without waiting for blinding polynomials + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + if err != nil { + return err + } + } else { + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitLagrangePolynomialOnGPU(s.polyL) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitLagrangePolynomialOnGPU(s.polyR) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitLagrangePolynomialOnGPU(s.polyO) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (no blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (no blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil + } + + // wait for blinding polynomials to be initialized or context to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chbp: + } + + if sequentialLRO { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + if err != nil { + return err + } + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + if err != nil { + return err + } + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + if err != nil { + return err + } + } else { + // Run the three commitments concurrently. + var g errgroup.Group + g.Go(func() error { + var err error + s.proof.LRO[0], err = s.commitToPolyAndBlinding(s.polyL, s.bp[id_Bl]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[1], err = s.commitToPolyAndBlinding(s.polyR, s.bp[id_Br]) + return err + }) + g.Go(func() error { + var err error + s.proof.LRO[2], err = s.commitToPolyAndBlinding(s.polyO, s.bp[id_Bo]) + return err + }) + if err := g.Wait(); err != nil { + return err + } + } + if isProfileMode { + l := logger.Logger() + if sequentialLRO { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: sequential commitments to L, R, O (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitLRO)).Msg("commitToLRO: concurrent commitments to L, R, O (with blinding)") + } + s.timings.record("commitToLRO (commit L,R,O)", time.Since(startCommitLRO)) + } + return nil +} + +// deriveGammaAndBeta (copy constraint) +func (s *instance) deriveGammaAndBeta() error { + wWitness, ok := s.fullWitness.Vector().(fr.Vector) + if !ok { + return witness.ErrInvalidWitness + } + + if err := bindPublicData(s.fs, "gamma", s.pk.VerifyingKey().(*plonk_bw6761.VerifyingKey), wWitness[:len(s.spr.Public)]); err != nil { + return err + } + + // wait for LRO to be committed + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + + gamma, err := deriveRandomness(s.fs, "gamma", &s.proof.LRO[0], &s.proof.LRO[1], &s.proof.LRO[2]) + if err != nil { + return err + } + + bbeta, err := s.fs.ComputeChallenge("beta") + if err != nil { + return err + } + s.gamma = gamma + s.beta.SetBytes(bbeta) + + close(s.chGammaBeta) + + return nil +} + +// commitToPolyAndBlinding computes the KZG commitment of a polynomial p +// in Lagrange form (large degree) +// and add the contribution of a blinding polynomial b (small degree) +// /!\ The polynomial p is supposed to be in Lagrange form. +// Only used when blinding is enabled (the default). +func (s *instance) commitToPolyAndBlinding(p, b *iop.Polynomial) (commit curve.G1Affine, err error) { + // Commit over the Lagrange SRS using shared device-resident polynomial data. + gpuCommit, err := s.commitLagrangePolynomialOnGPU(p) + if err != nil { + return curve.G1Affine{}, err + } + + // add CPU blinding contribution (two MSMs on canonical SRS) + n := int(s.domain0.Cardinality) + cb := commitBlindingFactor(n, b, s.pk.Kzg) + gpuCommit.Add(&gpuCommit, &cb) + return gpuCommit, nil +} + +func (s *instance) commitLagrangePolynomialOnGPU(p *iop.Polynomial) (curve.G1Affine, error) { + if p == nil { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: nil polynomial") + } + if p.Basis != iop.Lagrange { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial must be in Lagrange basis, got %v", p.Basis) + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{p}) + if err != nil { + return curve.G1Affine{}, err + } + idx, ok := gpuState.polyToIdx[p] + if !ok || idx < 0 || idx >= len(gpuState.deviceSlices) { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: polynomial is missing from shared GPU state") + } + if gpuState.deviceSlices[idx].IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitLagrangePolynomialOnGPU: empty device slice for polynomial") + } + + // Keep the polynomial in its native Lagrange basis; large MSMs are split + // into device-side chunks inside commitOnGPULagrangeDevice. + return commitOnGPULagrangeDevice(gpuState.deviceSlices[idx], &s.device, s.pk) +} + +func (s *instance) deriveAlpha() (err error) { + alphaDeps := make([]*curve.G1Affine, len(s.proof.Bsb22Commitments)+1) + for i := range s.proof.Bsb22Commitments { + alphaDeps[i] = &s.proof.Bsb22Commitments[i] + } + alphaDeps[len(alphaDeps)-1] = &s.proof.Z + s.alpha, err = deriveRandomness(s.fs, "alpha", alphaDeps...) + return err +} + +func (s *instance) deriveZeta() (err error) { + s.zeta, err = deriveRandomness(s.fs, "zeta", &s.proof.H[0], &s.proof.H[1], &s.proof.H[2]) + return +} + +func (s *instance) computeQuotient() (err error) { + // wait for solver to be done + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLRO: + } + if isProfileMode { + var startComputeQuotient time.Time + startComputeQuotient = time.Now() + defer func() { + l := logger.Logger() + if err != nil { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Err(err).Msg("computeQuotient: total (with error)") + } else { + l.Debug().Dur("took", time.Since(startComputeQuotient)).Msg("computeQuotient: total") + } + s.timings.record("computeQuotient (total)", time.Since(startComputeQuotient)) + }() + } + + // wait for Z to be committed or context done + doneWaitZ := profileStep("computeQuotient: wait Z commit") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZ: + } + doneWaitZ() + + // derive alpha + if err = s.deriveAlpha(); err != nil { + return err + } + + if err := s.waitForComputeNumeratorQk(); err != nil { + return err + } + if s.polyQk == nil { + return fmt.Errorf("computeQuotient: missing completed Qk polynomial") + } + + doneEnsureGPUState := profileStep("computeQuotient: ensure shared GPU state") + gpuState, err := s.ensurePolysOnSharedGPU(s.buildComputeNumeratorGPUBatch()) + if err != nil { + return err + } + doneEnsureGPUState() + + // compute Z shifted by one for copy-constraint terms. + if s.polyZ == nil { + return fmt.Errorf("computeQuotient: missing Z polynomial") + } + s.polyZS = s.polyZ.ShallowClone().Shift(1) + + var numeratorGPU *gpuNumeratorPolynomial + var quotientGPU *gpuQuotientPolynomial + var e error + doneComputeNumerator := profileStep("computeQuotient: computeNumerator") + numeratorGPU, e = s.computeNumerator(gpuState) + if e != nil { + return e + } + doneComputeNumerator() + + doneDivideByZH := profileStep("computeQuotient: divideByZHOnGPU") + quotientGPU, e = s.divideByZHOnGPU(numeratorGPU, [2]*fft.Domain{s.domain0, s.domain1}) + if e != nil { + return e + } + doneDivideByZH() + s.hGPU = quotientGPU + + // Shared state slices were mutated during numerator coset iterations and are no + // longer needed now; computeLinearizedPolynomial uses the immutable snapshot. + s.releaseSharedGPUState() + close(s.chRestoreLRO) + + doneCommitH := profileStep("computeQuotient: commit H from device") + if err := s.commitToQuotientGPUFromDevice(s.hGPU); err != nil { + return err + } + doneCommitH() + + if err := s.deriveZeta(); err != nil { + return err + } + + donePrepareLinearizedEval := profileStep("computeQuotient: prepare linearized eval GPU state") + if err := s.prepareLinearizedEvalGPUStateFromHost(); err != nil { + return fmt.Errorf("computeQuotient: prepare linearized eval GPU state failed: %w", err) + } + donePrepareLinearizedEval() + + close(s.chH) + + return nil +} + +func (s *instance) buildRatioCopyConstraint() (err error) { + // wait for gamma and beta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chGammaBeta: + } + + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("buildRatioCopyConstraint: missing L/R/O polynomials") + } + + gpuState, err := s.ensurePolysOnSharedGPU([]*iop.Polynomial{s.polyL, s.polyR, s.polyO}) + if err != nil { + return err + } + dL, err := getStateDeviceSlice(gpuState, s.polyL, "L") + if err != nil { + return err + } + dR, err := getStateDeviceSlice(gpuState, s.polyR, "R") + if err != nil { + return err + } + dO, err := getStateDeviceSlice(gpuState, s.polyO, "O") + if err != nil { + return err + } + + var startBuildRatioCopyConstraintIcicle time.Time + if isProfileMode { + startBuildRatioCopyConstraintIcicle = time.Now() + } + s.polyZ, err = s.BuildRatioCopyConstraintIcicle( + []icicle_core.DeviceSlice{dL, dR, dO}, + s.trace.S, + s.beta, + s.gamma, + s.domain0, + gpuState, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startBuildRatioCopyConstraintIcicle)).Msg("buildRatioCopyConstraint: BuildRatioCopyConstraintIcicle") + s.timings.record("buildRatioCopyConstraint (perm Z)", time.Since(startBuildRatioCopyConstraintIcicle)) + } + + dZ, err := getStateDeviceSlice(gpuState, s.polyZ, "Z") + if err != nil { + return err + } + copyDone := make(chan error, 1) + var dPersist icicle_core.DeviceSlice + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + var allocErr error + dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) + if allocErr != nil { + finish(fmt.Errorf("buildRatioCopyConstraint: alloc persist Z failed: %w", allocErr)) + return + } + if e := copyDeviceSliceIntoOnCurrentDevice(dPersist, dZ, cfg); e != icicle_runtime.Success { + _ = dPersist.Free() + dPersist = icicle_core.DeviceSlice{} + finish(fmt.Errorf("buildRatioCopyConstraint: persist Z copy failed: %s", e.AsString())) + return + } + finish(nil) + }) + if err := <-copyDone; err != nil { + return err + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + s.polyZLagrangeGPU = dPersist + + // commit to Z (with or without blinding) + var startCommitZ time.Time + if isProfileMode { + startCommitZ = time.Now() + } + if useBlinding { + s.proof.Z, err = s.commitToPolyAndBlinding(s.polyZ, s.bp[id_Bz]) + } else { + s.proof.Z, err = s.commitLagrangePolynomialOnGPU(s.polyZ) + } + if isProfileMode { + l := logger.Logger() + if useBlinding { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (with blinding)") + } else { + l.Debug().Dur("took", time.Since(startCommitZ)).Msg("buildRatioCopyConstraint: commit Z (no blinding)") + } + } + s.freeIdleTempGPUMemoryOnDevice() + + close(s.chZ) + + return +} + +// open Z (blinded) at ωζ +func (s *instance) openZ() (err error) { + // wait for H to be committed and zeta to be derived (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + + var zetaShifted fr.Element + zetaShifted.Mul(&s.zeta, &s.pk.Vk.Generator) + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("openZ: missing GPU Z polynomial") + } + + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + + dZLagrange := s.polyZLagrangeGPU + if dZLagrange.Len() <= 1 { + return fmt.Errorf("openZ: invalid Z size %d", dZLagrange.Len()) + } + + blindSize := order_blinding_Z + 1 + if useBlinding { + if len(s.bp) <= id_Bz || s.bp[id_Bz] == nil { + return fmt.Errorf("openZ: missing Z blinding polynomial") + } + blindSize = len(s.bp[id_Bz].Coefficients()) + if blindSize == 0 { + return fmt.Errorf("openZ: empty Z blinding polynomial") + } + } + + var dBlindedCanonical icicle_core.DeviceSlice + buildDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + buildDone <- cfgErr + return + } + finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + // Async boundary for canonicalization/blinding before exposing output. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + buildDone <- runErr + } + + n := dZLagrange.Len() + dZCanonical := s.getTempDeviceSlice(n) + if err := copyDeviceSliceIntoOnCurrentDevice(dZCanonical, dZLagrange, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy Z to canonical buffer failed: %s", err.AsString()), dZCanonical, false) + return + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + cfgNtt.Ordering = icicle_core.KNN // regular lagrange -> regular canonical + if err := icicle_ntt.Ntt(dZCanonical, icicle_core.KInverse, &cfgNtt, dZCanonical); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: inverse NTT on Z failed: %s", err.AsString()), dZCanonical, false) + return + } + + dBlindedCanonical = s.getTempDeviceSlice(n + blindSize) + dBlindedPrefix := (&dBlindedCanonical).Range(0, n, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedPrefix, dZCanonical, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: copy canonical Z into blinded buffer failed: %s", err.AsString()), dZCanonical, true) + return + } + + if useBlinding { + dBp := uploadVector(s.bp[id_Bz].Coefficients()) + dBlindedHead := (&dBlindedPrefix).Range(0, blindSize, false) + if err := icicle_vecops.VecOp(dBlindedHead, dBp, dBlindedHead, cfgVec, icicle_core.Sub); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: subtract Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := copyDeviceSliceIntoOnCurrentDevice(dBlindedTail, dBp, cfgVec); err != icicle_runtime.Success { + _ = dBp.FreeAsync(stream) + finalize(fmt.Errorf("openZ: append Z blinding polynomial failed: %s", err.AsString()), dZCanonical, true) + return + } + _ = dBp.FreeAsync(stream) + } else { + dBlindedTail := (&dBlindedCanonical).Range(n, n+blindSize, false) + if err := zeroDeviceSliceOnCurrentDevice(dBlindedTail, cfgVec); err != icicle_runtime.Success { + finalize(fmt.Errorf("openZ: zero-pad non-blinded Z failed: %s", err.AsString()), dZCanonical, true) + return + } + } + + finalize(nil, dZCanonical, false) + }) + if err := <-buildDone; err != nil { + return err + } + s.blindedZCanonicalGPU = dBlindedCanonical + + // open z at zeta*w. + var startKzgOpen time.Time + if isProfileMode { + startKzgOpen = time.Now() + } + s.proof.ZShiftedOpening, err = s.openPolynomialOnGPUCanonicalDevice(s.blindedZCanonicalGPU, zetaShifted) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startKzgOpen)).Msg("openZ: open polynomial on GPU") + s.timings.record("openZ (KZG open on GPU)", time.Since(startKzgOpen)) + } + if err != nil { + return err + } + close(s.chZOpening) + return nil +} + +func (s *instance) h1() []fr.Element { + var h1 []fr.Element + if !s.opt.StatisticalZK { + h1 = s.h.Coefficients()[:s.domain0.Cardinality+2] + } else { + h1 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h1, s.h.Coefficients()[:s.domain0.Cardinality+2]) + h1[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[0]) + } + return h1 +} + +func (s *instance) h2() []fr.Element { + var h2 []fr.Element + if !s.opt.StatisticalZK { + h2 = s.h.Coefficients()[s.domain0.Cardinality+2 : 2*(s.domain0.Cardinality+2)] + } else { + h2 = make([]fr.Element, s.domain0.Cardinality+3) + copy(h2, s.h.Coefficients()[s.domain0.Cardinality+2:2*(s.domain0.Cardinality+2)]) + h2[0].Sub(&h2[0], &s.quotientShardsRandomizers[0]) + h2[s.domain0.Cardinality+2].Set(&s.quotientShardsRandomizers[1]) + } + return h2 +} + +func (s *instance) h3() []fr.Element { + var h3 []fr.Element + if !s.opt.StatisticalZK { + h3 = s.h.Coefficients()[2*(s.domain0.Cardinality+2) : 3*(s.domain0.Cardinality+2)] + } else { + h3 = make([]fr.Element, s.domain0.Cardinality+2) + copy(h3, s.h.Coefficients()[2*(s.domain0.Cardinality+2):3*(s.domain0.Cardinality+2)]) + h3[0].Sub(&h3[0], &s.quotientShardsRandomizers[1]) + } + return h3 +} + +// witnessEvalAtZeta holds the scalar evaluations of witness and constraint +// polynomials at the challenge point zeta, as needed by the linearized +// polynomial computation. +type witnessEvalAtZeta struct { + blzeta, brzeta, bozeta fr.Element + s1zeta, s2zeta fr.Element + qcpzeta []fr.Element +} + +type linearizedSelectorScales struct { + s3, ql, qr, qm, qo, qk fr.Element + qcp []fr.Element +} + +// evaluateWitnessPolynomialsAtZeta evaluates L, R, O (with optional blinding), +// S1, S2, and all Qcp polynomials at the point zeta using the GPU-resident +// polynomial state. +func (s *instance) evaluateWitnessPolynomialsAtZeta( + evalGPUState *gpuPolysState, + zeta fr.Element, +) (witnessEvalAtZeta, error) { + doneTotal := profileStep("evaluateWitnessPolynomialsAtZeta: total") + defer doneTotal() + + var result witnessEvalAtZeta + var err error + + result.qcpzeta = make([]fr.Element, len(s.commitmentInfo)) + var startQcp time.Time + if isProfileMode { + startQcp = time.Now() + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i >= len(s.trace.Qcp) || s.trace.Qcp[i] == nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: missing Qcp polynomial at index %d", i) + } + var startQcpItem time.Time + if isProfileMode { + startQcpItem = time.Now() + } + result.qcpzeta[i], err = s.evalPolynomialInCurrentFormOnGPU(s.trace.Qcp[i], evalGPUState, zeta) + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: qcp[%d] GPU evaluation failed: %w", i, err) + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("idx", i).Dur("took", time.Since(startQcpItem)).Msg("evaluateWitnessPolynomialsAtZeta: qcp eval item") + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startQcp)).Msg("evaluateWitnessPolynomialsAtZeta: qcpZeta evaluate on GPU") + } + + if useBlinding { + result.blzeta, err = s.evaluateBlindedOnGPU(s.polyL, s.bp[id_Bl], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: L evaluate on GPU") + result.blzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyL, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: blzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.brzeta, err = s.evaluateBlindedOnGPU(s.polyR, s.bp[id_Br], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: R evaluate on GPU") + result.brzeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyR, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: brzeta GPU evaluation failed: %w", err) + } + + if useBlinding { + result.bozeta, err = s.evaluateBlindedOnGPU(s.polyO, s.bp[id_Bo], evalGPUState, zeta) + } else { + done := profileStep("evaluateWitnessPolynomialsAtZeta: O evaluate on GPU") + result.bozeta, err = s.evalPolynomialInCurrentFormOnGPU(s.polyO, evalGPUState, zeta) + done() + } + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: bozeta GPU evaluation failed: %w", err) + } + + doneS1 := profileStep("evaluateWitnessPolynomialsAtZeta: S1 evaluate on GPU") + result.s1zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S1, evalGPUState, zeta) + doneS1() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s1(zeta) GPU evaluation failed: %w", err) + } + doneS2 := profileStep("evaluateWitnessPolynomialsAtZeta: S2 evaluate on GPU") + result.s2zeta, err = s.evalPolynomialInCurrentFormOnGPU(s.trace.S2, evalGPUState, zeta) + doneS2() + if err != nil { + return witnessEvalAtZeta{}, fmt.Errorf("evaluateWitnessPolynomialsAtZeta: s2(zeta) GPU evaluation failed: %w", err) + } + + return result, nil +} + +func (s *instance) computeLinearizedPolynomial() error { + + // wait for H to be committed and zeta to be derived (or ctx.Done()) + var startWaitH time.Time + if isProfileMode { + startWaitH = time.Now() + } + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chH: + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startWaitH)).Msg("computeLinearizedPolynomial: wait H and zeta") + s.timings.record("computeLinearizedPoly (wait H+zeta, overlaps)", time.Since(startWaitH)) + } + if s.opt.StatisticalZK { + return fmt.Errorf("computeLinearizedPolynomial: GPU-only opening path does not support StatisticalZK=true") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil || s.polyZ == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing required polynomials") + } + + // Reuse the immutable snapshot prepared in computeQuotient before numerator + // coset iterations mutate shared state slices. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("computeLinearizedPolynomial: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized eval prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: required polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("computeLinearizedPolynomial: missing blinding polynomials for GPU evaluation") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "computeLinearized blinding prereq"); e != nil { + return fmt.Errorf("computeLinearizedPolynomial: blinding polynomial is not on GPU: %w", e) + } + } + } + + doneEvaluate := profileStep("computeLinearizedPolynomial: evaluate witness polynomials") + evals, err := s.evaluateWitnessPolynomialsAtZeta(evalGPUState, s.zeta) + doneEvaluate() + if err != nil { + return err + } + + // wait for Z to be opened at zeta (or ctx.Done()) + doneWaitZOpening := profileStep("computeLinearizedPolynomial: wait Z opening") + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chZOpening: + } + doneWaitZOpening() + if s.blindedZCanonicalGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing canonical blinded Z on GPU") + } + if s.polyZLagrangeGPU.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing lagrange Z on GPU") + } + defer func() { + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + }() + bzuzeta := s.proof.ZShiftedOpening.ClaimedValue + + if s.hGPU == nil || s.hGPU.coeffs.IsEmpty() { + return fmt.Errorf("computeLinearizedPolynomial: missing GPU quotient polynomial") + } + + doneBuild := profileStep("computeLinearizedPolynomial: build selector terms on GPU") + dLin, err := s.buildLinearizedSelectorTermsOnGPU(evals, bzuzeta, s.blindedZCanonicalGPU.Len()) + doneBuild() + if err != nil { + return err + } + + doneAddZ := profileStep("computeLinearizedPolynomial: add Z contribution on GPU") + err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) + doneAddZ() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + freeSliceOnDevice(&dLin, &s.device) + return err + } + s.linearizedPolynomialGPU = dLin + + doneEvalClaim := profileStep("computeLinearizedPolynomial: evaluate linearized claim") + claim, err := s.evalDevicePolynomialAtPoint(dLin, s.zeta) + doneEvalClaim() + if err != nil { + return err + } + s.linearizedPolynomialClaim = claim + + // Commit the linearized polynomial over the canonical SRS. + var startMSM time.Time + if isProfileMode { + startMSM = time.Now() + } + s.linearizedPolynomialDigest, err = commitOnGPUCanonicalDevice(dLin, &s.device, s.pk) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startMSM)).Msg("computeLinearizedPolynomial: KZG commit") + s.timings.record("computeLinearizedPoly (KZG commit)", time.Since(startMSM)) + } + close(s.chLinearizedPolynomial) + return nil +} + +func (s *instance) batchOpening() error { + + // wait for linearizedPolynomial to be computed (or ctx.Done()) + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chLinearizedPolynomial: + } + + defer func() { + freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } + }() + + if s.linearizedPolynomialGPU.IsEmpty() { + return fmt.Errorf("batchOpening: missing GPU linearized polynomial") + } + if s.polyL == nil || s.polyR == nil || s.polyO == nil { + return fmt.Errorf("batchOpening: missing L/R/O polynomials") + } + + // Reuse immutable GPU snapshot prepared before quotient iterations. + s.gpuStateMu.Lock() + evalGPUState := s.linearizedEvalGPUState + s.gpuStateMu.Unlock() + if evalGPUState == nil { + return fmt.Errorf("batchOpening: linearized eval GPU state is missing") + } + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening eval prereq"); e != nil { + return fmt.Errorf("batchOpening: required polynomial is not on GPU: %w", e) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] == nil { + return fmt.Errorf("batchOpening: missing Qcp polynomial at index %d", i) + } + if _, e := getStateDeviceSlice(evalGPUState, s.trace.Qcp[i], "batchOpening qcp prereq"); e != nil { + return fmt.Errorf("batchOpening: Qcp polynomial is not on GPU: %w", e) + } + } + if useBlinding { + if len(s.bp) <= id_Bo || s.bp[id_Bl] == nil || s.bp[id_Br] == nil || s.bp[id_Bo] == nil { + return fmt.Errorf("batchOpening: missing blinding polynomials") + } + for _, p := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if _, e := getStateDeviceSlice(evalGPUState, p, "batchOpening blinding prereq"); e != nil { + return fmt.Errorf("batchOpening: blinding polynomial is not on GPU: %w", e) + } + } + } + + devicePolys, ownedPolys, claimed, err := s.prepareBatchOpeningPolynomialsOnGPU(evalGPUState, s.zeta) + if err != nil { + return err + } + defer func() { + for i := 0; i < len(devicePolys); i++ { + if i < len(ownedPolys) && ownedPolys[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + } + } + s.releaseLinearizedEvalGPUState() + }() + + digestsToOpen := make([]curve.G1Affine, len(s.pk.Vk.Qcp)+6) + copy(digestsToOpen[6:], s.pk.Vk.Qcp) + digestsToOpen[0] = s.linearizedPolynomialDigest + digestsToOpen[1] = s.proof.LRO[0] + digestsToOpen[2] = s.proof.LRO[1] + digestsToOpen[3] = s.proof.LRO[2] + digestsToOpen[4] = s.pk.Vk.S[0] + digestsToOpen[5] = s.pk.Vk.S[1] + if len(claimed) != len(digestsToOpen) { + return fmt.Errorf("batchOpening: claimed size mismatch (%d != %d)", len(claimed), len(digestsToOpen)) + } + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + var dFold icicle_core.DeviceSlice + foldDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening fold", foldDone) + + dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: copy linearized polynomial failed: %s", e.AsString())) + return + } + + gammaPow := gamma + for i := 1; i < len(devicePolys); i++ { + dPoly := devicePolys[i] + if dPoly.IsEmpty() { + gammaPow.Mul(&gammaPow, &gamma) + continue + } + dScaled := s.getTempDeviceSlice(dPoly.Len()) + dGammaStd := uploadScalarStdOnCurrentDevice(gammaPow, cfg) + eMul := icicle_vecops.ScalarMulVec(dGammaStd, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dGammaStd.FreeAsync(cfg.StreamHandle) + } else { + _ = dGammaStd.Free() + } + if eMul != icicle_runtime.Success { + if cfg.IsAsync { + _ = icicle_runtime.SynchronizeStream(cfg.StreamHandle) + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: scale polynomial %d failed: %s", i, eMul.AsString())) + return + } + + dPrefix := (&dFold).Range(0, dPoly.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add) + if cfg.IsAsync { + // dScaled is recycled each iteration; wait before returning to pool. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dScaled, dPoly.Len()) + finish(fmt.Errorf("batchOpening: synchronize stream failed: %s", eSync.AsString())) + return + } + } + s.putTempDeviceSlice(dScaled, dPoly.Len()) + if eAdd != icicle_runtime.Success { + finish(fmt.Errorf("batchOpening: fold add polynomial %d failed: %s", i, eAdd.AsString())) + return + } + gammaPow.Mul(&gammaPow, &gamma) + } + finish(nil) + }) + if err := <-foldDone; err != nil { + if !dFold.IsEmpty() { + s.putTempDeviceSlice(dFold, dFold.Len()) + } + return err + } + var dWitness icicle_core.DeviceSlice + divDone := make(chan error, 1) + witnessSize := dFold.Len() - 1 + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dFold, dPoint, dWitness, cfg) + if cfg.IsAsync { + _ = dPoint.FreeAsync(cfg.StreamHandle) + } else { + _ = dPoint.Free() + } + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + finish(fmt.Errorf("batchOpening: divide by (x-zeta) failed: %s", eDiv.AsString())) + return + } + finish(nil) + }) + if err := <-divDone; err != nil { + s.putTempDeviceSlice(dFold, dFold.Len()) + return err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dFold, dFold.Len()) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return err + } + + s.proof.BatchedProof = kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + } + if err := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); err != nil { + l := logger.Logger() + l.Warn().Err(err).Msg("batchOpening: GPU folded opening failed raw-G2 validation; falling back to host fold with CPU KZG commitment") + fallbackProof, fallbackErr := s.batchOpeningHostFoldGPUCommitFromDevicePolys(devicePolys, digestsToOpen) + if fallbackErr != nil { + return fallbackErr + } + s.proof.BatchedProof = fallbackProof + if verifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); verifyErr != nil { + nativeProof, nativeErr := s.batchOpeningNativeCPUFromDevicePolys(devicePolys, digestsToOpen) + if nativeErr != nil { + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed: %v", verifyErr, nativeErr) + } + s.proof.BatchedProof = nativeProof + if nativeVerifyErr := s.verifyBatchOpeningWithRecomputedLines(digestsToOpen); nativeVerifyErr != nil { + diagnostic, diagnosticErr := s.diagnoseBatchOpeningDevicePolynomials(devicePolys, digestsToOpen, claimed) + if diagnosticErr != nil { + diagnostic = fmt.Sprintf("batch opening diagnostic failed: %v", diagnosticErr) + } + return fmt.Errorf("batchOpening: fallback folded opening failed raw-G2 validation: %w; native CPU batch opening also failed raw-G2 validation: %v; %s", verifyErr, nativeVerifyErr, diagnostic) + } + l.Warn().Msg("batchOpening: native CPU KZG fallback produced a valid proof after host-fold fallback failed") + } + } + _ = foldedEval // kept for parity with kzg.BatchOpenSinglePoint flow. + return nil +} + +func (s *instance) verifyBatchOpeningWithRecomputedLines(digestsToOpen []curve.G1Affine) error { + vk := s.pk.Vk.Kzg + vk.Lines[0] = curve.PrecomputeLines(vk.G2[0]) + vk.Lines[1] = curve.PrecomputeLines(vk.G2[1]) + return kzg.BatchVerifySinglePoint( + digestsToOpen, + &s.proof.BatchedProof, + s.zeta, + s.kzgFoldingHash, + vk, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningHostFoldGPUCommitFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningNativeCPUFromDevicePolys( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "batchOpeningNativeCPUFromDevicePolys", + ) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return kzg.BatchOpenSinglePoint( + polysToOpen, + digestsToOpen, + s.zeta, + s.kzgFoldingHash, + s.pk.Kzg, + s.proof.ZShiftedOpening.ClaimedValue.Marshal(), + ) +} + +func (s *instance) diagnoseBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + digestsToOpen []curve.G1Affine, + gpuClaimed []fr.Element, +) (string, error) { + polysToOpen, err := s.downloadBatchOpeningDevicePolynomials( + devicePolys, + len(digestsToOpen), + "diagnoseBatchOpeningDevicePolynomials", + ) + if err != nil { + return "", err + } + if len(gpuClaimed) != len(polysToOpen) { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: claimed/polynomial mismatch (%d != %d)", len(gpuClaimed), len(polysToOpen)) + } + + l := logger.Logger() + claimMismatches := make([]string, 0) + commitMismatches := make([]string, 0) + for i := range polysToOpen { + label := batchOpeningPolynomialLabel(i) + cpuClaim := evalCanonicalAtPoint(polysToOpen[i], s.zeta) + if !cpuClaim.Equal(&gpuClaimed[i]) { + claimMismatches = append(claimMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("gpuClaim", frFingerprint(gpuClaimed[i])). + Str("cpuClaim", frFingerprint(cpuClaim)). + Msg("batchOpening diagnostic: GPU claim differs from CPU evaluation") + } + + cpuDigest, err := kzg.Commit(polysToOpen[i], s.pk.Kzg) + if err != nil { + return "", fmt.Errorf("diagnoseBatchOpeningDevicePolynomials: commit %s: %w", label, err) + } + if !cpuDigest.Equal(&digestsToOpen[i]) { + commitMismatches = append(commitMismatches, label) + l.Warn(). + Int("index", i). + Str("label", label). + Int("len", len(polysToOpen[i])). + Str("expectedDigest", g1Fingerprint(digestsToOpen[i])). + Str("cpuDigest", g1Fingerprint(cpuDigest)). + Msg("batchOpening diagnostic: CPU commitment differs from proof digest") + } + } + + if len(claimMismatches) == 0 && len(commitMismatches) == 0 { + return "batch opening diagnostic found no per-polynomial claim or commitment mismatch", nil + } + return fmt.Sprintf( + "batch opening diagnostic claim mismatches=[%s] commitment mismatches=[%s]", + strings.Join(claimMismatches, ","), + strings.Join(commitMismatches, ","), + ), nil +} + +func batchOpeningPolynomialLabel(index int) string { + switch index { + case 0: + return "linearized" + case 1: + return "L" + case 2: + return "R" + case 3: + return "O" + case 4: + return "S1" + case 5: + return "S2" + default: + return fmt.Sprintf("Qcp[%d]", index-6) + } +} + +func frFingerprint(v fr.Element) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func g1Fingerprint(v curve.G1Affine) string { + b := v.Marshal() + if len(b) > 8 { + b = b[:8] + } + return fmt.Sprintf("%x", b) +} + +func (s *instance) downloadBatchOpeningDevicePolynomials( + devicePolys []icicle_core.DeviceSlice, + expectedDigests int, + label string, +) ([][]fr.Element, error) { + if len(devicePolys) != expectedDigests { + return nil, fmt.Errorf("%s: polynomial/digest mismatch (%d != %d)", label, len(devicePolys), expectedDigests) + } + + polysToOpen := make([][]fr.Element, len(devicePolys)) + for i := range devicePolys { + var err error + polysToOpen[i], err = s.downloadCanonicalDeviceCoefficients( + devicePolys[i], + fmt.Sprintf("%s[%d]", label, i), + ) + if err != nil { + return nil, err + } + } + return polysToOpen, nil +} + +func (s *instance) batchOpeningHostFoldGPUCommit(digestsToOpen []curve.G1Affine) (kzg.BatchOpeningProof, error) { + polysToOpen, err := s.batchOpeningHostPolynomials() + if err != nil { + return kzg.BatchOpeningProof{}, err + } + return s.batchOpeningHostFoldGPUCommitFromPolynomials(polysToOpen, digestsToOpen) +} + +func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( + polysToOpen [][]fr.Element, + digestsToOpen []curve.G1Affine, +) (kzg.BatchOpeningProof, error) { + if len(polysToOpen) != len(digestsToOpen) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: polynomial/digest mismatch (%d != %d)", len(polysToOpen), len(digestsToOpen)) + } + + largestPoly := 0 + for i := range polysToOpen { + if len(polysToOpen[i]) == 0 || len(polysToOpen[i]) > len(s.pk.Kzg.G1) { + return kzg.BatchOpeningProof{}, fmt.Errorf("batchOpeningHostFoldGPUCommitFromPolynomials: invalid polynomial %d size %d", i, len(polysToOpen[i])) + } + if len(polysToOpen[i]) > largestPoly { + largestPoly = len(polysToOpen[i]) + } + } + + claimed := make([]fr.Element, len(polysToOpen)) + utils.Parallelize(len(polysToOpen), func(start, end int) { + for i := start; i < end; i++ { + claimed[i] = evalCanonicalAtPoint(polysToOpen[i], s.zeta) + } + }) + + gamma, err := deriveBatchOpeningGamma(s.zeta, digestsToOpen, claimed, s.kzgFoldingHash, s.proof.ZShiftedOpening.ClaimedValue.Marshal()) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + foldedEval := claimed[len(claimed)-1] + for i := len(claimed) - 2; i >= 0; i-- { + foldedEval.Mul(&foldedEval, &gamma).Add(&foldedEval, &claimed[i]) + } + + foldedPolynomials := make([]fr.Element, largestPoly) + copy(foldedPolynomials, polysToOpen[0]) + gammaPow := gamma + for i := 1; i < len(polysToOpen); i++ { + poly := polysToOpen[i] + scale := gammaPow + utils.Parallelize(len(poly), func(start, end int) { + var term fr.Element + for j := start; j < end; j++ { + term.Mul(&poly[j], &scale) + foldedPolynomials[j].Add(&foldedPolynomials[j], &term) + } + }) + gammaPow.Mul(&gammaPow, &gamma) + } + + hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) + var dWitness icicle_core.DeviceSlice + uploadDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + dWitness = uploadVector(hCoeffs) + close(uploadDone) + }) + <-uploadDone + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + freeSliceOnDevice(&dWitness, &s.device) + if err != nil { + return kzg.BatchOpeningProof{}, err + } + + return kzg.BatchOpeningProof{ + H: h, + ClaimedValues: claimed, + }, nil +} + +func (s *instance) batchOpeningHostPolynomials() ([][]fr.Element, error) { + total := 6 + len(s.trace.Qcp) + polysToOpen := make([][]fr.Element, total) + + var err error + polysToOpen[0], err = s.downloadCanonicalDeviceCoefficients( + s.linearizedPolynomialGPU, + "batchOpeningHostPolynomials linearized", + ) + if err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare linearized: %w", err) + } + + prepareLRO := func(p, bp *iop.Polynomial, blindingOrder int) ([]fr.Element, error) { + base, err := canonicalRegularCoefficientsCopy(p, s.domain0) + if err != nil { + return nil, err + } + if useBlinding { + if bp == nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomial") + } + blind, err := canonicalRegularCoefficientsCopy(bp, s.domain0) + if err != nil { + return nil, err + } + out := make([]fr.Element, len(base)+len(blind)) + copy(out, base) + copy(out[len(base):], blind) + for i := range blind { + out[i].Sub(&out[i], &blind[i]) + } + return out, nil + } + out := make([]fr.Element, len(base)+blindingOrder+1) + copy(out, base) + return out, nil + } + + var bpL, bpR, bpO *iop.Polynomial + if useBlinding { + if len(s.bp) <= id_Bo { + return nil, fmt.Errorf("batchOpeningHostPolynomials: missing blinding polynomials") + } + bpL, bpR, bpO = s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo] + } + + if polysToOpen[1], err = prepareLRO(s.polyL, bpL, order_blinding_L); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare L: %w", err) + } + if polysToOpen[2], err = prepareLRO(s.polyR, bpR, order_blinding_R); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare R: %w", err) + } + if polysToOpen[3], err = prepareLRO(s.polyO, bpO, order_blinding_O); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare O: %w", err) + } + if polysToOpen[4], err = canonicalRegularCoefficientsCopy(s.trace.S1, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S1: %w", err) + } + if polysToOpen[5], err = canonicalRegularCoefficientsCopy(s.trace.S2, s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare S2: %w", err) + } + for i := range s.trace.Qcp { + if polysToOpen[6+i], err = canonicalRegularCoefficientsCopy(s.trace.Qcp[i], s.domain0); err != nil { + return nil, fmt.Errorf("batchOpeningHostPolynomials: prepare Qcp[%d]: %w", i, err) + } + } + return polysToOpen, nil +} + +func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceSlice, label string) ([]fr.Element, error) { + if dPoly.IsEmpty() { + return nil, fmt.Errorf("%s: empty device polynomial", label) + } + + coeffs := make([]fr.Element, dPoly.Len()) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + }) + if err := <-done; err != nil { + return nil, err + } + return coeffs, nil +} + +func canonicalRegularCoefficientsCopy(p *iop.Polynomial, domain *fft.Domain) ([]fr.Element, error) { + if p == nil { + return nil, fmt.Errorf("nil polynomial") + } + cp := p.Clone() + cp.ToCanonical(domain).ToRegular() + coeffs := cp.Coefficients() + out := make([]fr.Element, len(coeffs)) + copy(out, coeffs) + return out, nil +} + +func dividePolyByXMinusAHost(f []fr.Element, fa, a fr.Element) []fr.Element { + f[0].Sub(&f[0], &fa) + var t fr.Element + for i := len(f) - 2; i >= 0; i-- { + t.Mul(&f[i+1], &a) + f[i].Add(&f[i], &t) + } + return f[1:] +} + +// evaluate the full set of constraints on the GPU-resident polynomial state. +type computeNumeratorLoopContext struct { + n int + rho int + mm uint64 + bn *big.Int + shifters []fr.Element + twiddles0 []fr.Element + dTwiddles0 icicle_core.DeviceSlice + dPrecomputedDenominators *icicle_core.DeviceSlice + scalingVector []fr.Element + scalingVectorRev []fr.Element + gpuState *gpuPolysState + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + one fr.Element + cs fr.Element + css fr.Element + nbBsbGates int + numeratorShards []icicle_core.DeviceSlice +} + +type gpuNumeratorPolynomial struct { + shards []icicle_core.DeviceSlice + n int + rho int + mm uint64 +} + +type gpuQuotientPolynomial struct { + coeffs icicle_core.DeviceSlice + size int +} + +func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolynomial, error) { + twiddles0, err := s.buildComputeNumeratorTwiddles() + if err != nil { + return nil, err + } + if err := s.validateComputeNumeratorGPUState(gpuState); err != nil { + return nil, err + } + + var startComputeNumerator time.Time + if isProfileMode { + startComputeNumerator = time.Now() + } + + n := s.domain0.Cardinality + nbBsbGates := len(s.proof.Bsb22Commitments) + + var cs, css fr.Element + cs.Set(&s.domain1.FrMultiplicativeGen) + css.Square(&cs) + + bn := big.NewInt(int64(n)) + + rho := int(s.domain1.Cardinality / n) + shifters := make([]fr.Element, rho) + shifters[0].Set(&s.domain1.FrMultiplicativeGen) + for i := 1; i < rho; i++ { + shifters[i].Set(&s.domain1.Generator) + } + + cosetTable, err := s.domain0.CosetTable() + if err != nil { + return nil, err + } + + // for the first iteration, the scalingVector is the coset table + scalingVector := cosetTable + scalingVectorRev := make([]fr.Element, len(cosetTable)) + copy(scalingVectorRev, cosetTable) + fft.BitReverse(scalingVectorRev) + + // pre-computed to compute the bit reverse index + // of the result polynomial + m := uint64(s.domain1.Cardinality) + mm := uint64(64 - bits.TrailingZeros64(m)) + + var dPrecomputedDenominators icicle_core.DeviceSlice + defer func() { + if !dPrecomputedDenominators.IsEmpty() { + freeDone := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + freeDone <- dPrecomputedDenominators.Free() + }) + if err := <-freeDone; err != icicle_runtime.Success { + panic(fmt.Sprintf("computeNumerator: failed to free dPrecomputedDenominators: %s", err.AsString())) + } + } + }() + + var coset, cosetExponentiatedToNMinusOne, one fr.Element + coset.SetOne() + one.SetOne() + + dTwiddles0, err := s.uploadComputeNumeratorTwiddles(twiddles0) + if err != nil { + return nil, err + } + + loopCtx := &computeNumeratorLoopContext{ + n: int(n), + rho: rho, + mm: mm, + bn: bn, + shifters: shifters, + twiddles0: twiddles0, + dTwiddles0: dTwiddles0, + dPrecomputedDenominators: &dPrecomputedDenominators, + scalingVector: scalingVector, + scalingVectorRev: scalingVectorRev, + gpuState: gpuState, + coset: coset, + cosetExponentiatedToNMinusOne: cosetExponentiatedToNMinusOne, + one: one, + cs: cs, + css: css, + nbBsbGates: nbBsbGates, + numeratorShards: make([]icicle_core.DeviceSlice, rho), + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: setup before iteration loop") + } + if err := s.executeComputeNumeratorCosetIterations(loopCtx); err != nil { + return nil, err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startComputeNumerator)).Msg("computeNumerator: main body (post-wait)") + } + + return &gpuNumeratorPolynomial{ + shards: loopCtx.numeratorShards, + n: loopCtx.n, + rho: loopCtx.rho, + mm: loopCtx.mm, + }, nil + +} + +func (s *instance) buildComputeNumeratorTwiddles() ([]fr.Element, error) { + n := s.domain0.Cardinality + var startTwiddles time.Time + if isProfileMode { + startTwiddles = time.Now() + } + twiddles0 := make([]fr.Element, n) + if n == 1 { + // edge case + twiddles0[0].SetOne() + } else { + twiddles, err := s.domain0.Twiddles() + if err != nil { + return nil, err + } + copy(twiddles0, twiddles[0]) + w := twiddles0[1] + for i := len(twiddles[0]); i < len(twiddles0); i++ { + twiddles0[i].Mul(&twiddles0[i-1], &w) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startTwiddles)).Msg("computeNumerator: build twiddles") + } + return twiddles0, nil +} + +func (s *instance) waitForComputeNumeratorQk() error { + select { + case <-s.ctx.Done(): + return errContextDone + case <-s.chQk: + } + return nil +} + +func (s *instance) buildLinearizedEvalGPUBatch() []*iop.Polynomial { + baseCap := 5 + len(s.trace.Qcp) + if useBlinding { + baseCap += 3 + } + polys := make([]*iop.Polynomial, 0, baseCap) + for _, p := range []*iop.Polynomial{s.polyL, s.polyR, s.polyO, s.trace.S1, s.trace.S2} { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.trace.Qcp); i++ { + if s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + } + if useBlinding && len(s.bp) > id_Bo { + for _, bpPoly := range []*iop.Polynomial{s.bp[id_Bl], s.bp[id_Br], s.bp[id_Bo]} { + if bpPoly != nil { + polys = append(polys, bpPoly) + } + } + } + return polys +} + +func (s *instance) buildComputeNumeratorGPUBatch() []*iop.Polynomial { + polys := make([]*iop.Polynomial, 0, 13+2*len(s.commitmentInfo)) + for _, p := range []*iop.Polynomial{ + s.polyL, s.polyR, s.polyO, s.polyZ, + s.trace.Ql, s.trace.Qr, s.trace.Qm, s.trace.Qo, s.polyQk, + s.trace.S1, s.trace.S2, s.trace.S3, + } { + if p != nil { + polys = append(polys, p) + } + } + for i := 0; i < len(s.commitmentInfo); i++ { + if i < len(s.trace.Qcp) && s.trace.Qcp[i] != nil { + polys = append(polys, s.trace.Qcp[i]) + } + if i < len(s.cCommitments) && s.cCommitments[i] != nil { + polys = append(polys, s.cCommitments[i]) + } + } + return polys +} + +func (s *instance) validateComputeNumeratorGPUState(state *gpuPolysState) error { + if state == nil { + return fmt.Errorf("computeNumerator: shared GPU state is nil") + } + required := s.buildComputeNumeratorGPUBatch() + if len(required) == 0 { + return fmt.Errorf("computeNumerator: no polynomials prepared for GPU batch") + } + for _, p := range required { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; !ok { + return fmt.Errorf("computeNumerator: polynomial ptr=%p is missing from shared GPU state", p) + } + } + return nil +} + +func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicle_core.DeviceSlice, error) { + var dTwiddles0 icicle_core.DeviceSlice + uploadTwiddlesDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + if s.tempGPUMemPool != nil { + s.tempGPUMemPool.FreeAll() + } + host := icicle_core.HostSliceFromElements(twiddles0) + var allocErr error + dTwiddles0, allocErr = allocDeviceUninitialized(len(twiddles0)) + if allocErr != nil { + uploadTwiddlesDone <- fmt.Errorf("uploadComputeNumeratorTwiddles: %w", allocErr) + return + } + host.CopyToDevice(&dTwiddles0, false) + uploadTwiddlesDone <- nil + }) + if err := <-uploadTwiddlesDone; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dTwiddles0, nil +} + +// executeComputeNumeratorCosetIterations runs the rho coset iterations. +func (s *instance) executeComputeNumeratorCosetIterations(loopCtx *computeNumeratorLoopContext) error { + var startIterLoop time.Time + if isProfileMode { + startIterLoop = time.Now() + } + + for i := 0; i < loopCtx.rho; i++ { + if err := s.computeNumeratorIteration(i, loopCtx); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + freeSliceOnDevice(&loopCtx.dTwiddles0, &s.device) + return err + } + } + + // Free twiddles0 device slice (uploaded once before the loop). + freeTwiddlesDone := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + loopCtx.dTwiddles0.Free() + close(freeTwiddlesDone) + }) + <-freeTwiddlesDone + + if useBlinding { + var startRestoreBlindingPolys time.Time + if isProfileMode { + startRestoreBlindingPolys = time.Now() + } + csInv := inverseShifterProduct(loopCtx.shifters) + if err := s.restoreBlindingPolynomials(csInv); err != nil { + s.freeNumeratorShards(loopCtx.numeratorShards) + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startRestoreBlindingPolys)).Msg("computeNumerator: restore blinding polys") + } + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startIterLoop)).Msg("computeNumerator: full iteration loop") + } + return nil +} + +func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoopContext) error { + loopCtx.coset.Mul(&loopCtx.coset, &loopCtx.shifters[i]) + loopCtx.cosetExponentiatedToNMinusOne.Exp(loopCtx.coset, loopCtx.bn). + Sub(&loopCtx.cosetExponentiatedToNMinusOne, &loopCtx.one) + + batchInvertDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) + }) + if err := <-batchInvertDone; err != nil { + return err + } + + s.applyNumeratorBlindingScale(i, loopCtx) + if i == 1 { + // We have to update the scalingVector; instead of scaling by + // cosets we scale by the twiddles of the large domain. + w := s.domain1.Generator + loopCtx.scalingVector = make([]fr.Element, loopCtx.n) + fft.BuildExpTable(w, loopCtx.scalingVector) + + // Reuse memory. + copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) + fft.BitReverse(loopCtx.scalingVectorRev) + } + + // We do **a lot** of FFT here, but on the small domain. + // Note that for all the polynomials in the proving key + // (Ql, Qr, Qm, Qo, S1, S2, S3, Qcp, Qc) and ID, LOne + // we could pre-compute these rho*2 FFTs and store them + // at the cost of a huge memory footprint. + var startGpuInverseScaleForward time.Time + if isProfileMode { + startGpuInverseScaleForward = time.Now() + } + + // Inverse NTT -> Scale -> Forward NTT all on GPU using persistent GPU memory. + if err := s.gpuNTTInverseScaleForwardOnDevice(loopCtx.gpuState, loopCtx.scalingVector, loopCtx.scalingVectorRev, s.pk); err != nil { + return err + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startGpuInverseScaleForward)).Msg("computeNumerator: gpuNTTInverseScaleForwardOnDevice") + } + + // Evaluate constraints on GPU. + constraintParams := gpuConstraintEvalParams{ + beta: s.beta, + gamma: s.gamma, + alpha: s.alpha, + coset: loopCtx.coset, + cosetExponentiatedToNMinusOne: loopCtx.cosetExponentiatedToNMinusOne, + cs: loopCtx.cs, + css: loopCtx.css, + cardinalityInv: s.domain0.CardinalityInv, + n: loopCtx.n, + nbBsbGates: loopCtx.nbBsbGates, + } + var startEvalConstraints time.Time + if isProfileMode { + startEvalConstraints = time.Now() + } + dNumeratorShard, err := s.gpuEvaluateConstraints( + loopCtx.gpuState, + constraintParams, + loopCtx.twiddles0, // CPU version for computeBlindingPolynomials + loopCtx.dTwiddles0, // GPU version for computeOrderingConstraint + *loopCtx.dPrecomputedDenominators, + s.bp, + nil, + ) + if err != nil { + return err + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startEvalConstraints)).Msg("computeNumerator: gpuEvaluateConstraints") + } + loopCtx.numeratorShards[i] = dNumeratorShard + + loopCtx.cosetExponentiatedToNMinusOne. + Inverse(&loopCtx.cosetExponentiatedToNMinusOne) + s.applyNumeratorBlindingUnscale(i, loopCtx) + return nil +} + +func (s *instance) buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx *computeNumeratorLoopContext) error { + if loopCtx == nil || loopCtx.dPrecomputedDenominators == nil { + return fmt.Errorf("computeNumerator: nil denominator device slice") + } + if loopCtx.dTwiddles0.IsEmpty() || loopCtx.dTwiddles0.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid twiddles device slice size %d, expected %d", loopCtx.dTwiddles0.Len(), loopCtx.n) + } + + if loopCtx.dPrecomputedDenominators.IsEmpty() { + dDenominators, err := allocDeviceUninitialized(loopCtx.n) + if err != nil { + return err + } + *loopCtx.dPrecomputedDenominators = dDenominators + } else if loopCtx.dPrecomputedDenominators.Len() != loopCtx.n { + return fmt.Errorf("computeNumerator: invalid denominator device slice size %d, expected %d", loopCtx.dPrecomputedDenominators.Len(), loopCtx.n) + } + + cfg := icicle_core.DefaultVecOpsConfig() + + // dTwiddles0 is a Montgomery scalar vector in domain0 regular order. + // ScalarMulVec expects the scalar in standard form and preserves the + // Montgomery representation of the vector result. + dCosetStd := uploadScalarStdOnCurrentDevice(loopCtx.coset, cfg) + defer dCosetStd.Free() + if err := icicle_vecops.ScalarMulVec( + dCosetStd, + loopCtx.dTwiddles0, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators coset*twiddles failed: %s", err.AsString()) + } + + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneMont := uploadScalarMontOnCurrentDevice(minusOne, cfg) + defer dMinusOneMont.Free() + if err := icicle_vecops.ScalarAddVec( + dMinusOneMont, + *loopCtx.dPrecomputedDenominators, + *loopCtx.dPrecomputedDenominators, + cfg, + ); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: build denominators subtract one failed: %s", err.AsString()) + } + + if err := s.batchInvertOnCurrentDevice(*loopCtx.dPrecomputedDenominators); err != icicle_runtime.Success { + return fmt.Errorf("computeNumerator: batchInvert failed: %s", err.AsString()) + } + return nil +} + +func (s *instance) applyNumeratorBlindingScale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindScale time.Time + if isProfileMode { + startBlindScale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + acc := loopCtx.cosetExponentiatedToNMinusOne + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &acc) + acc.Mul(&acc, &loopCtx.shifters[i]) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindScale)).Msg("computeNumerator: scale blinding polys") + } +} + +func (s *instance) applyNumeratorBlindingUnscale(i int, loopCtx *computeNumeratorLoopContext) { + if !useBlinding { + return + } + + var startBlindUnscale time.Time + if isProfileMode { + startBlindUnscale = time.Now() + } + for _, q := range s.bp { + cq := q.Coefficients() + for j := 0; j < len(cq); j++ { + cq[j].Mul(&cq[j], &loopCtx.cosetExponentiatedToNMinusOne) + } + } + if isProfileMode { + l := logger.Logger() + l.Debug().Int("iter", i).Dur("took", time.Since(startBlindUnscale)).Msg("computeNumerator: unscale blinding polys") + } +} + +func (s *instance) restoreBlindingPolynomials(csInv fr.Element) error { + for _, q := range s.bp { + if q == nil { + continue + } + cp := q.Coefficients() + if len(cp) == 0 { + continue + } + var acc fr.Element + acc.SetOne() + for i := 0; i < len(cp); i++ { + cp[i].Mul(&cp[i], &acc) + acc.Mul(&acc, &csInv) + } + } + return nil +} + +func inverseShifterProduct(shifters []fr.Element) fr.Element { + var acc fr.Element + acc.SetOne() + for i := 0; i < len(shifters); i++ { + acc.Mul(&acc, &shifters[i]) + } + acc.Inverse(&acc) + return acc +} + +func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial) (_ *iop.Polynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("downloadNumeratorFromGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("downloadNumeratorFromGPU: invalid dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard count mismatch: got %d, expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + + defer func() { + if err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + } + }() + + for i := 0; i < gpuNumerator.rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("downloadNumeratorFromGPU: shard %d is empty", i) + } + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + var dMerged icicle_core.DeviceSlice + + mergeDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + return + } + // Async boundary before merged slice is consumed by host copy. + mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + }) + if mergeErr := <-mergeDone; mergeErr != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, mergeErr + } + + cres := make([]fr.Element, totalSize) + cresHost := icicle_core.HostSliceFromElements(cres) + downloadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + }) + if err := <-downloadDone; err != nil { + s.putTempDeviceSlice(dMerged, totalSize) + return nil, err + } + s.putTempDeviceSlice(dMerged, totalSize) + + s.freeNumeratorShards(gpuNumerator.shards) + return iop.NewPolynomial(&cres, iop.Form{Basis: iop.LagrangeCoset, Layout: iop.BitReverse}), nil +} + +func (s *instance) freeNumeratorShards(shards []icicle_core.DeviceSlice) { + if len(shards) == 0 { + return + } + for i := 0; i < len(shards); i++ { + if !shards[i].IsEmpty() { + s.putTempDeviceSlice(shards[i], shards[i].Len()) + shards[i] = icicle_core.DeviceSlice{} + } + } +} + +func (s *instance) batchInvert(dVec icicle_core.DeviceSlice) { + if dVec.Len() == 0 { + return + } + + done := make(chan icicle_runtime.EIcicleError, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + done <- s.batchInvertOnCurrentDevice(dVec) + }) + if err := <-done; err != icicle_runtime.Success { + panic(fmt.Sprintf("batchInvert: BatchInverseVec failed: %s", err.AsString())) + } +} + +// batchInvertOnCurrentDevice assumes caller already runs on the active device thread. +func (s *instance) batchInvertOnCurrentDevice(dVec icicle_core.DeviceSlice) icicle_runtime.EIcicleError { + if dVec.Len() == 0 { + return icicle_runtime.Success + } + err := icicle_bw6761.FromMontgomery(dVec) + if err == icicle_runtime.Success { + cfg := icicle_core.DefaultVecOpsConfig() + err = icicle_vecops.BatchInverseVec(dVec, dVec, cfg) + } + if err == icicle_runtime.Success { + err = icicle_bw6761.ToMontgomery(dVec) + } + return err +} + +// gpuPolysState holds GPU-resident polynomial data to avoid repeated CPU-GPU transfers. +// Use ensurePolysOnSharedGPU to populate/reuse and freeGPUPolys to release GPU memory. +type gpuPolysState struct { + deviceSlices []icicle_core.DeviceSlice + hostSlices []icicle_core.HostSlice[fr.Element] + polys []*iop.Polynomial + originalForm []iop.Form + polyToIdx map[*iop.Polynomial]int +} + +func (s *instance) sharedGPUStateInitialCap(extra int) int { + base := 16 + len(s.bp) + 2*len(s.commitmentInfo) + if extra > 0 { + base += extra + } + return base +} + +func (s *instance) initSharedGPUState() { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + if s.sharedGPUState != nil { + return + } + initialCap := s.sharedGPUStateInitialCap(0) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } +} + +func (s *instance) releaseSharedGPUState() { + s.gpuStateMu.Lock() + state := s.sharedGPUState + s.sharedGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) releaseLinearizedEvalGPUState() { + s.gpuStateMu.Lock() + state := s.linearizedEvalGPUState + s.linearizedEvalGPUState = nil + s.gpuStateMu.Unlock() + s.freeGPUPolys(state) +} + +func (s *instance) freeIdleTempGPUMemoryOnDevice() { + if s == nil || s.tempGPUMemPool == nil { + return + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +func (s *instance) prepareLinearizedEvalGPUState(source *gpuPolysState) error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.clonePolysOnGPUFromState(source, s.buildLinearizedEvalGPUBatch()) + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) prepareLinearizedEvalGPUStateFromHost() error { + s.freeIdleTempGPUMemoryOnDevice() + snapshot, err := s.uploadPolysToGPUState(s.buildLinearizedEvalGPUBatch(), "prepareLinearizedEvalGPUStateFromHost") + if err != nil { + return err + } + + s.gpuStateMu.Lock() + old := s.linearizedEvalGPUState + s.linearizedEvalGPUState = snapshot + s.gpuStateMu.Unlock() + s.freeGPUPolys(old) + return nil +} + +func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop.Polynomial) (*gpuPolysState, error) { + if source == nil { + return nil, fmt.Errorf("clonePolysOnGPUFromState: nil source state") + } + + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("clonePolysOnGPUFromState: empty polynomial batch") + } + + srcSlices := make([]icicle_core.DeviceSlice, len(unique)) + useSource := make([]bool, len(unique)) + for i, p := range unique { + idx, ok := source.polyToIdx[p] + if ok && idx >= 0 && idx < len(source.deviceSlices) && !source.deviceSlices[idx].IsEmpty() { + srcSlices[i] = source.deviceSlices[idx] + useSource[i] = true + } + } + + snapshot := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + snapshot.polys[i] = p + snapshot.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + snapshot.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("clonePolysOnGPUFromState") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for snapshot cloning before handing state to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i := range srcSlices { + if useSource[i] { + dst, allocErr := allocDeviceUninitialized(srcSlices[i].Len()) + if allocErr != nil { + runErr = fmt.Errorf("clonePolysOnGPUFromState: alloc failed at index %d: %w", i, allocErr) + return + } + if err := copyDeviceSliceIntoOnCurrentDevice(dst, srcSlices[i], cfg); err != icicle_runtime.Success { + _ = dst.Free() + runErr = fmt.Errorf("clonePolysOnGPUFromState: device copy failed at index %d: %s", i, err.AsString()) + return + } + snapshot.deviceSlices[i] = dst + continue + } + + coeffs := unique[i].Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("clonePolysOnGPUFromState: empty host coefficients at index %d", i) + return + } + host := icicle_core.HostSliceFromElements(coeffs) + var dst icicle_core.DeviceSlice + host.CopyToDeviceAsync(&dst, cfg.StreamHandle, true) + if dst.IsEmpty() { + runErr = fmt.Errorf("clonePolysOnGPUFromState: host upload failed at index %d", i) + return + } + snapshot.deviceSlices[i] = dst + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(snapshot) + return nil, err + } + return snapshot, nil +} + +func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) (*gpuPolysState, error) { + unique := make([]*iop.Polynomial, 0, len(polys)) + seen := make(map[*iop.Polynomial]struct{}, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + unique = append(unique, p) + } + if len(unique) == 0 { + return nil, fmt.Errorf("%s: empty polynomial batch", label) + } + + state := &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, len(unique)), + hostSlices: make([]icicle_core.HostSlice[fr.Element], len(unique)), + polys: make([]*iop.Polynomial, len(unique)), + originalForm: make([]iop.Form, len(unique)), + polyToIdx: make(map[*iop.Polynomial]int, len(unique)), + } + for i, p := range unique { + state.polys[i] = p + state.originalForm[i] = iop.Form{Basis: p.Basis, Layout: p.Layout} + state.polyToIdx[p] = i + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label) + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + + for i, p := range unique { + coeffs := p.Coefficients() + if len(coeffs) == 0 { + runErr = fmt.Errorf("%s: empty host coefficients at index %d", label, i) + return + } + state.hostSlices[i] = icicle_core.HostSliceFromElements(coeffs) + state.hostSlices[i].CopyToDeviceAsync(&state.deviceSlices[i], cfg.StreamHandle, true) + if state.deviceSlices[i].IsEmpty() { + runErr = fmt.Errorf("%s: host upload failed at index %d", label, i) + return + } + } + }) + if err := <-done; err != nil { + s.freeGPUPolys(state) + return nil, err + } + return state, nil +} + +func (s *instance) getTempDeviceSlice(n int) icicle_core.DeviceSlice { + if s == nil { + panic("getTempDeviceSlice: nil instance") + } + if s.tempGPUMemPool == nil { + panic("getTempDeviceSlice: temp GPU memory pool is not initialized") + } + return s.tempGPUMemPool.Get(n) +} + +func (s *instance) putTempDeviceSlice(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + if s != nil && s.tempGPUMemPool != nil { + s.tempGPUMemPool.Put(ds, n) + return + } + _ = ds.Free() +} + +func (s *instance) releaseTempGPUMemoryPool() { + if s == nil || s.tempGPUMemPool == nil { + return + } + freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) + if !s.blindedZCanonicalGPU.IsEmpty() { + s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) + s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} + } + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + s.tempGPUMemPool.FreeAll() + close(done) + }) + <-done +} + +// ensurePolysOnSharedGPU uploads missing polynomials once and reuses already-uploaded slices. +func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysState, error) { + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if s.sharedGPUState == nil { + initialCap := s.sharedGPUStateInitialCap(len(polys)) + s.sharedGPUState = &gpuPolysState{ + deviceSlices: make([]icicle_core.DeviceSlice, 0, initialCap), + hostSlices: make([]icicle_core.HostSlice[fr.Element], 0, initialCap), + polys: make([]*iop.Polynomial, 0, initialCap), + originalForm: make([]iop.Form, 0, initialCap), + polyToIdx: make(map[*iop.Polynomial]int, initialCap), + } + } + state := s.sharedGPUState + if state == nil { + return nil, fmt.Errorf("ensurePolysOnSharedGPU: shared GPU state is nil") + } + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int, len(polys)) + } + + newIndices := make([]int, 0, len(polys)) + for _, p := range polys { + if p == nil { + continue + } + if _, ok := state.polyToIdx[p]; ok { + continue + } + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, icicle_core.DeviceSlice{}) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + newIndices = append(newIndices, idx) + } + if len(newIndices) == 0 { + return state, nil + } + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("ensurePolysOnSharedGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + // Async boundary for GPU uploads in ensurePolysOnSharedGPU. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + }() + for _, idx := range newIndices { + p := state.polys[idx] + if p == nil { + continue + } + cp := p.Coefficients() + state.hostSlices[idx] = icicle_core.HostSliceFromElements(cp) + state.hostSlices[idx].CopyToDeviceAsync(&state.deviceSlices[idx], cfg.StreamHandle, true) + } + }) + if err := <-done; err != nil { + return nil, err + } + return state, nil +} + +func getStateDeviceSlice(state *gpuPolysState, p *iop.Polynomial, label string) (icicle_core.DeviceSlice, error) { + if state == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil GPU state", label) + } + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): nil polynomial", label) + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): polynomial is not registered on GPU", label) + } + ds := state.deviceSlices[idx] + if ds.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("getStateDeviceSlice(%s): empty device slice", label) + } + return ds, nil +} + +// gpuNTTInverseScaleForwardOnDevice performs inverse NTT → scale → forward NTT +// on GPU-resident polynomial data without CPU-GPU transfers for polynomial data. +// The scaling vectors are uploaded each call (they may change between iterations). +func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scalingVector, scalingVectorRev []fr.Element, pk *ProvingKey) error { + if state == nil || len(state.polys) == 0 { + return nil + } + + device := &s.device + var scalingVectorDevice, scalingVectorRevDevice icicle_core.DeviceSlice + + // Upload scaling vectors to GPU + uploadDone := make(chan error, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) + scalingHost := icicle_core.HostSliceFromElements(scalingVector) + scalingHost.CopyToDevice(&scalingVectorDevice, false) + scalingVectorRevDevice = s.getTempDeviceSlice(len(scalingVectorRev)) + scalingRevHost := icicle_core.HostSliceFromElements(scalingVectorRev) + scalingRevHost.CopyToDevice(&scalingVectorRevDevice, false) + + // Convert scaling vectors from Montgomery form to standard form + if err := icicle_bw6761.FromMontgomery(scalingVectorDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVector failed: %s", err.AsString()) + return + } + if err := icicle_bw6761.FromMontgomery(scalingVectorRevDevice); err != icicle_runtime.Success { + uploadDone <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: FromMontgomery scalingVectorRev failed: %s", err.AsString()) + return + } + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return err + } + + doneChans := make([]chan error, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + + done := make(chan error, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Step 1: Inverse NTT + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + return + } + + // Step 2: Scale by vector using GPU vecOps + vecCfg := icicle_core.DefaultVecOpsConfig() + vecCfg.StreamHandle = stream + vecCfg.IsAsync = true + + var scaleDevice icicle_core.DeviceSlice + if layout == iop.Regular { + // After KNR inverse, output is BitReverse → use scalingVectorRev + scaleDevice = scalingVectorRevDevice + } else { + // After KRN inverse, output is Regular → use scalingVector + scaleDevice = scalingVectorDevice + } + + if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + return + } + + // Step 3: Forward NTT to Lagrange + one := icicle_ntt.GetDefaultNttConfig().CosetGen + cfg.CosetGen = one + if layout == iop.Regular { + cfg.Ordering = icicle_core.KRN // BitReverse → Regular + } else { + cfg.Ordering = icicle_core.KNR // Regular → BitReverse + } + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + return + } + done <- nil + }) + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + if err := <-doneChans[i]; err != nil { + return err + } + } + } + + // Update polynomial metadata: final result is in Lagrange, same layout as original + for _, p := range state.polys { + if p != nil { + p.Basis = iop.Lagrange + } + } + + // Free scaling vectors from device + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + return nil +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +// +// NOTE: The prover hot path should use the state-based API and keep data on device. +// This wrapper exists for compatibility/testing where callers expect host coefficients +// to be materialized after the transform. +func (s *instance) gpuNTTInverseBatch(polys []*iop.Polynomial, pk *ProvingKey) { + if len(polys) == 0 { + return + } + state, err := s.ensurePolysOnSharedGPU(polys) + if err != nil { + panic(fmt.Sprintf("gpuNTTInverseBatch: ensurePolysOnSharedGPU failed: %v", err)) + } + + s.gpuNTTInverseBatchOnState(state, pk) + + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + for _, p := range polys { + if p == nil { + continue + } + idx, ok := state.polyToIdx[p] + if !ok || idx < 0 || idx >= len(state.deviceSlices) || state.deviceSlices[idx].IsEmpty() { + continue + } + cp := p.Coefficients() + host := icicle_core.HostSliceFromElements(cp) + host.CopyFromDevice(&state.deviceSlices[idx]) + copy(cp, ([]fr.Element)(host)) + } + close(done) + }) + <-done +} + +// gpuNTTInverseBatchOnState performs inverse NTT directly on GPU-resident polynomial data +// without CPU-GPU transfers. The polynomials must already be on GPU in gpuState.deviceSlices. +// This updates the polynomial metadata (basis, layout) after the operation. +func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKey) { + if state == nil || len(state.polys) == 0 { + return + } + + device := &s.device + doneChans := make([]chan struct{}, len(state.polys)) + + for i := range state.polys { + p := state.polys[i] + if p == nil { + continue + } + + switch p.Basis { + case iop.Canonical: + continue // already in canonical form + case iop.Lagrange, iop.LagrangeCoset: + // Schedule GPU work + done := make(chan struct{}, 1) + doneChans[i] = done + scalarsDevice := state.deviceSlices[i] + layout := p.Layout + basis := p.Basis + + icicle_runtime.RunOnDevice(device, func(args ...any) { + cfg := icicle_ntt.GetDefaultNttConfig() + stream, _ := icicle_runtime.CreateStream() + cfg.StreamHandle = stream + cfg.IsAsync = true + + // Select ordering and coset generator depending on basis and input layout + if basis == iop.LagrangeCoset { + cfg.CosetGen = pk.CosetGenerator + } + + // Base-domain inverse: + // - Regular input → KNR (output BitReverse) + // - BitReverse input → KRN (output Regular) + if layout == iop.Regular { + cfg.Ordering = icicle_core.KNR + } else { + cfg.Ordering = icicle_core.KRN + } + + // Run NTT inverse directly on the existing device slice (in-place) + if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) + } + icicle_runtime.SynchronizeStream(stream) + + // Update metadata inside closure to avoid race + p.Basis = iop.Canonical + if layout == iop.Regular { + p.Layout = iop.BitReverse + } else { + p.Layout = iop.Regular + } + close(done) + }) + default: + panic("unsupported polynomial basis") + } + } + + // Wait for all scheduled tasks + for i := range doneChans { + if doneChans[i] != nil { + <-doneChans[i] + } + } +} + +// freeGPUPolys releases GPU memory for polynomial data. +func (s *instance) freeGPUPolys(state *gpuPolysState) { + if state == nil { + return + } + + device := &s.device + freeDone := make(chan struct{}) + + icicle_runtime.RunOnDevice(device, func(args ...any) { + for i := range state.polys { + if state.deviceSlices[i].IsEmpty() { + continue + } + _ = state.deviceSlices[i].Free() + } + close(freeDone) + }) + <-freeDone +} + +// gpuMemoryPool manages a pool of reusable device slices to avoid repeated allocations. +// Must be used within RunOnDevice context to ensure thread safety per device. +type gpuMemoryPool struct { + freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice + mu sync.Mutex +} + +// newGPUMemoryPool creates a new GPU memory pool. +func newGPUMemoryPool() *gpuMemoryPool { + return &gpuMemoryPool{ + freeSlices: make(map[int][]icicle_core.DeviceSlice), + } +} + +// Get returns a device slice of the specified size, either from the pool or newly allocated. +func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { + p.mu.Lock() + defer p.mu.Unlock() + + // Check if we have a free slice of this size + if slices, ok := p.freeSlices[n]; ok && len(slices) > 0 { + // Reuse the last slice + slice := slices[len(slices)-1] + p.freeSlices[n] = slices[:len(slices)-1] + return slice + } + + // No free slice available, allocate a new one. + // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + // Free all currently cached (idle) slices to reduce memory pressure. + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + + if ds, err := allocDeviceUninitialized(n); err == nil { + return ds + } + + panic(fmt.Sprintf("gpuMemoryPool.Get: allocation failed for size %d after clearing idle cache", n)) +} + +// Put returns a device slice to the pool for reuse instead of freeing it. +func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { + if ds.IsEmpty() { + return + } + + p.mu.Lock() + defer p.mu.Unlock() + + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// FreeAll releases all pooled device slices. +func (p *gpuMemoryPool) FreeAll() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) +} + +// allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. +// Use when the destination is fully overwritten by a kernel. +func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { + var ds icicle_core.DeviceSlice + if _, err := ds.Malloc(int(unsafe.Sizeof(fr.Element{})), n); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("allocDeviceUninitialized: malloc failed for size %d: %s", n, err.AsString()) + } + return ds, nil +} + +// mustAllocDeviceUninitialized is like allocDeviceUninitialized but panics on failure. +// Use only in contexts where error propagation is impractical (e.g. upload helpers). +func mustAllocDeviceUninitialized(n int) icicle_core.DeviceSlice { + ds, err := allocDeviceUninitialized(n) + if err != nil { + panic(err) + } + return ds +} + +// freeDeviceSlice frees a device slice if non-empty and zeroes the pointer. +// Use for directly-allocated slices, NOT pool-allocated ones (use putTempDeviceSlice for those). +func freeDeviceSlice(ds *icicle_core.DeviceSlice) { + if ds != nil && !ds.IsEmpty() { + _ = ds.Free() + *ds = icicle_core.DeviceSlice{} + } +} + +// freeSliceOnDevice frees a device slice on the specified device and blocks +// until complete. Use outside RunOnDevice closures. Zeroes the slice after freeing. +func freeSliceOnDevice(ds *icicle_core.DeviceSlice, device *icicle_runtime.Device) { + if ds == nil || ds.IsEmpty() { + return + } + d := *ds + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + _ = d.Free() + close(done) + }) + <-done + *ds = icicle_core.DeviceSlice{} +} + +// copyDeviceSliceIntoOnCurrentDevice copies src into dst entirely on GPU. +func copyDeviceSliceIntoOnCurrentDevice( + dst, src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if src.IsEmpty() || src.Len() <= 0 || dst.IsEmpty() || dst.Len() < src.Len() { + return icicle_runtime.InvalidArgument + } + src.CheckDevice() + dst.CheckDevice() + + srcElemSize := src.SizeOfElement() + dstElemSize := dst.SizeOfElement() + if srcElemSize <= 0 || dstElemSize <= 0 || srcElemSize != dstElemSize { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(src.Len() * srcElemSize) + if cfg.IsAsync { + return icicle_runtime.CopyAsync(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen, cfg.StreamHandle) + } + _, err := icicle_runtime.Copy(dst.AsUnsafePointer(), src.AsUnsafePointer(), byteLen) + return err +} + +// zeroDeviceSliceOnCurrentDevice zero-fills dst entirely on GPU. +func zeroDeviceSliceOnCurrentDevice( + dst icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) icicle_runtime.EIcicleError { + if dst.IsEmpty() || dst.Len() <= 0 { + return icicle_runtime.InvalidArgument + } + dst.CheckDevice() + + elemSize := dst.SizeOfElement() + if elemSize <= 0 { + return icicle_runtime.InvalidArgument + } + + byteLen := uint(dst.Len() * elemSize) + if cfg.IsAsync { + return icicle_runtime.MemSetAsync(dst.AsUnsafePointer(), 0, byteLen, cfg.StreamHandle) + } + return icicle_runtime.MemSet(dst.AsUnsafePointer(), 0, byteLen) +} + +func createAsyncVecOpsConfigOnCurrentDevice(label string) (icicle_core.VecOpsConfig, icicle_runtime.Stream, error) { + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + return icicle_core.VecOpsConfig{}, nil, fmt.Errorf("%s: create stream failed: %s", label, eStream.AsString()) + } + cfg := icicle_core.DefaultVecOpsConfig() + cfg.StreamHandle = stream + cfg.IsAsync = true + return cfg, stream, nil +} + +func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label string) error { + if stream == nil { + return nil + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + _ = icicle_runtime.DestroyStream(stream) + return fmt.Errorf("%s: synchronize stream failed: %s", label, eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success { + return fmt.Errorf("%s: destroy stream failed: %s", label, eDestroy.AsString()) + } + return nil +} + +// makeFinisher returns a closure that synchronizes and destroys the stream, +// then sends the (possibly merged) error to done. Use inside RunOnDevice closures. +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { + return func(runErr error) { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { + runErr = syncErr + } + done <- runErr + } +} + +// uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. +// Use this for additions where the vector is already in Montgomery form. +func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + return ds +} + +func uploadScalarMontOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + if cfg.IsAsync { + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&ds, false) + } + return ds +} + +// uploadScalarStd uploads a scalar in STANDARD form (not Montgomery) as a single-element device slice. +// For use with ScalarMulVec: (a*R) * b_std = (a*b)*R +func uploadScalarStd(scalar fr.Element) icicle_core.DeviceSlice { + vec := []fr.Element{scalar} + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(1) + host.CopyToDevice(&ds, false) + // Convert from Montgomery form to standard form + if err := icicle_bw6761.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +func uploadScalarStdOnCurrentDevice(scalar fr.Element, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + ds := uploadScalarMontOnCurrentDevice(scalar, cfg) + // Fallback to sync conversion for compatibility with ICICLE wrappers + // that do not expose *_WithConfig APIs. + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bw6761.FromMontgomery(ds); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } + return ds +} + +// uploadVectorStd uploads a vector and converts to standard form. +func uploadVectorStd(vec []fr.Element) icicle_core.DeviceSlice { + ds := mustAllocDeviceUninitialized(len(vec)) + uploadVectorStdInto(&ds, vec) + return ds +} + +// uploadVectorStdInto uploads vec into an existing device slice and converts it to standard form. +// The destination must already be allocated with enough capacity for len(vec) elements. +func uploadVectorStdInto(dst *icicle_core.DeviceSlice, vec []fr.Element) { + cfg := icicle_core.DefaultVecOpsConfig() + uploadVectorStdIntoOnCurrentDevice(dst, vec, cfg) +} + +// uploadVectorStdIntoOnCurrentDevice uploads vec into an existing device slice and converts it +// to standard form while honoring the provided vector-op config/stream. +func uploadVectorStdIntoOnCurrentDevice( + dst *icicle_core.DeviceSlice, + vec []fr.Element, + cfg icicle_core.VecOpsConfig, +) { + host := icicle_core.HostSliceFromElements(vec) + if cfg.IsAsync { + host.CopyToDeviceAsync(dst, cfg.StreamHandle, false) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } else { + host.CopyToDevice(dst, false) + } + // Convert from Montgomery form to standard form + if err := icicle_bw6761.FromMontgomery(*dst); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +// uploadVector uploads a vector (keeps Montgomery form for additions). +func uploadVector(vec []fr.Element) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + ds := mustAllocDeviceUninitialized(len(vec)) + host.CopyToDevice(&ds, false) + return ds +} + +// uploadInt64Vector uploads int64 indices to a device slice. +func uploadInt64Vector(vec []int64) icicle_core.DeviceSlice { + cfg := icicle_core.DefaultVecOpsConfig() + return uploadInt64VectorOnCurrentDevice(vec, cfg) +} + +func uploadInt64VectorOnCurrentDevice(vec []int64, cfg icicle_core.VecOpsConfig) icicle_core.DeviceSlice { + host := icicle_core.HostSliceFromElements(vec) + var ds icicle_core.DeviceSlice + if cfg.IsAsync { + if _, err := ds.MallocAsync(int(unsafe.Sizeof(int64(0))), len(vec), cfg.StreamHandle); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc async failed: %s", err.AsString())) + } + host.CopyToDeviceAsync(&ds, cfg.StreamHandle, false) + return ds + } + if _, err := ds.Malloc(int(unsafe.Sizeof(int64(0))), len(vec)); err != icicle_runtime.Success { + panic(fmt.Sprintf("uploadInt64Vector: malloc failed: %s", err.AsString())) + } + host.CopyToDevice(&ds, false) + return ds +} + +// toStandardFormInPlace converts a device slice to standard form in-place (modifies the source). +// Use this for temporary vectors that won't be needed in Montgomery form. +func toStandardFormInPlace(src icicle_core.DeviceSlice) { + cfg := icicle_core.DefaultVecOpsConfig() + toStandardFormInPlaceWithCfg(src, cfg) +} + +func toStandardFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bw6761.FromMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("FromMontgomery failed: %s", err.AsString())) + } +} + +func toMontgomeryFormInPlaceWithCfg(src icicle_core.DeviceSlice, cfg icicle_core.VecOpsConfig) { + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("SynchronizeStream failed: %s", eSync.AsString())) + } + } + if err := icicle_bw6761.ToMontgomery(src); err != icicle_runtime.Success { + panic(fmt.Sprintf("ToMontgomery failed: %s", err.AsString())) + } +} + +// multiplyMontgomerySlices multiplies two device slices that are both in Montgomery form. +// It creates a copy of dSlice1Mont, converts the copy to standard form, and then multiplies +// it with dSlice2Mont (which remains in Montgomery form). The result is stored in dResult +// and will be in Montgomery form. +// +// Parameters: +// - dSlice1Mont: first device slice in Montgomery form (not modified) +// - dSlice2Mont: second device slice in Montgomery form (not modified) +// - dResult: destination device slice for the result (must be pre-allocated) +// - state: GPU state with memory pool and vector configuration +// - n: size of the slices +func multiplyMontgomerySlices( + dSlice1Mont, dSlice2Mont icicle_core.DeviceSlice, + dResult icicle_core.DeviceSlice, + state *gpuConstraintEvalState, + vecCfg icicle_core.VecOpsConfig, + n int, +) error { + // Copy dSlice1Mont to standard form + dSlice1Std := state.getTempDeviceSlice(n) + defer state.putTempDeviceSlice(dSlice1Std, n) + + if err := copyDeviceSliceIntoOnCurrentDevice(dSlice1Std, dSlice1Mont, vecCfg); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: device copy failed: %s", err.AsString()) + } + toStandardFormInPlace(dSlice1Std) + + // Multiply: dSlice1Std (standard) * dSlice2Mont (Montgomery) = dResult (Montgomery) + if err := icicle_vecops.VecOp(dSlice1Std, dSlice2Mont, dResult, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("multiplyMontgomerySlices: VecOp multiplication failed: %s", err.AsString()) + } + return nil +} + +// gpuConstraintEvalParams holds parameters for GPU constraint evaluation +type gpuConstraintEvalParams struct { + beta fr.Element + gamma fr.Element + alpha fr.Element + coset fr.Element + cosetExponentiatedToNMinusOne fr.Element + cs fr.Element // domain1.FrMultiplicativeGen + css fr.Element // cs^2 + cardinalityInv fr.Element + n int + nbBsbGates int +} + +// gpuConstraintEvalState holds intermediate state during constraint evaluation +type gpuConstraintEvalState struct { + // Polynomial device slices (may point to gpuState or allocated buffers) + dL, dR, dO, dZ, dZS icicle_core.DeviceSlice + dQl, dQr, dQm, dQo, dQk icicle_core.DeviceSlice + dS1, dS2, dS3 icicle_core.DeviceSlice + // Intermediate results + dGate, dOrdering, dLocal, dResult icicle_core.DeviceSlice + // Scalar device slices + dGammaScalar icicle_core.DeviceSlice + // Configuration + vecCfg icicle_core.VecOpsConfig + // Helper function to get device slices + getDeviceSlice func(int) icicle_core.DeviceSlice + // Shared prover-level temporary GPU memory pool accessors + getTempDeviceSlice func(int) icicle_core.DeviceSlice + putTempDeviceSlice func(icicle_core.DeviceSlice, int) + // Track allocated polynomial buffers for automatic cleanup + allocatedPolyBuffers []struct { + slice icicle_core.DeviceSlice + size int + } +} + +// allocate allocates a new device slice from the memory pool and tracks it for automatic cleanup. +// Returns the allocated device slice. +func (s *gpuConstraintEvalState) allocate(size int) icicle_core.DeviceSlice { + slice := s.getTempDeviceSlice(size) + s.allocatedPolyBuffers = append(s.allocatedPolyBuffers, struct { + slice icicle_core.DeviceSlice + size int + }{slice, size}) + return slice +} + +// freeAllocatedPolyBuffers returns all allocated polynomial buffers to the memory pool. +// This should be called during cleanup to free all buffers allocated via allocate(). +func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { + for _, buf := range s.allocatedPolyBuffers { + s.putTempDeviceSlice(buf.slice, buf.size) + } + s.allocatedPolyBuffers = nil +} + +// computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. +// Returns device slices for the blinding polynomials. +func computeBlindingPolynomials( + n int, + twiddles0 []fr.Element, + bp []*iop.Polynomial, +) (dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice) { + blindL := make([]fr.Element, n) + blindR := make([]fr.Element, n) + blindO := make([]fr.Element, n) + blindZ := make([]fr.Element, n) + blindZS := make([]fr.Element, n) // ZS uses shifted index + + // TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove this. Since we don't normally use blindings, we will + // not make this optimization. + utils.Parallelize(n, func(start, end int) { + for i := start; i < end; i++ { + blindL[i] = bp[id_Bl].Evaluate(twiddles0[i]) + blindR[i] = bp[id_Br].Evaluate(twiddles0[i]) + blindO[i] = bp[id_Bo].Evaluate(twiddles0[i]) + blindZ[i] = bp[id_Bz].Evaluate(twiddles0[i]) + blindZS[i] = bp[id_Bz].Evaluate(twiddles0[(i+1)%n]) + } + }) + + dBlindL = uploadVector(blindL) + dBlindR = uploadVector(blindR) + dBlindO = uploadVector(blindO) + dBlindZ = uploadVector(blindZ) + dBlindZS = uploadVector(blindZS) + + return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS +} + +// applyBlindingToPolynomials applies blinding to polynomials L, R, O, Z, ZS. +// Allocates new buffers for L, R, O, Z (tracked for cleanup) and modifies ZS in-place. +// The original slices in gpuState remain unchanged. +func applyBlindingToPolynomials( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS icicle_core.DeviceSlice, +) error { + // L' = L + blindL (allocate new buffer, tracked for cleanup) + dLBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dL, dBlindL, dLBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add L failed: %s", err.AsString()) + } + state.dL = dLBlinded + + // R' = R + blindR (allocate new buffer, tracked for cleanup) + dRBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dR, dBlindR, dRBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add R failed: %s", err.AsString()) + } + state.dR = dRBlinded + + // O' = O + blindO (allocate new buffer, tracked for cleanup) + dOBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dO, dBlindO, dOBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add O failed: %s", err.AsString()) + } + state.dO = dOBlinded + + // Z' = Z + blindZ (allocate new buffer, tracked for cleanup) + dZBlinded := state.allocate(params.n) + if err := icicle_vecops.VecOp(state.dZ, dBlindZ, dZBlinded, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add Z failed: %s", err.AsString()) + } + state.dZ = dZBlinded + + // ZS' = ZS + blindZS + // Note: dZS is a temporary buffer created inside gpuEvaluateConstraints, + // so it's safe to modify it in-place. + if err := icicle_vecops.VecOp(state.dZS, dBlindZS, state.dZS, state.vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) + } + + // Free blinding vectors - no longer needed after creating blinded polynomials + dBlindL.Free() + dBlindR.Free() + dBlindO.Free() + dBlindZ.Free() + dBlindZS.Free() + return nil +} + +// scaleSVectorsByBeta scales S1, S2, S3 by beta. +// Allocates new buffers for S1, S2, S3 (tracked for cleanup). +func scaleSVectorsByBeta( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, +) error { + // S1' = S1 * beta (need to scale S1, S2, S3 by beta for ordering constraint) + // Use standard form for beta so: S1_mont * beta_std = (S1*beta)_mont + dBetaStd := uploadScalarStd(params.beta) + + // S1' = S1 * beta (allocate new buffer, tracked for cleanup) + dS1Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dS1Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S1 failed: %s", err.AsString()) + } + state.dS1 = dS1Scaled + + // S2' = S2 * beta (allocate new buffer, tracked for cleanup) + dS2Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dS2Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S2 failed: %s", err.AsString()) + } + state.dS2 = dS2Scaled + + // S3' = S3 * beta (allocate new buffer, tracked for cleanup) + dS3Scaled := state.allocate(params.n) + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dS3Scaled, state.vecCfg); err != icicle_runtime.Success { + dBetaStd.Free() + return fmt.Errorf("scaleSVectorsByBeta: ScalarMulVec S3 failed: %s", err.AsString()) + } + state.dS3 = dS3Scaled + + // Free dBetaStd - no longer needed after scaling S vectors + dBetaStd.Free() + return nil +} + +// computeGateConstraint computes the gate constraint. +// Returns dGate device slice. +func computeGateConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // gate = Ql*L' + Qr*R' + Qm*L'*R' + Qo*O' + Qk + sum(Qci*Pi) + // We use multiplyMontgomerySlices for all poly×poly multiplications. + // Note: dL, dR, dO are used later in ordering constraint, so we preserve them. + + dGate := state.getTempDeviceSlice(params.n) + dTmp := state.getTempDeviceSlice(params.n) + + // Ql * L' + if err := multiplyMontgomerySlices(state.dQl, state.dL, dGate, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Ql*L: %w", err) + } + + // + Qr * R' + if err := multiplyMontgomerySlices(state.dQr, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qr*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qr*R failed: %s", err.AsString()) + } + + // + Qm * L' * R' (need two multiplications) + // First: Qm * L' = dTmp (Montgomery) + if err := multiplyMontgomerySlices(state.dQm, state.dL, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qm*L: %w", err) + } + // Second: dTmp (Montgomery) * R' (Montgomery) + if err := multiplyMontgomerySlices(dTmp, state.dR, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: (Qm*L)*R: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qm*L*R failed: %s", err.AsString()) + } + + // + Qo * O' + if err := multiplyMontgomerySlices(state.dQo, state.dO, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qo*O: %w", err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qo*O failed: %s", err.AsString()) + } + + // + Qk + if err := icicle_vecops.VecOp(dGate, state.dQk, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qk failed: %s", err.AsString()) + } + + // + BSB gates: sum(Qci[2*i] * Qci[2*i+1]) + for i := 0; i < params.nbBsbGates; i++ { + origQci0 := state.getDeviceSlice(id_Qci + 2*i) + origQci1 := state.getDeviceSlice(id_Qci + 2*i + 1) + if !origQci0.IsEmpty() && !origQci1.IsEmpty() { + // Use helper to multiply Qci0 * Qci1 without modifying original values + if err := multiplyMontgomerySlices(origQci0, origQci1, dTmp, state, vecCfg, params.n); err != nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: Qci[%d]: %w", i, err) + } + if err := icicle_vecops.VecOp(dGate, dTmp, dGate, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeGateConstraint: VecOp add Qci[%d] failed: %s", i, err.AsString()) + } + } + } + + // Return temporary buffer to pool - no longer needed after Step 3 + state.putTempDeviceSlice(dTmp, params.n) + + return dGate, nil +} + +// computeOrderingConstraint computes the ordering constraint. +// Returns dOrdering device slice. +func computeOrderingConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dTwiddles0 icicle_core.DeviceSlice, // twiddles0 already on GPU + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // This is complex: involves ID computation, gamma, beta, Z, ZS, S1, S2, S3 + // id = twiddles[i] * coset * beta + // a = gamma + L' + id + // b = gamma + R' + id*cs + // c = gamma + O' + id*css + // r = a * b * c * Z' + // + // a2 = gamma + L' + S1*beta + // b2 = gamma + R' + S2*beta + // c2 = gamma + O' + S3*beta + // l = a2 * b2 * c2 * ZS' + // + // ordering = l - r + + // Compute ID vector: twiddles * coset * beta (computed on GPU) + // dTwiddles0 is already on GPU (passed as parameter, don't free it here) + + // Compute coset * beta on CPU, then upload as scalar in standard form + var cosetTimesBeta fr.Element + cosetTimesBeta.Mul(¶ms.coset, ¶ms.beta) + dCosetTimesBetaStd := uploadScalarStd(cosetTimesBeta) + dBetaStd := uploadScalarStd(params.beta) + + // Multiply twiddles0 by cosetTimesBeta on GPU: dID = (cosetTimesBeta * twiddles0) * R (Montgomery form) + dID := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dCosetTimesBetaStd, dTwiddles0, dID, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec coset*beta*twiddles failed: %s", err.AsString()) + } + + // Free temporary device slice (dTwiddles0 is owned by caller, don't free it) + dCosetTimesBetaStd.Free() + + // id * cs - use standard form for cs + dIDcs := state.getTempDeviceSlice(params.n) + dCsStd := uploadScalarStd(params.cs) + if err := icicle_vecops.ScalarMulVec(dCsStd, dID, dIDcs, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*cs failed: %s", err.AsString()) + } + + // id * css - use standard form for css + dIDcss := state.getTempDeviceSlice(params.n) + dCssStd := uploadScalarStd(params.css) + if err := icicle_vecops.ScalarMulVec(dCssStd, dID, dIDcss, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec id*css failed: %s", err.AsString()) + } + + // a = gamma + L' + id (dL now contains L' after in-place blinding) + dA := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dID, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a+id failed: %s", err.AsString()) + } + + // b = gamma + R' + id*cs (dR now contains R' after in-place blinding) + dB := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dIDcs, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b+id*cs failed: %s", err.AsString()) + } + + // c = gamma + O' + id*css (dO now contains O' after in-place blinding) + dC := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dIDcss, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c+id*css failed: %s", err.AsString()) + } + + // Return to pool: dID, dIDcs, dIDcss - no longer needed after computing a, b, c + state.putTempDeviceSlice(dID, params.n) + state.putTempDeviceSlice(dIDcs, params.n) + state.putTempDeviceSlice(dIDcss, params.n) + dCsStd.Free() + dCssStd.Free() + + // r = a * b * c * Z' (dZ now contains Z' after in-place blinding) + // For chain multiplication, convert operands to std form in-place when possible + dR_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a*b failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dR_ord, dC, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*c failed: %s", err.AsString()) + } + // Convert dR_ord to standard form in-place (temporary result, dZ needs to be preserved) + toStandardFormInPlace(dR_ord) + if err := icicle_vecops.VecOp(dR_ord, state.dZ, dR_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp r*Z failed: %s", err.AsString()) + } + + // Reuse dA, dB, dC for a2, b2, c2 instead of freeing and reallocating. + // To reduce peak memory, we scale S vectors by beta on-demand through a single temp buffer. + dScaledS := state.getTempDeviceSlice(params.n) + + // a2 = gamma + L' + S1*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dL, dA, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+L (a2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS1, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S1 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dA, dScaledS, dA, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2+S1*beta failed: %s", err.AsString()) + } + + // b2 = gamma + R' + S2*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dR, dB, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+R (b2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS2, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S2 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dB, dScaledS, dB, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp b2+S2*beta failed: %s", err.AsString()) + } + + // c2 = gamma + O' + S3*beta + if err := icicle_vecops.ScalarAddVec(state.dGammaScalar, state.dO, dC, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarAddVec gamma+O (c2) failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarMulVec(dBetaStd, state.dS3, dScaledS, vecCfg); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: ScalarMulVec beta*S3 failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dC, dScaledS, dC, vecCfg, icicle_core.Add); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp c2+S3*beta failed: %s", err.AsString()) + } + + // Free dGammaScalar - no longer needed after computing a2, b2, c2 + // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. + state.dGammaScalar.Free() + state.putTempDeviceSlice(dScaledS, params.n) + dBetaStd.Free() + + // l = a2 * b2 * c2 * ZS' (dZS now contains ZS' after in-place blinding) + dL_ord := state.getTempDeviceSlice(params.n) + // Convert dB to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dB) + if err := icicle_vecops.VecOp(dA, dB, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp a2*b2 failed: %s", err.AsString()) + } + // Convert dC to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(dC) + if err := icicle_vecops.VecOp(dL_ord, dC, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*c2 failed: %s", err.AsString()) + } + // Convert dZS to standard form in-place (temporary vector, no longer needed in Montgomery form) + toStandardFormInPlace(state.dZS) + if err := icicle_vecops.VecOp(dL_ord, state.dZS, dL_ord, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l*ZS failed: %s", err.AsString()) + } + + // Return temporary buffers to pool - no longer needed after computing l + state.putTempDeviceSlice(dA, params.n) + state.putTempDeviceSlice(dB, params.n) + state.putTempDeviceSlice(dC, params.n) + state.putTempDeviceSlice(state.dZS, params.n) + state.dZS = icicle_core.DeviceSlice{} + + // ordering = l - r, reuse dL_ord as the final ordering vector + if err := icicle_vecops.VecOp(dL_ord, dR_ord, dL_ord, vecCfg, icicle_core.Sub); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeOrderingConstraint: VecOp l-r failed: %s", err.AsString()) + } + + // Return dR_ord to pool - no longer needed after computing ordering + state.putTempDeviceSlice(dR_ord, params.n) + + // Return dL_ord as ordering (caller is responsible for freeing) + return dL_ord, nil +} + +// computeLocalConstraint computes the local constraint. +// Returns dLocal device slice. +func computeLocalConstraint( + state *gpuConstraintEvalState, + params gpuConstraintEvalParams, + dPrecomputedDenominators icicle_core.DeviceSlice, + vecCfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + // local = (Z' - 1) * LagrangeOne + // where LagrangeOne[i] = cosetExpMinusOne * cardinalityInv / (coset*twiddles0[i] - 1) + + if dPrecomputedDenominators.IsEmpty() || dPrecomputedDenominators.Len() < params.n { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: invalid denominator device slice size %d, expected at least %d", dPrecomputedDenominators.Len(), params.n) + } + + // Compute LagrangeOne on device. dPrecomputedDenominators is already in + // Montgomery form after batch inversion; ScalarMulVec expects the scalar in + // standard form and preserves a Montgomery vector result. + var lagrangeCoeff fr.Element + lagrangeCoeff.Mul(¶ms.cosetExponentiatedToNMinusOne, ¶ms.cardinalityInv) + dLagrangeCoeffStd := uploadScalarStdOnCurrentDevice(lagrangeCoeff, vecCfg) + dLagrangeOneStd := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dLagrangeCoeffStd, dPrecomputedDenominators, dLagrangeOneStd, vecCfg); err != icicle_runtime.Success { + dLagrangeCoeffStd.Free() + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarMulVec lagrangeOne failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dLagrangeOneStd, vecCfg) + dLagrangeCoeffStd.Free() + + // Z' - 1 using ScalarAddVec with minus one + var minusOne fr.Element + minusOne.SetOne() + minusOne.Neg(&minusOne) + dMinusOneScalar := uploadScalarMont(minusOne) + + dZMinusOne := state.getTempDeviceSlice(params.n) + // dZ now contains Z' after in-place blinding + if err := icicle_vecops.ScalarAddVec(dMinusOneScalar, state.dZ, dZMinusOne, vecCfg); err != icicle_runtime.Success { + dMinusOneScalar.Free() + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: ScalarAddVec Z-1 failed: %s", err.AsString()) + } + + // Free dMinusOneScalar - no longer needed after computing Z' - 1 + // Note: dZ is part of gpuState and will be freed later + dMinusOneScalar.Free() + + // local = (Z' - 1) * LagrangeOne_std + dLocal := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.VecOp(dZMinusOne, dLagrangeOneStd, dLocal, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { + return icicle_core.DeviceSlice{}, fmt.Errorf("computeLocalConstraint: VecOp (Z-1)*LagrangeOne failed: %s", err.AsString()) + } + + // Return temporary buffers to pool + state.putTempDeviceSlice(dZMinusOne, params.n) + state.putTempDeviceSlice(dLagrangeOneStd, params.n) + + return dLocal, nil +} + +// createGetDeviceSliceFunc creates a function to get device slices for polynomials. +// It returns a function that maps polynomial indices to their device slices. +func (s *instance) polyByID(polyIdx int) *iop.Polynomial { + switch polyIdx { + case id_L: + return s.polyL + case id_R: + return s.polyR + case id_O: + return s.polyO + case id_Z: + return s.polyZ + case id_ZS: + return s.polyZS + case id_Ql: + return s.trace.Ql + case id_Qr: + return s.trace.Qr + case id_Qm: + return s.trace.Qm + case id_Qo: + return s.trace.Qo + case id_Qk: + return s.polyQk + case id_S1: + return s.trace.S1 + case id_S2: + return s.trace.S2 + case id_S3: + return s.trace.S3 + default: + if polyIdx < id_Qci { + return nil + } + offset := polyIdx - id_Qci + i := offset / 2 + if i < 0 { + return nil + } + if offset%2 == 0 { + if i < len(s.trace.Qcp) { + return s.trace.Qcp[i] + } + return nil + } + if i < len(s.cCommitments) { + return s.cCommitments[i] + } + return nil + } +} + +func createGetDeviceSliceFunc( + gpuState *gpuPolysState, + polyToIdx map[*iop.Polynomial]int, + resolvePoly func(int) *iop.Polynomial, +) func(int) icicle_core.DeviceSlice { + return func(polyIdx int) icicle_core.DeviceSlice { + p := resolvePoly(polyIdx) + if p == nil { + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d is nil", polyIdx)) + } + if idx, ok := polyToIdx[p]; ok { + return gpuState.deviceSlices[idx] + } + panic(fmt.Sprintf("getDeviceSlice: polynomial id %d (ptr=%p) not found in polyToIdx (map has %d entries)", polyIdx, p, len(polyToIdx))) + } +} + +// initializeConstraintEvalState initializes the GPU constraint evaluation state. +// It sets up all device slices. +// The slices in gpuState are treated as read-only; helper functions will allocate +// separate working buffers whenever they need to modify data. +func initializeConstraintEvalState( + getDeviceSlice func(int) icicle_core.DeviceSlice, + getTempDeviceSlice func(int) icicle_core.DeviceSlice, + putTempDeviceSlice func(icicle_core.DeviceSlice, int), +) *gpuConstraintEvalState { + vecCfg := icicle_core.DefaultVecOpsConfig() + + state := &gpuConstraintEvalState{ + dL: getDeviceSlice(id_L), + dR: getDeviceSlice(id_R), + dO: getDeviceSlice(id_O), + dZ: getDeviceSlice(id_Z), + dQl: getDeviceSlice(id_Ql), + dQr: getDeviceSlice(id_Qr), + dQm: getDeviceSlice(id_Qm), + dQo: getDeviceSlice(id_Qo), + dQk: getDeviceSlice(id_Qk), + dS1: getDeviceSlice(id_S1), + dS2: getDeviceSlice(id_S2), + dS3: getDeviceSlice(id_S3), + vecCfg: vecCfg, + getDeviceSlice: getDeviceSlice, + getTempDeviceSlice: getTempDeviceSlice, + putTempDeviceSlice: putTempDeviceSlice, + } + + return state +} + +// gpuEvaluateConstraints evaluates all PLONK constraints on GPU. +// It takes polynomials already on GPU (via gpuState), computes blinding polynomial evaluations, +// and evaluates gate, ordering, and local constraints entirely on GPU. +// If result is non-nil, it downloads into result and returns an empty device slice. +// If result is nil, it returns a persistent device slice with the result. +// TODO(martun): Once we have a GPU kernel to evaluate polynomials, we can remove 'twiddles0'. Since we don't +// normally use blindings, we will not make this optimization. +func (s *instance) gpuEvaluateConstraints( + gpuState *gpuPolysState, + params gpuConstraintEvalParams, + twiddles0 []fr.Element, // CPU vector for computeBlindingPolynomials + dTwiddles0 icicle_core.DeviceSlice, // GPU vector for computeOrderingConstraint + dPrecomputedDenominators icicle_core.DeviceSlice, + bp []*iop.Polynomial, // blinding polynomials (already scaled for this iteration) + result []fr.Element, +) (icicle_core.DeviceSlice, error) { + if gpuState == nil || len(gpuState.polys) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("gpuState is nil or empty") + } + + n := params.n + device := &s.device + + // Create a map from polynomial to its index in gpuState + polyToIdx := make(map[*iop.Polynomial]int) + for i, p := range gpuState.polys { + if p != nil { + polyToIdx[p] = i + } + } + + // Get device slices for the polynomials we need. + getDeviceSlice := createGetDeviceSliceFunc(gpuState, polyToIdx, s.polyByID) + + done := make(chan error, 1) + var resultOnDevice icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(device, func(args ...any) { + state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). + // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). + state.dGammaScalar = uploadScalarMont(params.gamma) + + state.dZS = state.getTempDeviceSlice(n) + if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + return + } + + // Step 1: Compute and apply blinding polynomial evaluations (if enabled) + if useBlinding { + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + done <- err + return + } + } + + // Step 2-4: Compute gate, ordering, and local constraints sequentially on a + // single synchronous stream, folding them into dResult as + // gate + alpha*ordering + alpha^2*local. Computing one family at a time + // keeps peak device allocation low, and sequential is not a compromise: + // the family kernels are memory-bandwidth-bound and each already saturates + // the device, so the parallel three-stream variant this replaces measured + // identical timings (111ms/iteration at n=2^23) — while racing on the + // shared temp-slice pool and lazily materialized inputs (it corrupted the + // numerator at every circuit size). + seqVecCfg := state.vecCfg + seqVecCfg.IsAsync = false + + // Compute ordering first to minimize peak memory before gate/local allocations. + var err error + state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + return + } + + // dResult = alpha * ordering + state.dResult = state.getTempDeviceSlice(params.n) + dAlphaStd := uploadScalarStd(params.alpha) + if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { + dAlphaStd.Free() + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + return + } + dAlphaStd.Free() + state.putTempDeviceSlice(state.dOrdering, params.n) + + // dResult += alpha^2 * local + state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + return + } + var alphaSquared fr.Element + alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) + dAlphaSquaredStd := uploadScalarStd(alphaSquared) + dTmp := state.getTempDeviceSlice(params.n) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + return + } + dAlphaSquaredStd.Free() + if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + state.putTempDeviceSlice(dTmp, params.n) + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dLocal, params.n) + state.putTempDeviceSlice(dTmp, params.n) + + // dResult += gate + state.dGate, err = computeGateConstraint(state, params, seqVecCfg) + if err != nil { + done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + return + } + if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + return + } + state.putTempDeviceSlice(state.dGate, params.n) + + // Step 5: materialize result either on host or as a persistent device slice. + if result != nil { + resultHost := icicle_core.HostSliceFromElements(result) + resultHost.CopyFromDevice(&state.dResult) + } else { + resultOnDevice = s.getTempDeviceSlice(params.n) + if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { + done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // Return dResult pool slice after materialization. + state.putTempDeviceSlice(state.dResult, params.n) + + // Return all allocated polynomial buffers to the pool. + // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). + // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are + // returned to pool inside computeGateConstraint. dQk and the original gpuState slices + // are owned by gpuState and will be freed separately. + state.freeAllocatedPolyBuffers() + + done <- nil + }) + + err := <-done + + if err != nil { + if !resultOnDevice.IsEmpty() { + s.putTempDeviceSlice(resultOnDevice, resultOnDevice.Len()) + } + return icicle_core.DeviceSlice{}, err + } + return resultOnDevice, nil +} + +func evaluateBlinded(p, bp *iop.Polynomial, zeta fr.Element) fr.Element { + // Get the size of the polynomial + n := big.NewInt(int64(p.Size())) + + var pEvaluatedAtZeta fr.Element + + // Evaluate the polynomial and blinded polynomial at zeta + chP := make(chan struct{}, 1) + go func() { + pEvaluatedAtZeta = p.Evaluate(zeta) + close(chP) + }() + + bpEvaluatedAtZeta := bp.Evaluate(zeta) + + // Multiply the evaluated blinded polynomial by tempElement + var t fr.Element + one := fr.One() + t.Exp(zeta, n).Sub(&t, &one) + bpEvaluatedAtZeta.Mul(&bpEvaluatedAtZeta, &t) + + // Add the evaluated polynomial and the evaluated blinded polynomial + <-chP + pEvaluatedAtZeta.Add(&pEvaluatedAtZeta, &bpEvaluatedAtZeta) + + // Return the result + return pEvaluatedAtZeta +} + +// /!\ modifies the size +func getBlindedCoefficients(p, bp *iop.Polynomial) []fr.Element { + cp := p.Coefficients() + cbp := bp.Coefficients() + cp = append(cp, cbp...) + for i := 0; i < len(cbp); i++ { + cp[i].Sub(&cp[i], &cbp[i]) + } + return cp +} + +// getNonBlindedCoefficients returns a padded copy of polynomial coefficients +// to match the size they would have with blinding enabled. +// The padding size is blindingOrder+1 (e.g., order 2 → 3 coefficients). +func getNonBlindedCoefficients(p *iop.Polynomial, blindingOrder int) []fr.Element { + cp := p.Coefficients() + padded := make([]fr.Element, len(cp)+blindingOrder+1) + copy(padded, cp) + return padded +} + +// commits to a polynomial of the form b*(Xⁿ-1) where b is of small degree +func commitBlindingFactor(n int, b *iop.Polynomial, key kzg.ProvingKey) curve.G1Affine { + cp := b.Coefficients() + np := b.Size() + + // lo + var tmp curve.G1Affine + tmp.MultiExp(key.G1[:np], cp, ecc.MultiExpConfig{}) + + // hi + var res curve.G1Affine + res.MultiExp(key.G1[n:n+np], cp, ecc.MultiExpConfig{}) + res.Sub(&res, &tmp) + return res +} + +// return a random polynomial of degree n, if n==-1 cancel the blinding +func getRandomPolynomial(n int) *iop.Polynomial { + var a []fr.Element + if n == -1 { + a := make([]fr.Element, 1) + a[0].SetZero() + } else { + a = make([]fr.Element, n+1) + for i := 0; i <= n; i++ { + a[i].SetRandom() + } + } + res := iop.NewPolynomial(&a, iop.Form{ + Basis: iop.Canonical, Layout: iop.Regular}) + return res +} + +func coefficients(p []*iop.Polynomial) [][]fr.Element { + res := make([][]fr.Element, len(p)) + for i, pI := range p { + res[i] = pI.Coefficients() + } + return res +} + +func (s *instance) freeGPUQuotient(quotient *gpuQuotientPolynomial) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return + } + s.putTempDeviceSlice(quotient.coeffs, quotient.coeffs.Len()) + quotient.coeffs = icicle_core.DeviceSlice{} + quotient.size = 0 +} + +// commitToQuotientGPUFromDevice commits H1/H2/H3 directly from device memory. +// For StatisticalZK=true we materialize adjusted device vectors for h1/h2/h3 +// and commit those without downloading quotient coefficients to host. +// prepareStatisticalZKQuotientShards constructs blinded quotient polynomial +// shards h1, h2, h3 on the GPU for the Statistical ZK path. Each shard is +// randomized so that the quotient split h = h1 + X^(n+2)*h2 + X^(2(n+2))*h3 +// hides the original polynomial. +// +// Caller is responsible for returning dH1, dH2, dH3 to the temp pool: +// - dH1 and dH2 have size nPlus2+1 +// - dH3 has size nPlus2 +func (s *instance) prepareStatisticalZKQuotientShards( + h1Device, h2Device, h3Device icicle_core.DeviceSlice, + nPlus2 int, +) (dH1, dH2, dH3 icicle_core.DeviceSlice, err error) { + nPlus3 := nPlus2 + 1 + + prepareDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + dH1 = s.getTempDeviceSlice(nPlus3) + dH2 = s.getTempDeviceSlice(nPlus3) + dH3 = s.getTempDeviceSlice(nPlus2) + + // h1 = base h1 with extra randomizer coefficient at degree n+2. + dH1Prefix := (&dH1).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH1Prefix, h1Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h1 failed: %s", e.AsString()) + return + } + dH1Tail := (&dH1).Range(nPlus2, nPlus3, false) + r0Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[0]}) + r0Host.CopyToDevice(&dH1Tail, false) + + // h2 = base h2 with first coefficient adjusted by -r0 and tail = r1. + dH2Prefix := (&dH2).Range(0, nPlus2, false) + if e := copyDeviceSliceIntoOnCurrentDevice(dH2Prefix, h2Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h2 failed: %s", e.AsString()) + return + } + dH2First := (&dH2).Range(0, 1, false) + var negR0 fr.Element + negR0.Neg(&s.quotientShardsRandomizers[0]) + dNegR0 := uploadScalarMont(negR0) + if e := icicle_vecops.ScalarAddVec(dNegR0, dH2First, dH2First, cfg); e != icicle_runtime.Success { + _ = dNegR0.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h2[0] failed: %s", e.AsString()) + return + } + _ = dNegR0.Free() + dH2Tail := (&dH2).Range(nPlus2, nPlus3, false) + r1Host := icicle_core.HostSliceFromElements([]fr.Element{s.quotientShardsRandomizers[1]}) + r1Host.CopyToDevice(&dH2Tail, false) + + // h3 = base h3 with first coefficient adjusted by -r1. + if e := copyDeviceSliceIntoOnCurrentDevice(dH3, h3Device, cfg); e != icicle_runtime.Success { + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: copy h3 failed: %s", e.AsString()) + return + } + dH3First := (&dH3).Range(0, 1, false) + var negR1 fr.Element + negR1.Neg(&s.quotientShardsRandomizers[1]) + dNegR1 := uploadScalarMont(negR1) + if e := icicle_vecops.ScalarAddVec(dNegR1, dH3First, dH3First, cfg); e != icicle_runtime.Success { + _ = dNegR1.Free() + prepareDone <- fmt.Errorf("prepareStatisticalZKQuotientShards: adjust h3[0] failed: %s", e.AsString()) + return + } + _ = dNegR1.Free() + prepareDone <- nil + }) + if err := <-prepareDone; err != nil { + if !dH1.IsEmpty() { + s.putTempDeviceSlice(dH1, nPlus3) + } + if !dH2.IsEmpty() { + s.putTempDeviceSlice(dH2, nPlus3) + } + if !dH3.IsEmpty() { + s.putTempDeviceSlice(dH3, nPlus2) + } + return icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, icicle_core.DeviceSlice{}, err + } + return dH1, dH2, dH3, nil +} + +func (s *instance) commitToQuotientGPUFromDevice(quotient *gpuQuotientPolynomial) error { + if quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("commitToQuotientGPUFromDevice: empty quotient") + } + + nPlus2 := int(s.domain0.Cardinality) + 2 + required := 3 * nPlus2 + if quotient.coeffs.Len() < required { + return fmt.Errorf("commitToQuotientGPUFromDevice: quotient too small: got %d need >= %d", quotient.coeffs.Len(), required) + } + + h1Device := ("ient.coeffs).Range(0, nPlus2, false) + h2Device := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3Device := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + if s.opt.StatisticalZK { + nPlus3 := nPlus2 + 1 + dH1, dH2, dH3, err := s.prepareStatisticalZKQuotientShards(h1Device, h2Device, h3Device, nPlus2) + if err != nil { + return err + } + defer s.putTempDeviceSlice(dH1, nPlus3) + defer s.putTempDeviceSlice(dH2, nPlus3) + defer s.putTempDeviceSlice(dH3, nPlus2) + + c0, err := commitOnGPUCanonicalDevice(dH1, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(dH2, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(dH3, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + return nil + } + + // Commit sequentially to avoid 3-way concurrent MSM memory spikes. + c0, err := commitOnGPUCanonicalDevice(h1Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[0] = c0 + + c1, err := commitOnGPUCanonicalDevice(h2Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[1] = c1 + + c2, err := commitOnGPUCanonicalDevice(h3Device, &s.device, s.pk) + if err != nil { + return err + } + s.proof.H[2] = c2 + + return nil +} + +func (s *instance) inverseAndMergeShards( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (icicle_core.DeviceSlice, error) { + if gpuNumerator == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: nil numerator") + } + n := gpuNumerator.n + rho := gpuNumerator.rho + if n <= 0 || rho <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: invalid n=%d rho=%d", n, rho) + } + if len(gpuNumerator.shards) != rho { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), rho) + } + for i := 0; i < rho; i++ { + if gpuNumerator.shards[i].IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("inverseAndMergeShards: shard %d is empty", i) + } + } + + expo := big.NewInt(int64(n)) + + // Per-shard cosets: c_i = c * g^i where c=FrMultiplicativeGen, g=Generator. + cosets := make([]fr.Element, rho) + cosets[0].Set(&domains[1].FrMultiplicativeGen) + for i := 1; i < rho; i++ { + cosets[i].Mul(&cosets[i-1], &domains[1].Generator) + } + invCosets := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + invCosets[i].Inverse(&cosets[i]) + } + + // ν = g^n is a rho-th root, used for the rho-point inverse DFT in combine. + var nu, nuInv fr.Element + nu.Exp(domains[1].Generator, expo) + nuInv.Inverse(&nu) + nuInvPowers := make([]fr.Element, rho) + nuInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + nuInvPowers[i].Mul(&nuInvPowers[i-1], &nuInv) + } + + // cN = c^n. Recover original coefficient blocks by scaling with cN^{-t}. + var cN, cNInv fr.Element + cN.Exp(domains[1].FrMultiplicativeGen, expo) + cNInv.Inverse(&cN) + cNInvPowers := make([]fr.Element, rho) + cNInvPowers[0].SetOne() + for i := 1; i < rho; i++ { + cNInvPowers[i].Mul(&cNInvPowers[i-1], &cNInv) + } + + // Each shard inverse contributes a 1/n factor; apply extra 1/rho. + var rhoFr, invRho fr.Element + rhoFr.SetUint64(uint64(rho)) + invRho.Inverse(&rhoFr) + combineScales := make([]fr.Element, rho) + for i := 0; i < rho; i++ { + combineScales[i].Mul(&invRho, &cNInvPowers[i]) + } + + totalSize := rho * n + done := make(chan error, 1) + var dMerged icicle_core.DeviceSlice + + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish := makeFinisher(stream, "inverseAndMergeShards", done) + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = true + cfgNtt.StreamHandle = stream + // KNR is often faster than KNN; we restore regular output explicitly + // by bit-reversing each shard after the inverse NTT. + cfgNtt.Ordering = icicle_core.KNR + + ext := config_extension.Create() + defer config_extension.Delete(ext) + alg := nttAlgorithmFromEnv("ICICLE_DIVIDE_BY_ZH_NTT_ALGO", icicle_core.MixedRadix) + ext.SetInt(icicle_core.CUDA_NTT_ALGORITHM, int(alg)) + cfgNtt.Ext = ext.AsUnsafePointer() + + // Step 1: inverse NTT each shard without coset, reorder to regular, + // then unscale by (c*g^i)^t to recover the coset-inverse equivalent. + nn := uint64(64 - bits.TrailingZeros64(uint64(n))) + invPowers := make([]fr.Element, n) + for i := 0; i < rho; i++ { + if nttErr := icicle_ntt.Ntt(gpuNumerator.shards[i], icicle_core.KInverse, &cfgNtt, gpuNumerator.shards[i]); nttErr != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: inverse NTT failed at shard %d: %s", i, nttErr.AsString())) + return + } + + // KNR outputs bit-reversed coefficients. Reorder back to regular. + dRegular := s.getTempDeviceSlice(n) + mergeErr := icicle_vecops.MergeShardsBitReverse( + []icicle_core.DeviceSlice{gpuNumerator.shards[i]}, + n, + nn, + dRegular, + cfgVec, + ) + if mergeErr != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: reorder failed at shard %d: %s", i, mergeErr.AsString())) + return + } + // gpuNumerator.shards[i] is returned to pool and replaced; wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(gpuNumerator.shards[i], n) + gpuNumerator.shards[i] = dRegular + + fft.BuildExpTable(invCosets[i], invPowers) + dInvPowers := s.getTempDeviceSlice(n) + uploadVectorStdIntoOnCurrentDevice(&dInvPowers, invPowers, cfgVec) + if e := icicle_vecops.VecOp(gpuNumerator.shards[i], dInvPowers, gpuNumerator.shards[i], cfgVec, icicle_core.Mul); e != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: normalize shard %d failed: %s", i, e.AsString())) + return + } + // dInvPowers is temporary and returned to pool each iteration. + if eSync := icicle_runtime.SynchronizeStream(cfgVec.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(dInvPowers, n) + finish(fmt.Errorf("inverseAndMergeShards: synchronize stream failed: %s", eSync.AsString())) + return + } + s.putTempDeviceSlice(dInvPowers, n) + } + + // Step 2: combine shard results via a size-rho inverse DFT per coefficient index. + dMerged = s.getTempDeviceSlice(totalSize) + keepMerged := false + defer func() { + if !keepMerged && !dMerged.IsEmpty() { + s.putTempDeviceSlice(dMerged, totalSize) + } + }() + + dTmp := s.getTempDeviceSlice(n) + defer func() { + if !dTmp.IsEmpty() { + s.putTempDeviceSlice(dTmp, n) + } + }() + + dNuWeights := make([]icicle_core.DeviceSlice, rho) + dCombineScales := make([]icicle_core.DeviceSlice, rho) + for i := 0; i < rho; i++ { + dNuWeights[i] = uploadScalarStdOnCurrentDevice(nuInvPowers[i], cfgVec) + dCombineScales[i] = uploadScalarStdOnCurrentDevice(combineScales[i], cfgVec) + } + defer func() { + for i := 0; i < rho; i++ { + if !dNuWeights[i].IsEmpty() { + _ = dNuWeights[i].Free() + } + if !dCombineScales[i].IsEmpty() { + _ = dCombineScales[i].Free() + } + } + }() + + for t := 0; t < rho; t++ { + outT := (&dMerged).Range(t*n, (t+1)*n, false) + if e := copyDeviceSliceIntoOnCurrentDevice(outT, gpuNumerator.shards[0], cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: init out[%d] failed: %s", t, e.AsString())) + return + } + for i := 1; i < rho; i++ { + weightIdx := (i * t) % rho + if weightIdx == 0 { + if e := icicle_vecops.VecOp(outT, gpuNumerator.shards[i], outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: add shard %d to out[%d] failed: %s", i, t, e.AsString())) + return + } + continue + } + if e := icicle_vecops.ScalarMulVec(dNuWeights[weightIdx], gpuNumerator.shards[i], dTmp, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: weight shard %d for out[%d] failed: %s", i, t, e.AsString())) + return + } + if e := icicle_vecops.VecOp(outT, dTmp, outT, cfgVec, icicle_core.Add); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: accumulate shard %d into out[%d] failed: %s", i, t, e.AsString())) + return + } + } + if e := icicle_vecops.ScalarMulVec(dCombineScales[t], outT, outT, cfgVec); e != icicle_runtime.Success { + finish(fmt.Errorf("inverseAndMergeShards: scale out[%d] failed: %s", t, e.AsString())) + return + } + } + keepMerged = true + finish(nil) + }) + + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dMerged, nil +} + +func (s *instance) divideByZHOnGPU( + gpuNumerator *gpuNumeratorPolynomial, + domains [2]*fft.Domain, +) (_ *gpuQuotientPolynomial, err error) { + if gpuNumerator == nil { + return nil, fmt.Errorf("divideByZHOnGPU: nil numerator") + } + if gpuNumerator.n <= 0 || gpuNumerator.rho <= 0 { + return nil, fmt.Errorf("divideByZHOnGPU: invalid numerator dimensions n=%d rho=%d", gpuNumerator.n, gpuNumerator.rho) + } + if len(gpuNumerator.shards) != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: shard count mismatch: got %d expected %d", len(gpuNumerator.shards), gpuNumerator.rho) + } + for i := range gpuNumerator.shards { + if gpuNumerator.shards[i].IsEmpty() { + return nil, fmt.Errorf("divideByZHOnGPU: shard %d is empty", i) + } + } + + rho := int(domains[1].Cardinality / domains[0].Cardinality) + if rho != gpuNumerator.rho { + return nil, fmt.Errorf("divideByZHOnGPU: rho mismatch domains=%d numerator=%d", rho, gpuNumerator.rho) + } + + // Evaluate 1/(X^n-1) over the large-domain coset values used by this quotient. + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + + // In bit-reversed merged layout, each shard maps to a fixed (iRev % rho) bucket. + // So we can divide by Z_H by scaling each shard with its corresponding inverse. + scaleDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + for i := 0; i < gpuNumerator.rho; i++ { + dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) + vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if vecErr != icicle_runtime.Success { + finish(fmt.Errorf("divideByZHOnGPU: shard scaling failed at %d: %s", i, vecErr.AsString())) + return + } + } + finish(nil) + }) + if err := <-scaleDone; err != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, err + } + + totalSize := gpuNumerator.n * gpuNumerator.rho + dMerged, splitErr := s.inverseAndMergeShards(gpuNumerator, domains) + if splitErr != nil { + s.freeNumeratorShards(gpuNumerator.shards) + return nil, splitErr + } + // Shards are not needed after split inverse+merge. + s.freeNumeratorShards(gpuNumerator.shards) + return &gpuQuotientPolynomial{coeffs: dMerged, size: totalSize}, nil +} + +func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*iop.Polynomial, error) { + if quotient == nil || quotient.coeffs.IsEmpty() { + return nil, fmt.Errorf("downloadQuotientFromGPU: empty quotient") + } + if quotient.size <= 0 { + return nil, fmt.Errorf("downloadQuotientFromGPU: invalid quotient size %d", quotient.size) + } + + coeffs := make([]fr.Element, quotient.size) + host := icicle_core.HostSliceFromElements(coeffs) + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + done <- cfgErr + return + } + host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) + // Async boundary for host materialization of quotient coefficients. + done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + }) + if err := <-done; err != nil { + return nil, err + } + + return iop.NewPolynomial(&coeffs, iop.Form{Basis: iop.Canonical, Layout: iop.Regular}), nil +} + +func commitOnGPUWithDeviceBases( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, +) (curve.G1Affine, error) { + return commitOnGPUWithDeviceBasesChunked(scalarsDevice, basesDevice, device, icicleMSMChunkSize()) +} + +func commitOnGPUWithDeviceBasesChunked( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + device *icicle_runtime.Device, + chunkSize int, +) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty scalar slice") + } + if basesDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: empty basis slice") + } + if scalarsDevice.Len() > basesDevice.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUWithDeviceBases: invalid scalar size %d", scalarsDevice.Len()) + } + if chunkSize <= 0 || chunkSize > scalarsDevice.Len() { + chunkSize = scalarsDevice.Len() + } + + var commit curve.G1Affine + var msmErr error + done := make(chan struct{}, 1) + icicle_runtime.RunOnDevice(device, func(args ...any) { + commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) + close(done) + }) + <-done + if msmErr != nil { + return curve.G1Affine{}, fmt.Errorf("icicle: MSM commit from device bases failed (%d scalars): %w", scalarsDevice.Len(), msmErr) + } + return commit, nil +} + +func commitOnGPUWithDeviceBasesChunkedOnCurrentDevice( + scalarsDevice icicle_core.DeviceSlice, + basesDevice icicle_core.DeviceSlice, + chunkSize int, +) (curve.G1Affine, error) { + var commit curve.G1Affine + for start := 0; start < scalarsDevice.Len(); start += chunkSize { + end := start + chunkSize + if end > scalarsDevice.Len() { + end = scalarsDevice.Len() + } + + // Each chunk must pair with exactly bases[start:end]: ICICLE treats a + // bases slice longer than the scalars as a batched MSM (and requires + // divisibility), so the full bases buffer cannot be passed as-is when + // it is longer than the scalar vector. + scalarsChunk := scalarsDevice + if start != 0 || end != scalarsDevice.Len() { + scalarsChunk = (&scalarsDevice).Range(start, end, false) + } + basesChunk := basesDevice + if start != 0 || end != basesDevice.Len() { + basesChunk = (&basesDevice).Range(start, end, false) + } + + res := make(icicle_core.HostSlice[icicle_bw6761.Projective], 1) + cfg := icicle_msm.GetDefaultMSMConfig() + cfg.AreBasesMontgomeryForm = true + cfg.AreScalarsMontgomeryForm = true + e := icicle_msm.Msm(scalarsChunk, basesChunk, &cfg, res) + if e != icicle_runtime.Success { + return curve.G1Affine{}, fmt.Errorf("icicle MSM failed for chunk [%d:%d]: %s", start, end, e.AsString()) + } + + chunkCommit, err := projectiveToGnarkAffine(res[0]) + if err != nil { + return curve.G1Affine{}, fmt.Errorf("convert chunk [%d:%d]: %w", start, end, err) + } + commit.Add(&commit, &chunkCommit) + } + return commit, nil +} + +func icicleMSMChunkSize() int { + // Production-sized MSMs still need chunking, but tiny chunks add thousands of + // ICICLE calls. 4M-point chunks passed the gnark replay profile; 8M did not. + const defaultChunkSize = 1 << 22 + v := strings.TrimSpace(os.Getenv("ICICLE_MSM_CHUNK_SIZE")) + if v == "" { + return defaultChunkSize + } + n, err := strconv.Atoi(v) + if err != nil || n <= 0 { + return defaultChunkSize + } + return n +} + +func commitOnGPUCanonicalDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: canonical SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPUCanonicalDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgDevice.G1, device) +} + +func evalCanonicalAtPoint(coeffs []fr.Element, point fr.Element) fr.Element { + var acc fr.Element + if len(coeffs) == 0 { + return acc + } + acc.Set(&coeffs[len(coeffs)-1]) + for i := len(coeffs) - 2; i >= 0; i-- { + acc.Mul(&acc, &point).Add(&acc, &coeffs[i]) + } + return acc +} + +func deriveBatchOpeningGamma( + point fr.Element, + digests []curve.G1Affine, + claimedValues []fr.Element, + hf hash.Hash, + dataTranscript ...[]byte, +) (fr.Element, error) { + fs := fiatshamir.NewTranscript(hf, "gamma") + if err := fs.Bind("gamma", point.Marshal()); err != nil { + return fr.Element{}, err + } + for i := range digests { + if err := fs.Bind("gamma", digests[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := range claimedValues { + if err := fs.Bind("gamma", claimedValues[i].Marshal()); err != nil { + return fr.Element{}, err + } + } + for i := 0; i < len(dataTranscript); i++ { + if err := fs.Bind("gamma", dataTranscript[i]); err != nil { + return fr.Element{}, err + } + } + gammaByte, err := fs.ComputeChallenge("gamma") + if err != nil { + return fr.Element{}, err + } + var gamma fr.Element + gamma.SetBytes(gammaByte) + return gamma, nil +} + +func (s *instance) evalDevicePolynomialAtPointOnCurrentDevice( + coeffsDevice icicle_core.DeviceSlice, + point fr.Element, + useBitReverse bool, + cfg icicle_core.VecOpsConfig, +) (fr.Element, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: empty coefficients") + } + + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStdOnCurrentDevice(point, cfg) + defer dPoint.Free() + + dOut := s.getTempDeviceSlice(1) + defer s.putTempDeviceSlice(dOut, 1) + + opName := "PolyEvalAt" + var eEval icicle_runtime.EIcicleError + if useBitReverse { + opName = "PolyEvalAtBitReverse" + mm := uint64(64 - bits.TrailingZeros64(uint64(coeffsDevice.Len()))) + eEval = icicle_vecops.PolyEvalAtBitReverse(coeffsDevice, dPoint, mm, dOut, cfg) + } else { + eEval = icicle_vecops.PolyEvalAt(coeffsDevice, dPoint, dOut, cfg) + } + if eEval != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: %s failed: %s", opName, eEval.AsString()) + } + + var out fr.Element + hostOut := icicle_core.HostSliceFromElements([]fr.Element{out}) + if cfg.IsAsync { + hostOut.CopyFromDeviceAsync(&dOut, cfg.StreamHandle) + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fr.Element{}, fmt.Errorf("evalDevicePolynomialAtPointOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } else { + hostOut.CopyFromDevice(&dOut) + } + return ([]fr.Element)(hostOut)[0], nil +} + +func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (fr.Element, error) { + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + + var runErr error + out, runErr = s.evalDevicePolynomialAtPointOnCurrentDevice(coeffsDevice, point, false, cfg) + done <- runErr + }) + return out, <-done +} + +func (s *instance) copyDeviceSliceOnCurrentDevice( + src icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, + label string, +) (icicle_core.DeviceSlice, error) { + if src.IsEmpty() || src.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: empty source slice", label) + } + dst := s.getTempDeviceSlice(src.Len()) + eCopy := copyDeviceSliceIntoOnCurrentDevice(dst, src, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dst, src.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("%s: copy failed: %s", label, eCopy.AsString()) + } + return dst, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromStateOnCurrentDevice( + p *iop.Polynomial, + state *gpuPolysState, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "batchOpening poly") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + dCanon, err := s.copyDeviceSliceOnCurrentDevice(dSrc, cfg, "materializePolynomialCanonicalRegularFromStateOnCurrentDevice") + if err != nil { + return icicle_core.DeviceSlice{}, err + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + // Canonical bit-reverse -> canonical regular. + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + dRegular := s.getTempDeviceSlice(dCanon.Len()) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNN // regular -> regular canonical + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular canonical + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromStateOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) buildBlindedCanonicalPolynomialOnCurrentDevice( + dBaseCanon, dBlindCanon icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBlindCanon.IsEmpty() { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: empty input") + } + n := dBaseCanon.Len() + blindLen := dBlindCanon.Len() + dOut := s.getTempDeviceSlice(n + blindLen) + + dPrefix := (&dOut).Range(0, n, false) + eCopyBase := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopyBase != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopyBase.AsString()) + } + + dTail := (&dOut).Range(n, n+blindLen, false) + eCopyBlind := copyDeviceSliceIntoOnCurrentDevice(dTail, dBlindCanon, cfg) + if eCopyBlind != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: copy tail failed: %s", eCopyBlind.AsString()) + } + + dHead := (&dOut).Range(0, blindLen, false) + if eSub := icicle_vecops.VecOp(dHead, dBlindCanon, dHead, cfg, icicle_core.Sub); eSub != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+blindLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildBlindedCanonicalPolynomialOnCurrentDevice: subtract blind from head failed: %s", eSub.AsString()) + } + return dOut, nil +} + +func (s *instance) buildPaddedCanonicalPolynomialOnCurrentDevice( + dBaseCanon icicle_core.DeviceSlice, + padLen int, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if dBaseCanon.IsEmpty() || dBaseCanon.Len() <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: empty base") + } + if padLen < 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: negative pad length %d", padLen) + } + n := dBaseCanon.Len() + dOut := s.getTempDeviceSlice(n + padLen) + dPrefix := (&dOut).Range(0, n, false) + + eCopy := copyDeviceSliceIntoOnCurrentDevice(dPrefix, dBaseCanon, cfg) + if eCopy != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: copy base failed: %s", eCopy.AsString()) + } + if padLen == 0 { + return dOut, nil + } + + dTail := (&dOut).Range(n, n+padLen, false) + eZero := zeroDeviceSliceOnCurrentDevice(dTail, cfg) + if eZero != icicle_runtime.Success { + s.putTempDeviceSlice(dOut, n+padLen) + return icicle_core.DeviceSlice{}, fmt.Errorf("buildPaddedCanonicalPolynomialOnCurrentDevice: zero tail failed: %s", eZero.AsString()) + } + return dOut, nil +} + +func (s *instance) prepareBatchOpeningPolynomialsOnGPU( + state *gpuPolysState, + point fr.Element, +) (devicePolys []icicle_core.DeviceSlice, owned []bool, claimed []fr.Element, err error) { + if state == nil { + return nil, nil, nil, fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: nil GPU state") + } + + total := 6 + len(s.trace.Qcp) + devicePolys = make([]icicle_core.DeviceSlice, total) + owned = make([]bool, total) + claimed = make([]fr.Element, total) + devicePolys[0] = s.linearizedPolynomialGPU + claimed[0] = s.linearizedPolynomialClaim + + prepDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + + cleanupOwned := func(from int) { + for i := from; i < len(devicePolys); i++ { + if owned[i] && !devicePolys[i].IsEmpty() { + s.putTempDeviceSlice(devicePolys[i], devicePolys[i].Len()) + devicePolys[i] = icicle_core.DeviceSlice{} + owned[i] = false + } + } + } + + prepareLRORow := func(dstIdx int, p, bp *iop.Polynomial, padLen int) error { + dBase, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + defer s.putTempDeviceSlice(dBase, dBase.Len()) + + var dFinal icicle_core.DeviceSlice + if useBlinding { + if bp == nil { + return fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: missing blinding polynomial") + } + dBlind, eBlind := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(bp, state, cfg) + if eBlind != nil { + return eBlind + } + defer s.putTempDeviceSlice(dBlind, dBlind.Len()) + dFinal, e = s.buildBlindedCanonicalPolynomialOnCurrentDevice(dBase, dBlind, cfg) + } else { + dFinal, e = s.buildPaddedCanonicalPolynomialOnCurrentDevice(dBase, padLen, cfg) + } + if e != nil { + return e + } + devicePolys[dstIdx] = dFinal + owned[dstIdx] = true + + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dFinal, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareLRORow(1, s.polyL, s.bp[id_Bl], order_blinding_L+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare L failed: %w", e)) + return + } + if e := prepareLRORow(2, s.polyR, s.bp[id_Br], order_blinding_R+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare R failed: %w", e)) + return + } + if e := prepareLRORow(3, s.polyO, s.bp[id_Bo], order_blinding_O+1); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare O failed: %w", e)) + return + } + + prepareDirect := func(dstIdx int, p *iop.Polynomial, label string) error { + dPoly, e := s.materializePolynomialCanonicalRegularFromStateOnCurrentDevice(p, state, cfg) + if e != nil { + return e + } + devicePolys[dstIdx] = dPoly + owned[dstIdx] = true + val, eEval := s.evalDevicePolynomialAtPointOnCurrentDevice(dPoly, point, false, cfg) + if eEval != nil { + return eEval + } + claimed[dstIdx] = val + return nil + } + + if e := prepareDirect(4, s.trace.S1, "S1"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S1 failed: %w", e)) + return + } + if e := prepareDirect(5, s.trace.S2, "S2"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare S2 failed: %w", e)) + return + } + + for i := 0; i < len(s.trace.Qcp); i++ { + idx := 6 + i + if e := prepareDirect(idx, s.trace.Qcp[i], "Qcp"); e != nil { + cleanupOwned(1) + finish(fmt.Errorf("prepareBatchOpeningPolynomialsOnGPU: prepare Qcp[%d] failed: %w", i, e)) + return + } + } + + finish(nil) + }) + if err := <-prepDone; err != nil { + return nil, nil, nil, err + } + return devicePolys, owned, claimed, nil +} + +type evalPolynomialInputPreparationResult struct { + dEval icicle_core.DeviceSlice + ownedLen int + useBitReverseEval bool +} + +func (s *instance) prepareEvalPolynomialInputOnCurrentDevice( + p *iop.Polynomial, + dSrc icicle_core.DeviceSlice, + cfgVec icicle_core.VecOpsConfig, +) (evalPolynomialInputPreparationResult, error) { + if p == nil { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: nil polynomial") + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: empty source polynomial") + } + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] prepareEvalInput begin n=%d basis=%v layout=%v step_profile=%q ntt_trace=%q ntt_profile_full=%q ntt_profile_arbitrary=%q\n", + dSrc.Len(), + p.Basis, + p.Layout, + os.Getenv("ICICLE_STEP_PROFILE"), + os.Getenv("ICICLE_NTT_TRACE"), + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + + result := evalPolynomialInputPreparationResult{ + dEval: dSrc, + } + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return result, nil + } + + releaseOwned := func() { + if result.ownedLen > 0 && !result.dEval.IsEmpty() { + s.putTempDeviceSlice(result.dEval, result.ownedLen) + result.dEval = icicle_core.DeviceSlice{} + result.ownedLen = 0 + } + } + + dWork := s.getTempDeviceSlice(dSrc.Len()) + if e := copyDeviceSliceIntoOnCurrentDevice(dWork, dSrc, cfgVec); e != icicle_runtime.Success { + s.putTempDeviceSlice(dWork, dSrc.Len()) + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: copy source polynomial failed: %s", e.AsString()) + } + + result.dEval = dWork + result.ownedLen = dSrc.Len() + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfgVec.IsAsync + cfgNtt.StreamHandle = cfgVec.StreamHandle + + var ownsNttStream bool + destroyOwnedNttStream := func() error { + if !ownsNttStream { + return nil + } + return syncAndDestroyStreamOnCurrentDevice(cfgNtt.StreamHandle, "prepareEvalPolynomialInputOnCurrentDevice") + } + + switch p.Basis { + case iop.Canonical: + // No transform required. + result.useBitReverseEval = p.Layout == iop.BitReverse + case iop.Lagrange, iop.LagrangeCoset: + if cfgNtt.StreamHandle == nil { + nttStream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: create NTT stream failed: %s", eStream.AsString()) + } + cfgNtt.StreamHandle = nttStream + cfgNtt.IsAsync = true + ownsNttStream = true + } + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + if p.Layout == iop.Regular { + cfgNtt.Ordering = icicle_core.KNR // regular -> bitreverse on inverse + result.useBitReverseEval = true + } else { + cfgNtt.Ordering = icicle_core.KRN // bitreverse -> regular on inverse + result.useBitReverseEval = false + } + startNtt := time.Now() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT launch n=%d ordering=%v has_coset=%t basis=%v layout=%v\n", + dSrc.Len(), + cfgNtt.Ordering, + p.Basis == iop.LagrangeCoset, + p.Basis, + p.Layout, + ) + } + + // Martun: This call to Ntt takes about 3 seconds, because Ntt is reusing NTT domain data that + // gets prepared during InitDomain (once per device), not re-derived every call. + // Inside ICICLE, InitDomain precomputes: domain.twiddles (main roots-of-unity table, N+1) + // internal_twiddles and basic_twiddles for mixed-radix kernels + // if fast mode is on (it is by default here), extra forward+inverse fast twiddle tables (fast_external/internal/basic and _inv) — comment says this costs ~4N extra memory + // CPU-side coset_index map (root -> index), then reused by later Ntt calls + eNtt := icicle_ntt.Ntt(result.dEval, icicle_core.KInverse, &cfgNtt, result.dEval) + nttElapsed := time.Since(startNtt) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] inverse NTT done status=%s took=%s\n", + eNtt.AsString(), + nttElapsed, + ) + } + if eNtt != icicle_runtime.Success { + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + l := logger.Logger() + l.Warn().Err(errDestroy).Msg("prepareEvalPolynomialInputOnCurrentDevice") + } + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + // Async boundary for this helper when it owns the stream. + if errDestroy := destroyOwnedNttStream(); errDestroy != nil { + releaseOwned() + return evalPolynomialInputPreparationResult{}, errDestroy + } + default: + releaseOwned() + return evalPolynomialInputPreparationResult{}, fmt.Errorf("prepareEvalPolynomialInputOnCurrentDevice: unsupported basis %v", p.Basis) + } + + return result, nil +} + +// evalPolynomialInCurrentFormOnGPU evaluates a polynomial at a point directly +// from the shared GPU state regardless of its current basis/layout by converting +// a temporary device copy to canonical/regular when needed. +func (s *instance) evalPolynomialInCurrentFormOnGPU( + p *iop.Polynomial, + state *gpuPolysState, + point fr.Element, +) (fr.Element, error) { + if p == nil { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: nil polynomial") + } + dSrc, err := getStateDeviceSlice(state, p, "eval") + if err != nil { + return fr.Element{}, err + } + if dSrc.IsEmpty() || dSrc.Len() <= 0 { + return fr.Element{}, fmt.Errorf("evalPolynomialInCurrentFormOnGPU: empty device polynomial") + } + + var out fr.Element + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + dEval := dSrc + ownedLen := 0 + releaseEval := func() { + if ownedLen > 0 && !dEval.IsEmpty() { + s.putTempDeviceSlice(dEval, ownedLen) + dEval = icicle_core.DeviceSlice{} + ownedLen = 0 + } + } + finish := func(runErr error) { + // Async boundary for eval path before handing result back to caller. + if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + releaseEval() + done <- runErr + } + + prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) + dEval = prepareResult.dEval + ownedLen = prepareResult.ownedLen + if prepErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", prepErr)) + return + } + + evalOut, evalErr := s.evalDevicePolynomialAtPointOnCurrentDevice(dEval, point, prepareResult.useBitReverseEval, cfgVec) + if evalErr != nil { + finish(fmt.Errorf("evalPolynomialInCurrentFormOnGPU: %w", evalErr)) + return + } + out = evalOut + finish(nil) + }) + return out, <-done +} + +func (s *instance) evaluateBlindedOnGPU( + p, bp *iop.Polynomial, + state *gpuPolysState, + zeta fr.Element, +) (fr.Element, error) { + if p == nil || bp == nil { + return fr.Element{}, fmt.Errorf("evaluateBlindedOnGPU: nil polynomial") + } + pAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(p, state, zeta) + if err != nil { + return fr.Element{}, err + } + bpAtZeta, err := s.evalPolynomialInCurrentFormOnGPU(bp, state, zeta) + if err != nil { + return fr.Element{}, err + } + + var t, one fr.Element + one.SetOne() + t.Exp(zeta, big.NewInt(int64(p.Size()))).Sub(&t, &one) + bpAtZeta.Mul(&bpAtZeta, &t) + pAtZeta.Add(&pAtZeta, &bpAtZeta) + return pAtZeta, nil +} + +func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.Element) (kzg.OpeningProof, error) { + if len(coeffs) < 2 || len(coeffs) > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonical: invalid polynomial size %d", len(coeffs)) + } + claimed := evalCanonicalAtPoint(coeffs, point) + + var dWitness icicle_core.DeviceSlice + witnessSize := len(coeffs) - 1 + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dCoeffs := uploadVector(coeffs) + // For Montgomery vectors, scalar multipliers must be standard-form. + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(dCoeffs, dPoint, dWitness, cfg) + _ = dCoeffs.Free() + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonical: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + if err := <-divDone; err != nil { + return kzg.OpeningProof{}, err + } + + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.DeviceSlice, point fr.Element) (kzg.OpeningProof, error) { + if coeffsDevice.IsEmpty() || coeffsDevice.Len() < 2 || coeffsDevice.Len() > len(s.pk.Kzg.G1) { + return kzg.OpeningProof{}, fmt.Errorf("openPolynomialOnGPUCanonicalDevice: invalid polynomial size %d", coeffsDevice.Len()) + } + n := coeffsDevice.Len() + + var startEval time.Time + if isProfileMode { + startEval = time.Now() + } + claimed, err := s.evalDevicePolynomialAtPoint(coeffsDevice, point) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startEval)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: evalDevicePolynomialAtPoint") + } + } + if err != nil { + return kzg.OpeningProof{}, err + } + + var dWitness icicle_core.DeviceSlice + witnessSize := coeffsDevice.Len() - 1 + var startDivideByXMinusA time.Time + if isProfileMode { + startDivideByXMinusA = time.Now() + } + divDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg := icicle_core.DefaultVecOpsConfig() + cfg.IsAsync = false + dPoint := uploadScalarStd(point) + dWitness = s.getTempDeviceSlice(witnessSize) + eDiv := icicle_vecops.DivideByXMinusA(coeffsDevice, dPoint, dWitness, cfg) + _ = dPoint.Free() + if eDiv != icicle_runtime.Success { + s.putTempDeviceSlice(dWitness, witnessSize) + divDone <- fmt.Errorf("openPolynomialOnGPUCanonicalDevice: divide by (x-a) failed: %s", eDiv.AsString()) + return + } + divDone <- nil + }) + divideErr := <-divDone + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startDivideByXMinusA)) + if divideErr != nil { + ev.Err(divideErr).Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: DivideByXMinusA") + } + } + if divideErr != nil { + return kzg.OpeningProof{}, divideErr + } + + var startCommit time.Time + if isProfileMode { + startCommit = time.Now() + } + h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) + if isProfileMode { + l := logger.Logger() + ev := l.Debug().Int("n", n).Dur("took", time.Since(startCommit)) + if err != nil { + ev.Err(err).Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice (with error)") + } else { + ev.Msg("openPolynomialOnGPUCanonicalDevice: commitOnGPUCanonicalDevice") + } + } + s.putTempDeviceSlice(dWitness, witnessSize) + if err != nil { + return kzg.OpeningProof{}, err + } + + return kzg.OpeningProof{ + H: h, + ClaimedValue: claimed, + }, nil +} + +func (s *instance) linearizedZContributionScale(lZeta, rZeta, oZeta fr.Element) fr.Element { + var s2, tmp fr.Element + var uzeta, uuzeta fr.Element + uzeta.Mul(&s.zeta, &s.pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &s.pk.Vk.CosetShift) + + s2.Mul(&s.beta, &s.zeta).Add(&s2, &lZeta).Add(&s2, &s.gamma) + tmp.Mul(&s.beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp) + tmp.Mul(&s.beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &s.gamma) + s2.Mul(&s2, &tmp).Neg(&s2).Mul(&s2, &s.alpha) + + var one, alphaSquareLagrangeZero, den fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&s.zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) + den.Sub(&s.zeta, &one).Inverse(&den) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) + + s2.Add(&s2, &alphaSquareLagrangeZero) + return s2 +} + +func (s *instance) linearizedSelectorScales(evals witnessEvalAtZeta, zu fr.Element) linearizedSelectorScales { + var scales linearizedSelectorScales + + // S3 scale: + // alpha * beta * Z(mu*zeta) * + // (L(zeta) + beta*S1(zeta) + gamma) * + // (R(zeta) + beta*S2(zeta) + gamma) + var tmp fr.Element + scales.s3.Mul(&evals.s1zeta, &s.beta).Add(&scales.s3, &evals.blzeta).Add(&scales.s3, &s.gamma) + tmp.Mul(&evals.s2zeta, &s.beta).Add(&tmp, &evals.brzeta).Add(&tmp, &s.gamma) + scales.s3.Mul(&scales.s3, &tmp).Mul(&scales.s3, &zu).Mul(&scales.s3, &s.beta).Mul(&scales.s3, &s.alpha) + + scales.ql.Set(&evals.blzeta) + scales.qr.Set(&evals.brzeta) + scales.qm.Mul(&evals.brzeta, &evals.blzeta) + scales.qo.Set(&evals.bozeta) + scales.qk.SetOne() + scales.qcp = append(scales.qcp, evals.qcpzeta...) + + return scales +} + +func (s *instance) buildLinearizedSelectorTermsOnGPU( + evals witnessEvalAtZeta, + zu fr.Element, + linearizedLen int, +) (icicle_core.DeviceSlice, error) { + if linearizedLen <= 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: invalid length %d", linearizedLen) + } + if len(evals.qcpzeta) > len(s.cCommitments) { + return icicle_core.DeviceSlice{}, fmt.Errorf("buildLinearizedSelectorTermsOnGPU: qcp/cCommitments mismatch (%d > %d)", len(evals.qcpzeta), len(s.cCommitments)) + } + + scales := s.linearizedSelectorScales(evals, zu) + + var dLinearized icicle_core.DeviceSlice + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildLinearizedSelectorTermsOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + defer func() { + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + if runErr != nil && !dLinearized.IsEmpty() { + s.putTempDeviceSlice(dLinearized, dLinearized.Len()) + dLinearized = icicle_core.DeviceSlice{} + } + done <- runErr + }() + + dLinearized = s.getTempDeviceSlice(linearizedLen) + if eZero := zeroDeviceSliceOnCurrentDevice(dLinearized, cfg); eZero != icicle_runtime.Success { + runErr = fmt.Errorf("buildLinearizedSelectorTermsOnGPU: zero output failed: %s", eZero.AsString()) + return + } + + addTerm := func(p *iop.Polynomial, scale fr.Element, label string) error { + if p == nil { + return fmt.Errorf("missing polynomial %s", label) + } + if scale.IsZero() { + return nil + } + + start := time.Now() + dPoly, err := s.materializePolynomialCanonicalRegularFromHostOnCurrentDevice(p, cfg) + if err != nil { + return fmt.Errorf("%s: %w", label, err) + } + defer s.putTempDeviceSlice(dPoly, dPoly.Len()) + if dPoly.Len() > dLinearized.Len() { + return fmt.Errorf("%s: polynomial too large (%d > %d)", label, dPoly.Len(), dLinearized.Len()) + } + + dScale := uploadScalarStdOnCurrentDevice(scale, cfg) + dScaled := s.getTempDeviceSlice(dPoly.Len()) + defer s.putTempDeviceSlice(dScaled, dScaled.Len()) + + eScale := icicle_vecops.ScalarMulVec(dScale, dPoly, dScaled, cfg) + if cfg.IsAsync { + _ = dScale.FreeAsync(cfg.StreamHandle) + } else { + _ = dScale.Free() + } + if eScale != icicle_runtime.Success { + return fmt.Errorf("%s: scale failed: %s", label, eScale.AsString()) + } + + dPrefix := (&dLinearized).Range(0, dPoly.Len(), false) + if eAdd := icicle_vecops.VecOp(dPrefix, dScaled, dPrefix, cfg, icicle_core.Add); eAdd != icicle_runtime.Success { + return fmt.Errorf("%s: add failed: %s", label, eAdd.AsString()) + } + + if isProfileMode { + l := logger.Logger() + l.Debug().Str("term", label).Int("n", dPoly.Len()).Dur("took", time.Since(start)).Msg("computeLinearizedPolynomial: add selector term on GPU") + } + return nil + } + + if runErr = addTerm(s.trace.S3, scales.s3, "S3"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Ql, scales.ql, "Ql"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qm, scales.qm, "Qm"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qr, scales.qr, "Qr"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qo, scales.qo, "Qo"); runErr != nil { + return + } + if runErr = addTerm(s.trace.Qk, scales.qk, "Qk"); runErr != nil { + return + } + for i := range scales.qcp { + if runErr = addTerm(s.cCommitments[i], scales.qcp[i], fmt.Sprintf("Qcp[%d]", i)); runErr != nil { + return + } + } + }) + if err := <-done; err != nil { + return icicle_core.DeviceSlice{}, err + } + return dLinearized, nil +} + +func (s *instance) materializePolynomialCanonicalRegularFromHostOnCurrentDevice( + p *iop.Polynomial, + cfg icicle_core.VecOpsConfig, +) (icicle_core.DeviceSlice, error) { + if p == nil { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: nil polynomial") + } + coeffs := p.Coefficients() + if len(coeffs) == 0 { + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: empty polynomial") + } + + dCanon := s.getTempDeviceSlice(len(coeffs)) + host := icicle_core.HostSliceFromElements(coeffs) + if cfg.IsAsync { + host.CopyToDeviceAsync(&dCanon, cfg.StreamHandle, false) + } else { + host.CopyToDevice(&dCanon, false) + } + if dCanon.IsEmpty() { + s.putTempDeviceSlice(dCanon, len(coeffs)) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: host upload failed") + } + + if p.Basis == iop.Canonical && p.Layout == iop.Regular { + return dCanon, nil + } + + cfgNtt := icicle_ntt.GetDefaultNttConfig() + cfgNtt.IsAsync = cfg.IsAsync + cfgNtt.StreamHandle = cfg.StreamHandle + + switch p.Basis { + case iop.Canonical: + if p.Layout != iop.BitReverse { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported canonical layout %v", p.Layout) + } + dRegular := s.getTempDeviceSlice(dCanon.Len()) + mm := uint64(64 - bits.TrailingZeros64(uint64(dCanon.Len()))) + eReorder := icicle_vecops.MergeShardsBitReverse([]icicle_core.DeviceSlice{dCanon}, dCanon.Len(), mm, dRegular, cfg) + s.putTempDeviceSlice(dCanon, dCanon.Len()) + if eReorder != icicle_runtime.Success { + s.putTempDeviceSlice(dRegular, dRegular.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: canonical reorder failed: %s", eReorder.AsString()) + } + return dRegular, nil + + case iop.Lagrange, iop.LagrangeCoset: + if p.Basis == iop.LagrangeCoset { + cfgNtt.CosetGen = s.pk.CosetGenerator + } + switch p.Layout { + case iop.Regular: + cfgNtt.Ordering = icicle_core.KNN + case iop.BitReverse: + cfgNtt.Ordering = icicle_core.KRN + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported layout %v", p.Layout) + } + if eNtt := icicle_ntt.Ntt(dCanon, icicle_core.KInverse, &cfgNtt, dCanon); eNtt != icicle_runtime.Success { + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: inverse NTT failed: %s", eNtt.AsString()) + } + return dCanon, nil + + default: + s.putTempDeviceSlice(dCanon, dCanon.Len()) + return icicle_core.DeviceSlice{}, fmt.Errorf("materializePolynomialCanonicalRegularFromHostOnCurrentDevice: unsupported basis %v", p.Basis) + } +} + +func (s *instance) addZContributionToLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + dBlindedZCanonical icicle_core.DeviceSlice, + lZeta, rZeta, oZeta fr.Element, +) error { + if dLinearized.IsEmpty() || dBlindedZCanonical.IsEmpty() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: empty input") + } + if dLinearized.Len() < dBlindedZCanonical.Len() { + return fmt.Errorf("addZContributionToLinearizedOnGPU: linearized too small (%d < %d)", dLinearized.Len(), dBlindedZCanonical.Len()) + } + + zScale := s.linearizedZContributionScale(lZeta, rZeta, oZeta) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("addZContributionToLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dScale icicle_core.DeviceSlice + var dScaledZ icicle_core.DeviceSlice + defer func() { + // Async boundary before returning temporary buffers to the pool. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dScale) + if !dScaledZ.IsEmpty() { + s.putTempDeviceSlice(dScaledZ, dScaledZ.Len()) + } + done <- runErr + }() + + dScale = uploadScalarStdOnCurrentDevice(zScale, cfg) + dScaledZ = s.getTempDeviceSlice(dBlindedZCanonical.Len()) + eMul := icicle_vecops.ScalarMulVec(dScale, dBlindedZCanonical, dScaledZ, cfg) + if eMul != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: scale Z failed: %s", eMul.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, dBlindedZCanonical.Len(), false) + eAdd := icicle_vecops.VecOp(dPrefix, dScaledZ, dPrefix, cfg, icicle_core.Add) + if eAdd != icicle_runtime.Success { + runErr = fmt.Errorf("addZContributionToLinearizedOnGPU: add scaled Z failed: %s", eAdd.AsString()) + return + } + }) + return <-done +} + +func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( + dLinearized icicle_core.DeviceSlice, + quotient *gpuQuotientPolynomial, +) error { + if dLinearized.IsEmpty() || quotient == nil || quotient.coeffs.IsEmpty() { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: empty input") + } + nPlus2 := int(s.domain0.Cardinality) + 2 + if dLinearized.Len() < nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: linearized too small") + } + if quotient.coeffs.Len() < 3*nPlus2 { + return fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: quotient too small") + } + + var one fr.Element + one.SetOne() + var zetaN, zetaNPlusTwo, zhZeta fr.Element + zetaN.Exp(s.zeta, big.NewInt(int64(s.domain0.Cardinality))) + zhZeta.Sub(&zetaN, &one) + zetaNPlusTwo.Mul(&zetaN, &s.zeta).Mul(&zetaNPlusTwo, &s.zeta) + + h1 := ("ient.coeffs).Range(0, nPlus2, false) + h2 := ("ient.coeffs).Range(nPlus2, 2*nPlus2, false) + h3 := ("ient.coeffs).Range(2*nPlus2, 3*nPlus2, false) + + done := make(chan error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("subtractQuotientContributionFromLinearizedOnGPU") + if cfgErr != nil { + done <- cfgErr + return + } + var runErr error + var dAcc icicle_core.DeviceSlice + var dZetaStd icicle_core.DeviceSlice + var dZhStd icicle_core.DeviceSlice + defer func() { + // Async boundary before reusing temporary quotient vectors. + if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { + runErr = syncErr + } + freeDeviceSlice(&dZhStd) + freeDeviceSlice(&dZetaStd) + if !dAcc.IsEmpty() { + s.putTempDeviceSlice(dAcc, dAcc.Len()) + } + done <- runErr + }() + + dAcc = s.getTempDeviceSlice(nPlus2) + dZetaStd = uploadScalarStdOnCurrentDevice(zetaNPlusTwo, cfg) + eMulH3 := icicle_vecops.ScalarMulVec(dZetaStd, h3, dAcc, cfg) + if eMulH3 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale h3 failed: %s", eMulH3.AsString()) + return + } + if eAddH2 := icicle_vecops.VecOp(dAcc, h2, dAcc, cfg, icicle_core.Add); eAddH2 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h2 failed: %s", eAddH2.AsString()) + return + } + if eMulPow := icicle_vecops.ScalarMulVec(dZetaStd, dAcc, dAcc, cfg); eMulPow != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by zeta^(n+2) failed: %s", eMulPow.AsString()) + return + } + if eAddH1 := icicle_vecops.VecOp(dAcc, h1, dAcc, cfg, icicle_core.Add); eAddH1 != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: add h1 failed: %s", eAddH1.AsString()) + return + } + + dZhStd = uploadScalarStdOnCurrentDevice(zhZeta, cfg) + eScaleZh := icicle_vecops.ScalarMulVec(dZhStd, dAcc, dAcc, cfg) + if eScaleZh != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: scale by Z_H(zeta) failed: %s", eScaleZh.AsString()) + return + } + + dPrefix := (&dLinearized).Range(0, nPlus2, false) + eSub := icicle_vecops.VecOp(dPrefix, dAcc, dPrefix, cfg, icicle_core.Sub) + if eSub != icicle_runtime.Success { + runErr = fmt.Errorf("subtractQuotientContributionFromLinearizedOnGPU: subtract term failed: %s", eSub.AsString()) + return + } + }) + return <-done +} + +// divideByZH +// The input must be in LagrangeCoset. +// The result is in Canonical Regular. (in place using a) +func (s *instance) divideByZH(a *iop.Polynomial, domains [2]*fft.Domain) (*iop.Polynomial, error) { + + // check that the basis is LagrangeCoset + if a.Basis != iop.LagrangeCoset || a.Layout != iop.BitReverse { + return nil, errors.New("invalid form") + } + + // prepare the evaluations of x^n-1 on the big domain's coset + var startEvaluateXnMinusOne time.Time + if isProfileMode { + startEvaluateXnMinusOne = time.Now() + } + xnMinusOneInverseLagrangeCoset := evaluateXnMinusOneDomainBigCoset(domains) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startEvaluateXnMinusOne)).Msg("divideByZH: evaluateXnMinusOneDomainBigCoset") + } + rho := int(domains[1].Cardinality / domains[0].Cardinality) + + r := a.Coefficients() + n := uint64(len(r)) + nn := uint64(64 - bits.TrailingZeros64(n)) + + var startParallelizeMul time.Time + if isProfileMode { + startParallelizeMul = time.Now() + } + utils.Parallelize(len(r), func(start, end int) { + for i := start; i < end; i++ { + iRev := bits.Reverse64(uint64(i)) >> nn + r[i].Mul(&r[i], &xnMinusOneInverseLagrangeCoset[int(iRev)%rho]) + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallelizeMul)).Msg("divideByZH: parallelize multiply coefficients") + } + + // Replace CPU FFT inverse by ICICLE NTT inverse. + var startGpuNTTInverse time.Time + if isProfileMode { + startGpuNTTInverse = time.Now() + } + // It's faster on CPU. + // s.gpuNTTInverse(a) + a.ToCanonical(domains[1]).ToRegular() + if isProfileMode { + l := logger.Logger() + l.Debug(). + Int("size", a.Size()). + Dur("took", time.Since(startGpuNTTInverse)). + Msg("divideByZH: gpuNTTInverse on input of size n") + } + + return a, nil +} + +// evaluateXnMinusOneDomainBigCoset evaluates Xᵐ-1 on DomainBig coset +func evaluateXnMinusOneDomainBigCoset(domains [2]*fft.Domain) []fr.Element { + + rho := domains[1].Cardinality / domains[0].Cardinality + + res := make([]fr.Element, rho) + + expo := big.NewInt(int64(domains[0].Cardinality)) + res[0].Exp(domains[1].FrMultiplicativeGen, expo) + + var t fr.Element + t.Exp(domains[1].Generator, expo) + + one := fr.One() + + for i := 1; i < int(rho); i++ { + res[i].Mul(&res[i-1], &t) + res[i-1].Sub(&res[i-1], &one) + } + res[len(res)-1].Sub(&res[len(res)-1], &one) + + res = fr.BatchInvert(res) + + return res +} + +// innerComputeLinearizedPoly computes the linearized polynomial in canonical basis. +// The purpose is to commit and open all in one ql, qr, qm, qo, qk. +// * lZeta, rZeta, oZeta are the evaluation of l, r, o at zeta +// * z is the permutation polynomial, zu is Z(μX), the shifted version of Z +// * pk is the proving key: the linearized polynomial is a linear combination of ql, qr, qm, qo, qk. +// +// The Linearized polynomial is: +// +// α²*L₁(ζ)*Z(X) +// + α*( (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*(β*s3(X))*Z(μζ) - Z(X)*(l(ζ)+β*id1(ζ)+γ)*(r(ζ)+β*id2(ζ)+γ)*(o(ζ)+β*id3(ζ)+γ)) +// + l(ζ)*Ql(X) + l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) +// - Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) +// +// /!\ blindedZCanonical is modified +func (s *instance) innerComputeLinearizedPoly( + lZeta, rZeta, oZeta, s1Zeta, s2Zeta, + alpha, beta, gamma, zeta, zu fr.Element, + qcpZeta, blindedZCanonical []fr.Element, + pi2Canonical [][]fr.Element, + pk *ProvingKey, +) []fr.Element { + + // l(ζ)r(ζ) + var rl fr.Element + rl.Mul(&rZeta, &lZeta) + + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + // the linearised polynomial is + // α²*L₁(ζ)*Z(X) + + // s1*s3(X)+s2*Z(X) + l(ζ)*Ql(X) + + // l(ζ)r(ζ)*Qm(X) + r(ζ)*Qr(X) + o(ζ)*Qo(X) + Qk(X) + ∑ᵢQcp_(ζ)Pi_(X) - + // Z_{H}(ζ)*((H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + var s1, s2, tmp fr.Element + s1.Mul(&s1Zeta, &beta).Add(&s1, &lZeta).Add(&s1, &gamma) // (l(ζ)+β*s1(ζ)+γ) + tmp.Mul(&s2Zeta, &beta).Add(&tmp, &rZeta).Add(&tmp, &gamma) + s1.Mul(&s1, &tmp).Mul(&s1, &zu).Mul(&s1, &beta).Mul(&s1, &alpha) // (l(ζ)+β*s1(ζ)+γ)*(r(ζ)+β*s2(ζ)+γ)*β*Z(μζ)*α + + var uzeta, uuzeta fr.Element + uzeta.Mul(&zeta, &pk.Vk.CosetShift) + uuzeta.Mul(&uzeta, &pk.Vk.CosetShift) + + s2.Mul(&beta, &zeta).Add(&s2, &lZeta).Add(&s2, &gamma) // (l(ζ)+β*ζ+γ) + tmp.Mul(&beta, &uzeta).Add(&tmp, &rZeta).Add(&tmp, &gamma) // (r(ζ)+β*u*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ) + tmp.Mul(&beta, &uuzeta).Add(&tmp, &oZeta).Add(&tmp, &gamma) // (o(ζ)+β*u²*ζ+γ) + s2.Mul(&s2, &tmp) // (l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + s2.Neg(&s2).Mul(&s2, &alpha) + + // Z_h(ζ), ζⁿ⁺², L₁(ζ)*α²*Z + var zhZeta, zetaNPlusTwo, alphaSquareLagrangeZero, one, den, frNbElmt fr.Element + one.SetOne() + nbElmt := int64(s.domain0.Cardinality) + alphaSquareLagrangeZero.Set(&zeta).Exp(alphaSquareLagrangeZero, big.NewInt(nbElmt)) // ζⁿ + zetaNPlusTwo.Mul(&alphaSquareLagrangeZero, &zeta).Mul(&zetaNPlusTwo, &zeta) // ζⁿ⁺² + alphaSquareLagrangeZero.Sub(&alphaSquareLagrangeZero, &one) // ζⁿ - 1 + zhZeta.Set(&alphaSquareLagrangeZero) // Z_h(ζ) = ζⁿ - 1 + frNbElmt.SetUint64(uint64(nbElmt)) + den.Sub(&zeta, &one).Inverse(&den) // 1/(ζ-1) + alphaSquareLagrangeZero.Mul(&alphaSquareLagrangeZero, &den). // L₁ = (ζⁿ - 1)/(ζ-1) + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &alpha). + Mul(&alphaSquareLagrangeZero, &s.domain0.CardinalityInv) // α²*L₁(ζ) + + s3canonical := s.trace.S3.Coefficients() + // Qk is prepared in canonical/regular form by computeLinearizedPolynomial. + + // len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 when Statistical ZK is activated + // len(h1)=len(h2)=len(h3)=len(blindedZCanonical)-1 when Statistical ZK is deactivated + h1 := s.h1() + h2 := s.h2() + h3 := s.h3() + + // at this stage we have + // s1 = α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ) + // s2 = -α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + startParallel := time.Now() + utils.Parallelize(len(blindedZCanonical), func(start, end int) { + + cql := s.trace.Ql.Coefficients() + cqr := s.trace.Qr.Coefficients() + cqm := s.trace.Qm.Coefficients() + cqo := s.trace.Qo.Coefficients() + cqk := s.trace.Qk.Coefficients() + + var t, t0, t1 fr.Element + + for i := start; i < end; i++ { + t.Mul(&blindedZCanonical[i], &s2) // -Z(X)*α*(l(ζ)+β*ζ+γ)*(r(ζ)+β*u*ζ+γ)*(o(ζ)+β*u²*ζ+γ) + if i < len(s3canonical) { + t0.Mul(&s3canonical[i], &s1) // α*(l(ζ)+β*s1(β)+γ)*(r(ζ)+β*s2(β)+γ)*β*Z(μζ)*β*s3(X) + t.Add(&t, &t0) + } + if i < len(cqm) { + t1.Mul(&cqm[i], &rl) // l(ζ)r(ζ)*Qm(X) + t.Add(&t, &t1) // linPol += l(ζ)r(ζ)*Qm(X) + t0.Mul(&cql[i], &lZeta) // l(ζ)Q_l(X) + t.Add(&t, &t0) // linPol += l(ζ)*Ql(X) + t0.Mul(&cqr[i], &rZeta) //r(ζ)*Qr(X) + t.Add(&t, &t0) // linPol += r(ζ)*Qr(X) + t0.Mul(&cqo[i], &oZeta) // o(ζ)*Qo(X) + t.Add(&t, &t0) // linPol += o(ζ)*Qo(X) + t.Add(&t, &cqk[i]) // linPol += Qk(X) + for j := range qcpZeta { // linPol += ∑ᵢQcp_(ζ)Pi_(X) + t0.Mul(&pi2Canonical[j][i], &qcpZeta[j]) + t.Add(&t, &t0) + } + } + + t0.Mul(&blindedZCanonical[i], &alphaSquareLagrangeZero) // α²L₁(ζ)Z(X) + blindedZCanonical[i].Add(&t, &t0) // linPol += α²L₁(ζ)Z(X) + + // if statistical zeroknowledge is deactivated, len(h1)=len(h2)=len(h3)=len(blindedZ)-1. + // Else len(h1)=len(h2)=len(blindedZCanonical)=len(h3)+1 + if i < len(h3) { + t.Mul(&h3[i], &zetaNPlusTwo). + Add(&t, &h2[i]). + Mul(&t, &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } else { + if s.opt.StatisticalZK { + t.Mul(&h2[i], &zetaNPlusTwo). + Add(&t, &h1[i]). + Mul(&t, &zhZeta) + blindedZCanonical[i].Sub(&blindedZCanonical[i], &t) // linPol -= Z_h(ζ)*(H₀(X) + ζᵐ⁺²*H₁(X) + ζ²⁽ᵐ⁺²⁾*H₂(X)) + } + } + } + }) + if isProfileMode { + l := logger.Logger() + l.Debug().Dur("took", time.Since(startParallel)).Msg("computeLinearizedPolynomial: inner parallel loop") + } + + return blindedZCanonical +} + +var errContextDone = errors.New("context done") + +// local copies of verification-time helpers used by prover transcript +func bindPublicData(fs *fiatshamir.Transcript, challenge string, vk *plonk_bw6761.VerifyingKey, publicInputs []fr.Element) error { + if err := fs.Bind(challenge, vk.S[0].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[1].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.S[2].Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Ql.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qr.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qm.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qo.Marshal()); err != nil { + return err + } + if err := fs.Bind(challenge, vk.Qk.Marshal()); err != nil { + return err + } + for i := range vk.Qcp { + if err := fs.Bind(challenge, vk.Qcp[i].Marshal()); err != nil { + return err + } + } + for i := 0; i < len(publicInputs); i++ { + if err := fs.Bind(challenge, publicInputs[i].Marshal()); err != nil { + return err + } + } + return nil +} + +func deriveRandomness(fs *fiatshamir.Transcript, challenge string, points ...*curve.G1Affine) (fr.Element, error) { + var buf [curve.SizeOfG1AffineUncompressed]byte + var r fr.Element + for _, p := range points { + buf = p.RawBytes() + if err := fs.Bind(challenge, buf[:]); err != nil { + return r, err + } + } + b, err := fs.ComputeChallenge(challenge) + if err != nil { + return r, err + } + r.SetBytes(b) + return r, nil +} + +// -------------------- GPU helpers and device setup -------------------- + +func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { + if pk.deviceInfo != nil { + return nil + } + pk.deviceInfo = &deviceInfo{} + + // Initialize ICICLE NTT domain (root of unity) and store coset generator for coset NTTs. + // ICICLE InitDomain expects a primitive root of unity; for coset transforms we use 𝔽ᵣ* generator. + + var gen fr.Element + var err error + if pk.Vk.Size < 6 { + gen, err = fft.Generator(8 * pk.Vk.Size) + if err != nil { + return err + } + } else { + gen, err = fft.Generator(4 * pk.Vk.Size) + if err != nil { + return err + } + } + genBits := gen.Bits() + limbs := icicle_core.ConvertUint64ArrToUint32Arr(genBits[:]) + // Initialize ICICLE NTT domain with root of unity + var rouIcicle icicle_bw6761.ScalarField + rouIcicle.FromLimbs(limbs) + + // Store coset generator = generator of 𝔽ᵣ* (matches CPU ToLagrangeCoset) + { + cosetGen := fft.GeneratorFullMultiplicativeGroup() + cosetBits := cosetGen.Bits() + cosetLimbs := icicle_core.ConvertUint64ArrToUint32Arr(cosetBits[:]) + copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) + } + + chInitDomain := make(chan struct{}) + initDomainQueuedAt := time.Now() + icicle_runtime.RunOnDevice(device, func(args ...any) { + initDomainStartedAt := time.Now() + initCfg := icicle_core.GetDefaultNTTInitDomainConfig() + ext := config_extension.Create() + defer config_extension.Delete(ext) + fastTwiddles := envEnabled("ICICLE_NTT_FAST_TWIDDLES", true) + ext.SetBool(icicle_core.CUDA_NTT_FAST_TWIDDLES_MODE, fastTwiddles) + initCfg.Ext = ext.AsUnsafePointer() + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain start vk_size=%d fast_twiddles=%t profile_full=%q profile_arbitrary=%q\n", + pk.Vk.Size, + fastTwiddles, + os.Getenv("ICICLE_NTT_PROFILE_FULL"), + os.Getenv("ICICLE_NTT_PROFILE_ARBITRARY_COSET"), + ) + } + e := icicle_ntt.InitDomain(rouIcicle, initCfg) + if isNttTrace { + fmt.Fprintf( + os.Stderr, + "[ICICLE_NTT_TRACE] InitDomain end status=%s call_took=%s\n", + e.AsString(), + time.Since(initDomainStartedAt), + ) + } + if e != icicle_runtime.Success { + panic("icicle: InitDomain failed") + } + close(chInitDomain) + }) + + <-chInitDomain + if isNttTrace { + fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) + } + + chLag := make(chan struct{}) + chCan := make(chan struct{}) + + if len(pk.KzgLagrange.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) + g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) + close(chLag) + }) + } else { + close(chLag) + } + + if len(pk.Kzg.G1) > 0 { + icicle_runtime.RunOnDevice(device, func(args ...any) { + g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) + g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) + close(chCan) + }) + } else { + close(chCan) + } + + <-chLag + <-chCan + return nil +} + +func projectiveToGnarkAffine(p icicle_bw6761.Projective) (curve.G1Affine, error) { + x, err := icicleBaseFieldToGnarkFp(p.X) + if err != nil { + return curve.G1Affine{}, err + } + y, err := icicleBaseFieldToGnarkFp(p.Y) + if err != nil { + return curve.G1Affine{}, err + } + z, err := icicleBaseFieldToGnarkFp(p.Z) + if err != nil { + return curve.G1Affine{}, err + } + if z.IsZero() { + return curve.G1Affine{}, nil + } + + var zInv fp.Element + zInv.Inverse(&z) + x.Mul(&x, &zInv) + y.Mul(&y, &zInv) + return curve.G1Affine{X: x, Y: y}, nil +} + +func icicleBaseFieldToGnarkFp(v icicle_bw6761.BaseField) (fp.Element, error) { + bytes := v.ToBytesLittleEndian() + if len(bytes) != fp.Bytes { + return fp.Element{}, fmt.Errorf("invalid ICICLE base field byte length %d", len(bytes)) + } + var buf [fp.Bytes]byte + copy(buf[:], bytes) + return fp.LittleEndian.Element(&buf) +} + +func commitOnGPULagrangeDevice(scalarsDevice icicle_core.DeviceSlice, device *icicle_runtime.Device, pk *ProvingKey) (curve.G1Affine, error) { + if scalarsDevice.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: empty scalar slice") + } + if pk == nil || pk.deviceInfo == nil || pk.deviceInfo.KzgLagrangeDevice.G1.IsEmpty() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: lagrange SRS is not available on device") + } + if scalarsDevice.Len() > pk.deviceInfo.KzgLagrangeDevice.G1.Len() { + return curve.G1Affine{}, fmt.Errorf("commitOnGPULagrangeDevice: invalid scalar size %d", scalarsDevice.Len()) + } + return commitOnGPUWithDeviceBases(scalarsDevice, pk.deviceInfo.KzgLagrangeDevice.G1, device) +} + +func (s *instance) registerDevicePolynomialInSharedState(state *gpuPolysState, p *iop.Polynomial, dSlice icicle_core.DeviceSlice) error { + if state == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil shared state") + } + if p == nil { + return fmt.Errorf("registerDevicePolynomialInSharedState: nil polynomial") + } + if dSlice.IsEmpty() { + return fmt.Errorf("registerDevicePolynomialInSharedState: empty device slice") + } + + s.gpuStateMu.Lock() + defer s.gpuStateMu.Unlock() + + if state.polyToIdx == nil { + state.polyToIdx = make(map[*iop.Polynomial]int) + } + if idx, ok := state.polyToIdx[p]; ok { + if idx < 0 || idx >= len(state.deviceSlices) { + return fmt.Errorf("registerDevicePolynomialInSharedState: invalid index %d", idx) + } + state.deviceSlices[idx] = dSlice + state.hostSlices[idx] = nil + state.originalForm[idx] = iop.Form{Basis: p.Basis, Layout: p.Layout} + return nil + } + + idx := len(state.polys) + state.polyToIdx[p] = idx + state.polys = append(state.polys, p) + state.deviceSlices = append(state.deviceSlices, dSlice) + state.hostSlices = append(state.hostSlices, nil) + state.originalForm = append(state.originalForm, iop.Form{Basis: p.Basis, Layout: p.Layout}) + return nil +} + +func (s *instance) gpuInclusivePrefixProductOnCurrentDevice( + dVec icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + if dVec.IsEmpty() || dVec.Len() <= 1 { + return nil + } + + n := dVec.Len() + for step := 1; step < n; step <<= 1 { + src := (&dVec).Range(0, n-step, false) + dst := (&dVec).Range(step, n, false) + + tmpStd := s.getTempDeviceSlice(n - step) + if err := copyDeviceSliceIntoOnCurrentDevice(tmpStd, src, cfg); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: copy stage failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(tmpStd, cfg) + if err := icicle_vecops.VecOp(tmpStd, dst, dst, cfg, icicle_core.Mul); err != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: multiply stage failed: %s", err.AsString()) + } + if cfg.IsAsync { + // tmpStd is returned to pool each stage, so we must wait before reuse. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + s.putTempDeviceSlice(tmpStd, n-step) + return fmt.Errorf("gpuInclusivePrefixProductOnCurrentDevice: synchronize stream failed: %s", eSync.AsString()) + } + } + s.putTempDeviceSlice(tmpStd, n-step) + } + return nil +} + +// buildPermutationGatherIndices prepares the subset of permutation indices that +// are consumed by the copy-constraint ratio loop (only rows [0, n-1) per copy). +func buildPermutationGatherIndices(permutation []int64, nbPolynomials, n, supportLen int) ([]int64, error) { + if n <= 1 { + return nil, nil + } + total := nbPolynomials * (n - 1) + indices := make([]int64, total) + + var permBuildErr error + var permBuildErrOnce sync.Once + utils.Parallelize(total, func(start, end int) { + for k := start; k < end; k++ { + j := k / (n - 1) + i := k % (n - 1) + base := j * n + permIdx := permutation[base+i] + if permIdx < 0 || int(permIdx) >= supportLen { + jj, ii, bad := j, i, permIdx + permBuildErrOnce.Do(func() { + permBuildErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: permutation index out of range at (%d,%d): %d", jj, ii, bad) + }) + continue + } + indices[k] = permIdx + } + }) + if permBuildErr != nil { + return nil, permBuildErr + } + return indices, nil +} + +func (s *instance) prepareCopyConstraintSupportsOnCurrentDevice( + n, nbPolynomials int, + domain *fft.Domain, + permGatherIndices []int64, + cfg icicle_core.VecOpsConfig, +) (dSupportFlat, dPermFlat icicle_core.DeviceSlice, err error) { + defer func() { + if err == nil { + return + } + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + }() + + if len(permGatherIndices) == 0 { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: empty permutation gather indices") + return + } + + dOmegaStd := uploadScalarStdOnCurrentDevice(domain.Generator, cfg) + defer dOmegaStd.Free() + dShiftStd := uploadScalarStdOnCurrentDevice(domain.FrMultiplicativeGen, cfg) + defer dShiftStd.Free() + + dSupportFlat, err = allocDeviceUninitialized(nbPolynomials * n) + if err != nil { + return + } + if e := icicle_vecops.SupportIdentity(dOmegaStd, dShiftStd, n, nbPolynomials, dSupportFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: generate identity support on GPU failed: %s", e.AsString()) + return + } + toMontgomeryFormInPlaceWithCfg(dSupportFlat, cfg) + + dPermIndicesDevice := uploadInt64VectorOnCurrentDevice(permGatherIndices, cfg) + defer dPermIndicesDevice.Free() + + dPermFlat, err = allocDeviceUninitialized(len(permGatherIndices)) + if err != nil { + return + } + if e := icicle_vecops.GatherByIndices(dSupportFlat, dPermIndicesDevice, dPermFlat, cfg); e != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: gather permutation support on GPU failed: %s", e.AsString()) + return + } + if cfg.IsAsync { + // Ensure temporary support/index slices are safe to free on return. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + err = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + return + } + } + + return +} + +func (s *instance) accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail icicle_core.DeviceSlice, + dBetaStd, dGammaMont icicle_core.DeviceSlice, + cfg icicle_core.VecOpsConfig, +) error { + nMinusOne := dEntryTail.Len() + dScaled := s.getTempDeviceSlice(nMinusOne) + dTerm := s.getTempDeviceSlice(nMinusOne) + defer func() { + s.putTempDeviceSlice(dScaled, nMinusOne) + s.putTempDeviceSlice(dTerm, nMinusOne) + }() + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dID, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale identity support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: numerator add beta*id failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dNumTail, dNumTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply numerator term failed: %s", err.AsString()) + } + + if err := icicle_vecops.ScalarMulVec(dBetaStd, dPerm, dScaled, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: scale permutation support failed: %s", err.AsString()) + } + if err := icicle_vecops.ScalarAddVec(dGammaMont, dEntryTail, dTerm, cfg); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add gamma failed: %s", err.AsString()) + } + if err := icicle_vecops.VecOp(dTerm, dScaled, dTerm, cfg, icicle_core.Add); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: denominator add beta*sigma failed: %s", err.AsString()) + } + toStandardFormInPlaceWithCfg(dTerm, cfg) + if err := icicle_vecops.VecOp(dTerm, dDenTail, dDenTail, cfg, icicle_core.Mul); err != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: multiply denominator term failed: %s", err.AsString()) + } + if cfg.IsAsync { + // Temp vectors are released at function exit, so ensure queued work is complete. + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + return fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize stream failed: %s", eSync.AsString()) + } + } + + return nil +} + +// validateDeviceEntries checks that all entries are non-empty and have consistent length. +// Returns the common length n. +func validateDeviceEntries(entries []icicle_core.DeviceSlice, label string) (int, error) { + if len(entries) == 0 { + return 0, fmt.Errorf("%s: no entries", label) + } + n := entries[0].Len() + if n == 0 { + return 0, fmt.Errorf("%s: empty device entry 0", label) + } + for i := range entries { + if entries[i].IsEmpty() { + return 0, fmt.Errorf("%s: empty device entry %d", label, i) + } + if entries[i].Len() != n { + return 0, fmt.Errorf("%s: inconsistent device entry size at %d (%d != %d)", label, i, entries[i].Len(), n) + } + } + return n, nil +} + +// BuildRatioCopyConstraintIcicle builds the accumulating ratio polynomial to prove that +// [P₁ ∥ .. ∥ P_{n—1}] is invariant by the permutation \sigma. +// Namely it returns the polynomial Z whose evaluation on the j-th root of unity is +// Z(ω^j) = Π_{i 1 { + dNumTail := (&dNum).Range(1, n, false) + dDenTail := (&dDen).Range(1, n, false) + var supportErr error + dSupportFlat, dPermFlat, supportErr = s.prepareCopyConstraintSupportsOnCurrentDevice(n, nbPolynomials, domain, permGatherIndices, cfg) + if supportErr != nil { + runErr = supportErr + return + } + + dBetaStd := uploadScalarStdOnCurrentDevice(beta, cfg) + dGammaMont := uploadScalarMontOnCurrentDevice(gamma, cfg) + + for j := 0; j < nbPolynomials; j++ { + dEntryTail := (&entriesDevice[j]).Range(0, n-1, false) + baseID := j * n + dID := (&dSupportFlat).Range(baseID, baseID+n-1, false) + basePerm := j * (n - 1) + dPerm := (&dPermFlat).Range(basePerm, basePerm+(n-1), false) + if err := s.accumulateCopyConstraintTermOnCurrentDevice( + dEntryTail, dID, dPerm, dNumTail, dDenTail, dBetaStd, dGammaMont, cfg, + ); err != nil { + runErr = err + return + } + } + if cfg.IsAsync { + if eSync := icicle_runtime.SynchronizeStream(cfg.StreamHandle); eSync != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: synchronize before releasing copy-constraint workspace failed: %s", eSync.AsString()) + return + } + } + _ = dBetaStd.Free() + _ = dGammaMont.Free() + + // Support vectors and loop temps are only needed for term accumulation. + // Free them before prefix products and batch inversion, whose ICICLE + // kernels allocate additional full-domain workspace internally. + freeDeviceSlice(&dPermFlat) + freeDeviceSlice(&dSupportFlat) + s.tempGPUMemPool.FreeAll() + } + + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dNum, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + if err := s.gpuInclusivePrefixProductOnCurrentDevice(dDen, cfg); err != nil { + runErr = err + return + } + s.tempGPUMemPool.FreeAll() + + if invErr := s.batchInvertOnCurrentDevice(dDen); invErr != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: GPU batch inversion failed: %s", invErr.AsString()) + return + } + + toStandardFormInPlace(dDen) + if err := icicle_vecops.VecOp(dDen, dNum, dNum, cfg, icicle_core.Mul); err != icicle_runtime.Success { + runErr = fmt.Errorf("BuildRatioCopyConstraintIcicle: final numerator*denominatorInv multiplication failed: %s", err.AsString()) + return + } + + dResult = dNum + dNum = icicle_core.DeviceSlice{} // transfer ownership to dResult + }) + if err := <-buildDone; err != nil { + return nil, err + } + + hostMirror := make([]fr.Element, n) + if len(hostMirror) > 0 { + hostMirror[0].SetOne() + } + res := iop.NewPolynomial(&hostMirror, iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}) + if err := s.registerDevicePolynomialInSharedState(gpuState, res, dResult); err != nil { + freeSliceOnDevice(&dResult, &s.device) + return nil, err + } + + return res, nil +} diff --git a/backend/accelerated/icicle/plonk/bw6-761/provingkey.go b/backend/accelerated/icicle/plonk/bw6-761/provingkey.go new file mode 100644 index 0000000000..d0f534c451 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bw6-761/provingkey.go @@ -0,0 +1,100 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +//go:build icicle + +package bw6761 + +import ( + "sync" + "time" + + "github.com/consensys/gnark-crypto/ecc/bw6-761/fr" + "github.com/consensys/gnark-crypto/ecc/bw6-761/fr/fft" + plonk_bw6761 "github.com/consensys/gnark/backend/plonk/bw6-761" + cs "github.com/consensys/gnark/constraint/bw6-761" + "github.com/consensys/gnark/logger" + icicle_core "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/core" +) + +// deviceInfo holds device-resident buffers for GPU acceleration. +type deviceInfo struct { + CosetGenerator [fr.Limbs * 2]uint32 + KzgDevice struct { + G1 icicle_core.DeviceSlice + } + KzgLagrangeDevice struct { + G1 icicle_core.DeviceSlice + } +} + +// hostSetup holds host-side, witness-independent prover state derived from the +// constraint system: the FFT domains and the PLONK trace (selector + +// permutation polynomials). Building the trace walks every constraint (~3s at +// 23M constraints), so it is computed once per proving key and shared across +// proofs. Everything here is read-only during proving: the prover clones Qk +// before patching public inputs into it, and every basis conversion of a trace +// polynomial copies first (see canonicalRegularCoefficientsCopy). +type hostSetup struct { + sizeSystem uint64 + domain0 *fft.Domain + domain1 *fft.Domain + trace *plonk_bw6761.Trace +} + +// ProvingKey wraps the native PLONK proving key with device-resident state +// (KZG bases, NTT domains, cached trace) that is uploaded once and reused +// across Prove calls. +// +// Concurrency: Prove calls sharing the same ProvingKey must be serialized by +// the caller. The device state hangs off the key and proofs share a single +// GPU; concurrent proves against the same key are not safe. +type ProvingKey struct { + plonk_bw6761.ProvingKey + *deviceInfo + hostSetupOnce sync.Once + hostSetup *hostSetup +} + +func buildHostSetup(spr *cs.SparseR1CS, sizeSystem uint64) *hostSetup { + domain0 := fft.NewDomain(sizeSystem) + + // h, the quotient polynomial is of degree 3(n+1)+2, so it's in a 3(n+2) dim + // vector space, the domain is the next power of 2 superior to 3(n+2). + // 4*domainNum is enough in all cases except when n<6. + var domain1 *fft.Domain + if sizeSystem < 6 { + domain1 = fft.NewDomain(8*sizeSystem, fft.WithoutPrecompute()) + } else { + domain1 = fft.NewDomain(4*sizeSystem, fft.WithoutPrecompute()) + } + + return &hostSetup{ + sizeSystem: sizeSystem, + domain0: domain0, + domain1: domain1, + trace: plonk_bw6761.NewTrace(spr, domain0), + } +} + +// hostSetupFor returns the FFT domains and trace for spr, building them on +// first use and caching them on the proving key. A PLONK proving key is bound +// to exactly one constraint system, so per-key caching is sound; as a +// defensive measure a system-size mismatch falls back to an uncached build +// rather than ever serving another circuit's trace. +func (pk *ProvingKey) hostSetupFor(spr *cs.SparseR1CS) *hostSetup { + nbConstraints := spr.GetNbConstraints() + sizeSystem := uint64(nbConstraints + len(spr.Public)) // len(spr.Public) is for the placeholder constraints + pk.hostSetupOnce.Do(func() { + start := time.Now() + pk.hostSetup = buildHostSetup(spr, sizeSystem) + log := logger.Logger() + log.Debug().Dur("took", time.Since(start)).Msg("built prover host setup (fft domains + trace)") + }) + if pk.hostSetup.sizeSystem != sizeSystem { + return buildHostSetup(spr, sizeSystem) + } + return pk.hostSetup +} diff --git a/backend/accelerated/icicle/plonk/e2e_test.go b/backend/accelerated/icicle/plonk/e2e_test.go new file mode 100644 index 0000000000..67d88e7699 --- /dev/null +++ b/backend/accelerated/icicle/plonk/e2e_test.go @@ -0,0 +1,88 @@ +//go:build icicle + +package plonk_test + +import ( + "errors" + "math/big" + "testing" + + "github.com/consensys/gnark-crypto/ecc" + accel_plonk "github.com/consensys/gnark/backend/accelerated/icicle/plonk" + "github.com/consensys/gnark/frontend" + "github.com/consensys/gnark/frontend/cs/scs" + "github.com/consensys/gnark/test" + "github.com/consensys/gnark/test/unsafekzg" +) + +const largeCircuitSize = 1 << 12 + +var errNoCommitter = errors.New("builder does not implement frontend.Committer") + +// largeCircuit is a sequential x = x² + a recurrence, padded with a BSB22 +// commitment so that the commitment path of the GPU prover is exercised +// (production circuits using std/rangecheck always carry one). +type largeCircuit struct { + A frontend.Variable `gnark:",public"` + Res frontend.Variable +} + +func (c *largeCircuit) Define(api frontend.API) error { + x := c.A + for i := 0; i < largeCircuitSize; i++ { + x = api.Add(api.Mul(x, x), c.A) + } + api.AssertIsEqual(x, c.Res) + committer, ok := api.(frontend.Committer) + if !ok { + return errNoCommitter + } + cm, err := committer.Commit(x, c.A) + if err != nil { + return err + } + api.AssertIsDifferent(cm, 0) + return nil +} + +// TestEndToEndLargeCircuit runs a full accelerated setup -> GPU prove -> +// verify cycle on every supported curve, with a circuit large enough +// (2^12+ constraints) to exercise the chunked-MSM and NTT paths. +func TestEndToEndLargeCircuit(t *testing.T) { + for _, curveID := range []ecc.ID{ecc.BN254, ecc.BLS12_377, ecc.BLS12_381, ecc.BW6_761} { + t.Run(curveID.String(), func(t *testing.T) { + assert := test.NewAssert(t) + + ccs, err := frontend.Compile(curveID.ScalarField(), scs.NewBuilder, &largeCircuit{}) + assert.NoError(err) + t.Logf("nb constraints: %d", ccs.GetNbConstraints()) + + srs, srsLagrange, err := unsafekzg.NewSRS(ccs) + assert.NoError(err) + + iciPK, iciVK, err := accel_plonk.Setup(ccs, srs, srsLagrange) + assert.NoError(err) + + // compute the expected result with the same recurrence + mod := curveID.ScalarField() + a := big.NewInt(3) + x := big.NewInt(3) + for i := 0; i < largeCircuitSize; i++ { + x.Mul(x, x) + x.Add(x, a) + x.Mod(x, mod) + } + + assignment := largeCircuit{A: a, Res: x} + w, err := frontend.NewWitness(&assignment, curveID.ScalarField()) + assert.NoError(err) + pw, err := w.Public() + assert.NoError(err) + + proof, err := accel_plonk.Prove(ccs, iciPK, w) + assert.NoError(err) + err = accel_plonk.Verify(proof, iciVK, pw) + assert.NoError(err) + }) + } +} diff --git a/backend/accelerated/icicle/plonk/fallback_test.go b/backend/accelerated/icicle/plonk/fallback_test.go new file mode 100644 index 0000000000..96f0f4988a --- /dev/null +++ b/backend/accelerated/icicle/plonk/fallback_test.go @@ -0,0 +1,51 @@ +//go:build !icicle + +package plonk_test + +import ( + "testing" + + "github.com/consensys/gnark-crypto/ecc" + accel_plonk "github.com/consensys/gnark/backend/accelerated/icicle/plonk" + "github.com/consensys/gnark/frontend" + "github.com/consensys/gnark/frontend/cs/scs" + "github.com/consensys/gnark/test" + "github.com/consensys/gnark/test/unsafekzg" +) + +type fallbackCircuit struct { + A, B frontend.Variable `gnark:",public"` + Res frontend.Variable +} + +func (c *fallbackCircuit) Define(api frontend.API) error { + api.AssertIsEqual(api.Mul(c.A, c.B), c.Res) + return nil +} + +// TestCPUFallback checks that, when compiled without the 'icicle' build tag, +// the package degrades gracefully to the native CPU PLONK prover instead of +// panicking. +func TestCPUFallback(t *testing.T) { + assert := test.NewAssert(t) + + ccs, err := frontend.Compile(ecc.BN254.ScalarField(), scs.NewBuilder, &fallbackCircuit{}) + assert.NoError(err) + + srs, srsLagrange, err := unsafekzg.NewSRS(ccs) + assert.NoError(err) + + pk, vk, err := accel_plonk.Setup(ccs, srs, srsLagrange) + assert.NoError(err) + + assignment := fallbackCircuit{A: 3, B: 5, Res: 15} + w, err := frontend.NewWitness(&assignment, ecc.BN254.ScalarField()) + assert.NoError(err) + pw, err := w.Public() + assert.NoError(err) + + proof, err := accel_plonk.Prove(ccs, pk, w) + assert.NoError(err) + err = accel_plonk.Verify(proof, vk, pw) + assert.NoError(err) +} diff --git a/backend/accelerated/icicle/plonk/marshal_test.go b/backend/accelerated/icicle/plonk/marshal_test.go new file mode 100644 index 0000000000..12ea0e261d --- /dev/null +++ b/backend/accelerated/icicle/plonk/marshal_test.go @@ -0,0 +1,105 @@ +//go:build icicle + +package plonk_test + +import ( + "bytes" + "testing" + + "github.com/consensys/gnark-crypto/ecc" + accel_plonk "github.com/consensys/gnark/backend/accelerated/icicle/plonk" + native_plonk "github.com/consensys/gnark/backend/plonk" + plonk_bn254 "github.com/consensys/gnark/backend/plonk/bn254" + cs_bn254 "github.com/consensys/gnark/constraint/bn254" + "github.com/consensys/gnark/frontend" + "github.com/consensys/gnark/frontend/cs/scs" + "github.com/consensys/gnark/test" + "github.com/consensys/gnark/test/unsafekzg" +) + +type circuit struct { + A, B frontend.Variable `gnark:",public"` + Res frontend.Variable +} + +func (c *circuit) Define(api frontend.API) error { + api.AssertIsEqual(api.Mul(c.A, c.B), c.Res) + return nil +} + +func TestMarshalBN254(t *testing.T) { + assert := test.NewAssert(t) + + ccs, err := frontend.Compile(ecc.BN254.ScalarField(), scs.NewBuilder, &circuit{}) + assert.NoError(err) + tCcs := ccs.(*cs_bn254.SparseR1CS) + + // Build an SRS suitable for this circuit + srs, srsLagrange, err := unsafekzg.NewSRS(tCcs) + assert.NoError(err) + + // Native setup + nativePK, nativeVK, err := native_plonk.Setup(tCcs, srs, srsLagrange) + assert.NoError(err) + + // Marshal native -> icicle bn254.ProvingKey + iciPK := accel_plonk.NewProvingKey(ecc.BN254) + buf := new(bytes.Buffer) + _, err = nativePK.WriteTo(buf) + assert.NoError(err) + _, err = iciPK.ReadFrom(buf) + assert.NoError(err) + + // Roundtrip back into native + buf.Reset() + _, err = iciPK.WriteTo(buf) + assert.NoError(err) + var nativePK2 plonk_bn254.ProvingKey + _, err = nativePK2.ReadFrom(buf) + assert.NoError(err) + + // Prove both ways + assignment := circuit{A: 3, B: 5, Res: 15} + w, err := frontend.NewWitness(&assignment, ecc.BN254.ScalarField()) + assert.NoError(err) + pw, err := w.Public() + assert.NoError(err) + + proofNative, err := native_plonk.Prove(tCcs, &nativePK2, w) + assert.NoError(err) + proofIcicle, err := accel_plonk.Prove(tCcs, iciPK, w) + assert.NoError(err) + + err = accel_plonk.Verify(proofNative, nativeVK, pw) + assert.NoError(err) + err = accel_plonk.Verify(proofIcicle, nativeVK, pw) + assert.NoError(err) +} + +func TestSetupBN254(t *testing.T) { + assert := test.NewAssert(t) + + ccs, err := frontend.Compile(ecc.BN254.ScalarField(), scs.NewBuilder, &circuit{}) + assert.NoError(err) + tCcs := ccs.(*cs_bn254.SparseR1CS) + + // Build an SRS suitable for this circuit + srs, srsLagrange, err := unsafekzg.NewSRS(tCcs) + assert.NoError(err) + + // Accelerated setup + iciPK, iciVK, err := accel_plonk.Setup(tCcs, srs, srsLagrange) + assert.NoError(err) + + // Prove/verify using accelerated Prove + assignment := circuit{A: 3, B: 5, Res: 15} + w, err := frontend.NewWitness(&assignment, ecc.BN254.ScalarField()) + assert.NoError(err) + pw, err := w.Public() + assert.NoError(err) + + proofIcicle, err := accel_plonk.Prove(tCcs, iciPK, w) + assert.NoError(err) + err = accel_plonk.Verify(proofIcicle, iciVK, pw) + assert.NoError(err) +} diff --git a/backend/accelerated/icicle/plonk/plonk_all.go b/backend/accelerated/icicle/plonk/plonk_all.go new file mode 100644 index 0000000000..c87e5cc44c --- /dev/null +++ b/backend/accelerated/icicle/plonk/plonk_all.go @@ -0,0 +1,30 @@ +// Package plonk provides wrappers for PLONK with ICICLE acceleration. +package plonk + +import ( + "github.com/consensys/gnark-crypto/ecc" + "github.com/consensys/gnark/backend" + native_plonk "github.com/consensys/gnark/backend/plonk" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" +) + +// Verify wraps native plonk.Verify for convenience. +func Verify(proof native_plonk.Proof, vk native_plonk.VerifyingKey, publicWitness witness.Witness, opts ...backend.VerifierOption) error { + return native_plonk.Verify(proof, vk, publicWitness, opts...) +} + +// NewVerifyingKey returns a new empty VerifyingKey compatible with native plonk.NewVerifyingKey. +func NewVerifyingKey(curveID ecc.ID) native_plonk.VerifyingKey { + return native_plonk.NewVerifyingKey(curveID) +} + +// NewProof returns a new empty Proof compatible with native plonk.NewProof. +func NewProof(curveID ecc.ID) native_plonk.Proof { + return native_plonk.NewProof(curveID) +} + +// NewCS returns a new typed CS compatible with native plonk.NewCS. +func NewCS(curveID ecc.ID) constraint.ConstraintSystem { + return native_plonk.NewCS(curveID) +} diff --git a/backend/accelerated/icicle/plonk/plonk_icicle.go b/backend/accelerated/icicle/plonk/plonk_icicle.go new file mode 100644 index 0000000000..c2638263a4 --- /dev/null +++ b/backend/accelerated/icicle/plonk/plonk_icicle.go @@ -0,0 +1,209 @@ +//go:build icicle + +package plonk + +import ( + "fmt" + "sync" + + "github.com/consensys/gnark-crypto/ecc" + kzg_bls12377 "github.com/consensys/gnark-crypto/ecc/bls12-377/kzg" + kzg_bls12381 "github.com/consensys/gnark-crypto/ecc/bls12-381/kzg" + kzg_bn254 "github.com/consensys/gnark-crypto/ecc/bn254/kzg" + kzg_bw6761 "github.com/consensys/gnark-crypto/ecc/bw6-761/kzg" + "github.com/consensys/gnark-crypto/kzg" + "github.com/consensys/gnark/backend" + native_plonk "github.com/consensys/gnark/backend/plonk" + plonk_bls12377 "github.com/consensys/gnark/backend/plonk/bls12-377" + plonk_bls12381 "github.com/consensys/gnark/backend/plonk/bls12-381" + plonk_bn254 "github.com/consensys/gnark/backend/plonk/bn254" + plonk_bw6761 "github.com/consensys/gnark/backend/plonk/bw6-761" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" + cs_bls12377 "github.com/consensys/gnark/constraint/bls12-377" + cs_bls12381 "github.com/consensys/gnark/constraint/bls12-381" + cs_bn254 "github.com/consensys/gnark/constraint/bn254" + cs_bw6761 "github.com/consensys/gnark/constraint/bw6-761" + + icicle_bls12377 "github.com/consensys/gnark/backend/accelerated/icicle/plonk/bls12-377" + icicle_bls12381 "github.com/consensys/gnark/backend/accelerated/icicle/plonk/bls12-381" + icicle_bn254 "github.com/consensys/gnark/backend/accelerated/icicle/plonk/bn254" + icicle_bw6761 "github.com/consensys/gnark/backend/accelerated/icicle/plonk/bw6-761" + "github.com/consensys/gnark/logger" + icicle_runtime "github.com/ingonyama-zk/icicle-gnark/v3/wrappers/golang/runtime" +) + +var onceWarmUpDevice sync.Once + +func warmUpDevice() { + onceWarmUpDevice.Do(func() { + log := logger.Logger() + err := icicle_runtime.LoadBackendFromEnvOrDefault() + if err != icicle_runtime.Success { + panic(fmt.Sprintf("ICICLE backend loading error: %s", err.AsString())) + } + + // PLONK currently proves on CUDA device 0; warm it once to reduce + // first-use latency spikes from allocator/runtime initialization. + device := icicle_runtime.CreateDevice("CUDA", 0) + warmDone := make(chan error, 1) + icicle_runtime.RunOnDevice(&device, func(args ...any) { + stream, streamErr := icicle_runtime.CreateStream() + if streamErr != icicle_runtime.Success { + warmDone <- fmt.Errorf("ICICLE create stream error: %s", streamErr.AsString()) + return + } + + var runErr error + defer func() { + if syncErr := icicle_runtime.SynchronizeStream(stream); syncErr != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("ICICLE device warmup synchronize error: %s", syncErr.AsString()) + } + if destroyErr := icicle_runtime.DestroyStream(stream); destroyErr != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("ICICLE destroy stream error: %s", destroyErr.AsString()) + } + warmDone <- runErr + }() + + if warmErr := icicle_runtime.WarmUpDevice(stream); warmErr != icicle_runtime.Success { + runErr = fmt.Errorf("ICICLE device warmup error: %s", warmErr.AsString()) + return + } + }) + + if warmErr := <-warmDone; warmErr != nil { + panic(warmErr) + } + log.Debug().Str("device", "CUDA:0").Msg("ICICLE backend initialized and warmed for PLONK") + }) +} + +// Prove runs the accelerated prover for supported curves. +func Prove(ccs constraint.ConstraintSystem, pk native_plonk.ProvingKey, fullWitness witness.Witness, opts ...backend.ProverOption) (native_plonk.Proof, error) { + warmUpDevice() + switch tccs := ccs.(type) { + case *cs_bn254.SparseR1CS: + // Accept both ICICLE-wrapped and native proving keys; wrap if needed. + var iciclePK *icicle_bn254.ProvingKey + switch t := pk.(type) { + case *icicle_bn254.ProvingKey: + iciclePK = t + case *plonk_bn254.ProvingKey: + // Wrap native proving key into ICICLE proving key; device buffers will be initialized lazily. + iciclePK = &icicle_bn254.ProvingKey{ProvingKey: *t} + default: + return nil, fmt.Errorf("icicle plonk: unsupported proving key type %T for BN254", pk) + } + return icicle_bn254.Prove(tccs, iciclePK, fullWitness, opts...) + case *cs_bls12377.SparseR1CS: + // Accept both ICICLE-wrapped and native proving keys; wrap if needed. + var iciclePK *icicle_bls12377.ProvingKey + switch t := pk.(type) { + case *icicle_bls12377.ProvingKey: + iciclePK = t + case *plonk_bls12377.ProvingKey: + // Wrap native proving key into ICICLE proving key; device buffers will be initialized lazily. + iciclePK = &icicle_bls12377.ProvingKey{ProvingKey: *t} + default: + return nil, fmt.Errorf("icicle plonk: unsupported proving key type %T for BLS12-377", pk) + } + return icicle_bls12377.Prove(tccs, iciclePK, fullWitness, opts...) + case *cs_bls12381.SparseR1CS: + // Accept both ICICLE-wrapped and native proving keys; wrap if needed. + var iciclePK *icicle_bls12381.ProvingKey + switch t := pk.(type) { + case *icicle_bls12381.ProvingKey: + iciclePK = t + case *plonk_bls12381.ProvingKey: + // Wrap native proving key into ICICLE proving key; device buffers will be initialized lazily. + iciclePK = &icicle_bls12381.ProvingKey{ProvingKey: *t} + default: + return nil, fmt.Errorf("icicle plonk: unsupported proving key type %T for BLS12-381", pk) + } + return icicle_bls12381.Prove(tccs, iciclePK, fullWitness, opts...) + case *cs_bw6761.SparseR1CS: + // Accept both ICICLE-wrapped and native proving keys; wrap if needed. + var iciclePK *icicle_bw6761.ProvingKey + switch t := pk.(type) { + case *icicle_bw6761.ProvingKey: + iciclePK = t + case *plonk_bw6761.ProvingKey: + // Wrap native proving key into ICICLE proving key; device buffers will be initialized lazily. + iciclePK = &icicle_bw6761.ProvingKey{ProvingKey: *t} + default: + return nil, fmt.Errorf("icicle plonk: unsupported proving key type %T for BW6-761", pk) + } + return icicle_bw6761.Prove(tccs, iciclePK, fullWitness, opts...) + default: + return nil, fmt.Errorf("icicle plonk: unsupported curve type") + } +} + +// Setup generates accelerated proving and verifying keys using the provided SRS. +func Setup(ccs constraint.ConstraintSystem, srs, srsLagrange kzg.SRS) (native_plonk.ProvingKey, native_plonk.VerifyingKey, error) { + warmUpDevice() + switch tccs := ccs.(type) { + case *cs_bn254.SparseR1CS: + // Mirror groth16: wrap native Setup into an ICICLE friendly ProvingKey + var pk icicle_bn254.ProvingKey + vk := new(plonk_bn254.VerifyingKey) + _nativePk, _vk, err := plonk_bn254.Setup(tccs, *srs.(*kzg_bn254.SRS), *srsLagrange.(*kzg_bn254.SRS)) + if err != nil { + return nil, nil, err + } + pk.ProvingKey = *_nativePk + *vk = *_vk + return &pk, vk, nil + case *cs_bls12377.SparseR1CS: + // Mirror groth16: wrap native Setup into an ICICLE friendly ProvingKey + var pk icicle_bls12377.ProvingKey + vk := new(plonk_bls12377.VerifyingKey) + _nativePk, _vk, err := plonk_bls12377.Setup(tccs, *srs.(*kzg_bls12377.SRS), *srsLagrange.(*kzg_bls12377.SRS)) + if err != nil { + return nil, nil, err + } + pk.ProvingKey = *_nativePk + *vk = *_vk + return &pk, vk, nil + case *cs_bls12381.SparseR1CS: + // Mirror groth16: wrap native Setup into an ICICLE friendly ProvingKey + var pk icicle_bls12381.ProvingKey + vk := new(plonk_bls12381.VerifyingKey) + _nativePk, _vk, err := plonk_bls12381.Setup(tccs, *srs.(*kzg_bls12381.SRS), *srsLagrange.(*kzg_bls12381.SRS)) + if err != nil { + return nil, nil, err + } + pk.ProvingKey = *_nativePk + *vk = *_vk + return &pk, vk, nil + case *cs_bw6761.SparseR1CS: + // Mirror groth16: wrap native Setup into an ICICLE friendly ProvingKey + var pk icicle_bw6761.ProvingKey + vk := new(plonk_bw6761.VerifyingKey) + _nativePk, _vk, err := plonk_bw6761.Setup(tccs, *srs.(*kzg_bw6761.SRS), *srsLagrange.(*kzg_bw6761.SRS)) + if err != nil { + return nil, nil, err + } + pk.ProvingKey = *_nativePk + *vk = *_vk + return &pk, vk, nil + default: + return nil, nil, fmt.Errorf("icicle plonk: unsupported curve type") + } +} + +// NewProvingKey creates an empty proving key for deserialization for supported curves. +func NewProvingKey(curveID ecc.ID) native_plonk.ProvingKey { + switch curveID { + case ecc.BN254: + return &icicle_bn254.ProvingKey{} + case ecc.BLS12_377: + return &icicle_bls12377.ProvingKey{} + case ecc.BLS12_381: + return &icicle_bls12381.ProvingKey{} + case ecc.BW6_761: + return &icicle_bw6761.ProvingKey{} + default: + panic("icicle plonk: unsupported curve") + } +} diff --git a/backend/accelerated/icicle/plonk/plonk_noicicle.go b/backend/accelerated/icicle/plonk/plonk_noicicle.go new file mode 100644 index 0000000000..a26e0f1eff --- /dev/null +++ b/backend/accelerated/icicle/plonk/plonk_noicicle.go @@ -0,0 +1,30 @@ +//go:build !icicle + +package plonk + +import ( + "github.com/consensys/gnark-crypto/ecc" + "github.com/consensys/gnark-crypto/kzg" + "github.com/consensys/gnark/backend" + native_plonk "github.com/consensys/gnark/backend/plonk" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" +) + +// Prove falls back to the native CPU PLONK prover when compiled without the +// 'icicle' build tag. +func Prove(ccs constraint.ConstraintSystem, pk native_plonk.ProvingKey, fullWitness witness.Witness, opts ...backend.ProverOption) (native_plonk.Proof, error) { + return native_plonk.Prove(ccs, pk, fullWitness, opts...) +} + +// Setup falls back to the native CPU PLONK setup when compiled without the +// 'icicle' build tag. +func Setup(ccs constraint.ConstraintSystem, srs, srsLagrange kzg.SRS) (native_plonk.ProvingKey, native_plonk.VerifyingKey, error) { + return native_plonk.Setup(ccs, srs, srsLagrange) +} + +// NewProvingKey falls back to the native PLONK proving key when compiled +// without the 'icicle' build tag. +func NewProvingKey(curveID ecc.ID) native_plonk.ProvingKey { + return native_plonk.NewProvingKey(curveID) +} diff --git a/backend/backend.go b/backend/backend.go index 5bc86fd142..856f5dadc1 100644 --- a/backend/backend.go +++ b/backend/backend.go @@ -57,11 +57,12 @@ type ProverOption func(*ProverConfig) error // ProverConfig is the configuration for the prover with the options applied. type ProverConfig struct { - SolverOpts []solver.Option - HashToFieldFn hash.Hash - ChallengeHash hash.Hash - KZGFoldingHash hash.Hash - StatisticalZK bool + SolverOpts []solver.Option + HashToFieldFn hash.Hash + ChallengeHash hash.Hash + KZGFoldingHash hash.Hash + StatisticalZK bool + SolutionCachePath string // if non-empty, path to cache/load solver solution } // NewProverConfig returns a default ProverConfig with given prover options opts @@ -137,6 +138,22 @@ func WithIcicleAcceleration() ProverOption { } } +// WithSolutionCachePath sets a file path used to cache the solver solution. +// It is currently honored only by the ICICLE-accelerated Groth16 backend +// (backend/accelerated/icicle/groth16). On the first prove call the solver +// runs normally and the result is written to the file. Subsequent calls with +// the same path load the cached solution (validated against the circuit's +// wire count) and skip the solver entirely. The caller is responsible for +// deleting the cache when the witness or circuit changes. +// Caching is automatically disabled when BSB22 commitments are present (the +// solver has side effects that cannot be replayed from cache). +func WithSolutionCachePath(path string) ProverOption { + return func(pc *ProverConfig) error { + pc.SolutionCachePath = path + return nil + } +} + // WithStatisticalZeroKnowledge ensures that statistical zero knowledgeness is achieved. // This option makes the prover more memory costly, as there are 3 more size n (size of the circuit) // allocations. diff --git a/backend/plonk/bench_impl_icicle_test.go b/backend/plonk/bench_impl_icicle_test.go new file mode 100644 index 0000000000..a595422381 --- /dev/null +++ b/backend/plonk/bench_impl_icicle_test.go @@ -0,0 +1,19 @@ +//go:build icicle + +package plonk_test + +import ( + "github.com/consensys/gnark-crypto/kzg" + accplonk "github.com/consensys/gnark/backend/accelerated/icicle/plonk" + "github.com/consensys/gnark/backend/plonk" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" +) + +func benchSetup(ccs constraint.ConstraintSystem, srs, srsLagrange kzg.SRS) (plonk.ProvingKey, plonk.VerifyingKey, error) { + return accplonk.Setup(ccs, srs, srsLagrange) +} + +func benchProve(ccs constraint.ConstraintSystem, pk plonk.ProvingKey, w witness.Witness) (plonk.Proof, error) { + return accplonk.Prove(ccs, pk, w) +} diff --git a/backend/plonk/bench_impl_native_test.go b/backend/plonk/bench_impl_native_test.go new file mode 100644 index 0000000000..a82793b543 --- /dev/null +++ b/backend/plonk/bench_impl_native_test.go @@ -0,0 +1,18 @@ +//go:build !icicle + +package plonk_test + +import ( + "github.com/consensys/gnark-crypto/kzg" + "github.com/consensys/gnark/backend/plonk" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" +) + +func benchSetup(ccs constraint.ConstraintSystem, srs, srsLagrange kzg.SRS) (plonk.ProvingKey, plonk.VerifyingKey, error) { + return plonk.Setup(ccs, srs, srsLagrange) +} + +func benchProve(ccs constraint.ConstraintSystem, pk plonk.ProvingKey, w witness.Witness) (plonk.Proof, error) { + return plonk.Prove(ccs, pk, w) +} diff --git a/backend/plonk/plonk_test.go b/backend/plonk/plonk_test.go index 4fc939dd58..873e1d559a 100644 --- a/backend/plonk/plonk_test.go +++ b/backend/plonk/plonk_test.go @@ -211,13 +211,13 @@ func BenchmarkProver(b *testing.B) { if err != nil { b.Fatal(err) } - pk, _, err := plonk.Setup(ccs, srs, srsLagrange) + pk, _, err := benchSetup(ccs, srs, srsLagrange) if err != nil { b.Fatal(err) } b.ResetTimer() for i := 0; i < b.N; i++ { - _, _ = plonk.Prove(ccs, pk, fullWitness) + _, _ = benchProve(ccs, pk, fullWitness) } }) } diff --git a/constraint/bls12-377/solution_cache.go b/constraint/bls12-377/solution_cache.go new file mode 100644 index 0000000000..ea6216de9a --- /dev/null +++ b/constraint/bls12-377/solution_cache.go @@ -0,0 +1,254 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +package cs + +// Package-level helpers for caching solver output to disk, used by the +// ICICLE-accelerated backends (backend/accelerated/icicle/...) to skip +// redundant solver runs when proving the same (circuit, witness) repeatedly. + +import ( + "encoding/binary" + "fmt" + "io" + "os" + "path/filepath" + "time" + + "github.com/consensys/gnark-crypto/ecc" + "github.com/consensys/gnark-crypto/ecc/bls12-377/fr" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" + csolver "github.com/consensys/gnark/constraint/solver" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + "unsafe" +) + +// evaluateLROSmallDomain extracts the solver l, r, o, and returns it in lagrange form. +// solver = [ public | secret | internal ] +func evaluateLROSmallDomain(cs *system, solution []fr.Element) ([]fr.Element, []fr.Element, []fr.Element) { + s := cs.GetNbConstraints() + len(cs.Public) // len(spr.Public) is for the placeholder constraints + s = int(ecc.NextPowerOfTwo(uint64(s))) + + var l, r, o []fr.Element + l = make([]fr.Element, s, s+4) // +4 to leave room for the blinding in plonk + r = make([]fr.Element, s, s+4) + o = make([]fr.Element, s, s+4) + s0 := solution[0] + + utils.Parallelize(len(cs.Public), func(start, end int) { + for i := start; i < end; i++ { // placeholders + l[i] = solution[i] + r[i] = s0 + o[i] = s0 + } + }) + offset := len(cs.Public) + nbConstraints := cs.GetNbConstraints() + + var sparseR1C constraint.SparseR1C + j := 0 + for _, inst := range cs.Instructions { + blueprint := cs.Blueprints[inst.BlueprintID] + if bc, ok := blueprint.(constraint.BlueprintSparseR1C); ok { + bc.DecompressSparseR1C(&sparseR1C, inst.Unpack(&cs.System)) + + l[offset+j] = solution[sparseR1C.XA] + r[offset+j] = solution[sparseR1C.XB] + o[offset+j] = solution[sparseR1C.XC] + j++ + } + } + + offset += nbConstraints + + padding := s - offset + utils.Parallelize(padding, func(start, end int) { + for i := start; i < end; i++ { // offset to reach 2**n constraints (where the id of l,r,o is 0, so we assign solver[0]) + j := offset + i + l[j] = s0 + r[j] = s0 + o[j] = s0 + } + }) + + return l, r, o +} + +// EvaluateLROSmallDomainFromValues derives the L, R, O Lagrange evaluations from a +// pre-computed full wire-value vector (as produced by SolveAndSaveRawValues / +// LoadRawSolverValues). It validates that the vector length matches the system. +func (cs *system) EvaluateLROSmallDomainFromValues(values []fr.Element) ([]fr.Element, []fr.Element, []fr.Element, error) { + expected := cs.GetNbPublicVariables() + cs.GetNbSecretVariables() + cs.GetNbInternalVariables() + if len(values) != expected { + return nil, nil, nil, fmt.Errorf("wire-value vector length mismatch: got %d, expected %d (stale cache?)", len(values), expected) + } + l, r, o := evaluateLROSmallDomain(cs, values) + return l, r, o, nil +} + +// SolveAndSaveRawValues behaves like Solve but, when rawCachePath is non-empty, +// additionally writes the solver's full wire-value vector to that path so that +// subsequent runs can skip the solver via LoadRawSolverValues + +// EvaluateLROSmallDomainFromValues. The caller is responsible for invalidating +// the cache when the circuit or witness changes. +func (cs *system) SolveAndSaveRawValues(witness witness.Witness, rawCachePath string, opts ...csolver.Option) (any, error) { + log := logger.Logger().With().Int("nbConstraints", cs.GetNbConstraints()).Logger() + start := time.Now() + + v := witness.Vector().(fr.Vector) + + // init the solver + solver, err := newSolver(cs, v, opts...) + if err != nil { + log.Err(err).Send() + return nil, err + } + + // reset the stateful blueprints + for i := range cs.Blueprints { + if b, ok := cs.Blueprints[i].(constraint.BlueprintStateful[constraint.U64]); ok { + b.Reset() + } + } + + // defer log printing once all solver.values are computed + // (or sooner, if a constraint is not satisfied) + defer solver.printLogs(cs.Logs) + + // run it. + if err := solver.run(); err != nil { + log.Err(err).Send() + return nil, err + } + + log.Debug().Dur("took", time.Since(start)).Msg("constraint system solver done") + + if rawCachePath != "" { + if err := SaveRawSolverValues(rawCachePath, solver.values); err != nil { + log.Warn().Err(err).Msg("failed to save raw solver values cache") + } else { + log.Debug().Str("file", rawCachePath).Int("wires", len(solver.values)).Msg("saved raw solver values cache") + } + } + + // format the solution + if cs.Type == constraint.SystemR1CS { + var res R1CSSolution + res.W = solver.values + res.A = solver.a + res.B = solver.b + res.C = solver.c + return &res, nil + } else { + // sparse R1CS: solver fills L,R,O during solving. + var res SparseR1CSSolution + res.L = solver.l + res.R = solver.r + res.O = solver.o + return &res, nil + } +} + +// LoadR1CSSolution reads a cached R1CSSolution (Groth16) from disk. +func LoadR1CSSolution(path string) (*R1CSSolution, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + var sol R1CSSolution + if _, err := sol.ReadFrom(f); err != nil { + return nil, fmt.Errorf("read cached Groth16 solution: %w", err) + } + return &sol, nil +} + +// SaveR1CSSolution writes a R1CSSolution (Groth16) to disk atomically. +func SaveR1CSSolution(path string, solution *R1CSSolution) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".groth16_solution_cache_*") + if err != nil { + return err + } + tmpName := tmp.Name() + if _, err := solution.WriteTo(tmp); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, +// no per-element conversion) with a little-endian uint64 length prefix. +// The format is not portable across architectures with different endianness. +func SaveRawSolverValues(path string, values []fr.Element) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".raw_solver_*") + if err != nil { + return err + } + tmpName := tmp.Name() + + nWires := uint64(len(values)) + if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if nWires > 0 { + // Bulk write: cast []fr.Element to []byte + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := tmp.Write(byteSlice); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. +// The declared length is validated against the file size before allocating. +func LoadRawSolverValues(path string) ([]fr.Element, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + fi, err := f.Stat() + if err != nil { + return nil, err + } + + var nWires uint64 + if err := binary.Read(f, binary.LittleEndian, &nWires); err != nil { + return nil, fmt.Errorf("read wire count: %w", err) + } + expectedSize := int64(8) + int64(nWires)*int64(fr.Bytes) + if fi.Size() != expectedSize { + return nil, fmt.Errorf("raw solver cache size mismatch: file is %d bytes, header declares %d wires (%d bytes)", + fi.Size(), nWires, expectedSize) + } + values := make([]fr.Element, nWires) + if nWires > 0 { + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := io.ReadFull(f, byteSlice); err != nil { + return nil, fmt.Errorf("read wires: %w", err) + } + } + return values, nil +} diff --git a/constraint/bls12-381/solution_cache.go b/constraint/bls12-381/solution_cache.go new file mode 100644 index 0000000000..8ae16167e6 --- /dev/null +++ b/constraint/bls12-381/solution_cache.go @@ -0,0 +1,254 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +package cs + +// Package-level helpers for caching solver output to disk, used by the +// ICICLE-accelerated backends (backend/accelerated/icicle/...) to skip +// redundant solver runs when proving the same (circuit, witness) repeatedly. + +import ( + "encoding/binary" + "fmt" + "io" + "os" + "path/filepath" + "time" + + "github.com/consensys/gnark-crypto/ecc" + "github.com/consensys/gnark-crypto/ecc/bls12-381/fr" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" + csolver "github.com/consensys/gnark/constraint/solver" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + "unsafe" +) + +// evaluateLROSmallDomain extracts the solver l, r, o, and returns it in lagrange form. +// solver = [ public | secret | internal ] +func evaluateLROSmallDomain(cs *system, solution []fr.Element) ([]fr.Element, []fr.Element, []fr.Element) { + s := cs.GetNbConstraints() + len(cs.Public) // len(spr.Public) is for the placeholder constraints + s = int(ecc.NextPowerOfTwo(uint64(s))) + + var l, r, o []fr.Element + l = make([]fr.Element, s, s+4) // +4 to leave room for the blinding in plonk + r = make([]fr.Element, s, s+4) + o = make([]fr.Element, s, s+4) + s0 := solution[0] + + utils.Parallelize(len(cs.Public), func(start, end int) { + for i := start; i < end; i++ { // placeholders + l[i] = solution[i] + r[i] = s0 + o[i] = s0 + } + }) + offset := len(cs.Public) + nbConstraints := cs.GetNbConstraints() + + var sparseR1C constraint.SparseR1C + j := 0 + for _, inst := range cs.Instructions { + blueprint := cs.Blueprints[inst.BlueprintID] + if bc, ok := blueprint.(constraint.BlueprintSparseR1C); ok { + bc.DecompressSparseR1C(&sparseR1C, inst.Unpack(&cs.System)) + + l[offset+j] = solution[sparseR1C.XA] + r[offset+j] = solution[sparseR1C.XB] + o[offset+j] = solution[sparseR1C.XC] + j++ + } + } + + offset += nbConstraints + + padding := s - offset + utils.Parallelize(padding, func(start, end int) { + for i := start; i < end; i++ { // offset to reach 2**n constraints (where the id of l,r,o is 0, so we assign solver[0]) + j := offset + i + l[j] = s0 + r[j] = s0 + o[j] = s0 + } + }) + + return l, r, o +} + +// EvaluateLROSmallDomainFromValues derives the L, R, O Lagrange evaluations from a +// pre-computed full wire-value vector (as produced by SolveAndSaveRawValues / +// LoadRawSolverValues). It validates that the vector length matches the system. +func (cs *system) EvaluateLROSmallDomainFromValues(values []fr.Element) ([]fr.Element, []fr.Element, []fr.Element, error) { + expected := cs.GetNbPublicVariables() + cs.GetNbSecretVariables() + cs.GetNbInternalVariables() + if len(values) != expected { + return nil, nil, nil, fmt.Errorf("wire-value vector length mismatch: got %d, expected %d (stale cache?)", len(values), expected) + } + l, r, o := evaluateLROSmallDomain(cs, values) + return l, r, o, nil +} + +// SolveAndSaveRawValues behaves like Solve but, when rawCachePath is non-empty, +// additionally writes the solver's full wire-value vector to that path so that +// subsequent runs can skip the solver via LoadRawSolverValues + +// EvaluateLROSmallDomainFromValues. The caller is responsible for invalidating +// the cache when the circuit or witness changes. +func (cs *system) SolveAndSaveRawValues(witness witness.Witness, rawCachePath string, opts ...csolver.Option) (any, error) { + log := logger.Logger().With().Int("nbConstraints", cs.GetNbConstraints()).Logger() + start := time.Now() + + v := witness.Vector().(fr.Vector) + + // init the solver + solver, err := newSolver(cs, v, opts...) + if err != nil { + log.Err(err).Send() + return nil, err + } + + // reset the stateful blueprints + for i := range cs.Blueprints { + if b, ok := cs.Blueprints[i].(constraint.BlueprintStateful[constraint.U64]); ok { + b.Reset() + } + } + + // defer log printing once all solver.values are computed + // (or sooner, if a constraint is not satisfied) + defer solver.printLogs(cs.Logs) + + // run it. + if err := solver.run(); err != nil { + log.Err(err).Send() + return nil, err + } + + log.Debug().Dur("took", time.Since(start)).Msg("constraint system solver done") + + if rawCachePath != "" { + if err := SaveRawSolverValues(rawCachePath, solver.values); err != nil { + log.Warn().Err(err).Msg("failed to save raw solver values cache") + } else { + log.Debug().Str("file", rawCachePath).Int("wires", len(solver.values)).Msg("saved raw solver values cache") + } + } + + // format the solution + if cs.Type == constraint.SystemR1CS { + var res R1CSSolution + res.W = solver.values + res.A = solver.a + res.B = solver.b + res.C = solver.c + return &res, nil + } else { + // sparse R1CS: solver fills L,R,O during solving. + var res SparseR1CSSolution + res.L = solver.l + res.R = solver.r + res.O = solver.o + return &res, nil + } +} + +// LoadR1CSSolution reads a cached R1CSSolution (Groth16) from disk. +func LoadR1CSSolution(path string) (*R1CSSolution, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + var sol R1CSSolution + if _, err := sol.ReadFrom(f); err != nil { + return nil, fmt.Errorf("read cached Groth16 solution: %w", err) + } + return &sol, nil +} + +// SaveR1CSSolution writes a R1CSSolution (Groth16) to disk atomically. +func SaveR1CSSolution(path string, solution *R1CSSolution) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".groth16_solution_cache_*") + if err != nil { + return err + } + tmpName := tmp.Name() + if _, err := solution.WriteTo(tmp); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, +// no per-element conversion) with a little-endian uint64 length prefix. +// The format is not portable across architectures with different endianness. +func SaveRawSolverValues(path string, values []fr.Element) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".raw_solver_*") + if err != nil { + return err + } + tmpName := tmp.Name() + + nWires := uint64(len(values)) + if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if nWires > 0 { + // Bulk write: cast []fr.Element to []byte + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := tmp.Write(byteSlice); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. +// The declared length is validated against the file size before allocating. +func LoadRawSolverValues(path string) ([]fr.Element, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + fi, err := f.Stat() + if err != nil { + return nil, err + } + + var nWires uint64 + if err := binary.Read(f, binary.LittleEndian, &nWires); err != nil { + return nil, fmt.Errorf("read wire count: %w", err) + } + expectedSize := int64(8) + int64(nWires)*int64(fr.Bytes) + if fi.Size() != expectedSize { + return nil, fmt.Errorf("raw solver cache size mismatch: file is %d bytes, header declares %d wires (%d bytes)", + fi.Size(), nWires, expectedSize) + } + values := make([]fr.Element, nWires) + if nWires > 0 { + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := io.ReadFull(f, byteSlice); err != nil { + return nil, fmt.Errorf("read wires: %w", err) + } + } + return values, nil +} diff --git a/constraint/bn254/solution_cache.go b/constraint/bn254/solution_cache.go new file mode 100644 index 0000000000..3052ad7295 --- /dev/null +++ b/constraint/bn254/solution_cache.go @@ -0,0 +1,254 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +package cs + +// Package-level helpers for caching solver output to disk, used by the +// ICICLE-accelerated backends (backend/accelerated/icicle/...) to skip +// redundant solver runs when proving the same (circuit, witness) repeatedly. + +import ( + "encoding/binary" + "fmt" + "io" + "os" + "path/filepath" + "time" + + "github.com/consensys/gnark-crypto/ecc" + "github.com/consensys/gnark-crypto/ecc/bn254/fr" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" + csolver "github.com/consensys/gnark/constraint/solver" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + "unsafe" +) + +// evaluateLROSmallDomain extracts the solver l, r, o, and returns it in lagrange form. +// solver = [ public | secret | internal ] +func evaluateLROSmallDomain(cs *system, solution []fr.Element) ([]fr.Element, []fr.Element, []fr.Element) { + s := cs.GetNbConstraints() + len(cs.Public) // len(spr.Public) is for the placeholder constraints + s = int(ecc.NextPowerOfTwo(uint64(s))) + + var l, r, o []fr.Element + l = make([]fr.Element, s, s+4) // +4 to leave room for the blinding in plonk + r = make([]fr.Element, s, s+4) + o = make([]fr.Element, s, s+4) + s0 := solution[0] + + utils.Parallelize(len(cs.Public), func(start, end int) { + for i := start; i < end; i++ { // placeholders + l[i] = solution[i] + r[i] = s0 + o[i] = s0 + } + }) + offset := len(cs.Public) + nbConstraints := cs.GetNbConstraints() + + var sparseR1C constraint.SparseR1C + j := 0 + for _, inst := range cs.Instructions { + blueprint := cs.Blueprints[inst.BlueprintID] + if bc, ok := blueprint.(constraint.BlueprintSparseR1C); ok { + bc.DecompressSparseR1C(&sparseR1C, inst.Unpack(&cs.System)) + + l[offset+j] = solution[sparseR1C.XA] + r[offset+j] = solution[sparseR1C.XB] + o[offset+j] = solution[sparseR1C.XC] + j++ + } + } + + offset += nbConstraints + + padding := s - offset + utils.Parallelize(padding, func(start, end int) { + for i := start; i < end; i++ { // offset to reach 2**n constraints (where the id of l,r,o is 0, so we assign solver[0]) + j := offset + i + l[j] = s0 + r[j] = s0 + o[j] = s0 + } + }) + + return l, r, o +} + +// EvaluateLROSmallDomainFromValues derives the L, R, O Lagrange evaluations from a +// pre-computed full wire-value vector (as produced by SolveAndSaveRawValues / +// LoadRawSolverValues). It validates that the vector length matches the system. +func (cs *system) EvaluateLROSmallDomainFromValues(values []fr.Element) ([]fr.Element, []fr.Element, []fr.Element, error) { + expected := cs.GetNbPublicVariables() + cs.GetNbSecretVariables() + cs.GetNbInternalVariables() + if len(values) != expected { + return nil, nil, nil, fmt.Errorf("wire-value vector length mismatch: got %d, expected %d (stale cache?)", len(values), expected) + } + l, r, o := evaluateLROSmallDomain(cs, values) + return l, r, o, nil +} + +// SolveAndSaveRawValues behaves like Solve but, when rawCachePath is non-empty, +// additionally writes the solver's full wire-value vector to that path so that +// subsequent runs can skip the solver via LoadRawSolverValues + +// EvaluateLROSmallDomainFromValues. The caller is responsible for invalidating +// the cache when the circuit or witness changes. +func (cs *system) SolveAndSaveRawValues(witness witness.Witness, rawCachePath string, opts ...csolver.Option) (any, error) { + log := logger.Logger().With().Int("nbConstraints", cs.GetNbConstraints()).Logger() + start := time.Now() + + v := witness.Vector().(fr.Vector) + + // init the solver + solver, err := newSolver(cs, v, opts...) + if err != nil { + log.Err(err).Send() + return nil, err + } + + // reset the stateful blueprints + for i := range cs.Blueprints { + if b, ok := cs.Blueprints[i].(constraint.BlueprintStateful[constraint.U64]); ok { + b.Reset() + } + } + + // defer log printing once all solver.values are computed + // (or sooner, if a constraint is not satisfied) + defer solver.printLogs(cs.Logs) + + // run it. + if err := solver.run(); err != nil { + log.Err(err).Send() + return nil, err + } + + log.Debug().Dur("took", time.Since(start)).Msg("constraint system solver done") + + if rawCachePath != "" { + if err := SaveRawSolverValues(rawCachePath, solver.values); err != nil { + log.Warn().Err(err).Msg("failed to save raw solver values cache") + } else { + log.Debug().Str("file", rawCachePath).Int("wires", len(solver.values)).Msg("saved raw solver values cache") + } + } + + // format the solution + if cs.Type == constraint.SystemR1CS { + var res R1CSSolution + res.W = solver.values + res.A = solver.a + res.B = solver.b + res.C = solver.c + return &res, nil + } else { + // sparse R1CS: solver fills L,R,O during solving. + var res SparseR1CSSolution + res.L = solver.l + res.R = solver.r + res.O = solver.o + return &res, nil + } +} + +// LoadR1CSSolution reads a cached R1CSSolution (Groth16) from disk. +func LoadR1CSSolution(path string) (*R1CSSolution, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + var sol R1CSSolution + if _, err := sol.ReadFrom(f); err != nil { + return nil, fmt.Errorf("read cached Groth16 solution: %w", err) + } + return &sol, nil +} + +// SaveR1CSSolution writes a R1CSSolution (Groth16) to disk atomically. +func SaveR1CSSolution(path string, solution *R1CSSolution) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".groth16_solution_cache_*") + if err != nil { + return err + } + tmpName := tmp.Name() + if _, err := solution.WriteTo(tmp); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, +// no per-element conversion) with a little-endian uint64 length prefix. +// The format is not portable across architectures with different endianness. +func SaveRawSolverValues(path string, values []fr.Element) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".raw_solver_*") + if err != nil { + return err + } + tmpName := tmp.Name() + + nWires := uint64(len(values)) + if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if nWires > 0 { + // Bulk write: cast []fr.Element to []byte + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := tmp.Write(byteSlice); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. +// The declared length is validated against the file size before allocating. +func LoadRawSolverValues(path string) ([]fr.Element, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + fi, err := f.Stat() + if err != nil { + return nil, err + } + + var nWires uint64 + if err := binary.Read(f, binary.LittleEndian, &nWires); err != nil { + return nil, fmt.Errorf("read wire count: %w", err) + } + expectedSize := int64(8) + int64(nWires)*int64(fr.Bytes) + if fi.Size() != expectedSize { + return nil, fmt.Errorf("raw solver cache size mismatch: file is %d bytes, header declares %d wires (%d bytes)", + fi.Size(), nWires, expectedSize) + } + values := make([]fr.Element, nWires) + if nWires > 0 { + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := io.ReadFull(f, byteSlice); err != nil { + return nil, fmt.Errorf("read wires: %w", err) + } + } + return values, nil +} diff --git a/constraint/bw6-761/solution_cache.go b/constraint/bw6-761/solution_cache.go new file mode 100644 index 0000000000..12e48e99eb --- /dev/null +++ b/constraint/bw6-761/solution_cache.go @@ -0,0 +1,254 @@ +// Copyright 2025-2026 Consensys Software Inc. +// Licensed under the Apache License, Version 2.0. See the LICENSE file for details. + +// Code generated by gnark DO NOT EDIT + +package cs + +// Package-level helpers for caching solver output to disk, used by the +// ICICLE-accelerated backends (backend/accelerated/icicle/...) to skip +// redundant solver runs when proving the same (circuit, witness) repeatedly. + +import ( + "encoding/binary" + "fmt" + "io" + "os" + "path/filepath" + "time" + + "github.com/consensys/gnark-crypto/ecc" + "github.com/consensys/gnark-crypto/ecc/bw6-761/fr" + "github.com/consensys/gnark/backend/witness" + "github.com/consensys/gnark/constraint" + csolver "github.com/consensys/gnark/constraint/solver" + "github.com/consensys/gnark/internal/utils" + "github.com/consensys/gnark/logger" + "unsafe" +) + +// evaluateLROSmallDomain extracts the solver l, r, o, and returns it in lagrange form. +// solver = [ public | secret | internal ] +func evaluateLROSmallDomain(cs *system, solution []fr.Element) ([]fr.Element, []fr.Element, []fr.Element) { + s := cs.GetNbConstraints() + len(cs.Public) // len(spr.Public) is for the placeholder constraints + s = int(ecc.NextPowerOfTwo(uint64(s))) + + var l, r, o []fr.Element + l = make([]fr.Element, s, s+4) // +4 to leave room for the blinding in plonk + r = make([]fr.Element, s, s+4) + o = make([]fr.Element, s, s+4) + s0 := solution[0] + + utils.Parallelize(len(cs.Public), func(start, end int) { + for i := start; i < end; i++ { // placeholders + l[i] = solution[i] + r[i] = s0 + o[i] = s0 + } + }) + offset := len(cs.Public) + nbConstraints := cs.GetNbConstraints() + + var sparseR1C constraint.SparseR1C + j := 0 + for _, inst := range cs.Instructions { + blueprint := cs.Blueprints[inst.BlueprintID] + if bc, ok := blueprint.(constraint.BlueprintSparseR1C); ok { + bc.DecompressSparseR1C(&sparseR1C, inst.Unpack(&cs.System)) + + l[offset+j] = solution[sparseR1C.XA] + r[offset+j] = solution[sparseR1C.XB] + o[offset+j] = solution[sparseR1C.XC] + j++ + } + } + + offset += nbConstraints + + padding := s - offset + utils.Parallelize(padding, func(start, end int) { + for i := start; i < end; i++ { // offset to reach 2**n constraints (where the id of l,r,o is 0, so we assign solver[0]) + j := offset + i + l[j] = s0 + r[j] = s0 + o[j] = s0 + } + }) + + return l, r, o +} + +// EvaluateLROSmallDomainFromValues derives the L, R, O Lagrange evaluations from a +// pre-computed full wire-value vector (as produced by SolveAndSaveRawValues / +// LoadRawSolverValues). It validates that the vector length matches the system. +func (cs *system) EvaluateLROSmallDomainFromValues(values []fr.Element) ([]fr.Element, []fr.Element, []fr.Element, error) { + expected := cs.GetNbPublicVariables() + cs.GetNbSecretVariables() + cs.GetNbInternalVariables() + if len(values) != expected { + return nil, nil, nil, fmt.Errorf("wire-value vector length mismatch: got %d, expected %d (stale cache?)", len(values), expected) + } + l, r, o := evaluateLROSmallDomain(cs, values) + return l, r, o, nil +} + +// SolveAndSaveRawValues behaves like Solve but, when rawCachePath is non-empty, +// additionally writes the solver's full wire-value vector to that path so that +// subsequent runs can skip the solver via LoadRawSolverValues + +// EvaluateLROSmallDomainFromValues. The caller is responsible for invalidating +// the cache when the circuit or witness changes. +func (cs *system) SolveAndSaveRawValues(witness witness.Witness, rawCachePath string, opts ...csolver.Option) (any, error) { + log := logger.Logger().With().Int("nbConstraints", cs.GetNbConstraints()).Logger() + start := time.Now() + + v := witness.Vector().(fr.Vector) + + // init the solver + solver, err := newSolver(cs, v, opts...) + if err != nil { + log.Err(err).Send() + return nil, err + } + + // reset the stateful blueprints + for i := range cs.Blueprints { + if b, ok := cs.Blueprints[i].(constraint.BlueprintStateful[constraint.U64]); ok { + b.Reset() + } + } + + // defer log printing once all solver.values are computed + // (or sooner, if a constraint is not satisfied) + defer solver.printLogs(cs.Logs) + + // run it. + if err := solver.run(); err != nil { + log.Err(err).Send() + return nil, err + } + + log.Debug().Dur("took", time.Since(start)).Msg("constraint system solver done") + + if rawCachePath != "" { + if err := SaveRawSolverValues(rawCachePath, solver.values); err != nil { + log.Warn().Err(err).Msg("failed to save raw solver values cache") + } else { + log.Debug().Str("file", rawCachePath).Int("wires", len(solver.values)).Msg("saved raw solver values cache") + } + } + + // format the solution + if cs.Type == constraint.SystemR1CS { + var res R1CSSolution + res.W = solver.values + res.A = solver.a + res.B = solver.b + res.C = solver.c + return &res, nil + } else { + // sparse R1CS: solver fills L,R,O during solving. + var res SparseR1CSSolution + res.L = solver.l + res.R = solver.r + res.O = solver.o + return &res, nil + } +} + +// LoadR1CSSolution reads a cached R1CSSolution (Groth16) from disk. +func LoadR1CSSolution(path string) (*R1CSSolution, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + var sol R1CSSolution + if _, err := sol.ReadFrom(f); err != nil { + return nil, fmt.Errorf("read cached Groth16 solution: %w", err) + } + return &sol, nil +} + +// SaveR1CSSolution writes a R1CSSolution (Groth16) to disk atomically. +func SaveR1CSSolution(path string, solution *R1CSSolution) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".groth16_solution_cache_*") + if err != nil { + return err + } + tmpName := tmp.Name() + if _, err := solution.WriteTo(tmp); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, +// no per-element conversion) with a little-endian uint64 length prefix. +// The format is not portable across architectures with different endianness. +func SaveRawSolverValues(path string, values []fr.Element) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".raw_solver_*") + if err != nil { + return err + } + tmpName := tmp.Name() + + nWires := uint64(len(values)) + if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if nWires > 0 { + // Bulk write: cast []fr.Element to []byte + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := tmp.Write(byteSlice); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + return os.Rename(tmpName, path) +} + +// LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. +// The declared length is validated against the file size before allocating. +func LoadRawSolverValues(path string) ([]fr.Element, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + fi, err := f.Stat() + if err != nil { + return nil, err + } + + var nWires uint64 + if err := binary.Read(f, binary.LittleEndian, &nWires); err != nil { + return nil, fmt.Errorf("read wire count: %w", err) + } + expectedSize := int64(8) + int64(nWires)*int64(fr.Bytes) + if fi.Size() != expectedSize { + return nil, fmt.Errorf("raw solver cache size mismatch: file is %d bytes, header declares %d wires (%d bytes)", + fi.Size(), nWires, expectedSize) + } + values := make([]fr.Element, nWires) + if nWires > 0 { + byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) + if _, err := io.ReadFull(f, byteSlice); err != nil { + return nil, fmt.Errorf("read wires: %w", err) + } + } + return values, nil +} diff --git a/go.mod b/go.mod index d00f03828f..fd847540a4 100644 --- a/go.mod +++ b/go.mod @@ -46,3 +46,10 @@ tool ( github.com/klauspost/asmfmt/cmd/asmfmt golang.org/x/tools/cmd/goimports ) + +// TODO(martun): temporary local-path replace — the accelerated backends need +// vec-ops that only exist in the extended icicle-gnark fork (see +// https://github.com/ingonyama-zk/icicle-gnark/pull/4). This breaks builds +// without a sibling checkout of that fork and MUST be replaced with a tagged +// icicle-gnark release before this branch can merge upstream. +replace github.com/ingonyama-zk/icicle-gnark/v3 => ../icicle-gnark-extended From b1d7c379faf698a8b179a41da1e32b6a4f682d96 Mon Sep 17 00:00:00 2001 From: Martun Karapetyan Date: Mon, 27 Jul 2026 19:46:20 +0400 Subject: [PATCH 2/5] fix(icicle/plonk): BSB22 cache ZK guard + GPU resource lifecycle Zero-knowledge: disable the raw solver cache (GNARK_RAW_SOLVER_CACHE) for circuits with BSB22 commitments unless GNARK_DISABLE_BLINDING is set. Replaying the cached commitment polynomials reuses their random blinding rows across proofs, making Bsb22Commitments linkable and leaking linear relations over the committed private wires once the polynomial has been opened at more than two distinct zetas. Re-randomizing on load is not possible because the commitment hash feeds back into a witness wire. GPU resource lifecycle: - Pair every CUDA stream in the NTT fan-outs with synchronize + destroy on all paths (previously ~50+ streams leaked per prove). - Drain all NTT worker channels before returning on error so the caller cannot free device slices sibling workers are still writing to; return the scaling-vector pool slices on every exit path. - Track outstanding pool slices (keyed by device pointer; CopyToDevice mutates the slice length field) and reclaim buffers abandoned by error paths in a new pool Shutdown at Prove teardown; mid-prove FreeAll keeps its idle-only semantics. Backstop hGPU / linearizedPolynomialGPU / blindedZCanonicalGPU / polyZLagrangeGPU there as well. - Return the linearized polynomial to the pool instead of direct-freeing it (it is a pool buffer; the bypass surfaced as a double free once outstanding tracking existed). - Convert device-goroutine and prover-stage panics (GPU allocation failure is the expected failure mode) into Prove errors via recover boundaries: devicePanicToError on error-channel closures, inline recover in runErr defers, a goStage wrapper around the errgroup stages, and error returns from setupDevicePointers. - gpuEvaluateConstraints frees all temporaries via a cleanup defer on every path; blinding vectors now come from the temp pool. Adds icicle-tagged regression tests: commitment-free cache round-trip, cache skipped for BSB22 circuits with fresh commitments per prove, and BSB22 replay under GNARK_DISABLE_BLINDING. Verified on RTX PRO 6000 / CUDA 13: full -tags=icicle plonk suite passes on all four curves; gofmt/go vet clean with and without the tag; the checked-in generated files match the generator output. Co-Authored-By: Claude Fable 5 --- README.md | 11 +- .../generator/templates/plonk.icicle.go.tmpl | 445 ++++++++++++++---- .../icicle/plonk/bls12-377/icicle.go | 445 ++++++++++++++---- .../icicle/plonk/bls12-381/icicle.go | 445 ++++++++++++++---- .../accelerated/icicle/plonk/bn254/icicle.go | 445 ++++++++++++++---- .../icicle/plonk/bw6-761/icicle.go | 445 ++++++++++++++---- .../accelerated/icicle/plonk/cache_test.go | 176 +++++++ 7 files changed, 1946 insertions(+), 466 deletions(-) create mode 100644 backend/accelerated/icicle/plonk/cache_test.go diff --git a/README.md b/README.md index ea89b3364c..6a34545275 100644 --- a/README.md +++ b/README.md @@ -171,11 +171,20 @@ entirely. 1. **First run** -- the solver runs normally and writes every wire value (Montgomery form, `[]fr.Element`) to a binary file. BSB22 commitment polynomials are saved alongside. -2. **Subsequent runs** -- the cache file is memory-mapped back, the L/R/O +2. **Subsequent runs** -- the cache file is read back, the L/R/O Lagrange evaluations are derived via `evaluateLROSmallDomain`, and the BSB22 commitment is recomputed from the cached committed-wire values. The solver is never invoked. +**Zero-knowledge caveat**: for circuits with BSB22 commitments (anything +using `std/rangecheck` and friends), the cached commitment polynomials +include their random blinding rows. Replaying them across proofs would +make the commitments linkable and progressively leak the committed private +wires, so in the default (blinding-on) mode the cache is **automatically +disabled** for such circuits, with a warning. It stays available for them +when `GNARK_DISABLE_BLINDING` is set, i.e. when zero-knowledge has already +been explicitly traded away. + ### Usage Set the `GNARK_RAW_SOLVER_CACHE` environment variable to a file path diff --git a/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl index 04a8b4b7b3..14b5597a4c 100644 --- a/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl +++ b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl @@ -257,32 +257,47 @@ func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts defer instance.releaseSharedGPUState() defer instance.releaseLinearizedEvalGPUState() + // goStage runs a prover stage on the errgroup, converting a panic into an + // error: several device helpers panic on GPU allocation failure, and an + // unrecovered panic in a stage goroutine would kill the embedding process + // instead of failing the Prove call. + goStage := func(name string, fn func() error) { + g.Go(func() (err error) { + defer func() { + if r := recover(); r != nil && err == nil { + err = fmt.Errorf("icicle plonk: stage %s panicked: %v", name, r) + } + }() + return fn() + }) + } + // solve constraints - g.Go(instance.solveConstraints) + goStage("solveConstraints", instance.solveConstraints) // complete qk - g.Go(instance.completeQk) + goStage("completeQk", instance.completeQk) // init blinding polynomials - g.Go(instance.initBlindingPolynomials) + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) // derive gamma, beta (copy constraint) - g.Go(instance.deriveGammaAndBeta) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) // compute accumulating ratio for the copy constraint - g.Go(instance.buildRatioCopyConstraint) + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) // compute h - g.Go(instance.computeQuotient) + goStage("computeQuotient", instance.computeQuotient) // open Z (blinded) at ωζ (proof.ZShiftedOpening) - g.Go(instance.openZ) + goStage("openZ", instance.openZ) // linearized polynomial - g.Go(instance.computeLinearizedPolynomial) + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) // Batch opening (no internal timer of its own — time the whole stage here) - g.Go(func() error { + goStage("batchOpening", func() error { startBatchOpening := time.Now() err := instance.batchOpening() if isProfileMode { @@ -495,6 +510,21 @@ func (s *instance) solveConstraints() error { // Try to load raw solver values from cache (fastest path) rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + // The raw solver cache stores the BSB22 commitment polynomials verbatim, + // including their two random blinding rows (see bsb22Hint). Replaying them + // across proofs makes Bsb22Commitments identical (linkable) and, once the + // committed polynomial has been opened at more than two distinct zetas, + // leaks linear relations over the committed private wires. Re-randomizing + // on load is not possible either: the commitment hash feeds back into a + // witness wire, so fresh blinding would invalidate every cached wire + // downstream of it. When zero-knowledge matters (blinding enabled, the + // default), the cache is therefore disabled for circuits with BSB22 + // commitments. With GNARK_DISABLE_BLINDING set, zero-knowledge is already + // explicitly forfeited and the replay leaks nothing new. + if rawCachePath != "" && len(s.commitmentInfo) > 0 && useBlinding { + log.Warn().Str("file", rawCachePath).Msg("GNARK_RAW_SOLVER_CACHE ignored: replaying cached BSB22 commitment blinding across proofs would break zero-knowledge (set GNARK_DISABLE_BLINDING to opt out of zero-knowledge and cache anyway)") + rawCachePath = "" + } if rawCachePath != "" { if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { // Reconstruct L, R, O from raw values @@ -981,6 +1011,7 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone := make(chan error, 1) var dPersist icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(copyDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") if cfgErr != nil { copyDone <- cfgErr @@ -1071,6 +1102,7 @@ func (s *instance) openZ() (err error) { var dBlindedCanonical icicle_core.DeviceSlice buildDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(buildDone) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") if cfgErr != nil { buildDone <- cfgErr @@ -1396,7 +1428,9 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) doneAddZ() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + // dLin is a pool buffer: return it via Put, not a direct free, so the + // pool's outstanding tracking stays consistent. + s.putTempDeviceSlice(dLin, dLin.Len()) return err } @@ -1404,7 +1438,7 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) doneSubtractH() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + s.putTempDeviceSlice(dLin, dLin.Len()) return err } s.linearizedPolynomialGPU = dLin @@ -1445,7 +1479,12 @@ func (s *instance) batchOpening() error { } defer func() { - freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + // The linearized polynomial is a pool buffer: return it via Put, not + // a direct free, so the pool's outstanding tracking stays consistent. + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } if s.hGPU != nil { s.freeGPUQuotient(s.hGPU) s.hGPU = nil @@ -1529,6 +1568,7 @@ func (s *instance) batchOpening() error { var dFold icicle_core.DeviceSlice foldDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(foldDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") if cfgErr != nil { foldDone <- cfgErr @@ -1595,6 +1635,7 @@ func (s *instance) batchOpening() error { divDone := make(chan error, 1) witnessSize := dFold.Len() - 1 icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") if cfgErr != nil { divDone <- cfgErr @@ -1891,12 +1932,15 @@ func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) var dWitness icicle_core.DeviceSlice - uploadDone := make(chan struct{}, 1) + uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) dWitness = uploadVector(hCoeffs) - close(uploadDone) + uploadDone <- nil }) - <-uploadDone + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) freeSliceOnDevice(&dWitness, &s.device) if err != nil { @@ -1988,6 +2032,7 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") if cfgErr != nil { done <- cfgErr @@ -2276,6 +2321,7 @@ func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicl var dTwiddles0 icicle_core.DeviceSlice uploadTwiddlesDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadTwiddlesDone) if s.tempGPUMemPool != nil { s.tempGPUMemPool.FreeAll() } @@ -2348,6 +2394,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo batchInvertDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(batchInvertDone) batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) }) if err := <-batchInvertDone; err != nil { @@ -2583,6 +2630,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(mergeDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") if cfgErr != nil { mergeDone <- cfgErr @@ -2613,6 +2661,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cresHost := icicle_core.HostSliceFromElements(cres) downloadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(downloadDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") if cfgErr != nil { downloadDone <- cfgErr @@ -2818,6 +2867,12 @@ func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop. } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for snapshot cloning before handing state to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { runErr = syncErr @@ -2902,6 +2957,12 @@ func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } @@ -2954,14 +3015,29 @@ func (s *instance) releaseTempGPUMemoryPool() { if s == nil || s.tempGPUMemPool == nil { return } + // Backstop for instance-owned device buffers whose owning stage may have + // exited (via error or ctx cancellation) before registering its own + // cleanup defer: openZ owns the two Z buffers, batchOpening owns the + // quotient and the linearized polynomial. On the success path these are + // already released and zeroed, making every free below a no-op. This must + // run before FreeAll so pool-owned buffers (blindedZCanonicalGPU, hGPU) + // are back in the pool when it frees everything. freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) if !s.blindedZCanonicalGPU.IsEmpty() { s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} } + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { - s.tempGPUMemPool.FreeAll() + s.tempGPUMemPool.Shutdown() close(done) }) <-done @@ -3019,6 +3095,12 @@ func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysSta } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for GPU uploads in ensurePolysOnSharedGPU. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -3073,6 +3155,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali // Upload scaling vectors to GPU uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(uploadDone) scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) scalingHost := icicle_core.HostSliceFromElements(scalingVector) scalingHost.CopyToDevice(&scalingVectorDevice, false) @@ -3091,10 +3174,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } uploadDone <- nil }) + // Return the scaling vectors to the pool on every exit path. This is safe + // because we never return before draining every per-polynomial done + // channel, and each worker synchronizes its stream before signalling, so + // no in-flight kernel can still reference the slices. + defer func() { + if !scalingVectorDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + } + if !scalingVectorRevDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + } + }() if err := <-uploadDone; err != nil { return err } + // Validate all polynomials before launching any GPU work so an invalid + // entry cannot abandon already-scheduled workers. + for _, p := range state.polys { + if p != nil && p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + } + doneChans := make([]chan error, len(state.polys)) for i := range state.polys { @@ -3103,10 +3206,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali continue } - if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { - return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) - } - done := make(chan error, 1) doneChans[i] = done scalarsDevice := state.deviceSlices[i] @@ -3115,7 +3214,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: CreateStream failed: %s", eStream.AsString()) + return + } + // Synchronize and destroy the stream on every path (including + // errors) before signalling done: the caller may free device + // buffers as soon as all workers have reported. + var runErr error + defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: DestroyStream failed: %s", eDestroy.AsString()) + } + done <- runErr + }() cfg.StreamHandle = stream cfg.IsAsync = true @@ -3129,7 +3251,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KRN } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) return } @@ -3148,7 +3270,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) return } @@ -3161,26 +3283,26 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KNR // Regular → BitReverse } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) - return - } - - if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) return } - done <- nil }) } - // Wait for all scheduled tasks + // Wait for ALL scheduled tasks even when one fails: returning on the + // first error would let the caller free device slices that sibling + // workers are still writing to. + var firstErr error for i := range doneChans { if doneChans[i] != nil { - if err := <-doneChans[i]; err != nil { - return err + if err := <-doneChans[i]; err != nil && firstErr == nil { + firstErr = err } } } + if firstErr != nil { + return firstErr + } // Update polynomial metadata: final result is in Lagrange, same layout as original for _, p := range state.polys { @@ -3188,10 +3310,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali p.Basis = iop.Lagrange } } - - // Free scaling vectors from device - s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) - s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) return nil } @@ -3263,7 +3381,11 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: CreateStream failed: %s", eStream.AsString())) + } + defer icicle_runtime.DestroyStream(stream) cfg.StreamHandle = stream cfg.IsAsync = true @@ -3285,7 +3407,9 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) } - icicle_runtime.SynchronizeStream(stream) + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: SynchronizeStream failed: %s", eSync.AsString())) + } // Update metadata inside closure to avoid race p.Basis = iop.Canonical @@ -3334,13 +3458,20 @@ func (s *instance) freeGPUPolys(state *gpuPolysState) { // Must be used within RunOnDevice context to ensure thread safety per device. type gpuMemoryPool struct { freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice - mu sync.Mutex + // outstanding tracks slices handed out by Get and not yet returned so + // Shutdown can reclaim buffers abandoned by error paths. Keyed by the + // device pointer, not the slice value: CopyToDevice mutates the slice's + // length field, so the value at Put time may differ from the one at Get + // time while the underlying buffer is the same. + outstanding map[unsafe.Pointer]icicle_core.DeviceSlice + mu sync.Mutex } // newGPUMemoryPool creates a new GPU memory pool. func newGPUMemoryPool() *gpuMemoryPool { return &gpuMemoryPool{ - freeSlices: make(map[int][]icicle_core.DeviceSlice), + freeSlices: make(map[int][]icicle_core.DeviceSlice), + outstanding: make(map[unsafe.Pointer]icicle_core.DeviceSlice), } } @@ -3354,12 +3485,14 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { // Reuse the last slice slice := slices[len(slices)-1] p.freeSlices[n] = slices[:len(slices)-1] + p.outstanding[slice.AsUnsafePointer()] = slice return slice } // No free slice available, allocate a new one. // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3372,6 +3505,7 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3387,11 +3521,14 @@ func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { p.mu.Lock() defer p.mu.Unlock() + delete(p.outstanding, ds.AsUnsafePointer()) // Add to the pool p.freeSlices[n] = append(p.freeSlices[n], ds) } -// FreeAll releases all pooled device slices. +// FreeAll releases all idle pooled device slices. Outstanding slices (handed +// out by Get and not yet returned) are left alone: FreeAll is also used +// mid-prove to relieve memory pressure while pool buffers are still live. func (p *gpuMemoryPool) FreeAll() { p.mu.Lock() defer p.mu.Unlock() @@ -3404,6 +3541,27 @@ func (p *gpuMemoryPool) FreeAll() { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) } +// Shutdown releases every pooled slice, including outstanding ones that were +// abandoned by error paths (a failed stage returns without Put-ing its +// temporaries; without this they would leak for the lifetime of the process). +// Only safe once no GPU work can still reference pool buffers, i.e. at the +// end of Prove after every stage goroutine has completed. +func (p *gpuMemoryPool) Shutdown() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + for _, ds := range p.outstanding { + _ = ds.Free() + } + p.outstanding = make(map[unsafe.Pointer]icicle_core.DeviceSlice) +} + // allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. // Use when the destination is fully overwritten by a kernel. func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { @@ -3532,6 +3690,26 @@ func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) } } +// devicePanicToError is deferred at the top of RunOnDevice closures that +// report completion by sending on an error channel. The pool and upload +// helpers panic on GPU allocation failure — the expected failure mode for a +// memory-hungry prover — and a panic in a goroutine cannot be recovered by +// the caller, so without this boundary it kills the whole embedding process +// instead of failing the Prove call. +func devicePanicToError(done chan<- error) { + if r := recover(); r != nil { + err := fmt.Errorf("icicle: device task panicked: %v", r) + select { + case done <- err: + default: + // The closure already reported success and the caller has moved + // on; all we can do is log. + log := logger.Logger() + log.Error().Err(err).Msg("icicle: panic on device goroutine after completion was signalled") + } + } +} + // uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. // Use this for additions where the vector is already in Montgomery form. func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { @@ -3772,8 +3950,11 @@ func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { } // computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. -// Returns device slices for the blinding polynomials. +// Returns device slices for the blinding polynomials. The buffers come from +// the shared temp pool (via state) so an aborted iteration cannot leak them +// past the end of Prove. func computeBlindingPolynomials( + state *gpuConstraintEvalState, n int, twiddles0 []fr.Element, bp []*iop.Polynomial, @@ -3796,11 +3977,17 @@ func computeBlindingPolynomials( } }) - dBlindL = uploadVector(blindL) - dBlindR = uploadVector(blindR) - dBlindO = uploadVector(blindO) - dBlindZ = uploadVector(blindZ) - dBlindZS = uploadVector(blindZS) + uploadFromPool := func(vec []fr.Element) icicle_core.DeviceSlice { + d := state.getTempDeviceSlice(n) + host := icicle_core.HostSliceFromElements(vec) + host.CopyToDevice(&d, false) + return d + } + dBlindL = uploadFromPool(blindL) + dBlindR = uploadFromPool(blindR) + dBlindO = uploadFromPool(blindO) + dBlindZ = uploadFromPool(blindZ) + dBlindZS = uploadFromPool(blindZS) return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS } @@ -3848,12 +4035,12 @@ func applyBlindingToPolynomials( return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) } - // Free blinding vectors - no longer needed after creating blinded polynomials - dBlindL.Free() - dBlindR.Free() - dBlindO.Free() - dBlindZ.Free() - dBlindZS.Free() + // Return blinding vectors to the pool - no longer needed after creating blinded polynomials + state.putTempDeviceSlice(dBlindL, params.n) + state.putTempDeviceSlice(dBlindR, params.n) + state.putTempDeviceSlice(dBlindO, params.n) + state.putTempDeviceSlice(dBlindZ, params.n) + state.putTempDeviceSlice(dBlindZS, params.n) return nil } @@ -4116,7 +4303,8 @@ func computeOrderingConstraint( // Free dGammaScalar - no longer needed after computing a2, b2, c2 // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. - state.dGammaScalar.Free() + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) state.putTempDeviceSlice(dScaledS, params.n) dBetaStd.Free() @@ -4359,21 +4547,45 @@ func (s *instance) gpuEvaluateConstraints( icicle_runtime.RunOnDevice(device, func(args ...any) { state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + var runErr error + var dTmp icicle_core.DeviceSlice + defer func() { + if r := recover(); r != nil && runErr == nil { + // The pool and upload helpers panic on GPU allocation + // failure; a panic in a goroutine cannot be recovered by the + // caller, so convert it into a prover error here instead of + // killing the embedding process. + runErr = fmt.Errorf("gpuEvaluateConstraints: device task panicked: %v", r) + } + // Free everything an aborted pipeline may have left behind. On + // the success path every release below is a no-op: state fields + // are zeroed when returned and the tracked list is emptied. + state.freeAllocatedPolyBuffers() + for _, t := range []*icicle_core.DeviceSlice{&state.dZS, &state.dOrdering, &state.dLocal, &state.dGate, &state.dResult, &dTmp} { + if !t.IsEmpty() { + state.putTempDeviceSlice(*t, t.Len()) + *t = icicle_core.DeviceSlice{} + } + } + freeDeviceSlice(&state.dGammaScalar) + done <- runErr + }() + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). state.dGammaScalar = uploadScalarMont(params.gamma) state.dZS = state.getTempDeviceSlice(n) if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + runErr = fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) return } // Step 1: Compute and apply blinding polynomial evaluations (if enabled) if useBlinding { - dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(state, n, twiddles0, bp) if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { - done <- err + runErr = err return } } @@ -4394,7 +4606,7 @@ func (s *instance) gpuEvaluateConstraints( var err error state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) return } @@ -4403,48 +4615,50 @@ func (s *instance) gpuEvaluateConstraints( dAlphaStd := uploadScalarStd(params.alpha) if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { dAlphaStd.Free() - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) return } dAlphaStd.Free() state.putTempDeviceSlice(state.dOrdering, params.n) + state.dOrdering = icicle_core.DeviceSlice{} // dResult += alpha^2 * local state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: local: %w", err) return } var alphaSquared fr.Element alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) dAlphaSquaredStd := uploadScalarStd(alphaSquared) - dTmp := state.getTempDeviceSlice(params.n) + dTmp = state.getTempDeviceSlice(params.n) if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { dAlphaSquaredStd.Free() - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) return } dAlphaSquaredStd.Free() if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dLocal, params.n) + state.dLocal = icicle_core.DeviceSlice{} state.putTempDeviceSlice(dTmp, params.n) + dTmp = icicle_core.DeviceSlice{} // dResult += gate state.dGate, err = computeGateConstraint(state, params, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) return } if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dGate, params.n) + state.dGate = icicle_core.DeviceSlice{} // Step 5: materialize result either on host or as a persistent device slice. if result != nil { @@ -4453,22 +4667,15 @@ func (s *instance) gpuEvaluateConstraints( } else { resultOnDevice = s.getTempDeviceSlice(params.n) if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) return } } - // Return dResult pool slice after materialization. - state.putTempDeviceSlice(state.dResult, params.n) - - // Return all allocated polynomial buffers to the pool. - // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). - // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are - // returned to pool inside computeGateConstraint. dQk and the original gpuState slices - // are owned by gpuState and will be freed separately. - state.freeAllocatedPolyBuffers() - - done <- nil + // dResult and the tracked polynomial buffers (blinded L, R, O, Z and + // scaled S1, S2, S3) are returned to the pool by the cleanup defer. + // dQk and the original gpuState slices are owned by gpuState and will + // be freed separately. }) err := <-done @@ -4601,6 +4808,7 @@ func (s *instance) prepareStatisticalZKQuotientShards( prepareDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepareDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -4805,6 +5013,7 @@ func (s *instance) inverseAndMergeShards( var dMerged icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") if cfgErr != nil { done <- cfgErr @@ -4978,6 +5187,7 @@ func (s *instance) divideByZHOnGPU( // So we can divide by Z_H by scaling each shard with its corresponding inverse. scaleDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(scaleDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") if cfgErr != nil { scaleDone <- cfgErr @@ -5027,6 +5237,7 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") if cfgErr != nil { done <- cfgErr @@ -5074,8 +5285,15 @@ func commitOnGPUWithDeviceBasesChunked( var msmErr error done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer func() { + if r := recover(); r != nil && msmErr == nil { + // Convert a device-goroutine panic into an error instead of + // killing the embedding process. + msmErr = fmt.Errorf("icicle: MSM device task panicked: %v", r) + } + close(done) + }() commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) - close(done) }) <-done if msmErr != nil { @@ -5249,6 +5467,7 @@ func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSl var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -5422,6 +5641,7 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") if cfgErr != nil { prepDone <- cfgErr @@ -5697,6 +5917,7 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") if cfgErr != nil { done <- cfgErr @@ -5774,6 +5995,7 @@ func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.El witnessSize := len(coeffs) - 1 divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dCoeffs := uploadVector(coeffs) @@ -5838,6 +6060,7 @@ func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.D } divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dPoint := uploadScalarStd(point) @@ -5963,6 +6186,12 @@ func (s *instance) buildLinearizedSelectorTermsOnGPU( } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr } @@ -6151,6 +6380,12 @@ func (s *instance) addZContributionToLinearizedOnGPU( var dScale icicle_core.DeviceSlice var dScaledZ icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before returning temporary buffers to the pool. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6218,6 +6453,12 @@ func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( var dZetaStd icicle_core.DeviceSlice var dZhStd icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before reusing temporary quotient vectors. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6588,9 +6829,10 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) } - chInitDomain := make(chan struct{}) + chInitDomain := make(chan error, 1) initDomainQueuedAt := time.Now() icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chInitDomain) initDomainStartedAt := time.Now() initCfg := icicle_core.GetDefaultNTTInitDomainConfig() ext := config_extension.Create() @@ -6618,41 +6860,52 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { ) } if e != icicle_runtime.Success { - panic("icicle: InitDomain failed") + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return } - close(chInitDomain) + chInitDomain <- nil }) - <-chInitDomain + if err := <-chInitDomain; err != nil { + return err + } if isNttTrace { fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) } - chLag := make(chan struct{}) - chCan := make(chan struct{}) + chLag := make(chan error, 1) + chCan := make(chan error, 1) if len(pk.KzgLagrange.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chLag) g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) - close(chLag) + chLag <- nil }) } else { - close(chLag) + chLag <- nil } if len(pk.Kzg.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chCan) g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) - close(chCan) + chCan <- nil }) } else { - close(chCan) + chCan <- nil } - <-chLag - <-chCan + errLag := <-chLag + errCan := <-chCan + if errLag != nil { + return fmt.Errorf("icicle: uploading Lagrange SRS to device failed: %w", errLag) + } + if errCan != nil { + return fmt.Errorf("icicle: uploading canonical SRS to device failed: %w", errCan) + } return nil } @@ -6982,6 +7235,12 @@ func (s *instance) BuildRatioCopyConstraintIcicle( var dNum, dDen icicle_core.DeviceSlice var dSupportFlat, dPermFlat icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for copy-constraint accumulation before cleanup. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "BuildRatioCopyConstraintIcicle"); syncErr != nil && runErr == nil { runErr = syncErr diff --git a/backend/accelerated/icicle/plonk/bls12-377/icicle.go b/backend/accelerated/icicle/plonk/bls12-377/icicle.go index b448e58a43..798fbdc9d1 100644 --- a/backend/accelerated/icicle/plonk/bls12-377/icicle.go +++ b/backend/accelerated/icicle/plonk/bls12-377/icicle.go @@ -264,32 +264,47 @@ func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts defer instance.releaseSharedGPUState() defer instance.releaseLinearizedEvalGPUState() + // goStage runs a prover stage on the errgroup, converting a panic into an + // error: several device helpers panic on GPU allocation failure, and an + // unrecovered panic in a stage goroutine would kill the embedding process + // instead of failing the Prove call. + goStage := func(name string, fn func() error) { + g.Go(func() (err error) { + defer func() { + if r := recover(); r != nil && err == nil { + err = fmt.Errorf("icicle plonk: stage %s panicked: %v", name, r) + } + }() + return fn() + }) + } + // solve constraints - g.Go(instance.solveConstraints) + goStage("solveConstraints", instance.solveConstraints) // complete qk - g.Go(instance.completeQk) + goStage("completeQk", instance.completeQk) // init blinding polynomials - g.Go(instance.initBlindingPolynomials) + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) // derive gamma, beta (copy constraint) - g.Go(instance.deriveGammaAndBeta) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) // compute accumulating ratio for the copy constraint - g.Go(instance.buildRatioCopyConstraint) + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) // compute h - g.Go(instance.computeQuotient) + goStage("computeQuotient", instance.computeQuotient) // open Z (blinded) at ωζ (proof.ZShiftedOpening) - g.Go(instance.openZ) + goStage("openZ", instance.openZ) // linearized polynomial - g.Go(instance.computeLinearizedPolynomial) + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) // Batch opening (no internal timer of its own — time the whole stage here) - g.Go(func() error { + goStage("batchOpening", func() error { startBatchOpening := time.Now() err := instance.batchOpening() if isProfileMode { @@ -502,6 +517,21 @@ func (s *instance) solveConstraints() error { // Try to load raw solver values from cache (fastest path) rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + // The raw solver cache stores the BSB22 commitment polynomials verbatim, + // including their two random blinding rows (see bsb22Hint). Replaying them + // across proofs makes Bsb22Commitments identical (linkable) and, once the + // committed polynomial has been opened at more than two distinct zetas, + // leaks linear relations over the committed private wires. Re-randomizing + // on load is not possible either: the commitment hash feeds back into a + // witness wire, so fresh blinding would invalidate every cached wire + // downstream of it. When zero-knowledge matters (blinding enabled, the + // default), the cache is therefore disabled for circuits with BSB22 + // commitments. With GNARK_DISABLE_BLINDING set, zero-knowledge is already + // explicitly forfeited and the replay leaks nothing new. + if rawCachePath != "" && len(s.commitmentInfo) > 0 && useBlinding { + log.Warn().Str("file", rawCachePath).Msg("GNARK_RAW_SOLVER_CACHE ignored: replaying cached BSB22 commitment blinding across proofs would break zero-knowledge (set GNARK_DISABLE_BLINDING to opt out of zero-knowledge and cache anyway)") + rawCachePath = "" + } if rawCachePath != "" { if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { // Reconstruct L, R, O from raw values @@ -988,6 +1018,7 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone := make(chan error, 1) var dPersist icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(copyDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") if cfgErr != nil { copyDone <- cfgErr @@ -1078,6 +1109,7 @@ func (s *instance) openZ() (err error) { var dBlindedCanonical icicle_core.DeviceSlice buildDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(buildDone) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") if cfgErr != nil { buildDone <- cfgErr @@ -1403,7 +1435,9 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) doneAddZ() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + // dLin is a pool buffer: return it via Put, not a direct free, so the + // pool's outstanding tracking stays consistent. + s.putTempDeviceSlice(dLin, dLin.Len()) return err } @@ -1411,7 +1445,7 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) doneSubtractH() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + s.putTempDeviceSlice(dLin, dLin.Len()) return err } s.linearizedPolynomialGPU = dLin @@ -1452,7 +1486,12 @@ func (s *instance) batchOpening() error { } defer func() { - freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + // The linearized polynomial is a pool buffer: return it via Put, not + // a direct free, so the pool's outstanding tracking stays consistent. + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } if s.hGPU != nil { s.freeGPUQuotient(s.hGPU) s.hGPU = nil @@ -1536,6 +1575,7 @@ func (s *instance) batchOpening() error { var dFold icicle_core.DeviceSlice foldDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(foldDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") if cfgErr != nil { foldDone <- cfgErr @@ -1602,6 +1642,7 @@ func (s *instance) batchOpening() error { divDone := make(chan error, 1) witnessSize := dFold.Len() - 1 icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") if cfgErr != nil { divDone <- cfgErr @@ -1898,12 +1939,15 @@ func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) var dWitness icicle_core.DeviceSlice - uploadDone := make(chan struct{}, 1) + uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) dWitness = uploadVector(hCoeffs) - close(uploadDone) + uploadDone <- nil }) - <-uploadDone + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) freeSliceOnDevice(&dWitness, &s.device) if err != nil { @@ -1995,6 +2039,7 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") if cfgErr != nil { done <- cfgErr @@ -2283,6 +2328,7 @@ func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicl var dTwiddles0 icicle_core.DeviceSlice uploadTwiddlesDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadTwiddlesDone) if s.tempGPUMemPool != nil { s.tempGPUMemPool.FreeAll() } @@ -2355,6 +2401,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo batchInvertDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(batchInvertDone) batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) }) if err := <-batchInvertDone; err != nil { @@ -2590,6 +2637,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(mergeDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") if cfgErr != nil { mergeDone <- cfgErr @@ -2620,6 +2668,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cresHost := icicle_core.HostSliceFromElements(cres) downloadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(downloadDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") if cfgErr != nil { downloadDone <- cfgErr @@ -2825,6 +2874,12 @@ func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop. } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for snapshot cloning before handing state to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { runErr = syncErr @@ -2909,6 +2964,12 @@ func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } @@ -2961,14 +3022,29 @@ func (s *instance) releaseTempGPUMemoryPool() { if s == nil || s.tempGPUMemPool == nil { return } + // Backstop for instance-owned device buffers whose owning stage may have + // exited (via error or ctx cancellation) before registering its own + // cleanup defer: openZ owns the two Z buffers, batchOpening owns the + // quotient and the linearized polynomial. On the success path these are + // already released and zeroed, making every free below a no-op. This must + // run before FreeAll so pool-owned buffers (blindedZCanonicalGPU, hGPU) + // are back in the pool when it frees everything. freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) if !s.blindedZCanonicalGPU.IsEmpty() { s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} } + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { - s.tempGPUMemPool.FreeAll() + s.tempGPUMemPool.Shutdown() close(done) }) <-done @@ -3026,6 +3102,12 @@ func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysSta } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for GPU uploads in ensurePolysOnSharedGPU. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -3080,6 +3162,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali // Upload scaling vectors to GPU uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(uploadDone) scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) scalingHost := icicle_core.HostSliceFromElements(scalingVector) scalingHost.CopyToDevice(&scalingVectorDevice, false) @@ -3098,10 +3181,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } uploadDone <- nil }) + // Return the scaling vectors to the pool on every exit path. This is safe + // because we never return before draining every per-polynomial done + // channel, and each worker synchronizes its stream before signalling, so + // no in-flight kernel can still reference the slices. + defer func() { + if !scalingVectorDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + } + if !scalingVectorRevDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + } + }() if err := <-uploadDone; err != nil { return err } + // Validate all polynomials before launching any GPU work so an invalid + // entry cannot abandon already-scheduled workers. + for _, p := range state.polys { + if p != nil && p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + } + doneChans := make([]chan error, len(state.polys)) for i := range state.polys { @@ -3110,10 +3213,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali continue } - if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { - return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) - } - done := make(chan error, 1) doneChans[i] = done scalarsDevice := state.deviceSlices[i] @@ -3122,7 +3221,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: CreateStream failed: %s", eStream.AsString()) + return + } + // Synchronize and destroy the stream on every path (including + // errors) before signalling done: the caller may free device + // buffers as soon as all workers have reported. + var runErr error + defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: DestroyStream failed: %s", eDestroy.AsString()) + } + done <- runErr + }() cfg.StreamHandle = stream cfg.IsAsync = true @@ -3136,7 +3258,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KRN } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) return } @@ -3155,7 +3277,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) return } @@ -3168,26 +3290,26 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KNR // Regular → BitReverse } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) - return - } - - if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) return } - done <- nil }) } - // Wait for all scheduled tasks + // Wait for ALL scheduled tasks even when one fails: returning on the + // first error would let the caller free device slices that sibling + // workers are still writing to. + var firstErr error for i := range doneChans { if doneChans[i] != nil { - if err := <-doneChans[i]; err != nil { - return err + if err := <-doneChans[i]; err != nil && firstErr == nil { + firstErr = err } } } + if firstErr != nil { + return firstErr + } // Update polynomial metadata: final result is in Lagrange, same layout as original for _, p := range state.polys { @@ -3195,10 +3317,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali p.Basis = iop.Lagrange } } - - // Free scaling vectors from device - s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) - s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) return nil } @@ -3270,7 +3388,11 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: CreateStream failed: %s", eStream.AsString())) + } + defer icicle_runtime.DestroyStream(stream) cfg.StreamHandle = stream cfg.IsAsync = true @@ -3292,7 +3414,9 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) } - icicle_runtime.SynchronizeStream(stream) + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: SynchronizeStream failed: %s", eSync.AsString())) + } // Update metadata inside closure to avoid race p.Basis = iop.Canonical @@ -3341,13 +3465,20 @@ func (s *instance) freeGPUPolys(state *gpuPolysState) { // Must be used within RunOnDevice context to ensure thread safety per device. type gpuMemoryPool struct { freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice - mu sync.Mutex + // outstanding tracks slices handed out by Get and not yet returned so + // Shutdown can reclaim buffers abandoned by error paths. Keyed by the + // device pointer, not the slice value: CopyToDevice mutates the slice's + // length field, so the value at Put time may differ from the one at Get + // time while the underlying buffer is the same. + outstanding map[unsafe.Pointer]icicle_core.DeviceSlice + mu sync.Mutex } // newGPUMemoryPool creates a new GPU memory pool. func newGPUMemoryPool() *gpuMemoryPool { return &gpuMemoryPool{ - freeSlices: make(map[int][]icicle_core.DeviceSlice), + freeSlices: make(map[int][]icicle_core.DeviceSlice), + outstanding: make(map[unsafe.Pointer]icicle_core.DeviceSlice), } } @@ -3361,12 +3492,14 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { // Reuse the last slice slice := slices[len(slices)-1] p.freeSlices[n] = slices[:len(slices)-1] + p.outstanding[slice.AsUnsafePointer()] = slice return slice } // No free slice available, allocate a new one. // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3379,6 +3512,7 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3394,11 +3528,14 @@ func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { p.mu.Lock() defer p.mu.Unlock() + delete(p.outstanding, ds.AsUnsafePointer()) // Add to the pool p.freeSlices[n] = append(p.freeSlices[n], ds) } -// FreeAll releases all pooled device slices. +// FreeAll releases all idle pooled device slices. Outstanding slices (handed +// out by Get and not yet returned) are left alone: FreeAll is also used +// mid-prove to relieve memory pressure while pool buffers are still live. func (p *gpuMemoryPool) FreeAll() { p.mu.Lock() defer p.mu.Unlock() @@ -3411,6 +3548,27 @@ func (p *gpuMemoryPool) FreeAll() { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) } +// Shutdown releases every pooled slice, including outstanding ones that were +// abandoned by error paths (a failed stage returns without Put-ing its +// temporaries; without this they would leak for the lifetime of the process). +// Only safe once no GPU work can still reference pool buffers, i.e. at the +// end of Prove after every stage goroutine has completed. +func (p *gpuMemoryPool) Shutdown() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + for _, ds := range p.outstanding { + _ = ds.Free() + } + p.outstanding = make(map[unsafe.Pointer]icicle_core.DeviceSlice) +} + // allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. // Use when the destination is fully overwritten by a kernel. func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { @@ -3539,6 +3697,26 @@ func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) } } +// devicePanicToError is deferred at the top of RunOnDevice closures that +// report completion by sending on an error channel. The pool and upload +// helpers panic on GPU allocation failure — the expected failure mode for a +// memory-hungry prover — and a panic in a goroutine cannot be recovered by +// the caller, so without this boundary it kills the whole embedding process +// instead of failing the Prove call. +func devicePanicToError(done chan<- error) { + if r := recover(); r != nil { + err := fmt.Errorf("icicle: device task panicked: %v", r) + select { + case done <- err: + default: + // The closure already reported success and the caller has moved + // on; all we can do is log. + log := logger.Logger() + log.Error().Err(err).Msg("icicle: panic on device goroutine after completion was signalled") + } + } +} + // uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. // Use this for additions where the vector is already in Montgomery form. func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { @@ -3779,8 +3957,11 @@ func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { } // computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. -// Returns device slices for the blinding polynomials. +// Returns device slices for the blinding polynomials. The buffers come from +// the shared temp pool (via state) so an aborted iteration cannot leak them +// past the end of Prove. func computeBlindingPolynomials( + state *gpuConstraintEvalState, n int, twiddles0 []fr.Element, bp []*iop.Polynomial, @@ -3803,11 +3984,17 @@ func computeBlindingPolynomials( } }) - dBlindL = uploadVector(blindL) - dBlindR = uploadVector(blindR) - dBlindO = uploadVector(blindO) - dBlindZ = uploadVector(blindZ) - dBlindZS = uploadVector(blindZS) + uploadFromPool := func(vec []fr.Element) icicle_core.DeviceSlice { + d := state.getTempDeviceSlice(n) + host := icicle_core.HostSliceFromElements(vec) + host.CopyToDevice(&d, false) + return d + } + dBlindL = uploadFromPool(blindL) + dBlindR = uploadFromPool(blindR) + dBlindO = uploadFromPool(blindO) + dBlindZ = uploadFromPool(blindZ) + dBlindZS = uploadFromPool(blindZS) return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS } @@ -3855,12 +4042,12 @@ func applyBlindingToPolynomials( return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) } - // Free blinding vectors - no longer needed after creating blinded polynomials - dBlindL.Free() - dBlindR.Free() - dBlindO.Free() - dBlindZ.Free() - dBlindZS.Free() + // Return blinding vectors to the pool - no longer needed after creating blinded polynomials + state.putTempDeviceSlice(dBlindL, params.n) + state.putTempDeviceSlice(dBlindR, params.n) + state.putTempDeviceSlice(dBlindO, params.n) + state.putTempDeviceSlice(dBlindZ, params.n) + state.putTempDeviceSlice(dBlindZS, params.n) return nil } @@ -4123,7 +4310,8 @@ func computeOrderingConstraint( // Free dGammaScalar - no longer needed after computing a2, b2, c2 // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. - state.dGammaScalar.Free() + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) state.putTempDeviceSlice(dScaledS, params.n) dBetaStd.Free() @@ -4366,21 +4554,45 @@ func (s *instance) gpuEvaluateConstraints( icicle_runtime.RunOnDevice(device, func(args ...any) { state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + var runErr error + var dTmp icicle_core.DeviceSlice + defer func() { + if r := recover(); r != nil && runErr == nil { + // The pool and upload helpers panic on GPU allocation + // failure; a panic in a goroutine cannot be recovered by the + // caller, so convert it into a prover error here instead of + // killing the embedding process. + runErr = fmt.Errorf("gpuEvaluateConstraints: device task panicked: %v", r) + } + // Free everything an aborted pipeline may have left behind. On + // the success path every release below is a no-op: state fields + // are zeroed when returned and the tracked list is emptied. + state.freeAllocatedPolyBuffers() + for _, t := range []*icicle_core.DeviceSlice{&state.dZS, &state.dOrdering, &state.dLocal, &state.dGate, &state.dResult, &dTmp} { + if !t.IsEmpty() { + state.putTempDeviceSlice(*t, t.Len()) + *t = icicle_core.DeviceSlice{} + } + } + freeDeviceSlice(&state.dGammaScalar) + done <- runErr + }() + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). state.dGammaScalar = uploadScalarMont(params.gamma) state.dZS = state.getTempDeviceSlice(n) if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + runErr = fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) return } // Step 1: Compute and apply blinding polynomial evaluations (if enabled) if useBlinding { - dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(state, n, twiddles0, bp) if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { - done <- err + runErr = err return } } @@ -4401,7 +4613,7 @@ func (s *instance) gpuEvaluateConstraints( var err error state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) return } @@ -4410,48 +4622,50 @@ func (s *instance) gpuEvaluateConstraints( dAlphaStd := uploadScalarStd(params.alpha) if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { dAlphaStd.Free() - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) return } dAlphaStd.Free() state.putTempDeviceSlice(state.dOrdering, params.n) + state.dOrdering = icicle_core.DeviceSlice{} // dResult += alpha^2 * local state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: local: %w", err) return } var alphaSquared fr.Element alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) dAlphaSquaredStd := uploadScalarStd(alphaSquared) - dTmp := state.getTempDeviceSlice(params.n) + dTmp = state.getTempDeviceSlice(params.n) if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { dAlphaSquaredStd.Free() - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) return } dAlphaSquaredStd.Free() if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dLocal, params.n) + state.dLocal = icicle_core.DeviceSlice{} state.putTempDeviceSlice(dTmp, params.n) + dTmp = icicle_core.DeviceSlice{} // dResult += gate state.dGate, err = computeGateConstraint(state, params, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) return } if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dGate, params.n) + state.dGate = icicle_core.DeviceSlice{} // Step 5: materialize result either on host or as a persistent device slice. if result != nil { @@ -4460,22 +4674,15 @@ func (s *instance) gpuEvaluateConstraints( } else { resultOnDevice = s.getTempDeviceSlice(params.n) if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) return } } - // Return dResult pool slice after materialization. - state.putTempDeviceSlice(state.dResult, params.n) - - // Return all allocated polynomial buffers to the pool. - // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). - // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are - // returned to pool inside computeGateConstraint. dQk and the original gpuState slices - // are owned by gpuState and will be freed separately. - state.freeAllocatedPolyBuffers() - - done <- nil + // dResult and the tracked polynomial buffers (blinded L, R, O, Z and + // scaled S1, S2, S3) are returned to the pool by the cleanup defer. + // dQk and the original gpuState slices are owned by gpuState and will + // be freed separately. }) err := <-done @@ -4608,6 +4815,7 @@ func (s *instance) prepareStatisticalZKQuotientShards( prepareDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepareDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -4812,6 +5020,7 @@ func (s *instance) inverseAndMergeShards( var dMerged icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") if cfgErr != nil { done <- cfgErr @@ -4985,6 +5194,7 @@ func (s *instance) divideByZHOnGPU( // So we can divide by Z_H by scaling each shard with its corresponding inverse. scaleDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(scaleDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") if cfgErr != nil { scaleDone <- cfgErr @@ -5034,6 +5244,7 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") if cfgErr != nil { done <- cfgErr @@ -5081,8 +5292,15 @@ func commitOnGPUWithDeviceBasesChunked( var msmErr error done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer func() { + if r := recover(); r != nil && msmErr == nil { + // Convert a device-goroutine panic into an error instead of + // killing the embedding process. + msmErr = fmt.Errorf("icicle: MSM device task panicked: %v", r) + } + close(done) + }() commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) - close(done) }) <-done if msmErr != nil { @@ -5256,6 +5474,7 @@ func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSl var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -5429,6 +5648,7 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") if cfgErr != nil { prepDone <- cfgErr @@ -5704,6 +5924,7 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") if cfgErr != nil { done <- cfgErr @@ -5781,6 +6002,7 @@ func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.El witnessSize := len(coeffs) - 1 divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dCoeffs := uploadVector(coeffs) @@ -5845,6 +6067,7 @@ func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.D } divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dPoint := uploadScalarStd(point) @@ -5970,6 +6193,12 @@ func (s *instance) buildLinearizedSelectorTermsOnGPU( } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr } @@ -6158,6 +6387,12 @@ func (s *instance) addZContributionToLinearizedOnGPU( var dScale icicle_core.DeviceSlice var dScaledZ icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before returning temporary buffers to the pool. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6225,6 +6460,12 @@ func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( var dZetaStd icicle_core.DeviceSlice var dZhStd icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before reusing temporary quotient vectors. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6595,9 +6836,10 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) } - chInitDomain := make(chan struct{}) + chInitDomain := make(chan error, 1) initDomainQueuedAt := time.Now() icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chInitDomain) initDomainStartedAt := time.Now() initCfg := icicle_core.GetDefaultNTTInitDomainConfig() ext := config_extension.Create() @@ -6625,41 +6867,52 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { ) } if e != icicle_runtime.Success { - panic("icicle: InitDomain failed") + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return } - close(chInitDomain) + chInitDomain <- nil }) - <-chInitDomain + if err := <-chInitDomain; err != nil { + return err + } if isNttTrace { fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) } - chLag := make(chan struct{}) - chCan := make(chan struct{}) + chLag := make(chan error, 1) + chCan := make(chan error, 1) if len(pk.KzgLagrange.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chLag) g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) - close(chLag) + chLag <- nil }) } else { - close(chLag) + chLag <- nil } if len(pk.Kzg.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chCan) g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) - close(chCan) + chCan <- nil }) } else { - close(chCan) + chCan <- nil } - <-chLag - <-chCan + errLag := <-chLag + errCan := <-chCan + if errLag != nil { + return fmt.Errorf("icicle: uploading Lagrange SRS to device failed: %w", errLag) + } + if errCan != nil { + return fmt.Errorf("icicle: uploading canonical SRS to device failed: %w", errCan) + } return nil } @@ -6989,6 +7242,12 @@ func (s *instance) BuildRatioCopyConstraintIcicle( var dNum, dDen icicle_core.DeviceSlice var dSupportFlat, dPermFlat icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for copy-constraint accumulation before cleanup. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "BuildRatioCopyConstraintIcicle"); syncErr != nil && runErr == nil { runErr = syncErr diff --git a/backend/accelerated/icicle/plonk/bls12-381/icicle.go b/backend/accelerated/icicle/plonk/bls12-381/icicle.go index 4914aa4051..5accd02547 100644 --- a/backend/accelerated/icicle/plonk/bls12-381/icicle.go +++ b/backend/accelerated/icicle/plonk/bls12-381/icicle.go @@ -264,32 +264,47 @@ func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts defer instance.releaseSharedGPUState() defer instance.releaseLinearizedEvalGPUState() + // goStage runs a prover stage on the errgroup, converting a panic into an + // error: several device helpers panic on GPU allocation failure, and an + // unrecovered panic in a stage goroutine would kill the embedding process + // instead of failing the Prove call. + goStage := func(name string, fn func() error) { + g.Go(func() (err error) { + defer func() { + if r := recover(); r != nil && err == nil { + err = fmt.Errorf("icicle plonk: stage %s panicked: %v", name, r) + } + }() + return fn() + }) + } + // solve constraints - g.Go(instance.solveConstraints) + goStage("solveConstraints", instance.solveConstraints) // complete qk - g.Go(instance.completeQk) + goStage("completeQk", instance.completeQk) // init blinding polynomials - g.Go(instance.initBlindingPolynomials) + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) // derive gamma, beta (copy constraint) - g.Go(instance.deriveGammaAndBeta) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) // compute accumulating ratio for the copy constraint - g.Go(instance.buildRatioCopyConstraint) + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) // compute h - g.Go(instance.computeQuotient) + goStage("computeQuotient", instance.computeQuotient) // open Z (blinded) at ωζ (proof.ZShiftedOpening) - g.Go(instance.openZ) + goStage("openZ", instance.openZ) // linearized polynomial - g.Go(instance.computeLinearizedPolynomial) + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) // Batch opening (no internal timer of its own — time the whole stage here) - g.Go(func() error { + goStage("batchOpening", func() error { startBatchOpening := time.Now() err := instance.batchOpening() if isProfileMode { @@ -502,6 +517,21 @@ func (s *instance) solveConstraints() error { // Try to load raw solver values from cache (fastest path) rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + // The raw solver cache stores the BSB22 commitment polynomials verbatim, + // including their two random blinding rows (see bsb22Hint). Replaying them + // across proofs makes Bsb22Commitments identical (linkable) and, once the + // committed polynomial has been opened at more than two distinct zetas, + // leaks linear relations over the committed private wires. Re-randomizing + // on load is not possible either: the commitment hash feeds back into a + // witness wire, so fresh blinding would invalidate every cached wire + // downstream of it. When zero-knowledge matters (blinding enabled, the + // default), the cache is therefore disabled for circuits with BSB22 + // commitments. With GNARK_DISABLE_BLINDING set, zero-knowledge is already + // explicitly forfeited and the replay leaks nothing new. + if rawCachePath != "" && len(s.commitmentInfo) > 0 && useBlinding { + log.Warn().Str("file", rawCachePath).Msg("GNARK_RAW_SOLVER_CACHE ignored: replaying cached BSB22 commitment blinding across proofs would break zero-knowledge (set GNARK_DISABLE_BLINDING to opt out of zero-knowledge and cache anyway)") + rawCachePath = "" + } if rawCachePath != "" { if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { // Reconstruct L, R, O from raw values @@ -988,6 +1018,7 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone := make(chan error, 1) var dPersist icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(copyDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") if cfgErr != nil { copyDone <- cfgErr @@ -1078,6 +1109,7 @@ func (s *instance) openZ() (err error) { var dBlindedCanonical icicle_core.DeviceSlice buildDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(buildDone) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") if cfgErr != nil { buildDone <- cfgErr @@ -1403,7 +1435,9 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) doneAddZ() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + // dLin is a pool buffer: return it via Put, not a direct free, so the + // pool's outstanding tracking stays consistent. + s.putTempDeviceSlice(dLin, dLin.Len()) return err } @@ -1411,7 +1445,7 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) doneSubtractH() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + s.putTempDeviceSlice(dLin, dLin.Len()) return err } s.linearizedPolynomialGPU = dLin @@ -1452,7 +1486,12 @@ func (s *instance) batchOpening() error { } defer func() { - freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + // The linearized polynomial is a pool buffer: return it via Put, not + // a direct free, so the pool's outstanding tracking stays consistent. + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } if s.hGPU != nil { s.freeGPUQuotient(s.hGPU) s.hGPU = nil @@ -1536,6 +1575,7 @@ func (s *instance) batchOpening() error { var dFold icicle_core.DeviceSlice foldDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(foldDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") if cfgErr != nil { foldDone <- cfgErr @@ -1602,6 +1642,7 @@ func (s *instance) batchOpening() error { divDone := make(chan error, 1) witnessSize := dFold.Len() - 1 icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") if cfgErr != nil { divDone <- cfgErr @@ -1898,12 +1939,15 @@ func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) var dWitness icicle_core.DeviceSlice - uploadDone := make(chan struct{}, 1) + uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) dWitness = uploadVector(hCoeffs) - close(uploadDone) + uploadDone <- nil }) - <-uploadDone + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) freeSliceOnDevice(&dWitness, &s.device) if err != nil { @@ -1995,6 +2039,7 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") if cfgErr != nil { done <- cfgErr @@ -2283,6 +2328,7 @@ func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicl var dTwiddles0 icicle_core.DeviceSlice uploadTwiddlesDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadTwiddlesDone) if s.tempGPUMemPool != nil { s.tempGPUMemPool.FreeAll() } @@ -2355,6 +2401,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo batchInvertDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(batchInvertDone) batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) }) if err := <-batchInvertDone; err != nil { @@ -2590,6 +2637,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(mergeDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") if cfgErr != nil { mergeDone <- cfgErr @@ -2620,6 +2668,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cresHost := icicle_core.HostSliceFromElements(cres) downloadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(downloadDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") if cfgErr != nil { downloadDone <- cfgErr @@ -2825,6 +2874,12 @@ func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop. } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for snapshot cloning before handing state to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { runErr = syncErr @@ -2909,6 +2964,12 @@ func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } @@ -2961,14 +3022,29 @@ func (s *instance) releaseTempGPUMemoryPool() { if s == nil || s.tempGPUMemPool == nil { return } + // Backstop for instance-owned device buffers whose owning stage may have + // exited (via error or ctx cancellation) before registering its own + // cleanup defer: openZ owns the two Z buffers, batchOpening owns the + // quotient and the linearized polynomial. On the success path these are + // already released and zeroed, making every free below a no-op. This must + // run before FreeAll so pool-owned buffers (blindedZCanonicalGPU, hGPU) + // are back in the pool when it frees everything. freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) if !s.blindedZCanonicalGPU.IsEmpty() { s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} } + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { - s.tempGPUMemPool.FreeAll() + s.tempGPUMemPool.Shutdown() close(done) }) <-done @@ -3026,6 +3102,12 @@ func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysSta } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for GPU uploads in ensurePolysOnSharedGPU. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -3080,6 +3162,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali // Upload scaling vectors to GPU uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(uploadDone) scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) scalingHost := icicle_core.HostSliceFromElements(scalingVector) scalingHost.CopyToDevice(&scalingVectorDevice, false) @@ -3098,10 +3181,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } uploadDone <- nil }) + // Return the scaling vectors to the pool on every exit path. This is safe + // because we never return before draining every per-polynomial done + // channel, and each worker synchronizes its stream before signalling, so + // no in-flight kernel can still reference the slices. + defer func() { + if !scalingVectorDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + } + if !scalingVectorRevDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + } + }() if err := <-uploadDone; err != nil { return err } + // Validate all polynomials before launching any GPU work so an invalid + // entry cannot abandon already-scheduled workers. + for _, p := range state.polys { + if p != nil && p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + } + doneChans := make([]chan error, len(state.polys)) for i := range state.polys { @@ -3110,10 +3213,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali continue } - if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { - return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) - } - done := make(chan error, 1) doneChans[i] = done scalarsDevice := state.deviceSlices[i] @@ -3122,7 +3221,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: CreateStream failed: %s", eStream.AsString()) + return + } + // Synchronize and destroy the stream on every path (including + // errors) before signalling done: the caller may free device + // buffers as soon as all workers have reported. + var runErr error + defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: DestroyStream failed: %s", eDestroy.AsString()) + } + done <- runErr + }() cfg.StreamHandle = stream cfg.IsAsync = true @@ -3136,7 +3258,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KRN } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) return } @@ -3155,7 +3277,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) return } @@ -3168,26 +3290,26 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KNR // Regular → BitReverse } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) - return - } - - if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) return } - done <- nil }) } - // Wait for all scheduled tasks + // Wait for ALL scheduled tasks even when one fails: returning on the + // first error would let the caller free device slices that sibling + // workers are still writing to. + var firstErr error for i := range doneChans { if doneChans[i] != nil { - if err := <-doneChans[i]; err != nil { - return err + if err := <-doneChans[i]; err != nil && firstErr == nil { + firstErr = err } } } + if firstErr != nil { + return firstErr + } // Update polynomial metadata: final result is in Lagrange, same layout as original for _, p := range state.polys { @@ -3195,10 +3317,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali p.Basis = iop.Lagrange } } - - // Free scaling vectors from device - s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) - s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) return nil } @@ -3270,7 +3388,11 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: CreateStream failed: %s", eStream.AsString())) + } + defer icicle_runtime.DestroyStream(stream) cfg.StreamHandle = stream cfg.IsAsync = true @@ -3292,7 +3414,9 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) } - icicle_runtime.SynchronizeStream(stream) + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: SynchronizeStream failed: %s", eSync.AsString())) + } // Update metadata inside closure to avoid race p.Basis = iop.Canonical @@ -3341,13 +3465,20 @@ func (s *instance) freeGPUPolys(state *gpuPolysState) { // Must be used within RunOnDevice context to ensure thread safety per device. type gpuMemoryPool struct { freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice - mu sync.Mutex + // outstanding tracks slices handed out by Get and not yet returned so + // Shutdown can reclaim buffers abandoned by error paths. Keyed by the + // device pointer, not the slice value: CopyToDevice mutates the slice's + // length field, so the value at Put time may differ from the one at Get + // time while the underlying buffer is the same. + outstanding map[unsafe.Pointer]icicle_core.DeviceSlice + mu sync.Mutex } // newGPUMemoryPool creates a new GPU memory pool. func newGPUMemoryPool() *gpuMemoryPool { return &gpuMemoryPool{ - freeSlices: make(map[int][]icicle_core.DeviceSlice), + freeSlices: make(map[int][]icicle_core.DeviceSlice), + outstanding: make(map[unsafe.Pointer]icicle_core.DeviceSlice), } } @@ -3361,12 +3492,14 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { // Reuse the last slice slice := slices[len(slices)-1] p.freeSlices[n] = slices[:len(slices)-1] + p.outstanding[slice.AsUnsafePointer()] = slice return slice } // No free slice available, allocate a new one. // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3379,6 +3512,7 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3394,11 +3528,14 @@ func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { p.mu.Lock() defer p.mu.Unlock() + delete(p.outstanding, ds.AsUnsafePointer()) // Add to the pool p.freeSlices[n] = append(p.freeSlices[n], ds) } -// FreeAll releases all pooled device slices. +// FreeAll releases all idle pooled device slices. Outstanding slices (handed +// out by Get and not yet returned) are left alone: FreeAll is also used +// mid-prove to relieve memory pressure while pool buffers are still live. func (p *gpuMemoryPool) FreeAll() { p.mu.Lock() defer p.mu.Unlock() @@ -3411,6 +3548,27 @@ func (p *gpuMemoryPool) FreeAll() { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) } +// Shutdown releases every pooled slice, including outstanding ones that were +// abandoned by error paths (a failed stage returns without Put-ing its +// temporaries; without this they would leak for the lifetime of the process). +// Only safe once no GPU work can still reference pool buffers, i.e. at the +// end of Prove after every stage goroutine has completed. +func (p *gpuMemoryPool) Shutdown() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + for _, ds := range p.outstanding { + _ = ds.Free() + } + p.outstanding = make(map[unsafe.Pointer]icicle_core.DeviceSlice) +} + // allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. // Use when the destination is fully overwritten by a kernel. func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { @@ -3539,6 +3697,26 @@ func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) } } +// devicePanicToError is deferred at the top of RunOnDevice closures that +// report completion by sending on an error channel. The pool and upload +// helpers panic on GPU allocation failure — the expected failure mode for a +// memory-hungry prover — and a panic in a goroutine cannot be recovered by +// the caller, so without this boundary it kills the whole embedding process +// instead of failing the Prove call. +func devicePanicToError(done chan<- error) { + if r := recover(); r != nil { + err := fmt.Errorf("icicle: device task panicked: %v", r) + select { + case done <- err: + default: + // The closure already reported success and the caller has moved + // on; all we can do is log. + log := logger.Logger() + log.Error().Err(err).Msg("icicle: panic on device goroutine after completion was signalled") + } + } +} + // uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. // Use this for additions where the vector is already in Montgomery form. func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { @@ -3779,8 +3957,11 @@ func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { } // computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. -// Returns device slices for the blinding polynomials. +// Returns device slices for the blinding polynomials. The buffers come from +// the shared temp pool (via state) so an aborted iteration cannot leak them +// past the end of Prove. func computeBlindingPolynomials( + state *gpuConstraintEvalState, n int, twiddles0 []fr.Element, bp []*iop.Polynomial, @@ -3803,11 +3984,17 @@ func computeBlindingPolynomials( } }) - dBlindL = uploadVector(blindL) - dBlindR = uploadVector(blindR) - dBlindO = uploadVector(blindO) - dBlindZ = uploadVector(blindZ) - dBlindZS = uploadVector(blindZS) + uploadFromPool := func(vec []fr.Element) icicle_core.DeviceSlice { + d := state.getTempDeviceSlice(n) + host := icicle_core.HostSliceFromElements(vec) + host.CopyToDevice(&d, false) + return d + } + dBlindL = uploadFromPool(blindL) + dBlindR = uploadFromPool(blindR) + dBlindO = uploadFromPool(blindO) + dBlindZ = uploadFromPool(blindZ) + dBlindZS = uploadFromPool(blindZS) return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS } @@ -3855,12 +4042,12 @@ func applyBlindingToPolynomials( return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) } - // Free blinding vectors - no longer needed after creating blinded polynomials - dBlindL.Free() - dBlindR.Free() - dBlindO.Free() - dBlindZ.Free() - dBlindZS.Free() + // Return blinding vectors to the pool - no longer needed after creating blinded polynomials + state.putTempDeviceSlice(dBlindL, params.n) + state.putTempDeviceSlice(dBlindR, params.n) + state.putTempDeviceSlice(dBlindO, params.n) + state.putTempDeviceSlice(dBlindZ, params.n) + state.putTempDeviceSlice(dBlindZS, params.n) return nil } @@ -4123,7 +4310,8 @@ func computeOrderingConstraint( // Free dGammaScalar - no longer needed after computing a2, b2, c2 // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. - state.dGammaScalar.Free() + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) state.putTempDeviceSlice(dScaledS, params.n) dBetaStd.Free() @@ -4366,21 +4554,45 @@ func (s *instance) gpuEvaluateConstraints( icicle_runtime.RunOnDevice(device, func(args ...any) { state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + var runErr error + var dTmp icicle_core.DeviceSlice + defer func() { + if r := recover(); r != nil && runErr == nil { + // The pool and upload helpers panic on GPU allocation + // failure; a panic in a goroutine cannot be recovered by the + // caller, so convert it into a prover error here instead of + // killing the embedding process. + runErr = fmt.Errorf("gpuEvaluateConstraints: device task panicked: %v", r) + } + // Free everything an aborted pipeline may have left behind. On + // the success path every release below is a no-op: state fields + // are zeroed when returned and the tracked list is emptied. + state.freeAllocatedPolyBuffers() + for _, t := range []*icicle_core.DeviceSlice{&state.dZS, &state.dOrdering, &state.dLocal, &state.dGate, &state.dResult, &dTmp} { + if !t.IsEmpty() { + state.putTempDeviceSlice(*t, t.Len()) + *t = icicle_core.DeviceSlice{} + } + } + freeDeviceSlice(&state.dGammaScalar) + done <- runErr + }() + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). state.dGammaScalar = uploadScalarMont(params.gamma) state.dZS = state.getTempDeviceSlice(n) if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + runErr = fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) return } // Step 1: Compute and apply blinding polynomial evaluations (if enabled) if useBlinding { - dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(state, n, twiddles0, bp) if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { - done <- err + runErr = err return } } @@ -4401,7 +4613,7 @@ func (s *instance) gpuEvaluateConstraints( var err error state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) return } @@ -4410,48 +4622,50 @@ func (s *instance) gpuEvaluateConstraints( dAlphaStd := uploadScalarStd(params.alpha) if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { dAlphaStd.Free() - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) return } dAlphaStd.Free() state.putTempDeviceSlice(state.dOrdering, params.n) + state.dOrdering = icicle_core.DeviceSlice{} // dResult += alpha^2 * local state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: local: %w", err) return } var alphaSquared fr.Element alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) dAlphaSquaredStd := uploadScalarStd(alphaSquared) - dTmp := state.getTempDeviceSlice(params.n) + dTmp = state.getTempDeviceSlice(params.n) if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { dAlphaSquaredStd.Free() - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) return } dAlphaSquaredStd.Free() if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dLocal, params.n) + state.dLocal = icicle_core.DeviceSlice{} state.putTempDeviceSlice(dTmp, params.n) + dTmp = icicle_core.DeviceSlice{} // dResult += gate state.dGate, err = computeGateConstraint(state, params, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) return } if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dGate, params.n) + state.dGate = icicle_core.DeviceSlice{} // Step 5: materialize result either on host or as a persistent device slice. if result != nil { @@ -4460,22 +4674,15 @@ func (s *instance) gpuEvaluateConstraints( } else { resultOnDevice = s.getTempDeviceSlice(params.n) if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) return } } - // Return dResult pool slice after materialization. - state.putTempDeviceSlice(state.dResult, params.n) - - // Return all allocated polynomial buffers to the pool. - // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). - // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are - // returned to pool inside computeGateConstraint. dQk and the original gpuState slices - // are owned by gpuState and will be freed separately. - state.freeAllocatedPolyBuffers() - - done <- nil + // dResult and the tracked polynomial buffers (blinded L, R, O, Z and + // scaled S1, S2, S3) are returned to the pool by the cleanup defer. + // dQk and the original gpuState slices are owned by gpuState and will + // be freed separately. }) err := <-done @@ -4608,6 +4815,7 @@ func (s *instance) prepareStatisticalZKQuotientShards( prepareDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepareDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -4812,6 +5020,7 @@ func (s *instance) inverseAndMergeShards( var dMerged icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") if cfgErr != nil { done <- cfgErr @@ -4985,6 +5194,7 @@ func (s *instance) divideByZHOnGPU( // So we can divide by Z_H by scaling each shard with its corresponding inverse. scaleDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(scaleDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") if cfgErr != nil { scaleDone <- cfgErr @@ -5034,6 +5244,7 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") if cfgErr != nil { done <- cfgErr @@ -5081,8 +5292,15 @@ func commitOnGPUWithDeviceBasesChunked( var msmErr error done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer func() { + if r := recover(); r != nil && msmErr == nil { + // Convert a device-goroutine panic into an error instead of + // killing the embedding process. + msmErr = fmt.Errorf("icicle: MSM device task panicked: %v", r) + } + close(done) + }() commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) - close(done) }) <-done if msmErr != nil { @@ -5256,6 +5474,7 @@ func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSl var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -5429,6 +5648,7 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") if cfgErr != nil { prepDone <- cfgErr @@ -5704,6 +5924,7 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") if cfgErr != nil { done <- cfgErr @@ -5781,6 +6002,7 @@ func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.El witnessSize := len(coeffs) - 1 divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dCoeffs := uploadVector(coeffs) @@ -5845,6 +6067,7 @@ func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.D } divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dPoint := uploadScalarStd(point) @@ -5970,6 +6193,12 @@ func (s *instance) buildLinearizedSelectorTermsOnGPU( } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr } @@ -6158,6 +6387,12 @@ func (s *instance) addZContributionToLinearizedOnGPU( var dScale icicle_core.DeviceSlice var dScaledZ icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before returning temporary buffers to the pool. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6225,6 +6460,12 @@ func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( var dZetaStd icicle_core.DeviceSlice var dZhStd icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before reusing temporary quotient vectors. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6595,9 +6836,10 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) } - chInitDomain := make(chan struct{}) + chInitDomain := make(chan error, 1) initDomainQueuedAt := time.Now() icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chInitDomain) initDomainStartedAt := time.Now() initCfg := icicle_core.GetDefaultNTTInitDomainConfig() ext := config_extension.Create() @@ -6625,41 +6867,52 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { ) } if e != icicle_runtime.Success { - panic("icicle: InitDomain failed") + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return } - close(chInitDomain) + chInitDomain <- nil }) - <-chInitDomain + if err := <-chInitDomain; err != nil { + return err + } if isNttTrace { fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) } - chLag := make(chan struct{}) - chCan := make(chan struct{}) + chLag := make(chan error, 1) + chCan := make(chan error, 1) if len(pk.KzgLagrange.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chLag) g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) - close(chLag) + chLag <- nil }) } else { - close(chLag) + chLag <- nil } if len(pk.Kzg.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chCan) g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) - close(chCan) + chCan <- nil }) } else { - close(chCan) + chCan <- nil } - <-chLag - <-chCan + errLag := <-chLag + errCan := <-chCan + if errLag != nil { + return fmt.Errorf("icicle: uploading Lagrange SRS to device failed: %w", errLag) + } + if errCan != nil { + return fmt.Errorf("icicle: uploading canonical SRS to device failed: %w", errCan) + } return nil } @@ -6989,6 +7242,12 @@ func (s *instance) BuildRatioCopyConstraintIcicle( var dNum, dDen icicle_core.DeviceSlice var dSupportFlat, dPermFlat icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for copy-constraint accumulation before cleanup. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "BuildRatioCopyConstraintIcicle"); syncErr != nil && runErr == nil { runErr = syncErr diff --git a/backend/accelerated/icicle/plonk/bn254/icicle.go b/backend/accelerated/icicle/plonk/bn254/icicle.go index 718ad8d37d..6d04b07146 100644 --- a/backend/accelerated/icicle/plonk/bn254/icicle.go +++ b/backend/accelerated/icicle/plonk/bn254/icicle.go @@ -264,32 +264,47 @@ func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts defer instance.releaseSharedGPUState() defer instance.releaseLinearizedEvalGPUState() + // goStage runs a prover stage on the errgroup, converting a panic into an + // error: several device helpers panic on GPU allocation failure, and an + // unrecovered panic in a stage goroutine would kill the embedding process + // instead of failing the Prove call. + goStage := func(name string, fn func() error) { + g.Go(func() (err error) { + defer func() { + if r := recover(); r != nil && err == nil { + err = fmt.Errorf("icicle plonk: stage %s panicked: %v", name, r) + } + }() + return fn() + }) + } + // solve constraints - g.Go(instance.solveConstraints) + goStage("solveConstraints", instance.solveConstraints) // complete qk - g.Go(instance.completeQk) + goStage("completeQk", instance.completeQk) // init blinding polynomials - g.Go(instance.initBlindingPolynomials) + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) // derive gamma, beta (copy constraint) - g.Go(instance.deriveGammaAndBeta) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) // compute accumulating ratio for the copy constraint - g.Go(instance.buildRatioCopyConstraint) + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) // compute h - g.Go(instance.computeQuotient) + goStage("computeQuotient", instance.computeQuotient) // open Z (blinded) at ωζ (proof.ZShiftedOpening) - g.Go(instance.openZ) + goStage("openZ", instance.openZ) // linearized polynomial - g.Go(instance.computeLinearizedPolynomial) + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) // Batch opening (no internal timer of its own — time the whole stage here) - g.Go(func() error { + goStage("batchOpening", func() error { startBatchOpening := time.Now() err := instance.batchOpening() if isProfileMode { @@ -502,6 +517,21 @@ func (s *instance) solveConstraints() error { // Try to load raw solver values from cache (fastest path) rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + // The raw solver cache stores the BSB22 commitment polynomials verbatim, + // including their two random blinding rows (see bsb22Hint). Replaying them + // across proofs makes Bsb22Commitments identical (linkable) and, once the + // committed polynomial has been opened at more than two distinct zetas, + // leaks linear relations over the committed private wires. Re-randomizing + // on load is not possible either: the commitment hash feeds back into a + // witness wire, so fresh blinding would invalidate every cached wire + // downstream of it. When zero-knowledge matters (blinding enabled, the + // default), the cache is therefore disabled for circuits with BSB22 + // commitments. With GNARK_DISABLE_BLINDING set, zero-knowledge is already + // explicitly forfeited and the replay leaks nothing new. + if rawCachePath != "" && len(s.commitmentInfo) > 0 && useBlinding { + log.Warn().Str("file", rawCachePath).Msg("GNARK_RAW_SOLVER_CACHE ignored: replaying cached BSB22 commitment blinding across proofs would break zero-knowledge (set GNARK_DISABLE_BLINDING to opt out of zero-knowledge and cache anyway)") + rawCachePath = "" + } if rawCachePath != "" { if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { // Reconstruct L, R, O from raw values @@ -988,6 +1018,7 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone := make(chan error, 1) var dPersist icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(copyDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") if cfgErr != nil { copyDone <- cfgErr @@ -1078,6 +1109,7 @@ func (s *instance) openZ() (err error) { var dBlindedCanonical icicle_core.DeviceSlice buildDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(buildDone) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") if cfgErr != nil { buildDone <- cfgErr @@ -1403,7 +1435,9 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) doneAddZ() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + // dLin is a pool buffer: return it via Put, not a direct free, so the + // pool's outstanding tracking stays consistent. + s.putTempDeviceSlice(dLin, dLin.Len()) return err } @@ -1411,7 +1445,7 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) doneSubtractH() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + s.putTempDeviceSlice(dLin, dLin.Len()) return err } s.linearizedPolynomialGPU = dLin @@ -1452,7 +1486,12 @@ func (s *instance) batchOpening() error { } defer func() { - freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + // The linearized polynomial is a pool buffer: return it via Put, not + // a direct free, so the pool's outstanding tracking stays consistent. + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } if s.hGPU != nil { s.freeGPUQuotient(s.hGPU) s.hGPU = nil @@ -1536,6 +1575,7 @@ func (s *instance) batchOpening() error { var dFold icicle_core.DeviceSlice foldDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(foldDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") if cfgErr != nil { foldDone <- cfgErr @@ -1602,6 +1642,7 @@ func (s *instance) batchOpening() error { divDone := make(chan error, 1) witnessSize := dFold.Len() - 1 icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") if cfgErr != nil { divDone <- cfgErr @@ -1898,12 +1939,15 @@ func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) var dWitness icicle_core.DeviceSlice - uploadDone := make(chan struct{}, 1) + uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) dWitness = uploadVector(hCoeffs) - close(uploadDone) + uploadDone <- nil }) - <-uploadDone + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) freeSliceOnDevice(&dWitness, &s.device) if err != nil { @@ -1995,6 +2039,7 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") if cfgErr != nil { done <- cfgErr @@ -2283,6 +2328,7 @@ func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicl var dTwiddles0 icicle_core.DeviceSlice uploadTwiddlesDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadTwiddlesDone) if s.tempGPUMemPool != nil { s.tempGPUMemPool.FreeAll() } @@ -2355,6 +2401,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo batchInvertDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(batchInvertDone) batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) }) if err := <-batchInvertDone; err != nil { @@ -2590,6 +2637,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(mergeDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") if cfgErr != nil { mergeDone <- cfgErr @@ -2620,6 +2668,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cresHost := icicle_core.HostSliceFromElements(cres) downloadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(downloadDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") if cfgErr != nil { downloadDone <- cfgErr @@ -2825,6 +2874,12 @@ func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop. } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for snapshot cloning before handing state to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { runErr = syncErr @@ -2909,6 +2964,12 @@ func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } @@ -2961,14 +3022,29 @@ func (s *instance) releaseTempGPUMemoryPool() { if s == nil || s.tempGPUMemPool == nil { return } + // Backstop for instance-owned device buffers whose owning stage may have + // exited (via error or ctx cancellation) before registering its own + // cleanup defer: openZ owns the two Z buffers, batchOpening owns the + // quotient and the linearized polynomial. On the success path these are + // already released and zeroed, making every free below a no-op. This must + // run before FreeAll so pool-owned buffers (blindedZCanonicalGPU, hGPU) + // are back in the pool when it frees everything. freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) if !s.blindedZCanonicalGPU.IsEmpty() { s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} } + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { - s.tempGPUMemPool.FreeAll() + s.tempGPUMemPool.Shutdown() close(done) }) <-done @@ -3026,6 +3102,12 @@ func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysSta } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for GPU uploads in ensurePolysOnSharedGPU. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -3080,6 +3162,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali // Upload scaling vectors to GPU uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(uploadDone) scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) scalingHost := icicle_core.HostSliceFromElements(scalingVector) scalingHost.CopyToDevice(&scalingVectorDevice, false) @@ -3098,10 +3181,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } uploadDone <- nil }) + // Return the scaling vectors to the pool on every exit path. This is safe + // because we never return before draining every per-polynomial done + // channel, and each worker synchronizes its stream before signalling, so + // no in-flight kernel can still reference the slices. + defer func() { + if !scalingVectorDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + } + if !scalingVectorRevDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + } + }() if err := <-uploadDone; err != nil { return err } + // Validate all polynomials before launching any GPU work so an invalid + // entry cannot abandon already-scheduled workers. + for _, p := range state.polys { + if p != nil && p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + } + doneChans := make([]chan error, len(state.polys)) for i := range state.polys { @@ -3110,10 +3213,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali continue } - if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { - return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) - } - done := make(chan error, 1) doneChans[i] = done scalarsDevice := state.deviceSlices[i] @@ -3122,7 +3221,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: CreateStream failed: %s", eStream.AsString()) + return + } + // Synchronize and destroy the stream on every path (including + // errors) before signalling done: the caller may free device + // buffers as soon as all workers have reported. + var runErr error + defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: DestroyStream failed: %s", eDestroy.AsString()) + } + done <- runErr + }() cfg.StreamHandle = stream cfg.IsAsync = true @@ -3136,7 +3258,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KRN } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) return } @@ -3155,7 +3277,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) return } @@ -3168,26 +3290,26 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KNR // Regular → BitReverse } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) - return - } - - if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) return } - done <- nil }) } - // Wait for all scheduled tasks + // Wait for ALL scheduled tasks even when one fails: returning on the + // first error would let the caller free device slices that sibling + // workers are still writing to. + var firstErr error for i := range doneChans { if doneChans[i] != nil { - if err := <-doneChans[i]; err != nil { - return err + if err := <-doneChans[i]; err != nil && firstErr == nil { + firstErr = err } } } + if firstErr != nil { + return firstErr + } // Update polynomial metadata: final result is in Lagrange, same layout as original for _, p := range state.polys { @@ -3195,10 +3317,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali p.Basis = iop.Lagrange } } - - // Free scaling vectors from device - s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) - s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) return nil } @@ -3270,7 +3388,11 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: CreateStream failed: %s", eStream.AsString())) + } + defer icicle_runtime.DestroyStream(stream) cfg.StreamHandle = stream cfg.IsAsync = true @@ -3292,7 +3414,9 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) } - icicle_runtime.SynchronizeStream(stream) + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: SynchronizeStream failed: %s", eSync.AsString())) + } // Update metadata inside closure to avoid race p.Basis = iop.Canonical @@ -3341,13 +3465,20 @@ func (s *instance) freeGPUPolys(state *gpuPolysState) { // Must be used within RunOnDevice context to ensure thread safety per device. type gpuMemoryPool struct { freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice - mu sync.Mutex + // outstanding tracks slices handed out by Get and not yet returned so + // Shutdown can reclaim buffers abandoned by error paths. Keyed by the + // device pointer, not the slice value: CopyToDevice mutates the slice's + // length field, so the value at Put time may differ from the one at Get + // time while the underlying buffer is the same. + outstanding map[unsafe.Pointer]icicle_core.DeviceSlice + mu sync.Mutex } // newGPUMemoryPool creates a new GPU memory pool. func newGPUMemoryPool() *gpuMemoryPool { return &gpuMemoryPool{ - freeSlices: make(map[int][]icicle_core.DeviceSlice), + freeSlices: make(map[int][]icicle_core.DeviceSlice), + outstanding: make(map[unsafe.Pointer]icicle_core.DeviceSlice), } } @@ -3361,12 +3492,14 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { // Reuse the last slice slice := slices[len(slices)-1] p.freeSlices[n] = slices[:len(slices)-1] + p.outstanding[slice.AsUnsafePointer()] = slice return slice } // No free slice available, allocate a new one. // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3379,6 +3512,7 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3394,11 +3528,14 @@ func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { p.mu.Lock() defer p.mu.Unlock() + delete(p.outstanding, ds.AsUnsafePointer()) // Add to the pool p.freeSlices[n] = append(p.freeSlices[n], ds) } -// FreeAll releases all pooled device slices. +// FreeAll releases all idle pooled device slices. Outstanding slices (handed +// out by Get and not yet returned) are left alone: FreeAll is also used +// mid-prove to relieve memory pressure while pool buffers are still live. func (p *gpuMemoryPool) FreeAll() { p.mu.Lock() defer p.mu.Unlock() @@ -3411,6 +3548,27 @@ func (p *gpuMemoryPool) FreeAll() { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) } +// Shutdown releases every pooled slice, including outstanding ones that were +// abandoned by error paths (a failed stage returns without Put-ing its +// temporaries; without this they would leak for the lifetime of the process). +// Only safe once no GPU work can still reference pool buffers, i.e. at the +// end of Prove after every stage goroutine has completed. +func (p *gpuMemoryPool) Shutdown() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + for _, ds := range p.outstanding { + _ = ds.Free() + } + p.outstanding = make(map[unsafe.Pointer]icicle_core.DeviceSlice) +} + // allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. // Use when the destination is fully overwritten by a kernel. func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { @@ -3539,6 +3697,26 @@ func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) } } +// devicePanicToError is deferred at the top of RunOnDevice closures that +// report completion by sending on an error channel. The pool and upload +// helpers panic on GPU allocation failure — the expected failure mode for a +// memory-hungry prover — and a panic in a goroutine cannot be recovered by +// the caller, so without this boundary it kills the whole embedding process +// instead of failing the Prove call. +func devicePanicToError(done chan<- error) { + if r := recover(); r != nil { + err := fmt.Errorf("icicle: device task panicked: %v", r) + select { + case done <- err: + default: + // The closure already reported success and the caller has moved + // on; all we can do is log. + log := logger.Logger() + log.Error().Err(err).Msg("icicle: panic on device goroutine after completion was signalled") + } + } +} + // uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. // Use this for additions where the vector is already in Montgomery form. func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { @@ -3779,8 +3957,11 @@ func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { } // computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. -// Returns device slices for the blinding polynomials. +// Returns device slices for the blinding polynomials. The buffers come from +// the shared temp pool (via state) so an aborted iteration cannot leak them +// past the end of Prove. func computeBlindingPolynomials( + state *gpuConstraintEvalState, n int, twiddles0 []fr.Element, bp []*iop.Polynomial, @@ -3803,11 +3984,17 @@ func computeBlindingPolynomials( } }) - dBlindL = uploadVector(blindL) - dBlindR = uploadVector(blindR) - dBlindO = uploadVector(blindO) - dBlindZ = uploadVector(blindZ) - dBlindZS = uploadVector(blindZS) + uploadFromPool := func(vec []fr.Element) icicle_core.DeviceSlice { + d := state.getTempDeviceSlice(n) + host := icicle_core.HostSliceFromElements(vec) + host.CopyToDevice(&d, false) + return d + } + dBlindL = uploadFromPool(blindL) + dBlindR = uploadFromPool(blindR) + dBlindO = uploadFromPool(blindO) + dBlindZ = uploadFromPool(blindZ) + dBlindZS = uploadFromPool(blindZS) return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS } @@ -3855,12 +4042,12 @@ func applyBlindingToPolynomials( return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) } - // Free blinding vectors - no longer needed after creating blinded polynomials - dBlindL.Free() - dBlindR.Free() - dBlindO.Free() - dBlindZ.Free() - dBlindZS.Free() + // Return blinding vectors to the pool - no longer needed after creating blinded polynomials + state.putTempDeviceSlice(dBlindL, params.n) + state.putTempDeviceSlice(dBlindR, params.n) + state.putTempDeviceSlice(dBlindO, params.n) + state.putTempDeviceSlice(dBlindZ, params.n) + state.putTempDeviceSlice(dBlindZS, params.n) return nil } @@ -4123,7 +4310,8 @@ func computeOrderingConstraint( // Free dGammaScalar - no longer needed after computing a2, b2, c2 // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. - state.dGammaScalar.Free() + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) state.putTempDeviceSlice(dScaledS, params.n) dBetaStd.Free() @@ -4366,21 +4554,45 @@ func (s *instance) gpuEvaluateConstraints( icicle_runtime.RunOnDevice(device, func(args ...any) { state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + var runErr error + var dTmp icicle_core.DeviceSlice + defer func() { + if r := recover(); r != nil && runErr == nil { + // The pool and upload helpers panic on GPU allocation + // failure; a panic in a goroutine cannot be recovered by the + // caller, so convert it into a prover error here instead of + // killing the embedding process. + runErr = fmt.Errorf("gpuEvaluateConstraints: device task panicked: %v", r) + } + // Free everything an aborted pipeline may have left behind. On + // the success path every release below is a no-op: state fields + // are zeroed when returned and the tracked list is emptied. + state.freeAllocatedPolyBuffers() + for _, t := range []*icicle_core.DeviceSlice{&state.dZS, &state.dOrdering, &state.dLocal, &state.dGate, &state.dResult, &dTmp} { + if !t.IsEmpty() { + state.putTempDeviceSlice(*t, t.Len()) + *t = icicle_core.DeviceSlice{} + } + } + freeDeviceSlice(&state.dGammaScalar) + done <- runErr + }() + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). state.dGammaScalar = uploadScalarMont(params.gamma) state.dZS = state.getTempDeviceSlice(n) if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + runErr = fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) return } // Step 1: Compute and apply blinding polynomial evaluations (if enabled) if useBlinding { - dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(state, n, twiddles0, bp) if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { - done <- err + runErr = err return } } @@ -4401,7 +4613,7 @@ func (s *instance) gpuEvaluateConstraints( var err error state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) return } @@ -4410,48 +4622,50 @@ func (s *instance) gpuEvaluateConstraints( dAlphaStd := uploadScalarStd(params.alpha) if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { dAlphaStd.Free() - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) return } dAlphaStd.Free() state.putTempDeviceSlice(state.dOrdering, params.n) + state.dOrdering = icicle_core.DeviceSlice{} // dResult += alpha^2 * local state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: local: %w", err) return } var alphaSquared fr.Element alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) dAlphaSquaredStd := uploadScalarStd(alphaSquared) - dTmp := state.getTempDeviceSlice(params.n) + dTmp = state.getTempDeviceSlice(params.n) if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { dAlphaSquaredStd.Free() - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) return } dAlphaSquaredStd.Free() if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dLocal, params.n) + state.dLocal = icicle_core.DeviceSlice{} state.putTempDeviceSlice(dTmp, params.n) + dTmp = icicle_core.DeviceSlice{} // dResult += gate state.dGate, err = computeGateConstraint(state, params, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) return } if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dGate, params.n) + state.dGate = icicle_core.DeviceSlice{} // Step 5: materialize result either on host or as a persistent device slice. if result != nil { @@ -4460,22 +4674,15 @@ func (s *instance) gpuEvaluateConstraints( } else { resultOnDevice = s.getTempDeviceSlice(params.n) if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) return } } - // Return dResult pool slice after materialization. - state.putTempDeviceSlice(state.dResult, params.n) - - // Return all allocated polynomial buffers to the pool. - // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). - // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are - // returned to pool inside computeGateConstraint. dQk and the original gpuState slices - // are owned by gpuState and will be freed separately. - state.freeAllocatedPolyBuffers() - - done <- nil + // dResult and the tracked polynomial buffers (blinded L, R, O, Z and + // scaled S1, S2, S3) are returned to the pool by the cleanup defer. + // dQk and the original gpuState slices are owned by gpuState and will + // be freed separately. }) err := <-done @@ -4608,6 +4815,7 @@ func (s *instance) prepareStatisticalZKQuotientShards( prepareDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepareDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -4812,6 +5020,7 @@ func (s *instance) inverseAndMergeShards( var dMerged icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") if cfgErr != nil { done <- cfgErr @@ -4985,6 +5194,7 @@ func (s *instance) divideByZHOnGPU( // So we can divide by Z_H by scaling each shard with its corresponding inverse. scaleDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(scaleDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") if cfgErr != nil { scaleDone <- cfgErr @@ -5034,6 +5244,7 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") if cfgErr != nil { done <- cfgErr @@ -5081,8 +5292,15 @@ func commitOnGPUWithDeviceBasesChunked( var msmErr error done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer func() { + if r := recover(); r != nil && msmErr == nil { + // Convert a device-goroutine panic into an error instead of + // killing the embedding process. + msmErr = fmt.Errorf("icicle: MSM device task panicked: %v", r) + } + close(done) + }() commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) - close(done) }) <-done if msmErr != nil { @@ -5256,6 +5474,7 @@ func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSl var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -5429,6 +5648,7 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") if cfgErr != nil { prepDone <- cfgErr @@ -5704,6 +5924,7 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") if cfgErr != nil { done <- cfgErr @@ -5781,6 +6002,7 @@ func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.El witnessSize := len(coeffs) - 1 divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dCoeffs := uploadVector(coeffs) @@ -5845,6 +6067,7 @@ func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.D } divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dPoint := uploadScalarStd(point) @@ -5970,6 +6193,12 @@ func (s *instance) buildLinearizedSelectorTermsOnGPU( } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr } @@ -6158,6 +6387,12 @@ func (s *instance) addZContributionToLinearizedOnGPU( var dScale icicle_core.DeviceSlice var dScaledZ icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before returning temporary buffers to the pool. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6225,6 +6460,12 @@ func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( var dZetaStd icicle_core.DeviceSlice var dZhStd icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before reusing temporary quotient vectors. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6595,9 +6836,10 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) } - chInitDomain := make(chan struct{}) + chInitDomain := make(chan error, 1) initDomainQueuedAt := time.Now() icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chInitDomain) initDomainStartedAt := time.Now() initCfg := icicle_core.GetDefaultNTTInitDomainConfig() ext := config_extension.Create() @@ -6625,41 +6867,52 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { ) } if e != icicle_runtime.Success { - panic("icicle: InitDomain failed") + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return } - close(chInitDomain) + chInitDomain <- nil }) - <-chInitDomain + if err := <-chInitDomain; err != nil { + return err + } if isNttTrace { fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) } - chLag := make(chan struct{}) - chCan := make(chan struct{}) + chLag := make(chan error, 1) + chCan := make(chan error, 1) if len(pk.KzgLagrange.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chLag) g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) - close(chLag) + chLag <- nil }) } else { - close(chLag) + chLag <- nil } if len(pk.Kzg.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chCan) g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) - close(chCan) + chCan <- nil }) } else { - close(chCan) + chCan <- nil } - <-chLag - <-chCan + errLag := <-chLag + errCan := <-chCan + if errLag != nil { + return fmt.Errorf("icicle: uploading Lagrange SRS to device failed: %w", errLag) + } + if errCan != nil { + return fmt.Errorf("icicle: uploading canonical SRS to device failed: %w", errCan) + } return nil } @@ -6989,6 +7242,12 @@ func (s *instance) BuildRatioCopyConstraintIcicle( var dNum, dDen icicle_core.DeviceSlice var dSupportFlat, dPermFlat icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for copy-constraint accumulation before cleanup. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "BuildRatioCopyConstraintIcicle"); syncErr != nil && runErr == nil { runErr = syncErr diff --git a/backend/accelerated/icicle/plonk/bw6-761/icicle.go b/backend/accelerated/icicle/plonk/bw6-761/icicle.go index 0796f293f0..2e510bf3e1 100644 --- a/backend/accelerated/icicle/plonk/bw6-761/icicle.go +++ b/backend/accelerated/icicle/plonk/bw6-761/icicle.go @@ -264,32 +264,47 @@ func Prove(spr *cs.SparseR1CS, pk *ProvingKey, fullWitness witness.Witness, opts defer instance.releaseSharedGPUState() defer instance.releaseLinearizedEvalGPUState() + // goStage runs a prover stage on the errgroup, converting a panic into an + // error: several device helpers panic on GPU allocation failure, and an + // unrecovered panic in a stage goroutine would kill the embedding process + // instead of failing the Prove call. + goStage := func(name string, fn func() error) { + g.Go(func() (err error) { + defer func() { + if r := recover(); r != nil && err == nil { + err = fmt.Errorf("icicle plonk: stage %s panicked: %v", name, r) + } + }() + return fn() + }) + } + // solve constraints - g.Go(instance.solveConstraints) + goStage("solveConstraints", instance.solveConstraints) // complete qk - g.Go(instance.completeQk) + goStage("completeQk", instance.completeQk) // init blinding polynomials - g.Go(instance.initBlindingPolynomials) + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) // derive gamma, beta (copy constraint) - g.Go(instance.deriveGammaAndBeta) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) // compute accumulating ratio for the copy constraint - g.Go(instance.buildRatioCopyConstraint) + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) // compute h - g.Go(instance.computeQuotient) + goStage("computeQuotient", instance.computeQuotient) // open Z (blinded) at ωζ (proof.ZShiftedOpening) - g.Go(instance.openZ) + goStage("openZ", instance.openZ) // linearized polynomial - g.Go(instance.computeLinearizedPolynomial) + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) // Batch opening (no internal timer of its own — time the whole stage here) - g.Go(func() error { + goStage("batchOpening", func() error { startBatchOpening := time.Now() err := instance.batchOpening() if isProfileMode { @@ -502,6 +517,21 @@ func (s *instance) solveConstraints() error { // Try to load raw solver values from cache (fastest path) rawCachePath := os.Getenv("GNARK_RAW_SOLVER_CACHE") + // The raw solver cache stores the BSB22 commitment polynomials verbatim, + // including their two random blinding rows (see bsb22Hint). Replaying them + // across proofs makes Bsb22Commitments identical (linkable) and, once the + // committed polynomial has been opened at more than two distinct zetas, + // leaks linear relations over the committed private wires. Re-randomizing + // on load is not possible either: the commitment hash feeds back into a + // witness wire, so fresh blinding would invalidate every cached wire + // downstream of it. When zero-knowledge matters (blinding enabled, the + // default), the cache is therefore disabled for circuits with BSB22 + // commitments. With GNARK_DISABLE_BLINDING set, zero-knowledge is already + // explicitly forfeited and the replay leaks nothing new. + if rawCachePath != "" && len(s.commitmentInfo) > 0 && useBlinding { + log.Warn().Str("file", rawCachePath).Msg("GNARK_RAW_SOLVER_CACHE ignored: replaying cached BSB22 commitment blinding across proofs would break zero-knowledge (set GNARK_DISABLE_BLINDING to opt out of zero-knowledge and cache anyway)") + rawCachePath = "" + } if rawCachePath != "" { if rawValues, err := cs.LoadRawSolverValues(rawCachePath); err == nil { // Reconstruct L, R, O from raw values @@ -988,6 +1018,7 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone := make(chan error, 1) var dPersist icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(copyDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") if cfgErr != nil { copyDone <- cfgErr @@ -1078,6 +1109,7 @@ func (s *instance) openZ() (err error) { var dBlindedCanonical icicle_core.DeviceSlice buildDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(buildDone) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") if cfgErr != nil { buildDone <- cfgErr @@ -1403,7 +1435,9 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.addZContributionToLinearizedOnGPU(dLin, s.blindedZCanonicalGPU, evals.blzeta, evals.brzeta, evals.bozeta) doneAddZ() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + // dLin is a pool buffer: return it via Put, not a direct free, so the + // pool's outstanding tracking stays consistent. + s.putTempDeviceSlice(dLin, dLin.Len()) return err } @@ -1411,7 +1445,7 @@ func (s *instance) computeLinearizedPolynomial() error { err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) doneSubtractH() if err != nil { - freeSliceOnDevice(&dLin, &s.device) + s.putTempDeviceSlice(dLin, dLin.Len()) return err } s.linearizedPolynomialGPU = dLin @@ -1452,7 +1486,12 @@ func (s *instance) batchOpening() error { } defer func() { - freeSliceOnDevice(&s.linearizedPolynomialGPU, &s.device) + // The linearized polynomial is a pool buffer: return it via Put, not + // a direct free, so the pool's outstanding tracking stays consistent. + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } if s.hGPU != nil { s.freeGPUQuotient(s.hGPU) s.hGPU = nil @@ -1536,6 +1575,7 @@ func (s *instance) batchOpening() error { var dFold icicle_core.DeviceSlice foldDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(foldDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") if cfgErr != nil { foldDone <- cfgErr @@ -1602,6 +1642,7 @@ func (s *instance) batchOpening() error { divDone := make(chan error, 1) witnessSize := dFold.Len() - 1 icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") if cfgErr != nil { divDone <- cfgErr @@ -1898,12 +1939,15 @@ func (s *instance) batchOpeningHostFoldGPUCommitFromPolynomials( hCoeffs := dividePolyByXMinusAHost(foldedPolynomials, foldedEval, s.zeta) var dWitness icicle_core.DeviceSlice - uploadDone := make(chan struct{}, 1) + uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) dWitness = uploadVector(hCoeffs) - close(uploadDone) + uploadDone <- nil }) - <-uploadDone + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } h, err := commitOnGPUCanonicalDevice(dWitness, &s.device, s.pk) freeSliceOnDevice(&dWitness, &s.device) if err != nil { @@ -1995,6 +2039,7 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") if cfgErr != nil { done <- cfgErr @@ -2283,6 +2328,7 @@ func (s *instance) uploadComputeNumeratorTwiddles(twiddles0 []fr.Element) (icicl var dTwiddles0 icicle_core.DeviceSlice uploadTwiddlesDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadTwiddlesDone) if s.tempGPUMemPool != nil { s.tempGPUMemPool.FreeAll() } @@ -2355,6 +2401,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo batchInvertDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(batchInvertDone) batchInvertDone <- s.buildAndInvertPrecomputedDenominatorsOnCurrentDevice(loopCtx) }) if err := <-batchInvertDone; err != nil { @@ -2590,6 +2637,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(mergeDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") if cfgErr != nil { mergeDone <- cfgErr @@ -2620,6 +2668,7 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cresHost := icicle_core.HostSliceFromElements(cres) downloadDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(downloadDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") if cfgErr != nil { downloadDone <- cfgErr @@ -2825,6 +2874,12 @@ func (s *instance) clonePolysOnGPUFromState(source *gpuPolysState, polys []*iop. } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for snapshot cloning before handing state to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "clonePolysOnGPUFromState"); syncErr != nil && runErr == nil { runErr = syncErr @@ -2909,6 +2964,12 @@ func (s *instance) uploadPolysToGPUState(polys []*iop.Polynomial, label string) } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } @@ -2961,14 +3022,29 @@ func (s *instance) releaseTempGPUMemoryPool() { if s == nil || s.tempGPUMemPool == nil { return } + // Backstop for instance-owned device buffers whose owning stage may have + // exited (via error or ctx cancellation) before registering its own + // cleanup defer: openZ owns the two Z buffers, batchOpening owns the + // quotient and the linearized polynomial. On the success path these are + // already released and zeroed, making every free below a no-op. This must + // run before FreeAll so pool-owned buffers (blindedZCanonicalGPU, hGPU) + // are back in the pool when it frees everything. freeSliceOnDevice(&s.polyZLagrangeGPU, &s.device) if !s.blindedZCanonicalGPU.IsEmpty() { s.putTempDeviceSlice(s.blindedZCanonicalGPU, s.blindedZCanonicalGPU.Len()) s.blindedZCanonicalGPU = icicle_core.DeviceSlice{} } + if !s.linearizedPolynomialGPU.IsEmpty() { + s.putTempDeviceSlice(s.linearizedPolynomialGPU, s.linearizedPolynomialGPU.Len()) + s.linearizedPolynomialGPU = icicle_core.DeviceSlice{} + } + if s.hGPU != nil { + s.freeGPUQuotient(s.hGPU) + s.hGPU = nil + } done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { - s.tempGPUMemPool.FreeAll() + s.tempGPUMemPool.Shutdown() close(done) }) <-done @@ -3026,6 +3102,12 @@ func (s *instance) ensurePolysOnSharedGPU(polys []*iop.Polynomial) (*gpuPolysSta } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for GPU uploads in ensurePolysOnSharedGPU. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "ensurePolysOnSharedGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -3080,6 +3162,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali // Upload scaling vectors to GPU uploadDone := make(chan error, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(uploadDone) scalingVectorDevice = s.getTempDeviceSlice(len(scalingVector)) scalingHost := icicle_core.HostSliceFromElements(scalingVector) scalingHost.CopyToDevice(&scalingVectorDevice, false) @@ -3098,10 +3181,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } uploadDone <- nil }) + // Return the scaling vectors to the pool on every exit path. This is safe + // because we never return before draining every per-polynomial done + // channel, and each worker synchronizes its stream before signalling, so + // no in-flight kernel can still reference the slices. + defer func() { + if !scalingVectorDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) + } + if !scalingVectorRevDevice.IsEmpty() { + s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) + } + }() if err := <-uploadDone; err != nil { return err } + // Validate all polynomials before launching any GPU work so an invalid + // entry cannot abandon already-scheduled workers. + for _, p := range state.polys { + if p != nil && p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { + return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) + } + } + doneChans := make([]chan error, len(state.polys)) for i := range state.polys { @@ -3110,10 +3213,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali continue } - if p.Basis != iop.Lagrange && p.Basis != iop.LagrangeCoset { - return fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: expected polynomial in Lagrange or LagrangeCoset form, got %v", p.Basis) - } - done := make(chan error, 1) doneChans[i] = done scalarsDevice := state.deviceSlices[i] @@ -3122,7 +3221,30 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: CreateStream failed: %s", eStream.AsString()) + return + } + // Synchronize and destroy the stream on every path (including + // errors) before signalling done: the caller may free device + // buffers as soon as all workers have reported. + var runErr error + defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + } + if eDestroy := icicle_runtime.DestroyStream(stream); eDestroy != icicle_runtime.Success && runErr == nil { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: DestroyStream failed: %s", eDestroy.AsString()) + } + done <- runErr + }() cfg.StreamHandle = stream cfg.IsAsync = true @@ -3136,7 +3258,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KRN } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: inverse NTT failed: %s", err.AsString()) return } @@ -3155,7 +3277,7 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali } if err := icicle_vecops.VecOp(scalarsDevice, scaleDevice, scalarsDevice, vecCfg, icicle_core.Mul); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: VecOp mul failed: %s", err.AsString()) return } @@ -3168,26 +3290,26 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali cfg.Ordering = icicle_core.KNR // Regular → BitReverse } if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KForward, &cfg, scalarsDevice); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) - return - } - - if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { - done <- fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: SynchronizeStream failed: %s", eSync.AsString()) + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) return } - done <- nil }) } - // Wait for all scheduled tasks + // Wait for ALL scheduled tasks even when one fails: returning on the + // first error would let the caller free device slices that sibling + // workers are still writing to. + var firstErr error for i := range doneChans { if doneChans[i] != nil { - if err := <-doneChans[i]; err != nil { - return err + if err := <-doneChans[i]; err != nil && firstErr == nil { + firstErr = err } } } + if firstErr != nil { + return firstErr + } // Update polynomial metadata: final result is in Lagrange, same layout as original for _, p := range state.polys { @@ -3195,10 +3317,6 @@ func (s *instance) gpuNTTInverseScaleForwardOnDevice(state *gpuPolysState, scali p.Basis = iop.Lagrange } } - - // Free scaling vectors from device - s.putTempDeviceSlice(scalingVectorDevice, scalingVectorDevice.Len()) - s.putTempDeviceSlice(scalingVectorRevDevice, scalingVectorRevDevice.Len()) return nil } @@ -3270,7 +3388,11 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe icicle_runtime.RunOnDevice(device, func(args ...any) { cfg := icicle_ntt.GetDefaultNttConfig() - stream, _ := icicle_runtime.CreateStream() + stream, eStream := icicle_runtime.CreateStream() + if eStream != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: CreateStream failed: %s", eStream.AsString())) + } + defer icicle_runtime.DestroyStream(stream) cfg.StreamHandle = stream cfg.IsAsync = true @@ -3292,7 +3414,9 @@ func (s *instance) gpuNTTInverseBatchOnState(state *gpuPolysState, pk *ProvingKe if err := icicle_ntt.Ntt(scalarsDevice, icicle_core.KInverse, &cfg, scalarsDevice); err != icicle_runtime.Success { panic(fmt.Sprintf("icicle: inverse NTT failed: %s", err.AsString())) } - icicle_runtime.SynchronizeStream(stream) + if eSync := icicle_runtime.SynchronizeStream(stream); eSync != icicle_runtime.Success { + panic(fmt.Sprintf("icicle: SynchronizeStream failed: %s", eSync.AsString())) + } // Update metadata inside closure to avoid race p.Basis = iop.Canonical @@ -3341,13 +3465,20 @@ func (s *instance) freeGPUPolys(state *gpuPolysState) { // Must be used within RunOnDevice context to ensure thread safety per device. type gpuMemoryPool struct { freeSlices map[int][]icicle_core.DeviceSlice // map[size][]slice - mu sync.Mutex + // outstanding tracks slices handed out by Get and not yet returned so + // Shutdown can reclaim buffers abandoned by error paths. Keyed by the + // device pointer, not the slice value: CopyToDevice mutates the slice's + // length field, so the value at Put time may differ from the one at Get + // time while the underlying buffer is the same. + outstanding map[unsafe.Pointer]icicle_core.DeviceSlice + mu sync.Mutex } // newGPUMemoryPool creates a new GPU memory pool. func newGPUMemoryPool() *gpuMemoryPool { return &gpuMemoryPool{ - freeSlices: make(map[int][]icicle_core.DeviceSlice), + freeSlices: make(map[int][]icicle_core.DeviceSlice), + outstanding: make(map[unsafe.Pointer]icicle_core.DeviceSlice), } } @@ -3361,12 +3492,14 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { // Reuse the last slice slice := slices[len(slices)-1] p.freeSlices[n] = slices[:len(slices)-1] + p.outstanding[slice.AsUnsafePointer()] = slice return slice } // No free slice available, allocate a new one. // If allocation fails (e.g. memory pressure), release idle cached slices and retry once. if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3379,6 +3512,7 @@ func (p *gpuMemoryPool) Get(n int) icicle_core.DeviceSlice { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) if ds, err := allocDeviceUninitialized(n); err == nil { + p.outstanding[ds.AsUnsafePointer()] = ds return ds } @@ -3394,11 +3528,14 @@ func (p *gpuMemoryPool) Put(ds icicle_core.DeviceSlice, n int) { p.mu.Lock() defer p.mu.Unlock() + delete(p.outstanding, ds.AsUnsafePointer()) // Add to the pool p.freeSlices[n] = append(p.freeSlices[n], ds) } -// FreeAll releases all pooled device slices. +// FreeAll releases all idle pooled device slices. Outstanding slices (handed +// out by Get and not yet returned) are left alone: FreeAll is also used +// mid-prove to relieve memory pressure while pool buffers are still live. func (p *gpuMemoryPool) FreeAll() { p.mu.Lock() defer p.mu.Unlock() @@ -3411,6 +3548,27 @@ func (p *gpuMemoryPool) FreeAll() { p.freeSlices = make(map[int][]icicle_core.DeviceSlice) } +// Shutdown releases every pooled slice, including outstanding ones that were +// abandoned by error paths (a failed stage returns without Put-ing its +// temporaries; without this they would leak for the lifetime of the process). +// Only safe once no GPU work can still reference pool buffers, i.e. at the +// end of Prove after every stage goroutine has completed. +func (p *gpuMemoryPool) Shutdown() { + p.mu.Lock() + defer p.mu.Unlock() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + p.freeSlices = make(map[int][]icicle_core.DeviceSlice) + for _, ds := range p.outstanding { + _ = ds.Free() + } + p.outstanding = make(map[unsafe.Pointer]icicle_core.DeviceSlice) +} + // allocDeviceUninitialized allocates device memory without uploading a zeroed host buffer. // Use when the destination is fully overwritten by a kernel. func allocDeviceUninitialized(n int) (icicle_core.DeviceSlice, error) { @@ -3539,6 +3697,26 @@ func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) } } +// devicePanicToError is deferred at the top of RunOnDevice closures that +// report completion by sending on an error channel. The pool and upload +// helpers panic on GPU allocation failure — the expected failure mode for a +// memory-hungry prover — and a panic in a goroutine cannot be recovered by +// the caller, so without this boundary it kills the whole embedding process +// instead of failing the Prove call. +func devicePanicToError(done chan<- error) { + if r := recover(); r != nil { + err := fmt.Errorf("icicle: device task panicked: %v", r) + select { + case done <- err: + default: + // The closure already reported success and the caller has moved + // on; all we can do is log. + log := logger.Logger() + log.Error().Err(err).Msg("icicle: panic on device goroutine after completion was signalled") + } + } +} + // uploadScalarMont uploads a scalar in MONTGOMERY form as a single-element device slice. // Use this for additions where the vector is already in Montgomery form. func uploadScalarMont(scalar fr.Element) icicle_core.DeviceSlice { @@ -3779,8 +3957,11 @@ func (s *gpuConstraintEvalState) freeAllocatedPolyBuffers() { } // computeBlindingPolynomials computes and uploads blinding polynomial evaluations to GPU. -// Returns device slices for the blinding polynomials. +// Returns device slices for the blinding polynomials. The buffers come from +// the shared temp pool (via state) so an aborted iteration cannot leak them +// past the end of Prove. func computeBlindingPolynomials( + state *gpuConstraintEvalState, n int, twiddles0 []fr.Element, bp []*iop.Polynomial, @@ -3803,11 +3984,17 @@ func computeBlindingPolynomials( } }) - dBlindL = uploadVector(blindL) - dBlindR = uploadVector(blindR) - dBlindO = uploadVector(blindO) - dBlindZ = uploadVector(blindZ) - dBlindZS = uploadVector(blindZS) + uploadFromPool := func(vec []fr.Element) icicle_core.DeviceSlice { + d := state.getTempDeviceSlice(n) + host := icicle_core.HostSliceFromElements(vec) + host.CopyToDevice(&d, false) + return d + } + dBlindL = uploadFromPool(blindL) + dBlindR = uploadFromPool(blindR) + dBlindO = uploadFromPool(blindO) + dBlindZ = uploadFromPool(blindZ) + dBlindZS = uploadFromPool(blindZS) return dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS } @@ -3855,12 +4042,12 @@ func applyBlindingToPolynomials( return fmt.Errorf("applyBlindingToPolynomials: VecOp add ZS failed: %s", err.AsString()) } - // Free blinding vectors - no longer needed after creating blinded polynomials - dBlindL.Free() - dBlindR.Free() - dBlindO.Free() - dBlindZ.Free() - dBlindZS.Free() + // Return blinding vectors to the pool - no longer needed after creating blinded polynomials + state.putTempDeviceSlice(dBlindL, params.n) + state.putTempDeviceSlice(dBlindR, params.n) + state.putTempDeviceSlice(dBlindO, params.n) + state.putTempDeviceSlice(dBlindZ, params.n) + state.putTempDeviceSlice(dBlindZS, params.n) return nil } @@ -4123,7 +4310,8 @@ func computeOrderingConstraint( // Free dGammaScalar - no longer needed after computing a2, b2, c2 // Note: dL, dR, dO, dS1, dS2, dS3 are part of gpuState and will be freed later. - state.dGammaScalar.Free() + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) state.putTempDeviceSlice(dScaledS, params.n) dBetaStd.Free() @@ -4366,21 +4554,45 @@ func (s *instance) gpuEvaluateConstraints( icicle_runtime.RunOnDevice(device, func(args ...any) { state := initializeConstraintEvalState(getDeviceSlice, s.getTempDeviceSlice, s.putTempDeviceSlice) + var runErr error + var dTmp icicle_core.DeviceSlice + defer func() { + if r := recover(); r != nil && runErr == nil { + // The pool and upload helpers panic on GPU allocation + // failure; a panic in a goroutine cannot be recovered by the + // caller, so convert it into a prover error here instead of + // killing the embedding process. + runErr = fmt.Errorf("gpuEvaluateConstraints: device task panicked: %v", r) + } + // Free everything an aborted pipeline may have left behind. On + // the success path every release below is a no-op: state fields + // are zeroed when returned and the tracked list is emptied. + state.freeAllocatedPolyBuffers() + for _, t := range []*icicle_core.DeviceSlice{&state.dZS, &state.dOrdering, &state.dLocal, &state.dGate, &state.dResult, &dTmp} { + if !t.IsEmpty() { + state.putTempDeviceSlice(*t, t.Len()) + *t = icicle_core.DeviceSlice{} + } + } + freeDeviceSlice(&state.dGammaScalar) + done <- runErr + }() + // Upload gamma as a scalar (inside RunOnDevice to ensure same device context). // For addition, we keep it in Montgomery form (unlike multiplication which needs standard form). state.dGammaScalar = uploadScalarMont(params.gamma) state.dZS = state.getTempDeviceSlice(n) if err := icicle_vecops.ShiftVec(state.dZ, state.dZS, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) + runErr = fmt.Errorf("ShiftVec failed while preparing ZS: %s", err.AsString()) return } // Step 1: Compute and apply blinding polynomial evaluations (if enabled) if useBlinding { - dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(n, twiddles0, bp) + dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS := computeBlindingPolynomials(state, n, twiddles0, bp) if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { - done <- err + runErr = err return } } @@ -4401,7 +4613,7 @@ func (s *instance) gpuEvaluateConstraints( var err error state.dOrdering, err = computeOrderingConstraint(state, params, dTwiddles0, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: ordering: %w", err) return } @@ -4410,48 +4622,50 @@ func (s *instance) gpuEvaluateConstraints( dAlphaStd := uploadScalarStd(params.alpha) if err := icicle_vecops.ScalarMulVec(dAlphaStd, state.dOrdering, state.dResult, seqVecCfg); err != icicle_runtime.Success { dAlphaStd.Free() - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha*ordering failed: %s", err.AsString()) return } dAlphaStd.Free() state.putTempDeviceSlice(state.dOrdering, params.n) + state.dOrdering = icicle_core.DeviceSlice{} // dResult += alpha^2 * local state.dLocal, err = computeLocalConstraint(state, params, dPrecomputedDenominators, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: local: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: local: %w", err) return } var alphaSquared fr.Element alphaSquared.Mul(¶ms.alpha, ¶ms.alpha) dAlphaSquaredStd := uploadScalarStd(alphaSquared) - dTmp := state.getTempDeviceSlice(params.n) + dTmp = state.getTempDeviceSlice(params.n) if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { dAlphaSquaredStd.Free() - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: combine alpha^2*local failed: %s", err.AsString()) return } dAlphaSquaredStd.Free() if err := icicle_vecops.VecOp(state.dResult, dTmp, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - state.putTempDeviceSlice(dTmp, params.n) - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add local failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dLocal, params.n) + state.dLocal = icicle_core.DeviceSlice{} state.putTempDeviceSlice(dTmp, params.n) + dTmp = icicle_core.DeviceSlice{} // dResult += gate state.dGate, err = computeGateConstraint(state, params, seqVecCfg) if err != nil { - done <- fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) + runErr = fmt.Errorf("gpuEvaluateConstraints: gate: %w", err) return } if err := icicle_vecops.VecOp(state.dResult, state.dGate, state.dResult, seqVecCfg, icicle_core.Add); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: VecOp add gate failed: %s", err.AsString()) return } state.putTempDeviceSlice(state.dGate, params.n) + state.dGate = icicle_core.DeviceSlice{} // Step 5: materialize result either on host or as a persistent device slice. if result != nil { @@ -4460,22 +4674,15 @@ func (s *instance) gpuEvaluateConstraints( } else { resultOnDevice = s.getTempDeviceSlice(params.n) if err := copyDeviceSliceIntoOnCurrentDevice(resultOnDevice, state.dResult, state.vecCfg); err != icicle_runtime.Success { - done <- fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) return } } - // Return dResult pool slice after materialization. - state.putTempDeviceSlice(state.dResult, params.n) - - // Return all allocated polynomial buffers to the pool. - // This automatically frees L, R, O, Z (if blinding was used) and S1, S2, S3 (always allocated). - // Note: dZS is freed in computeOrderingConstraint, and Q* working copies are - // returned to pool inside computeGateConstraint. dQk and the original gpuState slices - // are owned by gpuState and will be freed separately. - state.freeAllocatedPolyBuffers() - - done <- nil + // dResult and the tracked polynomial buffers (blinded L, R, O, Z and + // scaled S1, S2, S3) are returned to the pool by the cleanup defer. + // dQk and the original gpuState slices are owned by gpuState and will + // be freed separately. }) err := <-done @@ -4608,6 +4815,7 @@ func (s *instance) prepareStatisticalZKQuotientShards( prepareDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepareDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -4812,6 +5020,7 @@ func (s *instance) inverseAndMergeShards( var dMerged icicle_core.DeviceSlice icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") if cfgErr != nil { done <- cfgErr @@ -4985,6 +5194,7 @@ func (s *instance) divideByZHOnGPU( // So we can divide by Z_H by scaling each shard with its corresponding inverse. scaleDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(scaleDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") if cfgErr != nil { scaleDone <- cfgErr @@ -5034,6 +5244,7 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io host := icicle_core.HostSliceFromElements(coeffs) done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") if cfgErr != nil { done <- cfgErr @@ -5081,8 +5292,15 @@ func commitOnGPUWithDeviceBasesChunked( var msmErr error done := make(chan struct{}, 1) icicle_runtime.RunOnDevice(device, func(args ...any) { + defer func() { + if r := recover(); r != nil && msmErr == nil { + // Convert a device-goroutine panic into an error instead of + // killing the embedding process. + msmErr = fmt.Errorf("icicle: MSM device task panicked: %v", r) + } + close(done) + }() commit, msmErr = commitOnGPUWithDeviceBasesChunkedOnCurrentDevice(scalarsDevice, basesDevice, chunkSize) - close(done) }) <-done if msmErr != nil { @@ -5256,6 +5474,7 @@ func (s *instance) evalDevicePolynomialAtPoint(coeffsDevice icicle_core.DeviceSl var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false @@ -5429,6 +5648,7 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(prepDone) cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") if cfgErr != nil { prepDone <- cfgErr @@ -5704,6 +5924,7 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( var out fr.Element done := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(done) cfgVec, evalStream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("evalPolynomialInCurrentFormOnGPU") if cfgErr != nil { done <- cfgErr @@ -5781,6 +6002,7 @@ func (s *instance) openPolynomialOnGPUCanonical(coeffs []fr.Element, point fr.El witnessSize := len(coeffs) - 1 divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dCoeffs := uploadVector(coeffs) @@ -5845,6 +6067,7 @@ func (s *instance) openPolynomialOnGPUCanonicalDevice(coeffsDevice icicle_core.D } divDone := make(chan error, 1) icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(divDone) cfg := icicle_core.DefaultVecOpsConfig() cfg.IsAsync = false dPoint := uploadScalarStd(point) @@ -5970,6 +6193,12 @@ func (s *instance) buildLinearizedSelectorTermsOnGPU( } var runErr error defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "buildLinearizedSelectorTermsOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr } @@ -6158,6 +6387,12 @@ func (s *instance) addZContributionToLinearizedOnGPU( var dScale icicle_core.DeviceSlice var dScaledZ icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before returning temporary buffers to the pool. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "addZContributionToLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6225,6 +6460,12 @@ func (s *instance) subtractQuotientContributionFromLinearizedOnGPU( var dZetaStd icicle_core.DeviceSlice var dZhStd icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary before reusing temporary quotient vectors. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "subtractQuotientContributionFromLinearizedOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -6595,9 +6836,10 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { copy(pk.deviceInfo.CosetGenerator[:], cosetLimbs[:fr.Limbs*2]) } - chInitDomain := make(chan struct{}) + chInitDomain := make(chan error, 1) initDomainQueuedAt := time.Now() icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chInitDomain) initDomainStartedAt := time.Now() initCfg := icicle_core.GetDefaultNTTInitDomainConfig() ext := config_extension.Create() @@ -6625,41 +6867,52 @@ func (pk *ProvingKey) setupDevicePointers(device *icicle_runtime.Device) error { ) } if e != icicle_runtime.Success { - panic("icicle: InitDomain failed") + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return } - close(chInitDomain) + chInitDomain <- nil }) - <-chInitDomain + if err := <-chInitDomain; err != nil { + return err + } if isNttTrace { fmt.Fprintf(os.Stderr, "[ICICLE_NTT_TRACE] InitDomain total_wait_took=%s\n", time.Since(initDomainQueuedAt)) } - chLag := make(chan struct{}) - chCan := make(chan struct{}) + chLag := make(chan error, 1) + chCan := make(chan error, 1) if len(pk.KzgLagrange.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chLag) g1LagHost := (icicle_core.HostSlice[curve.G1Affine])(pk.KzgLagrange.G1) g1LagHost.CopyToDevice(&pk.deviceInfo.KzgLagrangeDevice.G1, true) - close(chLag) + chLag <- nil }) } else { - close(chLag) + chLag <- nil } if len(pk.Kzg.G1) > 0 { icicle_runtime.RunOnDevice(device, func(args ...any) { + defer devicePanicToError(chCan) g1CanHost := (icicle_core.HostSlice[curve.G1Affine])(pk.Kzg.G1) g1CanHost.CopyToDevice(&pk.deviceInfo.KzgDevice.G1, true) - close(chCan) + chCan <- nil }) } else { - close(chCan) + chCan <- nil } - <-chLag - <-chCan + errLag := <-chLag + errCan := <-chCan + if errLag != nil { + return fmt.Errorf("icicle: uploading Lagrange SRS to device failed: %w", errLag) + } + if errCan != nil { + return fmt.Errorf("icicle: uploading canonical SRS to device failed: %w", errCan) + } return nil } @@ -6989,6 +7242,12 @@ func (s *instance) BuildRatioCopyConstraintIcicle( var dNum, dDen icicle_core.DeviceSlice var dSupportFlat, dPermFlat icicle_core.DeviceSlice defer func() { + if r := recover(); r != nil && runErr == nil { + // Convert a device-goroutine panic (e.g. GPU allocation + // failure in the pool/upload helpers) into a prover error; + // unrecovered it would kill the embedding process. + runErr = fmt.Errorf("icicle: device task panicked: %v", r) + } // Async boundary for copy-constraint accumulation before cleanup. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "BuildRatioCopyConstraintIcicle"); syncErr != nil && runErr == nil { runErr = syncErr diff --git a/backend/accelerated/icicle/plonk/cache_test.go b/backend/accelerated/icicle/plonk/cache_test.go new file mode 100644 index 0000000000..a84406a114 --- /dev/null +++ b/backend/accelerated/icicle/plonk/cache_test.go @@ -0,0 +1,176 @@ +//go:build icicle + +package plonk_test + +import ( + "math/big" + "os" + "path/filepath" + "testing" + + "github.com/consensys/gnark-crypto/ecc" + accel_plonk "github.com/consensys/gnark/backend/accelerated/icicle/plonk" + plonk_bn254 "github.com/consensys/gnark/backend/plonk/bn254" + "github.com/consensys/gnark/frontend" + "github.com/consensys/gnark/frontend/cs/scs" + "github.com/consensys/gnark/test" + "github.com/consensys/gnark/test/unsafekzg" +) + +// cacheCircuit is largeCircuit without the BSB22 commitment, so the raw +// solver cache stays enabled in the default (blinding-on) mode. +type cacheCircuit struct { + A frontend.Variable `gnark:",public"` + Res frontend.Variable +} + +func (c *cacheCircuit) Define(api frontend.API) error { + x := c.A + for i := 0; i < largeCircuitSize; i++ { + x = api.Add(api.Mul(x, x), c.A) + } + api.AssertIsEqual(x, c.Res) + return nil +} + +// recurrenceResult computes the expected largeCircuit/cacheCircuit output. +func recurrenceResult(mod *big.Int) (a, x *big.Int) { + a = big.NewInt(3) + x = big.NewInt(3) + for i := 0; i < largeCircuitSize; i++ { + x.Mul(x, x) + x.Add(x, a) + x.Mod(x, mod) + } + return a, x +} + +// TestRawSolverCacheRoundTrip proves a commitment-free circuit twice with +// GNARK_RAW_SOLVER_CACHE set: the first prove writes the cache, the second +// loads it, and both proofs verify. +func TestRawSolverCacheRoundTrip(t *testing.T) { + assert := test.NewAssert(t) + curveID := ecc.BN254 + + cachePath := filepath.Join(t.TempDir(), "raw_solver.bin") + t.Setenv("GNARK_RAW_SOLVER_CACHE", cachePath) + + ccs, err := frontend.Compile(curveID.ScalarField(), scs.NewBuilder, &cacheCircuit{}) + assert.NoError(err) + srs, srsLagrange, err := unsafekzg.NewSRS(ccs) + assert.NoError(err) + pk, vk, err := accel_plonk.Setup(ccs, srs, srsLagrange) + assert.NoError(err) + + a, x := recurrenceResult(curveID.ScalarField()) + w, err := frontend.NewWitness(&cacheCircuit{A: a, Res: x}, curveID.ScalarField()) + assert.NoError(err) + pw, err := w.Public() + assert.NoError(err) + + proof, err := accel_plonk.Prove(ccs, pk, w) + assert.NoError(err) + assert.NoError(accel_plonk.Verify(proof, vk, pw)) + _, err = os.Stat(cachePath) + assert.NoError(err, "first prove must write the raw solver cache") + + proof, err = accel_plonk.Prove(ccs, pk, w) + assert.NoError(err) + assert.NoError(accel_plonk.Verify(proof, vk, pw)) +} + +// TestRawSolverCacheBSB22ReplayNoBlinding checks the path that stays allowed +// for BSB22 circuits: with GNARK_DISABLE_BLINDING set (zero-knowledge +// explicitly traded away), the cache and its commitment sidecars are written +// on the first prove and replayed on the second, and both proofs verify. +// The flag is read once in the package init, so this test can only run when +// the environment variable is set for the whole test process; it is skipped +// otherwise. +func TestRawSolverCacheBSB22ReplayNoBlinding(t *testing.T) { + if _, ok := os.LookupEnv("GNARK_DISABLE_BLINDING"); !ok { + t.Skip("requires GNARK_DISABLE_BLINDING (BSB22 cache replay is disabled in the default blinding-on mode)") + } + assert := test.NewAssert(t) + curveID := ecc.BN254 + + cacheDir := t.TempDir() + cachePath := filepath.Join(cacheDir, "raw_solver.bin") + t.Setenv("GNARK_RAW_SOLVER_CACHE", cachePath) + + ccs, err := frontend.Compile(curveID.ScalarField(), scs.NewBuilder, &largeCircuit{}) + assert.NoError(err) + srs, srsLagrange, err := unsafekzg.NewSRS(ccs) + assert.NoError(err) + pk, vk, err := accel_plonk.Setup(ccs, srs, srsLagrange) + assert.NoError(err) + + a, x := recurrenceResult(curveID.ScalarField()) + w, err := frontend.NewWitness(&largeCircuit{A: a, Res: x}, curveID.ScalarField()) + assert.NoError(err) + pw, err := w.Public() + assert.NoError(err) + + proof, err := accel_plonk.Prove(ccs, pk, w) + assert.NoError(err) + assert.NoError(accel_plonk.Verify(proof, vk, pw)) + _, err = os.Stat(cachePath) + assert.NoError(err, "first prove must write the raw solver cache") + sidecars, err := filepath.Glob(filepath.Join(cacheDir, "bsb22_commit_*.bin")) + assert.NoError(err) + assert.True(len(sidecars) > 0, "first prove must write the BSB22 sidecar files") + + // second prove replays the cache (same pk, so the recomputed commitment + // matches the cached wires) and must still produce a valid proof + proof, err = accel_plonk.Prove(ccs, pk, w) + assert.NoError(err) + assert.NoError(accel_plonk.Verify(proof, vk, pw)) +} + +// TestRawSolverCacheSkippedForBSB22 checks that in the default (blinding-on) +// mode the raw solver cache is disabled for circuits with BSB22 commitments: +// replaying the cached commitment blinding across proofs would make the +// commitments linkable and leak the committed private wires. No cache file +// may be written, and two proofs of the same witness must carry distinct +// (freshly blinded) Bsb22Commitments. +func TestRawSolverCacheSkippedForBSB22(t *testing.T) { + assert := test.NewAssert(t) + curveID := ecc.BN254 + + cacheDir := t.TempDir() + cachePath := filepath.Join(cacheDir, "raw_solver.bin") + t.Setenv("GNARK_RAW_SOLVER_CACHE", cachePath) + + ccs, err := frontend.Compile(curveID.ScalarField(), scs.NewBuilder, &largeCircuit{}) + assert.NoError(err) + srs, srsLagrange, err := unsafekzg.NewSRS(ccs) + assert.NoError(err) + pk, vk, err := accel_plonk.Setup(ccs, srs, srsLagrange) + assert.NoError(err) + + a, x := recurrenceResult(curveID.ScalarField()) + w, err := frontend.NewWitness(&largeCircuit{A: a, Res: x}, curveID.ScalarField()) + assert.NoError(err) + pw, err := w.Public() + assert.NoError(err) + + proof1, err := accel_plonk.Prove(ccs, pk, w) + assert.NoError(err) + assert.NoError(accel_plonk.Verify(proof1, vk, pw)) + proof2, err := accel_plonk.Prove(ccs, pk, w) + assert.NoError(err) + assert.NoError(accel_plonk.Verify(proof2, vk, pw)) + + _, err = os.Stat(cachePath) + assert.True(os.IsNotExist(err), "raw solver cache must not be written for BSB22 circuits with blinding enabled") + sidecars, err := filepath.Glob(filepath.Join(cacheDir, "bsb22_commit_*.bin")) + assert.NoError(err) + assert.Empty(sidecars, "BSB22 sidecar files must not be written with blinding enabled") + + p1 := proof1.(*plonk_bn254.Proof) + p2 := proof2.(*plonk_bn254.Proof) + assert.True(len(p1.Bsb22Commitments) > 0, "test circuit must carry a BSB22 commitment") + for i := range p1.Bsb22Commitments { + assert.False(p1.Bsb22Commitments[i].Equal(&p2.Bsb22Commitments[i]), + "Bsb22Commitments must be freshly blinded on every prove") + } +} From 09fa248ec406af0270c759b39790b50779c46a8a Mon Sep 17 00:00:00 2001 From: Martun Karapetyan Date: Mon, 27 Jul 2026 19:57:20 +0400 Subject: [PATCH 3/5] chore(icicle/plonk): make golangci-lint clean - annotate the atomic cache writes with #nosec G703: the temp file comes from os.CreateTemp in the caller-chosen cache directory and the rename target is the caller-provided cache location (gosec taint analysis flags them as path traversal) - replace deprecated fft.BitReverse with gnark-crypto utils.BitReverse - merge a variable declaration with its assignment (S1021) golangci-lint v2.10.1 (CI's version/config) now reports 0 issues on the default build and 0 issues on ./backend/accelerated/icicle/plonk/... + ./constraint/... with -tags=icicle. Remaining icicle-tagged findings are in the pre-existing accelerated groth16 code, untouched by this PR. Full icicle plonk suite re-verified on GPU after the BitReverse swap. Co-Authored-By: Claude Fable 5 --- .../templates/constraint.solution_cache.go.tmpl | 14 +++++++------- .../generator/templates/plonk.icicle.go.tmpl | 9 +++++---- .../accelerated/icicle/plonk/bls12-377/icicle.go | 9 +++++---- .../accelerated/icicle/plonk/bls12-381/icicle.go | 9 +++++---- backend/accelerated/icicle/plonk/bn254/icicle.go | 9 +++++---- backend/accelerated/icicle/plonk/bw6-761/icicle.go | 9 +++++---- constraint/bls12-377/solution_cache.go | 14 +++++++------- constraint/bls12-381/solution_cache.go | 14 +++++++------- constraint/bn254/solution_cache.go | 14 +++++++------- constraint/bw6-761/solution_cache.go | 14 +++++++------- 10 files changed, 60 insertions(+), 55 deletions(-) diff --git a/backend/accelerated/icicle/internal/generator/templates/constraint.solution_cache.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/constraint.solution_cache.go.tmpl index e9388dae74..7b8402824d 100644 --- a/backend/accelerated/icicle/internal/generator/templates/constraint.solution_cache.go.tmpl +++ b/backend/accelerated/icicle/internal/generator/templates/constraint.solution_cache.go.tmpl @@ -170,14 +170,14 @@ func SaveR1CSSolution(path string, solution *R1CSSolution) error { tmpName := tmp.Name() if _, err := solution.WriteTo(tmp); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, @@ -194,7 +194,7 @@ func SaveRawSolverValues(path string, values []fr.Element) error { nWires := uint64(len(values)) if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if nWires > 0 { @@ -202,15 +202,15 @@ func SaveRawSolverValues(path string, values []fr.Element) error { byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) if _, err := tmp.Write(byteSlice); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. diff --git a/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl index 14b5597a4c..bf74400f2c 100644 --- a/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl +++ b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl @@ -27,6 +27,8 @@ import ( "github.com/consensys/gnark/constraint/solver" fcs "github.com/consensys/gnark/frontend/cs" "github.com/consensys/gnark/internal/utils" + + crypto_utils "github.com/consensys/gnark-crypto/utils" "github.com/consensys/gnark/logger" "github.com/consensys/gnark-crypto/ecc" @@ -863,8 +865,7 @@ func (s *instance) computeQuotient() (err error) { case <-s.chLRO: } if isProfileMode { - var startComputeQuotient time.Time - startComputeQuotient = time.Now() + startComputeQuotient := time.Now() defer func() { l := logger.Logger() if err != nil { @@ -2142,7 +2143,7 @@ func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolyn scalingVector := cosetTable scalingVectorRev := make([]fr.Element, len(cosetTable)) copy(scalingVectorRev, cosetTable) - fft.BitReverse(scalingVectorRev) + crypto_utils.BitReverse(scalingVectorRev) // pre-computed to compute the bit reverse index // of the result polynomial @@ -2411,7 +2412,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo // Reuse memory. copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) - fft.BitReverse(loopCtx.scalingVectorRev) + crypto_utils.BitReverse(loopCtx.scalingVectorRev) } // We do **a lot** of FFT here, but on the small domain. diff --git a/backend/accelerated/icicle/plonk/bls12-377/icicle.go b/backend/accelerated/icicle/plonk/bls12-377/icicle.go index 798fbdc9d1..3ac0274854 100644 --- a/backend/accelerated/icicle/plonk/bls12-377/icicle.go +++ b/backend/accelerated/icicle/plonk/bls12-377/icicle.go @@ -34,6 +34,8 @@ import ( "github.com/consensys/gnark/constraint/solver" fcs "github.com/consensys/gnark/frontend/cs" "github.com/consensys/gnark/internal/utils" + + crypto_utils "github.com/consensys/gnark-crypto/utils" "github.com/consensys/gnark/logger" "github.com/consensys/gnark-crypto/ecc" @@ -870,8 +872,7 @@ func (s *instance) computeQuotient() (err error) { case <-s.chLRO: } if isProfileMode { - var startComputeQuotient time.Time - startComputeQuotient = time.Now() + startComputeQuotient := time.Now() defer func() { l := logger.Logger() if err != nil { @@ -2149,7 +2150,7 @@ func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolyn scalingVector := cosetTable scalingVectorRev := make([]fr.Element, len(cosetTable)) copy(scalingVectorRev, cosetTable) - fft.BitReverse(scalingVectorRev) + crypto_utils.BitReverse(scalingVectorRev) // pre-computed to compute the bit reverse index // of the result polynomial @@ -2418,7 +2419,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo // Reuse memory. copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) - fft.BitReverse(loopCtx.scalingVectorRev) + crypto_utils.BitReverse(loopCtx.scalingVectorRev) } // We do **a lot** of FFT here, but on the small domain. diff --git a/backend/accelerated/icicle/plonk/bls12-381/icicle.go b/backend/accelerated/icicle/plonk/bls12-381/icicle.go index 5accd02547..ecf9b174ef 100644 --- a/backend/accelerated/icicle/plonk/bls12-381/icicle.go +++ b/backend/accelerated/icicle/plonk/bls12-381/icicle.go @@ -34,6 +34,8 @@ import ( "github.com/consensys/gnark/constraint/solver" fcs "github.com/consensys/gnark/frontend/cs" "github.com/consensys/gnark/internal/utils" + + crypto_utils "github.com/consensys/gnark-crypto/utils" "github.com/consensys/gnark/logger" "github.com/consensys/gnark-crypto/ecc" @@ -870,8 +872,7 @@ func (s *instance) computeQuotient() (err error) { case <-s.chLRO: } if isProfileMode { - var startComputeQuotient time.Time - startComputeQuotient = time.Now() + startComputeQuotient := time.Now() defer func() { l := logger.Logger() if err != nil { @@ -2149,7 +2150,7 @@ func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolyn scalingVector := cosetTable scalingVectorRev := make([]fr.Element, len(cosetTable)) copy(scalingVectorRev, cosetTable) - fft.BitReverse(scalingVectorRev) + crypto_utils.BitReverse(scalingVectorRev) // pre-computed to compute the bit reverse index // of the result polynomial @@ -2418,7 +2419,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo // Reuse memory. copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) - fft.BitReverse(loopCtx.scalingVectorRev) + crypto_utils.BitReverse(loopCtx.scalingVectorRev) } // We do **a lot** of FFT here, but on the small domain. diff --git a/backend/accelerated/icicle/plonk/bn254/icicle.go b/backend/accelerated/icicle/plonk/bn254/icicle.go index 6d04b07146..260b7263dd 100644 --- a/backend/accelerated/icicle/plonk/bn254/icicle.go +++ b/backend/accelerated/icicle/plonk/bn254/icicle.go @@ -34,6 +34,8 @@ import ( "github.com/consensys/gnark/constraint/solver" fcs "github.com/consensys/gnark/frontend/cs" "github.com/consensys/gnark/internal/utils" + + crypto_utils "github.com/consensys/gnark-crypto/utils" "github.com/consensys/gnark/logger" "github.com/consensys/gnark-crypto/ecc" @@ -870,8 +872,7 @@ func (s *instance) computeQuotient() (err error) { case <-s.chLRO: } if isProfileMode { - var startComputeQuotient time.Time - startComputeQuotient = time.Now() + startComputeQuotient := time.Now() defer func() { l := logger.Logger() if err != nil { @@ -2149,7 +2150,7 @@ func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolyn scalingVector := cosetTable scalingVectorRev := make([]fr.Element, len(cosetTable)) copy(scalingVectorRev, cosetTable) - fft.BitReverse(scalingVectorRev) + crypto_utils.BitReverse(scalingVectorRev) // pre-computed to compute the bit reverse index // of the result polynomial @@ -2418,7 +2419,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo // Reuse memory. copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) - fft.BitReverse(loopCtx.scalingVectorRev) + crypto_utils.BitReverse(loopCtx.scalingVectorRev) } // We do **a lot** of FFT here, but on the small domain. diff --git a/backend/accelerated/icicle/plonk/bw6-761/icicle.go b/backend/accelerated/icicle/plonk/bw6-761/icicle.go index 2e510bf3e1..c26c2214cc 100644 --- a/backend/accelerated/icicle/plonk/bw6-761/icicle.go +++ b/backend/accelerated/icicle/plonk/bw6-761/icicle.go @@ -34,6 +34,8 @@ import ( "github.com/consensys/gnark/constraint/solver" fcs "github.com/consensys/gnark/frontend/cs" "github.com/consensys/gnark/internal/utils" + + crypto_utils "github.com/consensys/gnark-crypto/utils" "github.com/consensys/gnark/logger" "github.com/consensys/gnark-crypto/ecc" @@ -870,8 +872,7 @@ func (s *instance) computeQuotient() (err error) { case <-s.chLRO: } if isProfileMode { - var startComputeQuotient time.Time - startComputeQuotient = time.Now() + startComputeQuotient := time.Now() defer func() { l := logger.Logger() if err != nil { @@ -2149,7 +2150,7 @@ func (s *instance) computeNumerator(gpuState *gpuPolysState) (*gpuNumeratorPolyn scalingVector := cosetTable scalingVectorRev := make([]fr.Element, len(cosetTable)) copy(scalingVectorRev, cosetTable) - fft.BitReverse(scalingVectorRev) + crypto_utils.BitReverse(scalingVectorRev) // pre-computed to compute the bit reverse index // of the result polynomial @@ -2418,7 +2419,7 @@ func (s *instance) computeNumeratorIteration(i int, loopCtx *computeNumeratorLoo // Reuse memory. copy(loopCtx.scalingVectorRev, loopCtx.scalingVector) - fft.BitReverse(loopCtx.scalingVectorRev) + crypto_utils.BitReverse(loopCtx.scalingVectorRev) } // We do **a lot** of FFT here, but on the small domain. diff --git a/constraint/bls12-377/solution_cache.go b/constraint/bls12-377/solution_cache.go index ea6216de9a..15887f565c 100644 --- a/constraint/bls12-377/solution_cache.go +++ b/constraint/bls12-377/solution_cache.go @@ -177,14 +177,14 @@ func SaveR1CSSolution(path string, solution *R1CSSolution) error { tmpName := tmp.Name() if _, err := solution.WriteTo(tmp); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, @@ -201,7 +201,7 @@ func SaveRawSolverValues(path string, values []fr.Element) error { nWires := uint64(len(values)) if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if nWires > 0 { @@ -209,15 +209,15 @@ func SaveRawSolverValues(path string, values []fr.Element) error { byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) if _, err := tmp.Write(byteSlice); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. diff --git a/constraint/bls12-381/solution_cache.go b/constraint/bls12-381/solution_cache.go index 8ae16167e6..8363e37ba6 100644 --- a/constraint/bls12-381/solution_cache.go +++ b/constraint/bls12-381/solution_cache.go @@ -177,14 +177,14 @@ func SaveR1CSSolution(path string, solution *R1CSSolution) error { tmpName := tmp.Name() if _, err := solution.WriteTo(tmp); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, @@ -201,7 +201,7 @@ func SaveRawSolverValues(path string, values []fr.Element) error { nWires := uint64(len(values)) if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if nWires > 0 { @@ -209,15 +209,15 @@ func SaveRawSolverValues(path string, values []fr.Element) error { byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) if _, err := tmp.Write(byteSlice); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. diff --git a/constraint/bn254/solution_cache.go b/constraint/bn254/solution_cache.go index 3052ad7295..aebfde0d02 100644 --- a/constraint/bn254/solution_cache.go +++ b/constraint/bn254/solution_cache.go @@ -177,14 +177,14 @@ func SaveR1CSSolution(path string, solution *R1CSSolution) error { tmpName := tmp.Name() if _, err := solution.WriteTo(tmp); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, @@ -201,7 +201,7 @@ func SaveRawSolverValues(path string, values []fr.Element) error { nWires := uint64(len(values)) if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if nWires > 0 { @@ -209,15 +209,15 @@ func SaveRawSolverValues(path string, values []fr.Element) error { byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) if _, err := tmp.Write(byteSlice); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. diff --git a/constraint/bw6-761/solution_cache.go b/constraint/bw6-761/solution_cache.go index 12e48e99eb..8c5f279919 100644 --- a/constraint/bw6-761/solution_cache.go +++ b/constraint/bw6-761/solution_cache.go @@ -177,14 +177,14 @@ func SaveR1CSSolution(path string, solution *R1CSSolution) error { tmpName := tmp.Name() if _, err := solution.WriteTo(tmp); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // SaveRawSolverValues writes a wire-value vector as raw bytes (Montgomery form, @@ -201,7 +201,7 @@ func SaveRawSolverValues(path string, values []fr.Element) error { nWires := uint64(len(values)) if err := binary.Write(tmp, binary.LittleEndian, nWires); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } if nWires > 0 { @@ -209,15 +209,15 @@ func SaveRawSolverValues(path string, values []fr.Element) error { byteSlice := unsafe.Slice((*byte)(unsafe.Pointer(&values[0])), len(values)*fr.Bytes) if _, err := tmp.Write(byteSlice); err != nil { tmp.Close() - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } } if err := tmp.Close(); err != nil { - os.Remove(tmpName) + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory return err } - return os.Rename(tmpName, path) + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location } // LoadRawSolverValues reads a wire-value vector written by SaveRawSolverValues. From d529758d7c1c8f8b467e5ce720ab0dee6ebe2a0c Mon Sep 17 00:00:00 2001 From: Martun Karapetyan Date: Mon, 27 Jul 2026 20:09:13 +0400 Subject: [PATCH 4/5] fix(icicle/plonk): destroy CUDA streams when a device task panics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit devicePanicToError converts GPU allocation panics into Prove errors, but closures that create a stream and only release it inside their finish / finalize callback skipped that cleanup when the panic hit between CreateStream and the callback — leaking one stream per recovered OOM, which in a long-lived process eventually exhausts streams and fails later proves. makeFinisher now also returns a guard to defer right after stream creation: if finish never ran, the guard synchronizes and destroys the stream (a no-op on all non-panic paths). The four download closures that released their stream at the tail are converted to the same finish/guard pattern, and openZ's finalize and the eval closure's local finish get equivalent flag-based guards (the latter also releasing its owned eval buffer). Verified on GPU: full -tags=icicle plonk suite 6/6 consecutive runs plus the GNARK_DISABLE_BLINDING replay test; golangci-lint clean (default and icicle-tagged); regenerated files match the generator output. Co-Authored-By: Claude Fable 5 --- .../generator/templates/plonk.icicle.go.tmpl | 74 +++++++++++++++---- .../icicle/plonk/bls12-377/icicle.go | 74 +++++++++++++++---- .../icicle/plonk/bls12-381/icicle.go | 74 +++++++++++++++---- .../accelerated/icicle/plonk/bn254/icicle.go | 74 +++++++++++++++---- .../icicle/plonk/bw6-761/icicle.go | 74 +++++++++++++++---- 5 files changed, 300 insertions(+), 70 deletions(-) diff --git a/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl index bf74400f2c..0c86a77357 100644 --- a/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl +++ b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl @@ -1018,7 +1018,8 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone <- cfgErr return } - finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() var allocErr error dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) if allocErr != nil { @@ -1109,7 +1110,9 @@ func (s *instance) openZ() (err error) { buildDone <- cfgErr return } + finalized := false finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + finalized = true // Async boundary for canonicalization/blinding before exposing output. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { runErr = syncErr @@ -1123,6 +1126,14 @@ func (s *instance) openZ() (err error) { } buildDone <- runErr } + // If the closure panics before finalize runs (recovered by + // devicePanicToError), still release the stream so the GPU-OOM path + // cannot leak it. + defer func() { + if !finalized { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "openZ") + } + }() n := dZLagrange.Len() dZCanonical := s.getTempDeviceSlice(n) @@ -1575,7 +1586,8 @@ func (s *instance) batchOpening() error { foldDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening fold", foldDone) + finish, guard := makeFinisher(stream, "batchOpening fold", foldDone) + defer guard() dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { @@ -1642,7 +1654,8 @@ func (s *instance) batchOpening() error { divDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + finish, guard := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + defer guard() // For Montgomery vectors, scalar multipliers must be standard-form. dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) dWitness = s.getTempDeviceSlice(witnessSize) @@ -2039,8 +2052,10 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS done <- cfgErr return } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) - done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -2637,6 +2652,8 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() dMerged = s.getTempDeviceSlice(totalSize) mergeErr := icicle_vecops.MergeShardsBitReverse( gpuNumerator.shards, @@ -2646,12 +2663,11 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cfg, ) if mergeErr != icicle_runtime.Success { - _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") - mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) return } // Async boundary before merged slice is consumed by host copy. - mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + finish(nil) }) if mergeErr := <-mergeDone; mergeErr != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -2668,8 +2684,10 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial downloadDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) - downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + finish(nil) }) if err := <-downloadDone; err != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -3682,13 +3700,25 @@ func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label str // makeFinisher returns a closure that synchronizes and destroys the stream, // then sends the (possibly merged) error to done. Use inside RunOnDevice closures. -func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { - return func(runErr error) { +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) (finish func(error), guard func()) { + finished := false + finish = func(runErr error) { + finished = true if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } done <- runErr } + // guard must be deferred right after the stream is created: if the + // closure panics before finish runs (the panic itself is recovered and + // reported by devicePanicToError), the guard still synchronizes and + // destroys the stream so the GPU-OOM path cannot leak it. + guard = func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(stream, label) + } + } + return finish, guard } // devicePanicToError is deferred at the top of RunOnDevice closures that @@ -5020,7 +5050,8 @@ func (s *instance) inverseAndMergeShards( done <- cfgErr return } - finish := makeFinisher(stream, "inverseAndMergeShards", done) + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() cfgNtt := icicle_ntt.GetDefaultNttConfig() cfgNtt.IsAsync = true @@ -5194,7 +5225,8 @@ func (s *instance) divideByZHOnGPU( scaleDone <- cfgErr return } - finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + finish, guard := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + defer guard() for i := 0; i < gpuNumerator.rho; i++ { dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) @@ -5244,9 +5276,11 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io done <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadQuotientFromGPU", done) + defer guard() host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) // Async boundary for host materialization of quotient coefficients. - done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -5648,7 +5682,8 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone <- cfgErr return } - finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() cleanupOwned := func(from int) { for i := from; i < len(devicePolys); i++ { @@ -5933,7 +5968,9 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( ownedLen = 0 } } + finished := false finish := func(runErr error) { + finished = true // Async boundary for eval path before handing result back to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -5941,6 +5978,15 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( releaseEval() done <- runErr } + // If the closure panics before finish runs (recovered by + // devicePanicToError), still release the stream and any owned eval + // buffer so the GPU-OOM path cannot leak them. + defer func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU") + releaseEval() + } + }() prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) dEval = prepareResult.dEval diff --git a/backend/accelerated/icicle/plonk/bls12-377/icicle.go b/backend/accelerated/icicle/plonk/bls12-377/icicle.go index 3ac0274854..f2fd78f8f2 100644 --- a/backend/accelerated/icicle/plonk/bls12-377/icicle.go +++ b/backend/accelerated/icicle/plonk/bls12-377/icicle.go @@ -1025,7 +1025,8 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone <- cfgErr return } - finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() var allocErr error dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) if allocErr != nil { @@ -1116,7 +1117,9 @@ func (s *instance) openZ() (err error) { buildDone <- cfgErr return } + finalized := false finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + finalized = true // Async boundary for canonicalization/blinding before exposing output. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { runErr = syncErr @@ -1130,6 +1133,14 @@ func (s *instance) openZ() (err error) { } buildDone <- runErr } + // If the closure panics before finalize runs (recovered by + // devicePanicToError), still release the stream so the GPU-OOM path + // cannot leak it. + defer func() { + if !finalized { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "openZ") + } + }() n := dZLagrange.Len() dZCanonical := s.getTempDeviceSlice(n) @@ -1582,7 +1593,8 @@ func (s *instance) batchOpening() error { foldDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening fold", foldDone) + finish, guard := makeFinisher(stream, "batchOpening fold", foldDone) + defer guard() dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { @@ -1649,7 +1661,8 @@ func (s *instance) batchOpening() error { divDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + finish, guard := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + defer guard() // For Montgomery vectors, scalar multipliers must be standard-form. dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) dWitness = s.getTempDeviceSlice(witnessSize) @@ -2046,8 +2059,10 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS done <- cfgErr return } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) - done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -2644,6 +2659,8 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() dMerged = s.getTempDeviceSlice(totalSize) mergeErr := icicle_vecops.MergeShardsBitReverse( gpuNumerator.shards, @@ -2653,12 +2670,11 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cfg, ) if mergeErr != icicle_runtime.Success { - _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") - mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) return } // Async boundary before merged slice is consumed by host copy. - mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + finish(nil) }) if mergeErr := <-mergeDone; mergeErr != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -2675,8 +2691,10 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial downloadDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) - downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + finish(nil) }) if err := <-downloadDone; err != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -3689,13 +3707,25 @@ func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label str // makeFinisher returns a closure that synchronizes and destroys the stream, // then sends the (possibly merged) error to done. Use inside RunOnDevice closures. -func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { - return func(runErr error) { +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) (finish func(error), guard func()) { + finished := false + finish = func(runErr error) { + finished = true if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } done <- runErr } + // guard must be deferred right after the stream is created: if the + // closure panics before finish runs (the panic itself is recovered and + // reported by devicePanicToError), the guard still synchronizes and + // destroys the stream so the GPU-OOM path cannot leak it. + guard = func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(stream, label) + } + } + return finish, guard } // devicePanicToError is deferred at the top of RunOnDevice closures that @@ -5027,7 +5057,8 @@ func (s *instance) inverseAndMergeShards( done <- cfgErr return } - finish := makeFinisher(stream, "inverseAndMergeShards", done) + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() cfgNtt := icicle_ntt.GetDefaultNttConfig() cfgNtt.IsAsync = true @@ -5201,7 +5232,8 @@ func (s *instance) divideByZHOnGPU( scaleDone <- cfgErr return } - finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + finish, guard := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + defer guard() for i := 0; i < gpuNumerator.rho; i++ { dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) @@ -5251,9 +5283,11 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io done <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadQuotientFromGPU", done) + defer guard() host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) // Async boundary for host materialization of quotient coefficients. - done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -5655,7 +5689,8 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone <- cfgErr return } - finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() cleanupOwned := func(from int) { for i := from; i < len(devicePolys); i++ { @@ -5940,7 +5975,9 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( ownedLen = 0 } } + finished := false finish := func(runErr error) { + finished = true // Async boundary for eval path before handing result back to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -5948,6 +5985,15 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( releaseEval() done <- runErr } + // If the closure panics before finish runs (recovered by + // devicePanicToError), still release the stream and any owned eval + // buffer so the GPU-OOM path cannot leak them. + defer func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU") + releaseEval() + } + }() prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) dEval = prepareResult.dEval diff --git a/backend/accelerated/icicle/plonk/bls12-381/icicle.go b/backend/accelerated/icicle/plonk/bls12-381/icicle.go index ecf9b174ef..aa8368b044 100644 --- a/backend/accelerated/icicle/plonk/bls12-381/icicle.go +++ b/backend/accelerated/icicle/plonk/bls12-381/icicle.go @@ -1025,7 +1025,8 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone <- cfgErr return } - finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() var allocErr error dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) if allocErr != nil { @@ -1116,7 +1117,9 @@ func (s *instance) openZ() (err error) { buildDone <- cfgErr return } + finalized := false finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + finalized = true // Async boundary for canonicalization/blinding before exposing output. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { runErr = syncErr @@ -1130,6 +1133,14 @@ func (s *instance) openZ() (err error) { } buildDone <- runErr } + // If the closure panics before finalize runs (recovered by + // devicePanicToError), still release the stream so the GPU-OOM path + // cannot leak it. + defer func() { + if !finalized { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "openZ") + } + }() n := dZLagrange.Len() dZCanonical := s.getTempDeviceSlice(n) @@ -1582,7 +1593,8 @@ func (s *instance) batchOpening() error { foldDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening fold", foldDone) + finish, guard := makeFinisher(stream, "batchOpening fold", foldDone) + defer guard() dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { @@ -1649,7 +1661,8 @@ func (s *instance) batchOpening() error { divDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + finish, guard := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + defer guard() // For Montgomery vectors, scalar multipliers must be standard-form. dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) dWitness = s.getTempDeviceSlice(witnessSize) @@ -2046,8 +2059,10 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS done <- cfgErr return } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) - done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -2644,6 +2659,8 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() dMerged = s.getTempDeviceSlice(totalSize) mergeErr := icicle_vecops.MergeShardsBitReverse( gpuNumerator.shards, @@ -2653,12 +2670,11 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cfg, ) if mergeErr != icicle_runtime.Success { - _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") - mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) return } // Async boundary before merged slice is consumed by host copy. - mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + finish(nil) }) if mergeErr := <-mergeDone; mergeErr != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -2675,8 +2691,10 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial downloadDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) - downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + finish(nil) }) if err := <-downloadDone; err != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -3689,13 +3707,25 @@ func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label str // makeFinisher returns a closure that synchronizes and destroys the stream, // then sends the (possibly merged) error to done. Use inside RunOnDevice closures. -func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { - return func(runErr error) { +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) (finish func(error), guard func()) { + finished := false + finish = func(runErr error) { + finished = true if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } done <- runErr } + // guard must be deferred right after the stream is created: if the + // closure panics before finish runs (the panic itself is recovered and + // reported by devicePanicToError), the guard still synchronizes and + // destroys the stream so the GPU-OOM path cannot leak it. + guard = func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(stream, label) + } + } + return finish, guard } // devicePanicToError is deferred at the top of RunOnDevice closures that @@ -5027,7 +5057,8 @@ func (s *instance) inverseAndMergeShards( done <- cfgErr return } - finish := makeFinisher(stream, "inverseAndMergeShards", done) + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() cfgNtt := icicle_ntt.GetDefaultNttConfig() cfgNtt.IsAsync = true @@ -5201,7 +5232,8 @@ func (s *instance) divideByZHOnGPU( scaleDone <- cfgErr return } - finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + finish, guard := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + defer guard() for i := 0; i < gpuNumerator.rho; i++ { dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) @@ -5251,9 +5283,11 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io done <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadQuotientFromGPU", done) + defer guard() host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) // Async boundary for host materialization of quotient coefficients. - done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -5655,7 +5689,8 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone <- cfgErr return } - finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() cleanupOwned := func(from int) { for i := from; i < len(devicePolys); i++ { @@ -5940,7 +5975,9 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( ownedLen = 0 } } + finished := false finish := func(runErr error) { + finished = true // Async boundary for eval path before handing result back to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -5948,6 +5985,15 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( releaseEval() done <- runErr } + // If the closure panics before finish runs (recovered by + // devicePanicToError), still release the stream and any owned eval + // buffer so the GPU-OOM path cannot leak them. + defer func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU") + releaseEval() + } + }() prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) dEval = prepareResult.dEval diff --git a/backend/accelerated/icicle/plonk/bn254/icicle.go b/backend/accelerated/icicle/plonk/bn254/icicle.go index 260b7263dd..ea03ae87e2 100644 --- a/backend/accelerated/icicle/plonk/bn254/icicle.go +++ b/backend/accelerated/icicle/plonk/bn254/icicle.go @@ -1025,7 +1025,8 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone <- cfgErr return } - finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() var allocErr error dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) if allocErr != nil { @@ -1116,7 +1117,9 @@ func (s *instance) openZ() (err error) { buildDone <- cfgErr return } + finalized := false finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + finalized = true // Async boundary for canonicalization/blinding before exposing output. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { runErr = syncErr @@ -1130,6 +1133,14 @@ func (s *instance) openZ() (err error) { } buildDone <- runErr } + // If the closure panics before finalize runs (recovered by + // devicePanicToError), still release the stream so the GPU-OOM path + // cannot leak it. + defer func() { + if !finalized { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "openZ") + } + }() n := dZLagrange.Len() dZCanonical := s.getTempDeviceSlice(n) @@ -1582,7 +1593,8 @@ func (s *instance) batchOpening() error { foldDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening fold", foldDone) + finish, guard := makeFinisher(stream, "batchOpening fold", foldDone) + defer guard() dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { @@ -1649,7 +1661,8 @@ func (s *instance) batchOpening() error { divDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + finish, guard := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + defer guard() // For Montgomery vectors, scalar multipliers must be standard-form. dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) dWitness = s.getTempDeviceSlice(witnessSize) @@ -2046,8 +2059,10 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS done <- cfgErr return } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) - done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -2644,6 +2659,8 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() dMerged = s.getTempDeviceSlice(totalSize) mergeErr := icicle_vecops.MergeShardsBitReverse( gpuNumerator.shards, @@ -2653,12 +2670,11 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cfg, ) if mergeErr != icicle_runtime.Success { - _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") - mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) return } // Async boundary before merged slice is consumed by host copy. - mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + finish(nil) }) if mergeErr := <-mergeDone; mergeErr != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -2675,8 +2691,10 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial downloadDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) - downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + finish(nil) }) if err := <-downloadDone; err != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -3689,13 +3707,25 @@ func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label str // makeFinisher returns a closure that synchronizes and destroys the stream, // then sends the (possibly merged) error to done. Use inside RunOnDevice closures. -func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { - return func(runErr error) { +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) (finish func(error), guard func()) { + finished := false + finish = func(runErr error) { + finished = true if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } done <- runErr } + // guard must be deferred right after the stream is created: if the + // closure panics before finish runs (the panic itself is recovered and + // reported by devicePanicToError), the guard still synchronizes and + // destroys the stream so the GPU-OOM path cannot leak it. + guard = func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(stream, label) + } + } + return finish, guard } // devicePanicToError is deferred at the top of RunOnDevice closures that @@ -5027,7 +5057,8 @@ func (s *instance) inverseAndMergeShards( done <- cfgErr return } - finish := makeFinisher(stream, "inverseAndMergeShards", done) + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() cfgNtt := icicle_ntt.GetDefaultNttConfig() cfgNtt.IsAsync = true @@ -5201,7 +5232,8 @@ func (s *instance) divideByZHOnGPU( scaleDone <- cfgErr return } - finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + finish, guard := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + defer guard() for i := 0; i < gpuNumerator.rho; i++ { dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) @@ -5251,9 +5283,11 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io done <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadQuotientFromGPU", done) + defer guard() host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) // Async boundary for host materialization of quotient coefficients. - done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -5655,7 +5689,8 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone <- cfgErr return } - finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() cleanupOwned := func(from int) { for i := from; i < len(devicePolys); i++ { @@ -5940,7 +5975,9 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( ownedLen = 0 } } + finished := false finish := func(runErr error) { + finished = true // Async boundary for eval path before handing result back to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -5948,6 +5985,15 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( releaseEval() done <- runErr } + // If the closure panics before finish runs (recovered by + // devicePanicToError), still release the stream and any owned eval + // buffer so the GPU-OOM path cannot leak them. + defer func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU") + releaseEval() + } + }() prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) dEval = prepareResult.dEval diff --git a/backend/accelerated/icicle/plonk/bw6-761/icicle.go b/backend/accelerated/icicle/plonk/bw6-761/icicle.go index c26c2214cc..9f8978b5a0 100644 --- a/backend/accelerated/icicle/plonk/bw6-761/icicle.go +++ b/backend/accelerated/icicle/plonk/bw6-761/icicle.go @@ -1025,7 +1025,8 @@ func (s *instance) buildRatioCopyConstraint() (err error) { copyDone <- cfgErr return } - finish := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() var allocErr error dPersist, allocErr = allocDeviceUninitialized(dZ.Len()) if allocErr != nil { @@ -1116,7 +1117,9 @@ func (s *instance) openZ() (err error) { buildDone <- cfgErr return } + finalized := false finalize := func(runErr error, dZCanonical icicle_core.DeviceSlice, releaseBlinded bool) { + finalized = true // Async boundary for canonicalization/blinding before exposing output. if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, "openZ"); syncErr != nil && runErr == nil { runErr = syncErr @@ -1130,6 +1133,14 @@ func (s *instance) openZ() (err error) { } buildDone <- runErr } + // If the closure panics before finalize runs (recovered by + // devicePanicToError), still release the stream so the GPU-OOM path + // cannot leak it. + defer func() { + if !finalized { + _ = syncAndDestroyStreamOnCurrentDevice(stream, "openZ") + } + }() n := dZLagrange.Len() dZCanonical := s.getTempDeviceSlice(n) @@ -1582,7 +1593,8 @@ func (s *instance) batchOpening() error { foldDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening fold", foldDone) + finish, guard := makeFinisher(stream, "batchOpening fold", foldDone) + defer guard() dFold = s.getTempDeviceSlice(s.linearizedPolynomialGPU.Len()) if e := copyDeviceSliceIntoOnCurrentDevice(dFold, s.linearizedPolynomialGPU, cfg); e != icicle_runtime.Success { @@ -1649,7 +1661,8 @@ func (s *instance) batchOpening() error { divDone <- cfgErr return } - finish := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + finish, guard := makeFinisher(stream, "batchOpening divideByXMinusA", divDone) + defer guard() // For Montgomery vectors, scalar multipliers must be standard-form. dPoint := uploadScalarStdOnCurrentDevice(s.zeta, cfg) dWitness = s.getTempDeviceSlice(witnessSize) @@ -2046,8 +2059,10 @@ func (s *instance) downloadCanonicalDeviceCoefficients(dPoly icicle_core.DeviceS done <- cfgErr return } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) - done <- syncAndDestroyStreamOnCurrentDevice(stream, label+" download") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -2644,6 +2659,8 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial mergeDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() dMerged = s.getTempDeviceSlice(totalSize) mergeErr := icicle_vecops.MergeShardsBitReverse( gpuNumerator.shards, @@ -2653,12 +2670,11 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial cfg, ) if mergeErr != icicle_runtime.Success { - _ = syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") - mergeDone <- fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString()) + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) return } // Async boundary before merged slice is consumed by host copy. - mergeDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU merge") + finish(nil) }) if mergeErr := <-mergeDone; mergeErr != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -2675,8 +2691,10 @@ func (s *instance) downloadNumeratorFromGPU(gpuNumerator *gpuNumeratorPolynomial downloadDone <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) - downloadDone <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadNumeratorFromGPU download") + finish(nil) }) if err := <-downloadDone; err != nil { s.putTempDeviceSlice(dMerged, totalSize) @@ -3689,13 +3707,25 @@ func syncAndDestroyStreamOnCurrentDevice(stream icicle_runtime.Stream, label str // makeFinisher returns a closure that synchronizes and destroys the stream, // then sends the (possibly merged) error to done. Use inside RunOnDevice closures. -func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) func(error) { - return func(runErr error) { +func makeFinisher(stream icicle_runtime.Stream, label string, done chan<- error) (finish func(error), guard func()) { + finished := false + finish = func(runErr error) { + finished = true if syncErr := syncAndDestroyStreamOnCurrentDevice(stream, label); syncErr != nil && runErr == nil { runErr = syncErr } done <- runErr } + // guard must be deferred right after the stream is created: if the + // closure panics before finish runs (the panic itself is recovered and + // reported by devicePanicToError), the guard still synchronizes and + // destroys the stream so the GPU-OOM path cannot leak it. + guard = func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(stream, label) + } + } + return finish, guard } // devicePanicToError is deferred at the top of RunOnDevice closures that @@ -5027,7 +5057,8 @@ func (s *instance) inverseAndMergeShards( done <- cfgErr return } - finish := makeFinisher(stream, "inverseAndMergeShards", done) + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() cfgNtt := icicle_ntt.GetDefaultNttConfig() cfgNtt.IsAsync = true @@ -5201,7 +5232,8 @@ func (s *instance) divideByZHOnGPU( scaleDone <- cfgErr return } - finish := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + finish, guard := makeFinisher(stream, "divideByZHOnGPU", scaleDone) + defer guard() for i := 0; i < gpuNumerator.rho; i++ { dScale := uploadScalarStdOnCurrentDevice(xnMinusOneInverseLagrangeCoset[i], cfg) vecErr := icicle_vecops.ScalarMulVec(dScale, gpuNumerator.shards[i], gpuNumerator.shards[i], cfg) @@ -5251,9 +5283,11 @@ func (s *instance) downloadQuotientFromGPU(quotient *gpuQuotientPolynomial) (*io done <- cfgErr return } + finish, guard := makeFinisher(stream, "downloadQuotientFromGPU", done) + defer guard() host.CopyFromDeviceAsync("ient.coeffs, cfg.StreamHandle) // Async boundary for host materialization of quotient coefficients. - done <- syncAndDestroyStreamOnCurrentDevice(stream, "downloadQuotientFromGPU") + finish(nil) }) if err := <-done; err != nil { return nil, err @@ -5655,7 +5689,8 @@ func (s *instance) prepareBatchOpeningPolynomialsOnGPU( prepDone <- cfgErr return } - finish := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() cleanupOwned := func(from int) { for i := from; i < len(devicePolys); i++ { @@ -5940,7 +5975,9 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( ownedLen = 0 } } + finished := false finish := func(runErr error) { + finished = true // Async boundary for eval path before handing result back to caller. if syncErr := syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU"); syncErr != nil && runErr == nil { runErr = syncErr @@ -5948,6 +5985,15 @@ func (s *instance) evalPolynomialInCurrentFormOnGPU( releaseEval() done <- runErr } + // If the closure panics before finish runs (recovered by + // devicePanicToError), still release the stream and any owned eval + // buffer so the GPU-OOM path cannot leak them. + defer func() { + if !finished { + _ = syncAndDestroyStreamOnCurrentDevice(evalStream, "evalPolynomialInCurrentFormOnGPU") + releaseEval() + } + }() prepareResult, prepErr := s.prepareEvalPolynomialInputOnCurrentDevice(p, dSrc, cfgVec) dEval = prepareResult.dEval From 2150115442840ea4c402f6ffb8175348c9b72815 Mon Sep 17 00:00:00 2001 From: Martun Karapetyan Date: Tue, 28 Jul 2026 11:30:54 +0400 Subject: [PATCH 5/5] fix(icicle/groth16): validate all cached solution dimensions, not just W The cached R1CSSolution was accepted after checking only the wire count, but a circuit edit can keep the wire count while changing the constraint count, leaving A/B/C the wrong size. That reaches computeH, whose padding of Domain.Cardinality-len(a) panics when the stale vectors are larger than the domain, or silently proves a stale assignment otherwise. The load path now also requires len(A) == len(B) == len(C) == GetNbConstraints() and falls back to solving with a warning on any mismatch. Co-Authored-By: Claude Fable 5 --- .../icicle/groth16/bls12-377/icicle.go | 17 ++++++++++++++--- .../icicle/groth16/bls12-381/icicle.go | 17 ++++++++++++++--- .../accelerated/icicle/groth16/bn254/icicle.go | 17 ++++++++++++++--- .../icicle/groth16/bw6-761/icicle.go | 17 ++++++++++++++--- .../generator/templates/groth16.icicle.go.tmpl | 17 ++++++++++++++--- 5 files changed, 70 insertions(+), 15 deletions(-) diff --git a/backend/accelerated/icicle/groth16/bls12-377/icicle.go b/backend/accelerated/icicle/groth16/bls12-377/icicle.go index cb6946b098..4ddb0ab795 100644 --- a/backend/accelerated/icicle/groth16/bls12-377/icicle.go +++ b/backend/accelerated/icicle/groth16/bls12-377/icicle.go @@ -895,10 +895,21 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic if canCache { if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + // Validate every dimension against the circuit, not just the wire + // count: a circuit edit can keep the wire count while changing + // the constraint count, and computeH pads A/B/C with + // Domain.Cardinality-len(a) — wrong-sized vectors panic there or + // produce a proof for a stale assignment. expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() - if len(cached.W) != expectedWires { - log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). - Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + expectedConstraints := r1cs.GetNbConstraints() + if len(cached.W) != expectedWires || + len(cached.A) != expectedConstraints || + len(cached.B) != expectedConstraints || + len(cached.C) != expectedConstraints { + log.Warn().Str("file", cachePath). + Int("wires", len(cached.W)).Int("expectedWires", expectedWires). + Int("constraints", len(cached.A)).Int("expectedConstraints", expectedConstraints). + Msg("ignoring cached Groth16 solution: dimension mismatch (stale cache?)") } else { solution = cached log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver") diff --git a/backend/accelerated/icicle/groth16/bls12-381/icicle.go b/backend/accelerated/icicle/groth16/bls12-381/icicle.go index f385fd1718..4005f5b965 100644 --- a/backend/accelerated/icicle/groth16/bls12-381/icicle.go +++ b/backend/accelerated/icicle/groth16/bls12-381/icicle.go @@ -895,10 +895,21 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic if canCache { if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + // Validate every dimension against the circuit, not just the wire + // count: a circuit edit can keep the wire count while changing + // the constraint count, and computeH pads A/B/C with + // Domain.Cardinality-len(a) — wrong-sized vectors panic there or + // produce a proof for a stale assignment. expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() - if len(cached.W) != expectedWires { - log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). - Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + expectedConstraints := r1cs.GetNbConstraints() + if len(cached.W) != expectedWires || + len(cached.A) != expectedConstraints || + len(cached.B) != expectedConstraints || + len(cached.C) != expectedConstraints { + log.Warn().Str("file", cachePath). + Int("wires", len(cached.W)).Int("expectedWires", expectedWires). + Int("constraints", len(cached.A)).Int("expectedConstraints", expectedConstraints). + Msg("ignoring cached Groth16 solution: dimension mismatch (stale cache?)") } else { solution = cached log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver") diff --git a/backend/accelerated/icicle/groth16/bn254/icicle.go b/backend/accelerated/icicle/groth16/bn254/icicle.go index 9b36fb8bcb..d5c21d2748 100644 --- a/backend/accelerated/icicle/groth16/bn254/icicle.go +++ b/backend/accelerated/icicle/groth16/bn254/icicle.go @@ -895,10 +895,21 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic if canCache { if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + // Validate every dimension against the circuit, not just the wire + // count: a circuit edit can keep the wire count while changing + // the constraint count, and computeH pads A/B/C with + // Domain.Cardinality-len(a) — wrong-sized vectors panic there or + // produce a proof for a stale assignment. expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() - if len(cached.W) != expectedWires { - log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). - Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + expectedConstraints := r1cs.GetNbConstraints() + if len(cached.W) != expectedWires || + len(cached.A) != expectedConstraints || + len(cached.B) != expectedConstraints || + len(cached.C) != expectedConstraints { + log.Warn().Str("file", cachePath). + Int("wires", len(cached.W)).Int("expectedWires", expectedWires). + Int("constraints", len(cached.A)).Int("expectedConstraints", expectedConstraints). + Msg("ignoring cached Groth16 solution: dimension mismatch (stale cache?)") } else { solution = cached log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver") diff --git a/backend/accelerated/icicle/groth16/bw6-761/icicle.go b/backend/accelerated/icicle/groth16/bw6-761/icicle.go index 05b2bb3eeb..02f8ac4b38 100644 --- a/backend/accelerated/icicle/groth16/bw6-761/icicle.go +++ b/backend/accelerated/icicle/groth16/bw6-761/icicle.go @@ -884,10 +884,21 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic if canCache { if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + // Validate every dimension against the circuit, not just the wire + // count: a circuit edit can keep the wire count while changing + // the constraint count, and computeH pads A/B/C with + // Domain.Cardinality-len(a) — wrong-sized vectors panic there or + // produce a proof for a stale assignment. expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() - if len(cached.W) != expectedWires { - log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). - Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + expectedConstraints := r1cs.GetNbConstraints() + if len(cached.W) != expectedWires || + len(cached.A) != expectedConstraints || + len(cached.B) != expectedConstraints || + len(cached.C) != expectedConstraints { + log.Warn().Str("file", cachePath). + Int("wires", len(cached.W)).Int("expectedWires", expectedWires). + Int("constraints", len(cached.A)).Int("expectedConstraints", expectedConstraints). + Msg("ignoring cached Groth16 solution: dimension mismatch (stale cache?)") } else { solution = cached log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver") diff --git a/backend/accelerated/icicle/internal/generator/templates/groth16.icicle.go.tmpl b/backend/accelerated/icicle/internal/generator/templates/groth16.icicle.go.tmpl index b4d2a80a07..7192f855fd 100644 --- a/backend/accelerated/icicle/internal/generator/templates/groth16.icicle.go.tmpl +++ b/backend/accelerated/icicle/internal/generator/templates/groth16.icicle.go.tmpl @@ -909,10 +909,21 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, cfg *icic if canCache { if cached, err := cs.LoadR1CSSolution(cachePath); err == nil { + // Validate every dimension against the circuit, not just the wire + // count: a circuit edit can keep the wire count while changing + // the constraint count, and computeH pads A/B/C with + // Domain.Cardinality-len(a) — wrong-sized vectors panic there or + // produce a proof for a stale assignment. expectedWires := r1cs.GetNbPublicVariables() + r1cs.GetNbSecretVariables() + r1cs.GetNbInternalVariables() - if len(cached.W) != expectedWires { - log.Warn().Str("file", cachePath).Int("got", len(cached.W)).Int("expected", expectedWires). - Msg("ignoring cached Groth16 solution: wire count mismatch (stale cache?)") + expectedConstraints := r1cs.GetNbConstraints() + if len(cached.W) != expectedWires || + len(cached.A) != expectedConstraints || + len(cached.B) != expectedConstraints || + len(cached.C) != expectedConstraints { + log.Warn().Str("file", cachePath). + Int("wires", len(cached.W)).Int("expectedWires", expectedWires). + Int("constraints", len(cached.A)).Int("expectedConstraints", expectedConstraints). + Msg("ignoring cached Groth16 solution: dimension mismatch (stale cache?)") } else { solution = cached log.Debug().Str("file", cachePath).Msg("loaded cached Groth16 solution, skipping solver")