diff --git a/README.md b/README.md index 50af91482e..6a34545275 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,68 @@ 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 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 +(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..4ddb0ab795 100644 --- a/backend/accelerated/icicle/groth16/bls12-377/icicle.go +++ b/backend/accelerated/icicle/groth16/bls12-377/icicle.go @@ -888,12 +888,50 @@ 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 { + // 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() + 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") + } + } } - 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..4005f5b965 100644 --- a/backend/accelerated/icicle/groth16/bls12-381/icicle.go +++ b/backend/accelerated/icicle/groth16/bls12-381/icicle.go @@ -888,12 +888,50 @@ 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 { + // 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() + 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") + } + } } - 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..d5c21d2748 100644 --- a/backend/accelerated/icicle/groth16/bn254/icicle.go +++ b/backend/accelerated/icicle/groth16/bn254/icicle.go @@ -888,12 +888,50 @@ 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 { + // 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() + 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") + } + } } - 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..02f8ac4b38 100644 --- a/backend/accelerated/icicle/groth16/bw6-761/icicle.go +++ b/backend/accelerated/icicle/groth16/bw6-761/icicle.go @@ -877,12 +877,50 @@ 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 { + // 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() + 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") + } + } } - 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..7b8402824d --- /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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + 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, +// 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location +} + +// 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..7192f855fd 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,50 @@ 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 { + // 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() + 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") + } + } } - 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..0c86a77357 --- /dev/null +++ b/backend/accelerated/icicle/internal/generator/templates/plonk.icicle.go.tmpl @@ -0,0 +1,7392 @@ +//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" + + crypto_utils "github.com/consensys/gnark-crypto/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() + + // 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 + goStage("solveConstraints", instance.solveConstraints) + + // complete qk + goStage("completeQk", instance.completeQk) + + // init blinding polynomials + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) + + // compute h + goStage("computeQuotient", instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + goStage("openZ", instance.openZ) + + // linearized polynomial + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + goStage("batchOpening", 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") + // 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 + 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 { + 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) { + defer devicePanicToError(copyDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() + 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) { + defer devicePanicToError(buildDone) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + 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 + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + 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) + 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 { + // 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 + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + s.putTempDeviceSlice(dLin, dLin.Len()) + 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() { + // 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 + } + }() + + 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) { + defer devicePanicToError(foldDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + 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 { + 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) { + defer devicePanicToError(divDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + 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) + 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 error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) + dWitness = uploadVector(hCoeffs) + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + finish(nil) + }) + 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) + crypto_utils.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) { + defer devicePanicToError(uploadTwiddlesDone) + 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) { + defer devicePanicToError(batchInvertDone) + 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) + crypto_utils.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) { + defer devicePanicToError(mergeDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) + return + } + // Async boundary before merged slice is consumed by host copy. + finish(nil) + }) + 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) { + defer devicePanicToError(downloadDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + finish(nil) + }) + 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() { + 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 + } + 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 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 + } + 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 + } + // 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.Shutdown() + 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() { + 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 + } + 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) { + defer devicePanicToError(uploadDone) + 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 + }) + // 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 { + p := state.polys[i] + if p == nil { + continue + } + + 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, 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 + + // 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 { + runErr = 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 { + runErr = 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 { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + }) + } + + // 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 && 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 { + if p != nil { + p.Basis = iop.Lagrange + } + } + 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, 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 + + // 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())) + } + 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 + 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 + // 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), + outstanding: make(map[unsafe.Pointer]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] + 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 + } + + // 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 { + p.outstanding[ds.AsUnsafePointer()] = ds + 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() + + delete(p.outstanding, ds.AsUnsafePointer()) + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// 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() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + 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) { + 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) (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 +// 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 { + 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. 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, +) (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]) + } + }) + + 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 +} + +// 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()) + } + + // 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 +} + +// 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. + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) + 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) + + 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 { + 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(state, n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + runErr = 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 { + runErr = 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() + 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 { + 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) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + 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 { + 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 { + 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 { + 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 { + 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 { + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // 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 + + 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) { + defer devicePanicToError(prepareDone) + 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) { + defer devicePanicToError(done) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() + + 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) { + defer devicePanicToError(scaleDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + 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) + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + 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. + finish(nil) + }) + 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) { + 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) + }) + <-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) { + defer devicePanicToError(done) + 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) { + defer devicePanicToError(prepDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() + + 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) { + defer devicePanicToError(done) + 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 + } + } + 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 + } + 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 + 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) { + defer devicePanicToError(divDone) + 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) { + defer devicePanicToError(divDone) + 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 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 + } + 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() { + 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 + } + 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() { + 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 + } + 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 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() + 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 { + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return + } + chInitDomain <- nil + }) + + 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 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) + chLag <- nil + }) + } else { + 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) + chCan <- nil + }) + } else { + chCan <- nil + } + + 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 +} + +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..f2fd78f8f2 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bls12-377/icicle.go @@ -0,0 +1,7399 @@ +// 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" + + crypto_utils "github.com/consensys/gnark-crypto/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() + + // 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 + goStage("solveConstraints", instance.solveConstraints) + + // complete qk + goStage("completeQk", instance.completeQk) + + // init blinding polynomials + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) + + // compute h + goStage("computeQuotient", instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + goStage("openZ", instance.openZ) + + // linearized polynomial + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + goStage("batchOpening", 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") + // 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 + 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 { + 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) { + defer devicePanicToError(copyDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() + 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) { + defer devicePanicToError(buildDone) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + 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 + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + 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) + 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 { + // 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 + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + s.putTempDeviceSlice(dLin, dLin.Len()) + 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() { + // 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 + } + }() + + 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) { + defer devicePanicToError(foldDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + 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 { + 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) { + defer devicePanicToError(divDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + 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) + 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 error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) + dWitness = uploadVector(hCoeffs) + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + finish(nil) + }) + 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) + crypto_utils.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) { + defer devicePanicToError(uploadTwiddlesDone) + 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) { + defer devicePanicToError(batchInvertDone) + 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) + crypto_utils.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) { + defer devicePanicToError(mergeDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) + return + } + // Async boundary before merged slice is consumed by host copy. + finish(nil) + }) + 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) { + defer devicePanicToError(downloadDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + finish(nil) + }) + 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() { + 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 + } + 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 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 + } + 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 + } + // 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.Shutdown() + 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() { + 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 + } + 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) { + defer devicePanicToError(uploadDone) + 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 + }) + // 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 { + p := state.polys[i] + if p == nil { + continue + } + + 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, 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 + + // 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 { + runErr = 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 { + runErr = 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 { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + }) + } + + // 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 && 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 { + if p != nil { + p.Basis = iop.Lagrange + } + } + 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, 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 + + // 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())) + } + 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 + 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 + // 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), + outstanding: make(map[unsafe.Pointer]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] + 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 + } + + // 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 { + p.outstanding[ds.AsUnsafePointer()] = ds + 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() + + delete(p.outstanding, ds.AsUnsafePointer()) + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// 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() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + 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) { + 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) (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 +// 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 { + 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. 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, +) (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]) + } + }) + + 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 +} + +// 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()) + } + + // 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 +} + +// 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. + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) + 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) + + 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 { + 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(state, n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + runErr = 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 { + runErr = 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() + 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 { + 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) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + 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 { + 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 { + 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 { + 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 { + 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 { + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // 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 + + 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) { + defer devicePanicToError(prepareDone) + 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) { + defer devicePanicToError(done) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() + + 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) { + defer devicePanicToError(scaleDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + 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) + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + 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. + finish(nil) + }) + 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) { + 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) + }) + <-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) { + defer devicePanicToError(done) + 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) { + defer devicePanicToError(prepDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() + + 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) { + defer devicePanicToError(done) + 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 + } + } + 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 + } + 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 + 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) { + defer devicePanicToError(divDone) + 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) { + defer devicePanicToError(divDone) + 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 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 + } + 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() { + 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 + } + 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() { + 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 + } + 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 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() + 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 { + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return + } + chInitDomain <- nil + }) + + 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 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) + chLag <- nil + }) + } else { + 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) + chCan <- nil + }) + } else { + chCan <- nil + } + + 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 +} + +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..aa8368b044 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bls12-381/icicle.go @@ -0,0 +1,7399 @@ +// 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" + + crypto_utils "github.com/consensys/gnark-crypto/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() + + // 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 + goStage("solveConstraints", instance.solveConstraints) + + // complete qk + goStage("completeQk", instance.completeQk) + + // init blinding polynomials + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) + + // compute h + goStage("computeQuotient", instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + goStage("openZ", instance.openZ) + + // linearized polynomial + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + goStage("batchOpening", 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") + // 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 + 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 { + 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) { + defer devicePanicToError(copyDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() + 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) { + defer devicePanicToError(buildDone) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + 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 + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + 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) + 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 { + // 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 + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + s.putTempDeviceSlice(dLin, dLin.Len()) + 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() { + // 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 + } + }() + + 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) { + defer devicePanicToError(foldDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + 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 { + 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) { + defer devicePanicToError(divDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + 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) + 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 error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) + dWitness = uploadVector(hCoeffs) + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + finish(nil) + }) + 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) + crypto_utils.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) { + defer devicePanicToError(uploadTwiddlesDone) + 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) { + defer devicePanicToError(batchInvertDone) + 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) + crypto_utils.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) { + defer devicePanicToError(mergeDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) + return + } + // Async boundary before merged slice is consumed by host copy. + finish(nil) + }) + 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) { + defer devicePanicToError(downloadDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + finish(nil) + }) + 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() { + 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 + } + 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 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 + } + 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 + } + // 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.Shutdown() + 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() { + 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 + } + 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) { + defer devicePanicToError(uploadDone) + 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 + }) + // 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 { + p := state.polys[i] + if p == nil { + continue + } + + 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, 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 + + // 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 { + runErr = 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 { + runErr = 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 { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + }) + } + + // 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 && 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 { + if p != nil { + p.Basis = iop.Lagrange + } + } + 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, 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 + + // 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())) + } + 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 + 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 + // 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), + outstanding: make(map[unsafe.Pointer]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] + 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 + } + + // 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 { + p.outstanding[ds.AsUnsafePointer()] = ds + 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() + + delete(p.outstanding, ds.AsUnsafePointer()) + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// 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() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + 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) { + 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) (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 +// 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 { + 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. 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, +) (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]) + } + }) + + 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 +} + +// 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()) + } + + // 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 +} + +// 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. + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) + 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) + + 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 { + 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(state, n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + runErr = 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 { + runErr = 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() + 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 { + 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) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + 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 { + 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 { + 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 { + 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 { + 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 { + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // 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 + + 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) { + defer devicePanicToError(prepareDone) + 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) { + defer devicePanicToError(done) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() + + 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) { + defer devicePanicToError(scaleDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + 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) + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + 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. + finish(nil) + }) + 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) { + 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) + }) + <-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) { + defer devicePanicToError(done) + 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) { + defer devicePanicToError(prepDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() + + 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) { + defer devicePanicToError(done) + 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 + } + } + 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 + } + 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 + 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) { + defer devicePanicToError(divDone) + 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) { + defer devicePanicToError(divDone) + 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 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 + } + 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() { + 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 + } + 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() { + 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 + } + 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 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() + 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 { + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return + } + chInitDomain <- nil + }) + + 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 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) + chLag <- nil + }) + } else { + 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) + chCan <- nil + }) + } else { + chCan <- nil + } + + 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 +} + +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..ea03ae87e2 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bn254/icicle.go @@ -0,0 +1,7399 @@ +// 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" + + crypto_utils "github.com/consensys/gnark-crypto/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() + + // 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 + goStage("solveConstraints", instance.solveConstraints) + + // complete qk + goStage("completeQk", instance.completeQk) + + // init blinding polynomials + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) + + // compute h + goStage("computeQuotient", instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + goStage("openZ", instance.openZ) + + // linearized polynomial + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + goStage("batchOpening", 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") + // 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 + 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 { + 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) { + defer devicePanicToError(copyDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() + 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) { + defer devicePanicToError(buildDone) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + 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 + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + 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) + 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 { + // 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 + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + s.putTempDeviceSlice(dLin, dLin.Len()) + 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() { + // 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 + } + }() + + 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) { + defer devicePanicToError(foldDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + 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 { + 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) { + defer devicePanicToError(divDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + 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) + 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 error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) + dWitness = uploadVector(hCoeffs) + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + finish(nil) + }) + 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) + crypto_utils.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) { + defer devicePanicToError(uploadTwiddlesDone) + 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) { + defer devicePanicToError(batchInvertDone) + 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) + crypto_utils.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) { + defer devicePanicToError(mergeDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) + return + } + // Async boundary before merged slice is consumed by host copy. + finish(nil) + }) + 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) { + defer devicePanicToError(downloadDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + finish(nil) + }) + 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() { + 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 + } + 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 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 + } + 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 + } + // 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.Shutdown() + 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() { + 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 + } + 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) { + defer devicePanicToError(uploadDone) + 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 + }) + // 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 { + p := state.polys[i] + if p == nil { + continue + } + + 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, 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 + + // 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 { + runErr = 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 { + runErr = 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 { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + }) + } + + // 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 && 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 { + if p != nil { + p.Basis = iop.Lagrange + } + } + 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, 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 + + // 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())) + } + 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 + 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 + // 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), + outstanding: make(map[unsafe.Pointer]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] + 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 + } + + // 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 { + p.outstanding[ds.AsUnsafePointer()] = ds + 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() + + delete(p.outstanding, ds.AsUnsafePointer()) + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// 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() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + 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) { + 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) (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 +// 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 { + 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. 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, +) (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]) + } + }) + + 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 +} + +// 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()) + } + + // 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 +} + +// 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. + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) + 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) + + 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 { + 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(state, n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + runErr = 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 { + runErr = 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() + 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 { + 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) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + 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 { + 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 { + 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 { + 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 { + 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 { + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // 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 + + 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) { + defer devicePanicToError(prepareDone) + 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) { + defer devicePanicToError(done) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() + + 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) { + defer devicePanicToError(scaleDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + 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) + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + 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. + finish(nil) + }) + 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) { + 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) + }) + <-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) { + defer devicePanicToError(done) + 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) { + defer devicePanicToError(prepDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() + + 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) { + defer devicePanicToError(done) + 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 + } + } + 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 + } + 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 + 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) { + defer devicePanicToError(divDone) + 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) { + defer devicePanicToError(divDone) + 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 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 + } + 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() { + 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 + } + 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() { + 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 + } + 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 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() + 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 { + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return + } + chInitDomain <- nil + }) + + 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 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) + chLag <- nil + }) + } else { + 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) + chCan <- nil + }) + } else { + chCan <- nil + } + + 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 +} + +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..9f8978b5a0 --- /dev/null +++ b/backend/accelerated/icicle/plonk/bw6-761/icicle.go @@ -0,0 +1,7399 @@ +// 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" + + crypto_utils "github.com/consensys/gnark-crypto/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() + + // 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 + goStage("solveConstraints", instance.solveConstraints) + + // complete qk + goStage("completeQk", instance.completeQk) + + // init blinding polynomials + goStage("initBlindingPolynomials", instance.initBlindingPolynomials) + + // derive gamma, beta (copy constraint) + goStage("deriveGammaAndBeta", instance.deriveGammaAndBeta) + + // compute accumulating ratio for the copy constraint + goStage("buildRatioCopyConstraint", instance.buildRatioCopyConstraint) + + // compute h + goStage("computeQuotient", instance.computeQuotient) + + // open Z (blinded) at ωζ (proof.ZShiftedOpening) + goStage("openZ", instance.openZ) + + // linearized polynomial + goStage("computeLinearizedPolynomial", instance.computeLinearizedPolynomial) + + // Batch opening (no internal timer of its own — time the whole stage here) + goStage("batchOpening", 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") + // 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 + 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 { + 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) { + defer devicePanicToError(copyDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("buildRatioCopyConstraint") + if cfgErr != nil { + copyDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "buildRatioCopyConstraint", copyDone) + defer guard() + 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) { + defer devicePanicToError(buildDone) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("openZ") + if cfgErr != nil { + 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 + } + if !dZCanonical.IsEmpty() { + s.putTempDeviceSlice(dZCanonical, dZCanonical.Len()) + } + if releaseBlinded && !dBlindedCanonical.IsEmpty() { + s.putTempDeviceSlice(dBlindedCanonical, dBlindedCanonical.Len()) + dBlindedCanonical = icicle_core.DeviceSlice{} + } + 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) + 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 { + // 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 + } + + doneSubtractH := profileStep("computeLinearizedPolynomial: subtract quotient contribution on GPU") + err = s.subtractQuotientContributionFromLinearizedOnGPU(dLin, s.hGPU) + doneSubtractH() + if err != nil { + s.putTempDeviceSlice(dLin, dLin.Len()) + 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() { + // 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 + } + }() + + 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) { + defer devicePanicToError(foldDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening fold") + if cfgErr != nil { + foldDone <- cfgErr + return + } + 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 { + 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) { + defer devicePanicToError(divDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("batchOpening divideByXMinusA") + if cfgErr != nil { + divDone <- cfgErr + return + } + 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) + 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 error, 1) + icicle_runtime.RunOnDevice(&s.device, func(args ...any) { + defer devicePanicToError(uploadDone) + dWitness = uploadVector(hCoeffs) + uploadDone <- nil + }) + if err := <-uploadDone; err != nil { + return kzg.BatchOpeningProof{}, err + } + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice(label + " download") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, label+" download", done) + defer guard() + host.CopyFromDeviceAsync(&dPoly, cfg.StreamHandle) + finish(nil) + }) + 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) + crypto_utils.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) { + defer devicePanicToError(uploadTwiddlesDone) + 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) { + defer devicePanicToError(batchInvertDone) + 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) + crypto_utils.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) { + defer devicePanicToError(mergeDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU merge") + if cfgErr != nil { + mergeDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU merge", mergeDone) + defer guard() + dMerged = s.getTempDeviceSlice(totalSize) + mergeErr := icicle_vecops.MergeShardsBitReverse( + gpuNumerator.shards, + gpuNumerator.n, + gpuNumerator.mm, + dMerged, + cfg, + ) + if mergeErr != icicle_runtime.Success { + finish(fmt.Errorf("downloadNumeratorFromGPU: merge shards kernel failed: %s", mergeErr.AsString())) + return + } + // Async boundary before merged slice is consumed by host copy. + finish(nil) + }) + 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) { + defer devicePanicToError(downloadDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadNumeratorFromGPU download") + if cfgErr != nil { + downloadDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "downloadNumeratorFromGPU download", downloadDone) + defer guard() + cresHost.CopyFromDeviceAsync(&dMerged, cfg.StreamHandle) + finish(nil) + }) + 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() { + 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 + } + 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 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 + } + 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 + } + // 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.Shutdown() + 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() { + 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 + } + 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) { + defer devicePanicToError(uploadDone) + 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 + }) + // 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 { + p := state.polys[i] + if p == nil { + continue + } + + 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, 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 + + // 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 { + runErr = 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 { + runErr = 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 { + runErr = fmt.Errorf("gpuNTTInverseScaleForwardOnDevice: forward NTT failed: %s", err.AsString()) + return + } + }) + } + + // 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 && 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 { + if p != nil { + p.Basis = iop.Lagrange + } + } + 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, 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 + + // 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())) + } + 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 + 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 + // 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), + outstanding: make(map[unsafe.Pointer]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] + 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 + } + + // 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 { + p.outstanding[ds.AsUnsafePointer()] = ds + 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() + + delete(p.outstanding, ds.AsUnsafePointer()) + // Add to the pool + p.freeSlices[n] = append(p.freeSlices[n], ds) +} + +// 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() + + for _, slices := range p.freeSlices { + for _, ds := range slices { + _ = ds.Free() + } + } + 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) { + 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) (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 +// 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 { + 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. 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, +) (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]) + } + }) + + 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 +} + +// 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()) + } + + // 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 +} + +// 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. + // Zero the field so the caller's cleanup defer does not double-free it. + freeDeviceSlice(&state.dGammaScalar) + 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) + + 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 { + 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(state, n, twiddles0, bp) + if err := applyBlindingToPolynomials(state, params, dBlindL, dBlindR, dBlindO, dBlindZ, dBlindZS); err != nil { + runErr = 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 { + runErr = 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() + 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 { + 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) + if err := icicle_vecops.ScalarMulVec(dAlphaSquaredStd, state.dLocal, dTmp, seqVecCfg); err != icicle_runtime.Success { + dAlphaSquaredStd.Free() + 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 { + 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 { + 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 { + 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 { + 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 { + runErr = fmt.Errorf("gpuEvaluateConstraints: failed to copy result to persistent device slice: %s", err.AsString()) + return + } + } + + // 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 + + 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) { + defer devicePanicToError(prepareDone) + 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) { + defer devicePanicToError(done) + cfgVec, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("inverseAndMergeShards") + if cfgErr != nil { + done <- cfgErr + return + } + finish, guard := makeFinisher(stream, "inverseAndMergeShards", done) + defer guard() + + 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) { + defer devicePanicToError(scaleDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("divideByZHOnGPU") + if cfgErr != nil { + scaleDone <- cfgErr + return + } + 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) + 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) { + defer devicePanicToError(done) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("downloadQuotientFromGPU") + if cfgErr != nil { + 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. + finish(nil) + }) + 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) { + 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) + }) + <-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) { + defer devicePanicToError(done) + 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) { + defer devicePanicToError(prepDone) + cfg, stream, cfgErr := createAsyncVecOpsConfigOnCurrentDevice("prepareBatchOpeningPolynomialsOnGPU") + if cfgErr != nil { + prepDone <- cfgErr + return + } + finish, guard := makeFinisher(stream, "prepareBatchOpeningPolynomialsOnGPU", prepDone) + defer guard() + + 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) { + defer devicePanicToError(done) + 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 + } + } + 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 + } + 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 + 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) { + defer devicePanicToError(divDone) + 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) { + defer devicePanicToError(divDone) + 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 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 + } + 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() { + 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 + } + 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() { + 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 + } + 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 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() + 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 { + chInitDomain <- fmt.Errorf("icicle: InitDomain failed: %s", e.AsString()) + return + } + chInitDomain <- nil + }) + + 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 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) + chLag <- nil + }) + } else { + 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) + chCan <- nil + }) + } else { + chCan <- nil + } + + 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 +} + +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/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") + } +} 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..15887f565c --- /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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + 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, +// 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location +} + +// 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..8363e37ba6 --- /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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + 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, +// 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location +} + +// 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..aebfde0d02 --- /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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + 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, +// 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location +} + +// 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..8c5f279919 --- /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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + 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, +// 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + 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) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) // #nosec G703 -- tmpName comes from os.CreateTemp in the caller-chosen cache directory + return err + } + return os.Rename(tmpName, path) // #nosec G703 -- path is the caller-provided cache location +} + +// 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