From 419da3969e7e1138881fbb55bcfafcb0709aae8e Mon Sep 17 00:00:00 2001 From: Thomas Piellard Date: Tue, 26 May 2026 17:31:03 +0200 Subject: [PATCH 1/8] feat: bench/main.go --- bench/main.go | 116 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) create mode 100644 bench/main.go diff --git a/bench/main.go b/bench/main.go new file mode 100644 index 0000000..f274ec3 --- /dev/null +++ b/bench/main.go @@ -0,0 +1,116 @@ +package main + +import ( + "fmt" + "os" + "runtime/pprof" + + "github.com/consensys/loom" + "github.com/consensys/loom/arguments" + "github.com/consensys/loom/board" + "github.com/consensys/loom/expr" + gnarkplonk "github.com/consensys/loom/integration_test/gnark_plonk" + "github.com/consensys/loom/prover" + "github.com/consensys/loom/setup" + "github.com/consensys/loom/trace" +) + +func main() { + + nbInstances := 40 + ns := make([]int, nbInstances) + for i := 0; i < nbInstances; i++ { + ns[i] = 1 << 15 + } + builder := board.NewBuilder() + traces := make([]trace.Trace, len(ns)) + for i, n := range ns { + tr, sigma, size, err := gnarkplonk.GetIthPlonkTrace(n, i) + if err != nil { + fmt.Println(err) + os.Exit(-1) + } + traces[i] = tr + builder.AddModule(gnarkplonk.PrepareIthPlonk(size, i)) + + lro := []expr.Expr{expr.Col(gnarkplonk.Ith(gnarkplonk.ID_L, i)), expr.Col(gnarkplonk.Ith(gnarkplonk.ID_R, i)), expr.Col(gnarkplonk.Ith(gnarkplonk.ID_O, i))} + sigmaGen := board.NewPermutationGen(sigma, gnarkplonk.Ith("plonk.S", i)) + if err := arguments.CopyConstraint(&builder, gnarkplonk.Ith("plonk", i), lro, sigmaGen); err != nil { + fmt.Println(err) + os.Exit(-1) + } + } + + fullTrace := prover.MergeTrace(traces[0], traces[1:]...) + program, err := board.Compile(&builder) + if err != nil { + fmt.Println(err) + os.Exit(-1) + } + + fresh := trace.New(len(fullTrace.Base)) + for k, v := range fullTrace.Base { + fresh.SetBase(k, v) + } + for k, v := range fullTrace.Ext { + fresh.SetExt(k, v) + } + + // runtime.GC() + // fBefore, err := os.Create("heap_before.prof") + // if err != nil { + // fmt.Println(err) + // os.Exit(-1) + // } + // fAfter, err := os.Create("heap_after.prof") + // if err != nil { + // fmt.Println(err) + // os.Exit(-1) + // } + // defer fBefore.Close() + // defer fAfter.Close() + // if err := pprof.WriteHeapProfile(fBefore); err != nil { + // fmt.Println(err) + // os.Exit(-1) + // } + f, err := os.Create("cpu_poseidon2.pprof") + if err != nil { + fmt.Println(err) + os.Exit(-1) + } + defer f.Close() + if err := pprof.StartCPUProfile(f); err != nil { + fmt.Println(err) + os.Exit(-1) + } + _, err = prover.Prove(fresh, setup.ProvingKey{}, nil, program, prover.WithHashBackend(loom.Poseidon2HashBackend())) + if err != nil { + fmt.Println(err) + os.Exit(-1) + } + pprof.StopCPUProfile() + // if err := pprof.WriteHeapProfile(fAfter); err != nil { + // fmt.Println(err) + // os.Exit(-1) + // } + + // f, err := os.Create("cpu.pprof") + // if err != nil { + // fmt.Println(err) + // os.Exit(-1) + // } + // defer f.Close() + // if err := pprof.StartCPUProfile(f); err != nil { + // fmt.Println(err) + // os.Exit(-1) + // } + // for i := 0; i < 100; i++ { + // err := verifier.Verify(nil, setup.VerificationKey{}, program, proof) + // if err != nil { + // fmt.Println(err) + // os.Exit(-1) + // } + // } + // pprof.StopCPUProfile() + +} From 00efd02ab4eac49b2d093da9e3a31d8e8a825130 Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Tue, 26 May 2026 16:49:32 +0000 Subject: [PATCH 2/8] perf(prover): parallelize ComputeAIRQuotients MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-module AIR quotient computation is now dispatched via parallel.Execute over modules sorted by name. On the bench workload (40 PLONK instances of size 2^17), Prove drops from 41.87s → 24.45s wall (1.71× speedup); par rises from 1.09× → 2.13× on a 32-core box. Correctness: - DomainCache is now thread-safe (sync.Mutex around lazy domain build). - Added dag.DAG.Clone(): per-goroutine deep copy of nodes + Leaf pointers, needed because ComputeQuotient* mutates Leaf.Idx for its local indexing and modules can share *expr.Leaf via cross-module arguments (Lookup, CopyConstraint). - Trace map writes for chunk results are batched under a single mutex. Memory: - New poly.Coset[Ext]LagrangeNormalToCanonicalWithCache helpers stop at canonical Normal form; computeAIRChunks uses them and skips the FFT → BitReverse → FFTInverse → BitReverse round-trip that previously bracketed the chunking step. Saves 2 size-bigSize FFTs per module. - Chunks are now sub-slices of the quotient backing array (three-index slice with capacity bound) instead of fresh make + copy. Saves numChunks·N·sizeof(E4|base) of allocation per module. - Net: 12.89 GiB → 12.27 GiB allocated per Prove (-5%), with the same end-to-end test/verify behavior. Bench (bench/main.go): - Fixed the parallelization metric: was using /cpu/classes/total which is GOMAXPROCS × wall (includes idle); now uses /cpu/classes/user + /cpu/classes/gc, giving real on-CPU time. - Defaults bumped to -instances 40 -log2-size 17 so Prove runs ~20-40s and short single-threaded sections become visible in pprof. Also adds optim_plan_ideas.md cataloguing the broader perf opportunities (this PR addresses #1 from that plan). Co-Authored-By: Claude Opus 4.7 (1M context) --- .gitignore | 2 + bench/main.go | 506 +++++++++++++++++++++++++----- internal/dag/dag.go | 36 +++ internal/poly/compute_quotient.go | 42 +-- internal/poly/domain_cache.go | 12 +- optim_plan_ideas.md | 388 +++++++++++++++++++++++ prover/prover.go | 151 ++++++--- 7 files changed, 989 insertions(+), 148 deletions(-) create mode 100644 optim_plan_ideas.md diff --git a/.gitignore b/.gitignore index 46e9dc5..23784f2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,8 @@ playground/ *.html *.prof +*.pprof +bench_profiles/ *.txt *.tmp *.csv diff --git a/bench/main.go b/bench/main.go index f274ec3..acc424b 100644 --- a/bench/main.go +++ b/bench/main.go @@ -1,9 +1,35 @@ +// loom bench: a beefy end-to-end benchmark that exercises the prover on an +// aggregated batch of PLONK instances and reports the metrics we care about +// when hunting for bottlenecks: +// +// - wall time per phase +// - CPU time per phase (user + GC), and the effective parallelization +// factor (CPU/wall) so we can tell single-threaded sections at a glance +// - bytes / objects allocated per phase +// - peak heap during each phase (sampled in the background) +// - GC count, GC CPU share, total CPU consumed +// +// pprof profiles (CPU, heap before/after Prove, allocations) are written to +// -profile-dir so we can drill down with `go tool pprof`. +// +// Run examples: +// +// go run ./bench +// go run ./bench -instances 80 -log2-size 16 -hash sha256 +// go run ./bench -skip-fri -profile-dir /tmp/loom-bench package main import ( + "flag" "fmt" "os" + "path/filepath" + "runtime" + "runtime/metrics" "runtime/pprof" + "sync/atomic" + "text/tabwriter" + "time" "github.com/consensys/loom" "github.com/consensys/loom/arguments" @@ -15,102 +41,430 @@ import ( "github.com/consensys/loom/trace" ) +// Defaults are sized so that Prove runs for ~15-30s on a 16-32 core box: large +// enough that any phase under measurement collects plenty of pprof samples and +// short single-threaded sections (which would be invisible in a 1s run) become +// obvious. Tune down with -instances / -log2-size for quick iteration. +var ( + nbInstances = flag.Int("instances", 40, "number of PLONK instances to aggregate") + log2Size = flag.Int("log2-size", 17, "log2 of each PLONK instance size") + hashName = flag.String("hash", "poseidon2", "hash backend: poseidon2 | sha256") + profileDir = flag.String("profile-dir", "bench_profiles", "directory to write pprof profiles") + skipFRI = flag.Bool("skip-fri", false, "skip the FRI / sampling phase of Prove") + gomaxprocs = flag.Int("gomaxprocs", 0, "override GOMAXPROCS (0 = leave default)") + sampleMs = flag.Int("sample-ms", 50, "heap sampling interval (ms) for peak-heap tracking") +) + func main() { + flag.Parse() - nbInstances := 40 - ns := make([]int, nbInstances) - for i := 0; i < nbInstances; i++ { - ns[i] = 1 << 15 + if *gomaxprocs > 0 { + runtime.GOMAXPROCS(*gomaxprocs) } + if err := os.MkdirAll(*profileDir, 0o755); err != nil { + fail("mkdir profile dir: %v", err) + } + + hashBackend := resolveHashBackend(*hashName) + procs := runtime.GOMAXPROCS(0) + n := 1 << *log2Size + + fmt.Printf("loom bench instances=%d size=2^%d (=%d) hash=%s GOMAXPROCS=%d NumCPU=%d\n", + *nbInstances, *log2Size, n, *hashName, procs, runtime.NumCPU()) + fmt.Printf("profiles -> %s\n\n", *profileDir) + + var phases []phaseReport + + // ---- Phase 1: build modules + generate per-instance traces ----------- + tr := newTracker("traces+modules", *sampleMs) builder := board.NewBuilder() - traces := make([]trace.Trace, len(ns)) - for i, n := range ns { - tr, sigma, size, err := gnarkplonk.GetIthPlonkTrace(n, i) + traces := make([]trace.Trace, *nbInstances) + for i := 0; i < *nbInstances; i++ { + t, sigma, size, err := gnarkplonk.GetIthPlonkTrace(n, i) if err != nil { - fmt.Println(err) - os.Exit(-1) + fail("GetIthPlonkTrace[%d]: %v", i, err) } - traces[i] = tr + traces[i] = t builder.AddModule(gnarkplonk.PrepareIthPlonk(size, i)) - lro := []expr.Expr{expr.Col(gnarkplonk.Ith(gnarkplonk.ID_L, i)), expr.Col(gnarkplonk.Ith(gnarkplonk.ID_R, i)), expr.Col(gnarkplonk.Ith(gnarkplonk.ID_O, i))} + lro := []expr.Expr{ + expr.Col(gnarkplonk.Ith(gnarkplonk.ID_L, i)), + expr.Col(gnarkplonk.Ith(gnarkplonk.ID_R, i)), + expr.Col(gnarkplonk.Ith(gnarkplonk.ID_O, i)), + } sigmaGen := board.NewPermutationGen(sigma, gnarkplonk.Ith("plonk.S", i)) if err := arguments.CopyConstraint(&builder, gnarkplonk.Ith("plonk", i), lro, sigmaGen); err != nil { - fmt.Println(err) - os.Exit(-1) + fail("CopyConstraint[%d]: %v", i, err) } } + phases = append(phases, tr.stop()) - fullTrace := prover.MergeTrace(traces[0], traces[1:]...) + // ---- Phase 2: compile program ---------------------------------------- + tr = newTracker("compile", *sampleMs) program, err := board.Compile(&builder) if err != nil { - fmt.Println(err) - os.Exit(-1) - } - - fresh := trace.New(len(fullTrace.Base)) - for k, v := range fullTrace.Base { - fresh.SetBase(k, v) - } - for k, v := range fullTrace.Ext { - fresh.SetExt(k, v) - } - - // runtime.GC() - // fBefore, err := os.Create("heap_before.prof") - // if err != nil { - // fmt.Println(err) - // os.Exit(-1) - // } - // fAfter, err := os.Create("heap_after.prof") - // if err != nil { - // fmt.Println(err) - // os.Exit(-1) - // } - // defer fBefore.Close() - // defer fAfter.Close() - // if err := pprof.WriteHeapProfile(fBefore); err != nil { - // fmt.Println(err) - // os.Exit(-1) - // } - f, err := os.Create("cpu_poseidon2.pprof") + fail("Compile: %v", err) + } + phases = append(phases, tr.stop()) + + // ---- Phase 3: merge per-instance traces into one --------------------- + tr = newTracker("merge-trace", *sampleMs) + fullTrace := prover.MergeTrace(traces[0], traces[1:]...) + traces = nil + phases = append(phases, tr.stop()) + + // Drop trace-generation junk so the Prove heap snapshot is clean. + runtime.GC() + dumpHeap(filepath.Join(*profileDir, "heap_before_prove.pprof")) + + // ---- Phase 4: prove -------------------------------------------------- + cpuFile := mustCreate(filepath.Join(*profileDir, "cpu_prove.pprof")) + if err := pprof.StartCPUProfile(cpuFile); err != nil { + fail("StartCPUProfile: %v", err) + } + + opts := []prover.Option{prover.WithHashBackend(hashBackend)} + if *skipFRI { + opts = append(opts, prover.SkipFRI()) + } + + tr = newTracker("prove", *sampleMs) + prf, err := prover.Prove(fullTrace, setup.ProvingKey{}, nil, program, opts...) if err != nil { - fmt.Println(err) - os.Exit(-1) + fail("Prove: %v", err) + } + phases = append(phases, tr.stop()) + + pprof.StopCPUProfile() + cpuFile.Close() + dumpHeap(filepath.Join(*profileDir, "heap_after_prove.pprof")) + dumpProfile("allocs", filepath.Join(*profileDir, "allocs_after_prove.pprof")) + + // ---- Report ---------------------------------------------------------- + fmt.Println() + printSummary(phases, procs) + + // Touch the proof so the compiler doesn't get clever; also give a sense of + // output size. + fmt.Printf("\nproof: %d commitments, %d FRI levels, %d query samplings\n", + len(prf.Commitments), len(prf.DeepQuotientCommitment), len(prf.PointSamplings)) +} + +// ----------------------------------------------------------------------------- +// phase tracker +// ----------------------------------------------------------------------------- + +type phaseReport struct { + name string + + wall time.Duration + + // cpuBusy is on-CPU time actually spent doing work: user goroutine time + // + GC. Excludes idle, so cpuBusy/wall is a real parallelization metric + // (cf. /cpu/classes/total which is just GOMAXPROCS × wall and tells you + // nothing). + cpuBusy time.Duration + cpuUser time.Duration + cpuGC time.Duration + + allocBytes uint64 + allocObjects uint64 + + heapStart uint64 + heapEnd uint64 + heapPeak uint64 + + gcCount uint32 +} + +type tracker struct { + name string + + wallStart time.Time + + cpuUserStart float64 + cpuGCStart float64 + + allocBytesStart uint64 + allocObjectsStart uint64 + + memStart runtime.MemStats + + peakHeap atomic.Uint64 + stopCh chan struct{} + doneCh chan struct{} +} + +func newTracker(name string, sampleMs int) *tracker { + t := &tracker{name: name} + t.cpuUserStart, t.cpuGCStart = readCPU() + t.allocBytesStart, t.allocObjectsStart = readAllocCounters() + runtime.ReadMemStats(&t.memStart) + t.peakHeap.Store(t.memStart.HeapAlloc) + t.stopCh = make(chan struct{}) + t.doneCh = make(chan struct{}) + go t.sampleLoop(time.Duration(sampleMs) * time.Millisecond) + t.wallStart = time.Now() + return t +} + +func (t *tracker) sampleLoop(interval time.Duration) { + defer close(t.doneCh) + if interval <= 0 { + <-t.stopCh + return + } + tk := time.NewTicker(interval) + defer tk.Stop() + var ms runtime.MemStats + for { + select { + case <-t.stopCh: + return + case <-tk.C: + runtime.ReadMemStats(&ms) + if cur := ms.HeapAlloc; cur > t.peakHeap.Load() { + t.peakHeap.Store(cur) + } + } + } +} + +func (t *tracker) stop() phaseReport { + wall := time.Since(t.wallStart) + close(t.stopCh) + <-t.doneCh + + cpuUser, cpuGC := readCPU() + allocBytes, allocObjects := readAllocCounters() + var memEnd runtime.MemStats + runtime.ReadMemStats(&memEnd) + + if memEnd.HeapAlloc > t.peakHeap.Load() { + t.peakHeap.Store(memEnd.HeapAlloc) + } + + dUser := secondsToDuration(cpuUser - t.cpuUserStart) + dGC := secondsToDuration(cpuGC - t.cpuGCStart) + + return phaseReport{ + name: t.name, + wall: wall, + cpuBusy: dUser + dGC, + cpuUser: dUser, + cpuGC: dGC, + allocBytes: allocBytes - t.allocBytesStart, + allocObjects: allocObjects - t.allocObjectsStart, + heapStart: t.memStart.HeapAlloc, + heapEnd: memEnd.HeapAlloc, + heapPeak: t.peakHeap.Load(), + gcCount: memEnd.NumGC - t.memStart.NumGC, } +} + +// ----------------------------------------------------------------------------- +// runtime/metrics helpers +// ----------------------------------------------------------------------------- + +var cpuMetricNames = []string{ + "/cpu/classes/user:cpu-seconds", + "/cpu/classes/gc/total:cpu-seconds", +} + +var allocMetricNames = []string{ + "/gc/heap/allocs:bytes", + "/gc/heap/allocs:objects", +} + +// readCPU returns cumulative on-CPU seconds (user goroutine, GC) since +// process start. We deliberately avoid /cpu/classes/total because it counts +// idle slots (GOMAXPROCS × wall) and so tells you nothing about real CPU use. +func readCPU() (user, gc float64) { + samples := make([]metrics.Sample, len(cpuMetricNames)) + for i, n := range cpuMetricNames { + samples[i].Name = n + } + metrics.Read(samples) + return samples[0].Value.Float64(), samples[1].Value.Float64() +} + +func readAllocCounters() (bytes, objects uint64) { + samples := make([]metrics.Sample, len(allocMetricNames)) + for i, n := range allocMetricNames { + samples[i].Name = n + } + metrics.Read(samples) + return samples[0].Value.Uint64(), samples[1].Value.Uint64() +} + +func secondsToDuration(s float64) time.Duration { + return time.Duration(s * float64(time.Second)) +} + +// ----------------------------------------------------------------------------- +// reporting +// ----------------------------------------------------------------------------- + +func printSummary(phases []phaseReport, procs int) { + var totalWall, totalCPU, totalGC time.Duration + var totalAllocBytes, totalAllocObjects uint64 + var peakHeap uint64 + var gcCount uint32 + for _, p := range phases { + totalWall += p.wall + totalCPU += p.cpuBusy + totalGC += p.cpuGC + totalAllocBytes += p.allocBytes + totalAllocObjects += p.allocObjects + gcCount += p.gcCount + if p.heapPeak > peakHeap { + peakHeap = p.heapPeak + } + } + + w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0) + fmt.Fprintln(w, "phase\twall\tcpu\tpar\tgc%\talloc\tobjs\tpeakHeap\tGCs") + fmt.Fprintln(w, "-----\t----\t---\t---\t---\t-----\t----\t--------\t---") + for _, p := range phases { + fmt.Fprintf(w, "%s\t%s\t%s\t%5.2fx\t%4.1f%%\t%s\t%s\t%s\t%d\n", + p.name, + fmtDur(p.wall), + fmtDur(p.cpuBusy), + parallelization(p.cpuBusy, p.wall), + gcShare(p.cpuGC, p.cpuBusy), + fmtBytes(p.allocBytes), + fmtCount(p.allocObjects), + fmtBytes(p.heapPeak), + p.gcCount, + ) + } + fmt.Fprintln(w, "-----\t----\t---\t---\t---\t-----\t----\t--------\t---") + fmt.Fprintf(w, "TOTAL\t%s\t%s\t%5.2fx\t%4.1f%%\t%s\t%s\t%s\t%d\n", + fmtDur(totalWall), + fmtDur(totalCPU), + parallelization(totalCPU, totalWall), + gcShare(totalGC, totalCPU), + fmtBytes(totalAllocBytes), + fmtCount(totalAllocObjects), + fmtBytes(peakHeap), + gcCount, + ) + w.Flush() + + fmt.Println() + fmt.Printf("cpu = on-CPU time (user goroutines + GC); excludes idle\n") + fmt.Printf("par = cpu / wall (ideal: %dx = %d cores fully busy; 1x = single-threaded)\n", procs, procs) + fmt.Printf("gc%% = GC CPU time / on-CPU time\n") + fmt.Println("peakHeap = max HeapAlloc observed during phase (sampled in background)") +} + +// ----------------------------------------------------------------------------- +// formatting +// ----------------------------------------------------------------------------- + +func fmtDur(d time.Duration) string { + switch { + case d >= time.Minute: + return fmt.Sprintf("%dm%02ds", int(d.Minutes()), int(d.Seconds())%60) + case d >= time.Second: + return fmt.Sprintf("%.2fs", d.Seconds()) + case d >= time.Millisecond: + return fmt.Sprintf("%.1fms", float64(d.Microseconds())/1000) + default: + return d.String() + } +} + +func fmtBytes(b uint64) string { + const ( + KiB = 1 << 10 + MiB = 1 << 20 + GiB = 1 << 30 + ) + switch { + case b >= GiB: + return fmt.Sprintf("%.2f GiB", float64(b)/GiB) + case b >= MiB: + return fmt.Sprintf("%.1f MiB", float64(b)/MiB) + case b >= KiB: + return fmt.Sprintf("%.1f KiB", float64(b)/KiB) + default: + return fmt.Sprintf("%d B", b) + } +} + +func fmtCount(n uint64) string { + switch { + case n >= 1_000_000_000: + return fmt.Sprintf("%.2fB", float64(n)/1e9) + case n >= 1_000_000: + return fmt.Sprintf("%.2fM", float64(n)/1e6) + case n >= 1_000: + return fmt.Sprintf("%.1fk", float64(n)/1e3) + default: + return fmt.Sprintf("%d", n) + } +} + +func parallelization(cpu, wall time.Duration) float64 { + if wall <= 0 { + return 0 + } + return float64(cpu) / float64(wall) +} + +func gcShare(gc, cpu time.Duration) float64 { + if cpu <= 0 { + return 0 + } + return 100 * float64(gc) / float64(cpu) +} + +// ----------------------------------------------------------------------------- +// misc helpers +// ----------------------------------------------------------------------------- + +func resolveHashBackend(name string) loom.HashBackend { + switch name { + case "poseidon2": + return loom.Poseidon2HashBackend() + case "sha256": + return loom.SHA256HashBackend() + default: + fail("unknown hash backend %q (want poseidon2 | sha256)", name) + return loom.HashBackend{} + } +} + +func dumpHeap(path string) { + f := mustCreate(path) defer f.Close() - if err := pprof.StartCPUProfile(f); err != nil { - fmt.Println(err) - os.Exit(-1) + if err := pprof.WriteHeapProfile(f); err != nil { + fail("WriteHeapProfile(%s): %v", path, err) + } +} + +func dumpProfile(name, path string) { + p := pprof.Lookup(name) + if p == nil { + fail("pprof.Lookup(%q) returned nil", name) } - _, err = prover.Prove(fresh, setup.ProvingKey{}, nil, program, prover.WithHashBackend(loom.Poseidon2HashBackend())) + f := mustCreate(path) + defer f.Close() + if err := p.WriteTo(f, 0); err != nil { + fail("write %s profile: %v", name, err) + } +} + +func mustCreate(path string) *os.File { + f, err := os.Create(path) if err != nil { - fmt.Println(err) - os.Exit(-1) + fail("create %s: %v", path, err) } - pprof.StopCPUProfile() - // if err := pprof.WriteHeapProfile(fAfter); err != nil { - // fmt.Println(err) - // os.Exit(-1) - // } - - // f, err := os.Create("cpu.pprof") - // if err != nil { - // fmt.Println(err) - // os.Exit(-1) - // } - // defer f.Close() - // if err := pprof.StartCPUProfile(f); err != nil { - // fmt.Println(err) - // os.Exit(-1) - // } - // for i := 0; i < 100; i++ { - // err := verifier.Verify(nil, setup.VerificationKey{}, program, proof) - // if err != nil { - // fmt.Println(err) - // os.Exit(-1) - // } - // } - // pprof.StopCPUProfile() + return f +} +func fail(format string, args ...any) { + fmt.Fprintf(os.Stderr, "bench: "+format+"\n", args...) + os.Exit(1) } diff --git a/internal/dag/dag.go b/internal/dag/dag.go index 6fa0e9b..3c7e3fe 100644 --- a/internal/dag/dag.go +++ b/internal/dag/dag.go @@ -66,6 +66,42 @@ type DAG struct { VarIndex map[string]int // leaf name → index in vars slice for EvalWithCacheVars } +// Clone returns a deep copy of the DAG: fresh DAGNodes and fresh *expr.Leaf for +// every leaf, with Children rewired to the new nodes. Use this to make per-call +// mutations (e.g. ComputeQuotient*'s assignment of Leaf.Idx) safe across +// goroutines that would otherwise share leaf pointers via shared sub-expressions +// upstream of ExprToDAG. +// +// VarIndex is shared (read-only after build). +func (d *DAG) Clone() *DAG { + if d == nil { + return nil + } + cloneOf := make(map[*DAGNode]*DAGNode, len(d.Nodes)) + nodes := make([]*DAGNode, len(d.Nodes)) + for i, n := range d.Nodes { + nn := *n + if n.Leaf != nil { + leafCp := *n.Leaf + nn.Leaf = &leafCp + } + if len(n.Children) > 0 { + nn.Children = make([]*DAGNode, len(n.Children)) + for j, c := range n.Children { + // children come before parents in topological order, so cloneOf[c] is populated + nn.Children[j] = cloneOf[c] + } + } + nodes[i] = &nn + cloneOf[n] = &nn + } + return &DAG{ + Root: cloneOf[d.Root], + Nodes: nodes, + VarIndex: d.VarIndex, + } +} + // ExprToDAG converts an Expr tree into a DAG by merging identical // sub-expressions into shared nodes. Commutativity is respected: (a+b) and // (b+a) produce the same node, as do (a*b) and (b*a). Sub is not commutative. diff --git a/internal/poly/compute_quotient.go b/internal/poly/compute_quotient.go index 3986ee2..3801008 100644 --- a/internal/poly/compute_quotient.go +++ b/internal/poly/compute_quotient.go @@ -382,16 +382,24 @@ func CosetLagrangeToLagrangeNormal(p Polynomial) { // CosetLagrangeToLagrangeNormalWithCache converts p in place using cache for // the FFT domain. func CosetLagrangeToLagrangeNormalWithCache(p Polynomial, cache *DomainCache) { - N := uint64(len(p)) - d := cache.Get(N) + CosetLagrangeNormalToCanonicalWithCache(p, cache) + cache.Get(uint64(len(p))).FFT(p, fft.DIF) + utils.BitReverse(p) +} - // Step 1: coset-Lagrange Normal → BitReversed IFFT (= c_k * FrGen^k, BitReversed) - d.FFTInverse(p, fft.DIF) +// CosetLagrangeNormalToCanonicalWithCache converts p in place from +// coset-Lagrange Normal form (evaluations on the coset {FrGen·ω^j}) to +// canonical Normal form (coefficients c_k in standard order). This is the +// first half of CosetLagrangeToLagrangeNormalWithCache; callers that +// subsequently need a different transform (e.g. AIR quotient chunking, which +// FFTs each N-sized chunk individually) avoid a redundant round-trip FFT. +func CosetLagrangeNormalToCanonicalWithCache(p Polynomial, cache *DomainCache) { + d := cache.Get(uint64(len(p))) - // Step 2: BitReverse → Normal order (= c_k * FrGen^k) - utils.BitReverse(p) + d.FFTInverse(p, fft.DIF) // coset-Lagrange Normal → coset-canonical BitReversed + utils.BitReverse(p) // → coset-canonical Normal - // Step 3: divide by FrGen^k to get canonical coefficients c_k + // Multiply by invFrGen^k to get standard canonical coefficients c_k. invFrGen := d.FrMultiplicativeGen invFrGen.Inverse(&invFrGen) acc := koalabear.One() @@ -399,12 +407,6 @@ func CosetLagrangeToLagrangeNormalWithCache(p Polynomial, cache *DomainCache) { p[k].Mul(&p[k], &acc) acc.Mul(&acc, &invFrGen) } - - // Step 4: canonical Normal → Lagrange BitReversed - d.FFT(p, fft.DIF) - - // Step 5: BitReverse → Lagrange Normal - utils.BitReverse(p) } // CosetExtLagrangeToLagrangeNormal converts an extension polynomial from @@ -416,8 +418,15 @@ func CosetExtLagrangeToLagrangeNormal(p ExtPolynomial) { // CosetExtLagrangeToLagrangeNormalWithCache converts p in place using cache for // the FFT domain. func CosetExtLagrangeToLagrangeNormalWithCache(p ExtPolynomial, cache *DomainCache) { - N := uint64(len(p)) - d := cache.Get(N) + CosetExtLagrangeNormalToCanonicalWithCache(p, cache) + cache.Get(uint64(len(p))).FFTExt(p, fft.DIF) + utils.BitReverse(p) +} + +// CosetExtLagrangeNormalToCanonicalWithCache is the ext-field counterpart of +// CosetLagrangeNormalToCanonicalWithCache. +func CosetExtLagrangeNormalToCanonicalWithCache(p ExtPolynomial, cache *DomainCache) { + d := cache.Get(uint64(len(p))) d.FFTInverseExt(p, fft.DIF) utils.BitReverse(p) @@ -429,7 +438,4 @@ func CosetExtLagrangeToLagrangeNormalWithCache(p ExtPolynomial, cache *DomainCac p[k].MulByElement(&p[k], &acc) acc.Mul(&acc, &invFrGen) } - - d.FFTExt(p, fft.DIF) - utils.BitReverse(p) } diff --git a/internal/poly/domain_cache.go b/internal/poly/domain_cache.go index 36f55c3..46534ff 100644 --- a/internal/poly/domain_cache.go +++ b/internal/poly/domain_cache.go @@ -13,11 +13,15 @@ package poly -import "github.com/consensys/gnark-crypto/field/koalabear/fft" +import ( + "sync" -// DomainCache memoizes FFT domains by cardinality. It is intended as -// single-run scratch state and is not safe for concurrent use. + "github.com/consensys/gnark-crypto/field/koalabear/fft" +) + +// DomainCache memoizes FFT domains by cardinality. Safe for concurrent use. type DomainCache struct { + mu sync.Mutex domains map[uint64]*fft.Domain } @@ -27,6 +31,8 @@ func (c *DomainCache) Get(n uint64) *fft.Domain { if c == nil { return fft.NewDomain(n) } + c.mu.Lock() + defer c.mu.Unlock() if c.domains == nil { c.domains = make(map[uint64]*fft.Domain) } diff --git a/optim_plan_ideas.md b/optim_plan_ideas.md new file mode 100644 index 0000000..3ef86e0 --- /dev/null +++ b/optim_plan_ideas.md @@ -0,0 +1,388 @@ +# Loom prover — optimisation plan & ideas + +Status: draft, 2026-05-26. All numbers in this document come from `bench/main.go` +(in this repo) running on a 32-core / 32-GOMAXPROCS Linux box with AVX-512, +gnark-crypto v0.20.1, Go 1.26. Workload: 40 aggregated PLONK instances of size +2^17, Poseidon2 leaf hashing. + +The goal of this file is to give a coding agent enough context to pick up any +one of these items, work it in isolation, and verify the impact. Each item +includes profile evidence, file:line pointers, the proposed change, and a +verification recipe. + +--- + +## 0. Baseline & reproduction + +``` +$ go build -o /tmp/loom-bench ./bench +$ /tmp/loom-bench # default workload, ~45s wall +$ /tmp/loom-bench -instances N -log2-size L # tune up/down +$ /tmp/loom-bench -skip-fri # to isolate non-FRI work +$ /tmp/loom-bench -hash sha256 # to swap Merkle hash backend +``` + +Default run produces: + +``` +phase wall cpu par gc% alloc peakHeap +traces+modules 2.02s 3.08s 1.52x 7.3% 5.40 GiB 942 MiB +compile 84.2ms 0s 0.00x 0.0% 244 MiB 920 MiB +merge-trace 77 µs 0s 0.00x 0.0% 6.3 KiB 920 MiB +prove 41.98s 45.99s 1.10x 0.1% 12.89 GiB 4.86 GiB +TOTAL 44.08s 49.07s 1.11x 0.6% 18.53 GiB 4.86 GiB +``` + +Read `par = cpu / wall` as "average cores busy". The headline finding is that +Prove uses **~1 core out of 32** despite parallel primitives existing. Almost +every optimization below either (a) raises that number or (b) reduces work that +would otherwise be parallelised. + +The pprof CPU profile referenced throughout lives at +`bench_profiles/cpu_prove.pprof` after a run. Re-render with: + +``` +go tool pprof -top -cum -nodecount 40 bench_profiles/cpu_prove.pprof +go tool pprof -list bench_profiles/cpu_prove.pprof +go tool pprof -focus -tree bench_profiles/cpu_prove.pprof +``` + +Heap before/after Prove and an allocations profile are written alongside. + +### Phase breakdown (cumulative CPU in the profile) + +| Phase | % of Prove CPU | Notes | +|-----------------------------|----------------|----------------------------------------| +| `ComputeAIRQuotients` | ~35% | per-module sequential outer loop | +| `ExecuteSteps` | ~12% | sequential per module, leaf hashing | +| `commitTraceRound` (in #2) | ~7% | Poseidon2 leaf hashing of trace polys | +| `ComputeDeepQuotient` | ~6% | sequential per size, per-elem inverse | +| `SampleEvaluations` | ~6% | sequential per query position | +| Everything else | ~34% | runtime / GC / minor paths | + +### Hottest leaf functions (flat CPU) + +| Function (flat %) | Where | +|--------------------------------------------------------|------------------------------------| +| `permutation16x24_avx512` 10.0% | Poseidon2 batch-16, already SIMD | +| `extensions.montReduce` 9.95% | Every E4.Mul; called from #1, #2 | +| `extensions.E4.Mul` 8.84% | Driven by `mulChildVectorIntoExt` | +| `bitReverseNaive[E4]` 5.28% (+0.58% mem swap) | See #3 | +| `vectorButterfly_avx512` 5.04% | gnark-crypto FFT, already SIMD | +| `koalabear.Element.Add` 4.34% | Base-field accumulation | + +--- + +## Top-3 optimisations (high confidence, real numbers) + +### #1 — Parallelise the per-module loop in `ComputeAIRQuotients` + +**Expected impact:** Up to ~4–5× Prove speedup (this phase is ~35% of CPU and +currently runs at ~1.1× parallelism with 40 independent modules on 32 cores). + +**Evidence:** +- `prover/prover.go:420` iterates `for moduleName, module := range pr.program.Modules` + sequentially. With 40 instances of the bench, there are 40+ modules whose + AIR quotients are independent. +- pprof tree under `ComputeAIRQuotients` (`go tool pprof -focus=ComputeAIRQuotients -tree`): + - `poly.ComputeQuotientMixed` — 24.3s cumulative, **26.6% of total Prove CPU**. + - `dag.EvalOnAllEntriesMixedInto` (`internal/dag/dag.go:1097`) — 24.5s cum. + - `dag.mulChildVectorIntoExt` (`internal/dag/dag.go:1401`) — 14.0s cum, **15% of total Prove CPU**, scalar Go loop over E4. +- `Total samples = 91.52s` over `Duration: 61.02s` on the 60×2^17 run = ~1.5x + effective parallelism for the whole Prove; ComputeAIRQuotients dominates. + +**Proposed change:** + +1. Materialise a deterministic module ordering once: + ```go + names := make([]string, 0, len(pr.program.Modules)) + for name := range pr.program.Modules { names = append(names, name) } + sort.Strings(names) + ``` +2. Wrap the work between `prover.go:420` and `prover.go:472` (i.e. the body + that produces per-module `airTrace` chunks + populates `chunkDomains`) in + `internal/parallel.Execute(len(names), func(start, end int) { ... })`. +3. Each goroutine writes into its own pre-allocated `[]ext.E4` / `[]koalabear.Element` + workspace and into a private `localChunkDomains` and `localAirChunks` map; + merge under a single mutex (or pre-size the global maps and write disjoint + keys) after the parallel section. +4. `pr.airTrace.SetBase` / `SetExt` are not safe for concurrent map writes — + either lock around them, or buffer per goroutine and merge serially in + <1% time. +5. The downstream "commit by size group" block (`prover.go:474–531`) stays + serial — it groups across modules and already uses a parallel committer + internally. + +**Watch out for:** +- `pr.domainCache` is shared (`poly.WithDomainCache(&pr.domainCache)`). Check + `internal/poly/domain_cache.go` for thread-safety; if it isn't safe, either + give each goroutine its own cache or add an `sync.RWMutex` around it. This + is the single biggest correctness risk for this change. +- The per-goroutine workspace pool inside `dag.EvalWorkspace` + (`internal/dag/dag.go:1118-1140`) is currently per-call. If you reuse a + workspace across modules per goroutine you should also confirm it's reset + between calls. + +**Verification:** +- `go test ./...` must stay green (especially `prover/...` and integration). +- `/tmp/loom-bench` should show `prove`'s `par` jump from ~1.1x toward + GOMAXPROCS. Even a modest 8x would translate to ~3× Prove speedup. +- `go tool pprof -focus=ComputeAIRQuotients -top` should show `sync.(*WaitGroup).Go.func1` + cum time rise significantly. + +--- + +### #2 — Batch-invert in `DeepQuotientExt` (and parallelise the per-size loop) + +**Expected impact:** 4–8% Prove speedup from the batch invert alone; another +2–4% from parallelising sizes (limited by the small number of distinct sizes, +typically 3–4). + +**Evidence:** +- `internal/poly/ext.go:130` — `DeepQuotientExt` does, in a tight loop of N + iterations: + ```go + for j := 0; j < N; j++ { + var num ext.E4 + num.Sub(&v, &p[j]) + var den ext.E4 + den.Sub(&z, &ω[j]) + var inv ext.E4 + inv.Inverse(&den) // ← one full E4 inversion per element + out[j].Mul(&num, &inv) + } + ``` +- An E4 inversion (Frobenius + base-field inverse + multiplies) is ~10–30× + the cost of an E4 multiply. The denominators form a vector ideal for + Montgomery batch inversion: 1 inversion + 3(N-1) multiplies for N elements. +- `ext.BatchInvertE4` already exists and is used at + `internal/poly/iop_utils.go:375` and `:555` — same pattern. +- `prover.go:612` — `for i, N := range sizes { ... }` is the outer loop over + distinct module sizes. Each iteration owns its own `deepQuotient` map entry, + its own RS encoder, its own FRI level tree → trivially independent. +- `prover.go:629` — `C_s := make(poly.ExtPolynomial, N)` is allocated **per + shift, per size**. For N=2^17, ~20 shifts, ~3 sizes that's ~96 MiB of E4 + churn per Prove (and a contributor to the 12.9 GiB of `prove`-phase + allocations). + +**Proposed change:** + +1. In `DeepQuotientExt`: + - Build `dens[j] = z - ω[j]` into a scratch `[]ext.E4`. + - Call `ext.BatchInvertE4(dens)` once. + - Final loop becomes `out[j].Mul(&num, &dens[j])` — single E4 multiply per + iteration instead of an inversion. + - Stream-friendly: reuse a scratch buffer across calls via an optional + argument or a `sync.Pool`. + +2. In `ComputeDeepQuotient` (`prover.go:596`): + - Convert the `for i, N := range sizes` loop body into a closure dispatched + via `parallel.Execute(len(sizes), ...)`. + - `deepQuotients` (a `map[int]poly.ExtPolynomial`) is written from disjoint + keys per iteration → preallocate to `len(sizes)` and write under + `pr.mu`, **or** swap to `[]poly.ExtPolynomial` indexed by `i` and assemble + after the parallel section. + - The inner per-shift work (`prover.go:620–668`) also accumulates into a + per-size `deepQuotient` — that stays serial within a size, fine. + - `pr.Proof.DeepQuotientCommitment` and the FRI proving block (`prover.go:709–737`) + stay serial since `fri.Prove` is invoked once after all level trees are + built. + +3. Bonus, low-effort: reuse `C_s` across shifts. Currently allocated fresh at + `prover.go:629`; the only operation that writes it is `C_s[x].Add(&C_s[x], &term)` + inside the `k`-loop, so a zero-on-entry policy + a single allocation per + size eliminates the per-shift alloc. + +**Verification:** +- Existing prover tests cover this end-to-end; if they pass, the batch invert + is correct. +- Microbench (write a `BenchmarkDeepQuotientExt` in `internal/poly/ext_test.go` + parameterised on `N`) should show ~5–10× speedup on the function alone. +- `bench_profiles/cpu_prove.pprof` should show `ext.E4.Inverse` cum time drop + by ~95% and `DeepQuotientExt` flat time shrink to a fraction. + +--- + +### #3 — Stop running `bitReverseNaive` on E4 / E2 polynomials + +**Expected impact:** ~5.9% of Prove CPU, ~50 LOC of bit-reverse code, or +**zero work** if you can elide the bit-reverses entirely (see option B). + +**Evidence:** +- `utils.bitReverseNaive[E4-shaped]` shows up at **4.83s flat / 5.36s cum = + 5.86% of total Prove CPU**. The function (`gnark-crypto/utils/bitreverse.go:29`) + is literally a swap-pairs loop — no algorithm. +- `gnark-crypto/utils/bitreverse.go:20` dispatches: + ```go + if runtime.GOARCH == "arm64" || len(v) < (1<<21) || unsafe.Sizeof(v[0]) < 8 { + bitReverseNaive(v) + } else { + bitReverseCobra(v) + } + ``` + Our quotient FFTs land on `bigSize = eDeg·N`, typically 2^19–2^20 (N=2^17, + eDeg up to 4–8) — just under the 2^21 cutoff. So even though gnark-crypto + ships a cache-friendly Cobra variant and specialised `bitReverseCobraInPlace_9_*` + unrolls, we never reach them. +- Call sites: `internal/poly/compute_quotient.go:128, 131, 301, 304, 392, + 407, 423, 434` and `prover/prover.go:436, 443, 460, 467`. + +**Option A — Local fast bit-reverse for E4 (and base):** + +Add `internal/poly/bitreverse.go` with a Cobra-style or block-transpose +bit-reverse specialised for `ext.E4` (and likely `koalabear.Element`). +Route all loom call sites through it. + +The Cobra threshold is set conservatively in gnark-crypto for arrays with +non-trivial element size — for E4 (16 bytes) and N=2^19 the cache benefit is +real, the threshold could safely come down. Crib the algorithm from +`gnark-crypto/utils/bitreverse.go:128-160` and the `_9_21` family. + +**Option B — Eliminate the bit-reverses:** + +Several call sites are of the form "FFT (DIF) → BitReverse → … → BitReverse → +FFT (DIT)". When two BitReverses bracket a region with only pointwise +operations, both can be deleted. Inspect: +- `prover.go:435-437` then `:442-443` — `FFTInverseExt(DIF)` → BitReverse → + chunk → `FFTExt(DIF)` → BitReverse. Either keep everything in bit-reversed + order or use DIT/DIF pairing carefully. +- `compute_quotient.go:391-407` — same pattern around the coset FFT. + +If you can keep polynomials in bit-reversed order between `ComputeQuotient*` +and the per-chunk FFT, you may delete 4 BitReverse calls outright. This is +the loomwide win. + +**Verification:** +- Cross-check by running prover_test.go and integration_test — these polys + feed into Merkle leaves that the verifier re-checks, so any bit-order + mistake fails the proof. +- Profile: `bitReverseNaive[E4]` should disappear from `pprof -top`. + +**Risks:** +- BitReverse interacts with `fft.DIF` vs `fft.DIT` directionality; getting + the algebra right requires reading `gnark-crypto`'s FFT API. Don't ship + option B without exhaustive tests. + +--- + +## Secondary opportunities (medium confidence) + +These are visible in the profile but harder to size without prototyping. + +### #4 — Vectorise `dag.mulChildVectorIntoExt` and friends + +**Evidence:** `internal/dag/dag.go:1401` is a scalar Go loop: +```go +for j := range N { + dst[j].Mul(&dst[j], &src[j]) +} +``` +This is **15% of Prove CPU** (cumulative including the called E4.Mul). The same +file has analogous `addChildVectorToExt`, `subChildVectorFromExt`, +`copyChildVectorToExt`. + +gnark-crypto exposes `extensions.Vector.Mul`, `Vector.MulByElement` (already +shows up at 5.7% via FFT), `Vector.Add`, `Vector.Butterfly` with AVX-512 +implementations. Most of those entry points need a third argument shape +(`Vector.Mul(dst, a, b)`); the loom DAG site is `Mul(dst, dst, src)`. Either +adapt the call site to use a scratch buffer, or add an in-place +`Vector.MulSelf(dst, src)` to gnark-crypto. + +If a vector entry point cuts this 15% in half, that's ~7% Prove on its own, +**stackable with #1** (parallelising modules just means each module's inner +loop runs on one core; a SIMD inner loop multiplies the per-core throughput). + +### #5 — Reuse `EvalWorkspace` across modules (after #1) + +`internal/dag/dag.go:1118-1140` has a per-call `basePool` / `extPool`. When #1 +lands and modules run in parallel, give each worker goroutine a persistent +`EvalWorkspace` instead of creating one per `EvalOnAllEntriesMixedInto`. This +keeps the slices in the same goroutine's local cache and dodges the allocator +on every coset. + +The current `make([]ext.E4, N)` at `dag.go:1138` is allocated `n_nodes × n_cosets` +times per module. With degree ~5 and ~10 cosets per module, that's hundreds +of N-sized E4 slabs in the churn budget per Prove. + +### #6 — Parallelise `SampleEvaluations` over query positions + +`prover/prover.go:771-786`: NUM_QUERIES (~128) FRI queries, each opening every +committed tree at one position. Currently `for q, s := range pr.queryPositions` +is sequential. Each query owns its own `pr.Proof.PointSamplings[q]` slot — no +data races. `parallel.Execute(NQ, ...)`. + +This phase is ~6% of CPU; parallelising should push it to single-digit ms. + +### #7 — Parallelise `ExecuteSteps`' `GenCol` loop + +`prover.go:365`: `for _, m := range pr.program.Modules { ... gen.Gen(pr.t, &mCopy) }`. +Generators may write to disjoint trace columns, but they share `pr.t` (a +`trace.Trace`); concurrent map writes will race. Requires a per-module +`trace.Trace` produced under the goroutine, then merged. Speculative — confirm +no cross-module data dependencies in `GenCol` before attempting. + +### #8 — Poseidon2 `WriteExtBatch` element-by-element absorb + +`internal/hash/poseidon2_batch.go`'s `WriteExtBatch` shows up at **7.0s cum = +7.7% of Prove**. It calls `WriteElementBatch` four times per E4 (one per +coefficient). Each `WriteElementBatch` does an `absorbFullBlock` if the sponge +is at capacity. There's probably a vector store / unrolled writer hiding here +— inspect whether you can stage 4 elements at once instead of looping. + +Lower confidence: this kernel may already be well-tuned and the time is +fundamental work. But "WriteExt does 4× WriteElement" smells loose. + +### #9 — Memory churn: reduce `make([]ext.E4, N)` in inner loops + +`prover.go:629`, `prover.go:613`, `prover.go:672` each allocate a per-iteration +N-sized E4 slice. With N=2^17 (= 2 MiB per E4 slice), Prove allocates ~12.9 +GiB in total. GC is at 0.1% so the *runtime* cost is negligible — but the +peak heap (4.86 GiB) is driven by these large transient buffers, and Linux +page-fault cost on first-touch is non-zero. Pool / reuse them per Prove run. + +This is mostly a memory-footprint win (helpful for running larger workloads +without OOM), not a CPU win. + +--- + +## Things NOT worth optimising (per profile) + +- **GC tuning.** 0.1–0.6% of CPU. Already invisible. +- **Poseidon2 permutation kernel.** Already AVX-512 batch-16, near the SIMD + ceiling. +- **FFT kernels.** `kerDIFNP_512Ext`, `vectorButterfly_avx512` are gnark-crypto + AVX-512 code — leave alone. +- **`ComputeEvaluationsAtZeta`** (`prover.go:556`). Doesn't even show in top + 40 of pprof. Skip. +- **`merge-trace`**, **`compile`**. 0% and ~0.2% of total wall. + +--- + +## Suggested ordering for a coding-agent campaign + +1. **Land #3 Option A** first — small, isolated, low risk, ~6% Prove. Good warmup. +2. **Land #2 batch-invert** — small isolated change in `internal/poly/ext.go`, + ~5% Prove. Verify with a microbench. +3. **Land #1** (parallelise `ComputeAIRQuotients`) — biggest payoff but most + careful work (domain-cache thread-safety + per-goroutine workspace). After + this, re-run the bench and re-profile — the percentage breakdowns will + shift substantially. +4. **Land #6** while you're in the prover (`SampleEvaluations`). +5. Re-profile. If E4 inner loops are still the leaf cost, do **#4** + (vector E4 ops in `dag`). Otherwise pick the next-biggest item from the + refreshed profile. + +After 1–4, target should be Prove `par` ≥ 8x on 32 cores and wall time ≤ +15s on the default workload. + +--- + +## Open questions for the team + +- Is `poly.DomainCache` safe for concurrent reads from multiple goroutines? + This blocks #1. +- Are `board.Module.GenCol` generators required to be sequential by design, + or just by current implementation? This blocks #7. +- Why is gnark-crypto's `bitReverseCobra` threshold set at 2^21? Could it be + lowered upstream for E4-shaped slices (16 bytes/elem) where the cache + benefit kicks in earlier? diff --git a/prover/prover.go b/prover/prover.go index 6bf088b..c7a297e 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -25,11 +25,13 @@ import ( "github.com/consensys/loom/board" "github.com/consensys/loom/expr" "github.com/consensys/loom/field" + "github.com/consensys/loom/internal/dag" "github.com/consensys/loom/internal/commitment" "github.com/consensys/loom/internal/constants" fiatshamir "github.com/consensys/loom/internal/fiat-shamir" "github.com/consensys/loom/internal/fri" "github.com/consensys/loom/internal/hash" + "github.com/consensys/loom/internal/parallel" "github.com/consensys/loom/internal/poly" "github.com/consensys/loom/internal/reedsolomon" "github.com/consensys/loom/proof" @@ -414,69 +416,116 @@ func (pr *proverRuntime) ExecuteSteps() error { return nil } +// computeExtAIRChunks computes the AIR quotient for a single ext-rooted module +// and returns its N-sized chunks in Lagrange form, ready to commit. Chunks +// alias non-overlapping windows of the quotient backing array (no per-chunk +// copy); after this returns the caller owns the chunks and the original +// quotient is unreachable. +func computeExtAIRChunks(piBase map[string]poly.Polynomial, piExt map[string]poly.ExtPolynomial, vrel *dag.DAG, N int, D *fft.Domain, cache *poly.DomainCache) ([]poly.ExtPolynomial, error) { + quotient, err := poly.ComputeQuotientMixed(piBase, piExt, *vrel, N, poly.WithDomainCache(cache)) + if err != nil { + return nil, err + } + // quotient is in coset-Lagrange Normal; convert directly to canonical + // Normal so chunking can sub-slice the backing array. + poly.CosetExtLagrangeNormalToCanonicalWithCache(quotient, cache) + + chunks := make([]poly.ExtPolynomial, len(quotient)/N) + for i := range chunks { + chunk := quotient[i*N : (i+1)*N : (i+1)*N] + D.FFTExt(chunk, fft.DIF) + utils.BitReverse(chunk) + chunks[i] = chunk + } + return chunks, nil +} + +// computeBaseAIRChunks is the base-field counterpart of computeExtAIRChunks. +func computeBaseAIRChunks(piBase map[string]poly.Polynomial, vrel *dag.DAG, N int, D *fft.Domain, cache *poly.DomainCache) ([]poly.Polynomial, error) { + quotient, err := poly.ComputeQuotient(piBase, *vrel, N, poly.WithDomainCache(cache)) + if err != nil { + return nil, err + } + poly.CosetLagrangeNormalToCanonicalWithCache(quotient, cache) + + chunks := make([]poly.Polynomial, len(quotient)/N) + for i := range chunks { + chunk := quotient[i*N : (i+1)*N : (i+1)*N] + D.FFT(chunk, fft.DIF) + utils.BitReverse(chunk) + chunks[i] = chunk + } + return chunks, nil +} + func (pr *proverRuntime) ComputeAIRQuotients() error { chunkDomains := make(map[string]*fft.Domain) - for moduleName, module := range pr.program.Modules { - if module.VanishingRelation.Degree() <= 0 { - continue - } + moduleNames := make([]string, 0, len(pr.program.Modules)) + for name := range pr.program.Modules { + moduleNames = append(moduleNames, name) + } + sort.Strings(moduleNames) - N := module.N - if module.VanishingRelation.Root.Field == field.Ext { - quotient, err := poly.ComputeQuotientMixed(pr.t.Base, pr.t.Ext, *module.VanishingRelation, N, poly.WithDomainCache(&pr.domainCache)) - if err != nil { - return err + // Per-module quotient computation is independent: each writes disjoint chunk + // names into airTrace, reads only from the shared pr.t. pr.domainCache is + // internally locked; trace map writes need an explicit lock. The DAG is + // cloned per goroutine because ComputeQuotient* mutates Leaf.Idx and modules + // can share *expr.Leaf pointers via cross-module arguments (Lookup, etc.). + var ( + writeMu sync.Mutex + errMu sync.Mutex + firstErr error + ) + parallel.Execute(len(moduleNames), func(start, end int) { + for idx := start; idx < end; idx++ { + moduleName := moduleNames[idx] + module := pr.program.Modules[moduleName] + if module.VanishingRelation.Degree() <= 0 { + continue } - poly.CosetExtLagrangeToLagrangeNormalWithCache(quotient, &pr.domainCache) - bigSize := len(quotient) - bigD := pr.domainCache.Get(uint64(bigSize)) - bigD.FFTInverseExt(quotient, fft.DIF) - utils.BitReverse(quotient) - - numChunks := bigSize / N - for i := 0; i < numChunks; i++ { - chunk := make(poly.ExtPolynomial, N) - copy(chunk, quotient[i*N:(i+1)*N]) - module.D.FFTExt(chunk, fft.DIF) - utils.BitReverse(chunk) - chunkName := constants.QuotientChunkName(moduleName, i) - pr.airTrace.SetExt(chunkName, chunk) - chunkDomains[chunkName] = module.D + N := module.N + vrel := module.VanishingRelation.Clone() + + var ( + baseChunks []poly.Polynomial + extChunks []poly.ExtPolynomial + err error + ) + if module.VanishingRelation.Root.Field == field.Ext { + extChunks, err = computeExtAIRChunks(pr.t.Base, pr.t.Ext, vrel, N, module.D, &pr.domainCache) + } else { + baseChunks, err = computeBaseAIRChunks(pr.t.Base, vrel, N, module.D, &pr.domainCache) + } + if err != nil { + errMu.Lock() + if firstErr == nil { + firstErr = err + } + errMu.Unlock() + return } - continue - } - - quotient, err := poly.ComputeQuotient(pr.t.Base, *module.VanishingRelation, N, poly.WithDomainCache(&pr.domainCache)) - if err != nil { - return err - } - - poly.CosetLagrangeToLagrangeNormalWithCache(quotient, &pr.domainCache) - bigSize := len(quotient) - bigD := pr.domainCache.Get(uint64(bigSize)) - bigD.FFTInverse(quotient, fft.DIF) - utils.BitReverse(quotient) - numChunks := bigSize / N - for i := 0; i < numChunks; i++ { - chunk := make(poly.Polynomial, N) - copy(chunk, quotient[i*N:(i+1)*N]) - module.D.FFT(chunk, fft.DIF) - utils.BitReverse(chunk) - chunkName := constants.QuotientChunkName(moduleName, i) - pr.airTrace.SetBase(chunkName, chunk) - chunkDomains[chunkName] = module.D + writeMu.Lock() + for i, c := range baseChunks { + name := constants.QuotientChunkName(moduleName, i) + pr.airTrace.SetBase(name, c) + chunkDomains[name] = module.D + } + for i, c := range extChunks { + name := constants.QuotientChunkName(moduleName, i) + pr.airTrace.SetExt(name, c) + chunkDomains[name] = module.D + } + writeMu.Unlock() } + }) + if firstErr != nil { + return firstErr } chunksByN := map[int]*mixedCommitGroup{} - moduleNames := make([]string, 0, len(pr.program.Modules)) - for name := range pr.program.Modules { - moduleNames = append(moduleNames, name) - } - sort.Strings(moduleNames) for _, moduleName := range moduleNames { module := pr.program.Modules[moduleName] N := module.N From 17bf51e6ad9adc03aa62044d9eb70be6a88deda3 Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Tue, 26 May 2026 18:30:16 +0000 Subject: [PATCH 3/8] perf(dag,poly): use ext/koalabear Vector ops in hot inner loops dag.{add,sub,mul}ChildVectorIntoExt and poly.{Ext,}EvaluateLagrangeWithWeights each ran scalar Go loops over N=2^17 elements on the proving hot path. Switching to gnark-crypto's `ext.Vector` / `koalabear.Vector` SIMD (AVX-512) entry points: - mulChildVectorIntoExt was ~15% of Prove CPU (per optim_plan_ideas.md #4); - (Ext)EvaluateLagrangeWithWeights now uses `InnerProduct[ByElement]` instead of an accumulating loop. The base-rail add/sub paths skip the lift-to-E4 indirection and touch only the B0.A0 coefficient directly, since the lifted value is zero in the other slots. Co-Authored-By: Claude Opus 4.7 (1M context) --- internal/dag/dag.go | 23 ++++++----------------- internal/poly/utils.go | 17 +++-------------- 2 files changed, 9 insertions(+), 31 deletions(-) diff --git a/internal/dag/dag.go b/internal/dag/dag.go index 3c7e3fe..cf05302 100644 --- a/internal/dag/dag.go +++ b/internal/dag/dag.go @@ -1403,15 +1403,12 @@ func copyChildVectorToExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.El func addChildVectorToExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.Element, extVec [][]ext.E4, N int) { if child.Field == field.Ext { src := extVec[child.Index] - for j := range N { - dst[j].Add(&dst[j], &src[j]) - } + ext.Vector(dst).Add(ext.Vector(dst), ext.Vector(src)) return } src := baseVec[child.Index] for j := range N { - rhs := liftBaseToE4(src[j]) - dst[j].Add(&dst[j], &rhs) + dst[j].B0.A0.Add(&dst[j].B0.A0, &src[j]) } } @@ -1420,15 +1417,12 @@ func addChildVectorToExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.Ele func subChildVectorFromExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.Element, extVec [][]ext.E4, N int) { if child.Field == field.Ext { src := extVec[child.Index] - for j := range N { - dst[j].Sub(&dst[j], &src[j]) - } + ext.Vector(dst).Sub(ext.Vector(dst), ext.Vector(src)) return } src := baseVec[child.Index] for j := range N { - rhs := liftBaseToE4(src[j]) - dst[j].Sub(&dst[j], &rhs) + dst[j].B0.A0.Sub(&dst[j].B0.A0, &src[j]) } } @@ -1437,16 +1431,11 @@ func subChildVectorFromExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.E func mulChildVectorIntoExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.Element, extVec [][]ext.E4, N int) { if child.Field == field.Ext { src := extVec[child.Index] - for j := range N { - dst[j].Mul(&dst[j], &src[j]) - } + ext.Vector(dst).Mul(ext.Vector(dst), ext.Vector(src)) return } src := baseVec[child.Index] - for j := range N { - rhs := liftBaseToE4(src[j]) - dst[j].Mul(&dst[j], &rhs) - } + ext.Vector(dst).MulByElement(ext.Vector(dst), koalabear.Vector(src)) } // EvalWithCache evaluates the DAG using the caller-supplied cache slice instead diff --git a/internal/poly/utils.go b/internal/poly/utils.go index b9ac51d..4506f79 100644 --- a/internal/poly/utils.go +++ b/internal/poly/utils.go @@ -149,13 +149,8 @@ func EvaluateLagrangeWithWeights(p Polynomial, weights []koalabear.Element) koal if len(p) != len(weights) { panic("EvaluateLagrangeWithWeights: length mismatch") } - var res koalabear.Element - for i := range p { - var term koalabear.Element - term.Mul(&p[i], &weights[i]) - res.Add(&res, &term) - } - return res + pv := koalabear.Vector(p) + return (&pv).InnerProduct(koalabear.Vector(weights)) } // ExtEvaluateLagrangeWithWeights is the extension-field counterpart of @@ -164,13 +159,7 @@ func ExtEvaluateLagrangeWithWeights(p ExtPolynomial, weights []koalabear.Element if len(p) != len(weights) { panic("ExtEvaluateLagrangeWithWeights: length mismatch") } - var res ext.E4 - for i := range p { - var term ext.E4 - term.MulByElement(&p[i], &weights[i]) - res.Add(&res, &term) - } - return res + return ext.Vector(p).InnerProductByElement(koalabear.Vector(weights)) } // EvaluateOnExtendedDomainRoot evaluates p, given in Lagrange form over d, at From cfad4873ac2cc33516d7edc5cd72964a47fa9635 Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Tue, 26 May 2026 18:30:22 +0000 Subject: [PATCH 4/8] perf(rs): drop two BitReverse passes per Encode via DIT FFT MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per optim_plan_ideas.md #3 option B: the Encode path was FFTInverse(DIF) → BitReverse → FFT(DIF) → BitReverse, where each BitReverse on N up to 2^20 fell through gnark-crypto's threshold to `bitReverseNaive` (~5.9% of Prove CPU). We instead scatter the n bit-reversed inverse-FFT coefficients directly into N bit-reversed positions (zero-padding the gaps) and run the forward FFT in DIT order, which consumes bit-reversed input and produces normal output — no explicit BitReverse needed. Co-Authored-By: Claude Opus 4.7 (1M context) --- internal/reedsolomon/rs.go | 44 +++++++++++++++++++++++++++----------- 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/internal/reedsolomon/rs.go b/internal/reedsolomon/rs.go index adc4029..08bb827 100644 --- a/internal/reedsolomon/rs.go +++ b/internal/reedsolomon/rs.go @@ -14,8 +14,9 @@ package reedsolomon import ( + "math/bits" + "github.com/consensys/gnark-crypto/field/koalabear/fft" - "github.com/consensys/gnark-crypto/utils" "github.com/consensys/loom/internal/poly" ) @@ -33,6 +34,28 @@ type Encoder struct { Domain *fft.Domain } +// scatterBitReversedCoeffs expands n-bit-reversed coefficients into the +// matching N-bit-reversed zero-padded slots, in place. +func scatterBitReversedCoeffs[T any](p []T, n, N int) { + if n <= 1 { + return + } + shift := bits.TrailingZeros64(uint64(N)) - bits.TrailingZeros64(uint64(n)) + stride := 1 << shift + for i := n - 1; i >= 0; i-- { + p[i< len(p)) // p is in Lagrange form // it returns a copy of p @@ -46,15 +69,12 @@ func (encoder *Encoder) Encode(p poly.Polynomial, d *fft.Domain) poly.Polynomial _p := make(poly.Polynomial, N) copy(_p, p) - // compute fftinv(_p[:n]) using d (d must be of the size of p) - // Lagrange normal → canonical bit-reversed (w.r.t. n); then un-reverse to canonical normal + // Lagrange normal to canonical bit-reversed (w.r.t. n). We place those + // coefficients directly in N-bit-reversed order and use a DIT FFT, avoiding + // the two explicit BitReverse passes previously needed for normal order. d.FFTInverse(_p[:n], fft.DIF) - utils.BitReverse(_p[:n]) - - // compute fft(_p) using the Encoder domain - // canonical normal (zero-padded to N) → Lagrange bit-reversed (w.r.t. N) → Lagrange normal - encoder.Domain.FFT(_p, fft.DIF) - utils.BitReverse(_p) + scatterBitReversedCoeffs(_p, n, int(N)) + encoder.Domain.FFT(_p, fft.DIT) // return _p return _p @@ -71,10 +91,8 @@ func (encoder *Encoder) EncodeExt(p poly.ExtPolynomial, d *fft.Domain) poly.ExtP copy(_p, p) d.FFTInverseExt(_p[:n], fft.DIF) - utils.BitReverse(_p[:n]) - - encoder.Domain.FFTExt(_p, fft.DIF) - utils.BitReverse(_p) + scatterBitReversedCoeffs(_p, n, int(N)) + encoder.Domain.FFTExt(_p, fft.DIT) return _p } From 9ce50dc0e5df668b2060fca78a584ba80b1364c2 Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Tue, 26 May 2026 18:30:31 +0000 Subject: [PATCH 5/8] perf(prover): parallelize zeta evaluations, vectorize DeepQuotient accumulation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ComputeAIRQuotients: the per-chunk evaluate-at-zeta loop (over both Base and Ext chunks) now runs through parallel.Execute on a deterministically sorted key list; the proof map is written serially after the parallel section. ComputeEvaluationsAtZeta: was a serial loop over modules. Now parallelized per-module; each goroutine buffers (key, val) pairs locally and the proof map is merged in the main goroutine. Errors are collected per-module without a mutex. ComputeDeepQuotient: the inner accumulation `C_s[x] += alpha · col[x]` (scaling each trace column by alphaAcc and folding into C_s) was a scalar loop. Now goes through `ext.Vector.MulAccByElement` (base columns) or `ScalarMul` + `Add` via a per-size scratch buffer (ext columns). Constant columns (len == 1) still use the broadcast scalar form. Scratch is allocated once per size and reused across all shifts. Co-Authored-By: Claude Opus 4.7 (1M context) --- prover/prover.go | 158 ++++++++++++++++++++++++++++++++--------------- 1 file changed, 107 insertions(+), 51 deletions(-) diff --git a/prover/prover.go b/prover/prover.go index c7a297e..eae4035 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -25,9 +25,9 @@ import ( "github.com/consensys/loom/board" "github.com/consensys/loom/expr" "github.com/consensys/loom/field" - "github.com/consensys/loom/internal/dag" "github.com/consensys/loom/internal/commitment" "github.com/consensys/loom/internal/constants" + "github.com/consensys/loom/internal/dag" fiatshamir "github.com/consensys/loom/internal/fiat-shamir" "github.com/consensys/loom/internal/fri" "github.com/consensys/loom/internal/hash" @@ -592,11 +592,28 @@ func (pr *proverRuntime) ComputeAIRQuotients() error { pr.zeta = hash.OutputToExt(zeta) } - for chunkName, chunkPoly := range pr.airTrace.Base { - pr.Proof.SetValueAtZetaExt(chunkName, poly.EvaluateAtExt(chunkPoly, chunkDomains[chunkName], pr.zeta)) + chunkNames := make([]string, 0, len(pr.airTrace.Base)+len(pr.airTrace.Ext)) + for name := range pr.airTrace.Base { + chunkNames = append(chunkNames, name) } - for chunkName, chunkPoly := range pr.airTrace.Ext { - pr.Proof.SetValueAtZetaExt(chunkName, poly.ExtEvaluateAtExt(chunkPoly, chunkDomains[chunkName], pr.zeta)) + for name := range pr.airTrace.Ext { + chunkNames = append(chunkNames, name) + } + sort.Strings(chunkNames) + + evals := make([]ext.E4, len(chunkNames)) + parallel.Execute(len(chunkNames), func(start, end int) { + for i := start; i < end; i++ { + name := chunkNames[i] + if p, ok := pr.airTrace.Base[name]; ok { + evals[i] = poly.EvaluateAtExt(p, chunkDomains[name], pr.zeta) + } else { + evals[i] = poly.ExtEvaluateAtExt(pr.airTrace.Ext[name], chunkDomains[name], pr.zeta) + } + } + }) + for i, name := range chunkNames { + pr.Proof.SetValueAtZetaExt(name, evals[i]) } return nil @@ -612,36 +629,90 @@ func (pr *proverRuntime) ComputeEvaluationsAtZeta() error { expr.WithoutPublicColumns(), ) - for _, module := range pr.program.Modules { - leaves := module.VanishingRelation.LeavesFull(config) - for _, leaf := range leaves { - evalPoint := pr.zeta - if leaf.Type == expr.RotatedColumn { - shift := ((leaf.Shift % module.N) + module.N) % module.N - var omegaPow koalabear.Element - omegaPow.SetOne() - for k := 0; k < shift; k++ { - omegaPow.Mul(&omegaPow, &module.D.Generator) + type zetaEval struct { + key string + val ext.E4 + } + + moduleNames := make([]string, 0, len(pr.program.Modules)) + for name := range pr.program.Modules { + moduleNames = append(moduleNames, name) + } + sort.Strings(moduleNames) + + results := make([][]zetaEval, len(moduleNames)) + errs := make([]error, len(moduleNames)) + parallel.Execute(len(moduleNames), func(start, end int) { + for idx := start; idx < end; idx++ { + module := pr.program.Modules[moduleNames[idx]] + leaves := module.VanishingRelation.LeavesFull(config) + local := make([]zetaEval, 0, len(leaves)) + for _, leaf := range leaves { + evalPoint := pr.zeta + if leaf.Type == expr.RotatedColumn { + shift := ((leaf.Shift % module.N) + module.N) % module.N + var omegaPow koalabear.Element + omegaPow.SetOne() + for k := 0; k < shift; k++ { + omegaPow.Mul(&omegaPow, &module.D.Generator) + } + evalPoint.MulByElement(&evalPoint, &omegaPow) } - evalPoint.MulByElement(&evalPoint, &omegaPow) - } - if p, ok := pr.t.Ext[leaf.Name]; ok { - val := poly.ExtEvaluateAtExt(p, module.D, evalPoint) - pr.Proof.SetValueAtZetaExt(leaf.String(), val) - continue - } - p, ok := pr.t.Base[leaf.Name] - if !ok { - return fmt.Errorf("ComputeEvaluationsAtZeta: column %q not found in trace", leaf.Name) + if p, ok := pr.t.Ext[leaf.Name]; ok { + val := poly.ExtEvaluateAtExt(p, module.D, evalPoint) + local = append(local, zetaEval{key: leaf.String(), val: val}) + continue + } + p, ok := pr.t.Base[leaf.Name] + if !ok { + errs[idx] = fmt.Errorf("ComputeEvaluationsAtZeta: column %q not found in trace", leaf.Name) + return + } + val := poly.EvaluateAtExt(p, module.D, evalPoint) + local = append(local, zetaEval{key: leaf.String(), val: val}) } - val := poly.EvaluateAtExt(p, module.D, evalPoint) - pr.Proof.SetValueAtZetaExt(leaf.String(), val) + results[idx] = local + } + }) + for _, err := range errs { + if err != nil { + return err + } + } + for _, local := range results { + for _, ev := range local { + pr.Proof.SetValueAtZetaExt(ev.key, ev.val) } } return nil } +func addScaledExtColumn(dst, col, scratch poly.ExtPolynomial, alpha *ext.E4) { + if len(col) == 1 { + var term ext.E4 + term.Mul(&col[0], alpha) + for i := range dst { + dst[i].Add(&dst[i], &term) + } + return + } + ext.Vector(scratch).ScalarMul(ext.Vector(col), alpha) + ext.Vector(dst).Add(ext.Vector(dst), ext.Vector(scratch)) +} + +func addScaledBaseColumn(dst poly.ExtPolynomial, col poly.Polynomial, alpha *ext.E4) { + if len(col) == 1 { + var term ext.E4 + term.MulByElement(alpha, &col[0]) + for i := range dst { + dst[i].Add(&dst[i], &term) + } + return + } + ext.Vector(dst).MulAccByElement(col, alpha) +} + func (pr *proverRuntime) ComputeDeepQuotient() error { dqLayout := BuildDeepQuotientLayout(pr.program) sizes := dqLayout.Sizes @@ -665,6 +736,7 @@ func (pr *proverRuntime) ComputeDeepQuotient() error { alphaAcc.SetOne() domainN := domainBySize[N] + scratch := make(poly.ExtPolynomial, N) for j, shift := range dqLayout.Shifts[i] { var omegaShift koalabear.Element @@ -689,21 +761,10 @@ func (pr *proverRuntime) ComputeDeepQuotient() error { if !hasExt && !hasBase { return fmt.Errorf("ComputeDeepQuotient: column %q not found in trace", names[k]) } - for x := 0; x < N; x++ { - var value, term ext.E4 - if hasExt { - if len(colExt) == 1 { - value.Set(&colExt[0]) - } else { - value.Set(&colExt[x]) - } - } else if len(colBase) == 1 { - value.Lift(&colBase[0]) - } else { - value.Lift(&colBase[x]) - } - term.Mul(&value, &alphaAcc) - C_s[x].Add(&C_s[x], &term) + if hasExt { + addScaledExtColumn(C_s, colExt, scratch, &alphaAcc) + } else { + addScaledBaseColumn(C_s, colBase, &alphaAcc) } var term ext.E4 term.Mul(&evalAtZ, &alphaAcc) @@ -730,15 +791,10 @@ func (pr *proverRuntime) ComputeDeepQuotient() error { if !hasExt && !hasBase { return fmt.Errorf("ComputeDeepQuotient: AIR chunk %q not found in trace", chunkName) } - for x := 0; x < N; x++ { - var value, term ext.E4 - if hasExt { - value.Set(&chunkExt[x]) - } else { - value.Lift(&chunkBase[x]) - } - term.Mul(&value, &alphaAcc) - C_s[x].Add(&C_s[x], &term) + if hasExt { + addScaledExtColumn(C_s, chunkExt, scratch, &alphaAcc) + } else { + addScaledBaseColumn(C_s, chunkBase, &alphaAcc) } var term ext.E4 term.Mul(&evalAtZ, &alphaAcc) From 724560bdc760e097b20d938013e6f923c18d4fc8 Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Tue, 26 May 2026 15:05:33 -0500 Subject: [PATCH 6/8] fix: remove .md file --- optim_plan_ideas.md | 388 -------------------------------------------- 1 file changed, 388 deletions(-) delete mode 100644 optim_plan_ideas.md diff --git a/optim_plan_ideas.md b/optim_plan_ideas.md deleted file mode 100644 index 3ef86e0..0000000 --- a/optim_plan_ideas.md +++ /dev/null @@ -1,388 +0,0 @@ -# Loom prover — optimisation plan & ideas - -Status: draft, 2026-05-26. All numbers in this document come from `bench/main.go` -(in this repo) running on a 32-core / 32-GOMAXPROCS Linux box with AVX-512, -gnark-crypto v0.20.1, Go 1.26. Workload: 40 aggregated PLONK instances of size -2^17, Poseidon2 leaf hashing. - -The goal of this file is to give a coding agent enough context to pick up any -one of these items, work it in isolation, and verify the impact. Each item -includes profile evidence, file:line pointers, the proposed change, and a -verification recipe. - ---- - -## 0. Baseline & reproduction - -``` -$ go build -o /tmp/loom-bench ./bench -$ /tmp/loom-bench # default workload, ~45s wall -$ /tmp/loom-bench -instances N -log2-size L # tune up/down -$ /tmp/loom-bench -skip-fri # to isolate non-FRI work -$ /tmp/loom-bench -hash sha256 # to swap Merkle hash backend -``` - -Default run produces: - -``` -phase wall cpu par gc% alloc peakHeap -traces+modules 2.02s 3.08s 1.52x 7.3% 5.40 GiB 942 MiB -compile 84.2ms 0s 0.00x 0.0% 244 MiB 920 MiB -merge-trace 77 µs 0s 0.00x 0.0% 6.3 KiB 920 MiB -prove 41.98s 45.99s 1.10x 0.1% 12.89 GiB 4.86 GiB -TOTAL 44.08s 49.07s 1.11x 0.6% 18.53 GiB 4.86 GiB -``` - -Read `par = cpu / wall` as "average cores busy". The headline finding is that -Prove uses **~1 core out of 32** despite parallel primitives existing. Almost -every optimization below either (a) raises that number or (b) reduces work that -would otherwise be parallelised. - -The pprof CPU profile referenced throughout lives at -`bench_profiles/cpu_prove.pprof` after a run. Re-render with: - -``` -go tool pprof -top -cum -nodecount 40 bench_profiles/cpu_prove.pprof -go tool pprof -list bench_profiles/cpu_prove.pprof -go tool pprof -focus -tree bench_profiles/cpu_prove.pprof -``` - -Heap before/after Prove and an allocations profile are written alongside. - -### Phase breakdown (cumulative CPU in the profile) - -| Phase | % of Prove CPU | Notes | -|-----------------------------|----------------|----------------------------------------| -| `ComputeAIRQuotients` | ~35% | per-module sequential outer loop | -| `ExecuteSteps` | ~12% | sequential per module, leaf hashing | -| `commitTraceRound` (in #2) | ~7% | Poseidon2 leaf hashing of trace polys | -| `ComputeDeepQuotient` | ~6% | sequential per size, per-elem inverse | -| `SampleEvaluations` | ~6% | sequential per query position | -| Everything else | ~34% | runtime / GC / minor paths | - -### Hottest leaf functions (flat CPU) - -| Function (flat %) | Where | -|--------------------------------------------------------|------------------------------------| -| `permutation16x24_avx512` 10.0% | Poseidon2 batch-16, already SIMD | -| `extensions.montReduce` 9.95% | Every E4.Mul; called from #1, #2 | -| `extensions.E4.Mul` 8.84% | Driven by `mulChildVectorIntoExt` | -| `bitReverseNaive[E4]` 5.28% (+0.58% mem swap) | See #3 | -| `vectorButterfly_avx512` 5.04% | gnark-crypto FFT, already SIMD | -| `koalabear.Element.Add` 4.34% | Base-field accumulation | - ---- - -## Top-3 optimisations (high confidence, real numbers) - -### #1 — Parallelise the per-module loop in `ComputeAIRQuotients` - -**Expected impact:** Up to ~4–5× Prove speedup (this phase is ~35% of CPU and -currently runs at ~1.1× parallelism with 40 independent modules on 32 cores). - -**Evidence:** -- `prover/prover.go:420` iterates `for moduleName, module := range pr.program.Modules` - sequentially. With 40 instances of the bench, there are 40+ modules whose - AIR quotients are independent. -- pprof tree under `ComputeAIRQuotients` (`go tool pprof -focus=ComputeAIRQuotients -tree`): - - `poly.ComputeQuotientMixed` — 24.3s cumulative, **26.6% of total Prove CPU**. - - `dag.EvalOnAllEntriesMixedInto` (`internal/dag/dag.go:1097`) — 24.5s cum. - - `dag.mulChildVectorIntoExt` (`internal/dag/dag.go:1401`) — 14.0s cum, **15% of total Prove CPU**, scalar Go loop over E4. -- `Total samples = 91.52s` over `Duration: 61.02s` on the 60×2^17 run = ~1.5x - effective parallelism for the whole Prove; ComputeAIRQuotients dominates. - -**Proposed change:** - -1. Materialise a deterministic module ordering once: - ```go - names := make([]string, 0, len(pr.program.Modules)) - for name := range pr.program.Modules { names = append(names, name) } - sort.Strings(names) - ``` -2. Wrap the work between `prover.go:420` and `prover.go:472` (i.e. the body - that produces per-module `airTrace` chunks + populates `chunkDomains`) in - `internal/parallel.Execute(len(names), func(start, end int) { ... })`. -3. Each goroutine writes into its own pre-allocated `[]ext.E4` / `[]koalabear.Element` - workspace and into a private `localChunkDomains` and `localAirChunks` map; - merge under a single mutex (or pre-size the global maps and write disjoint - keys) after the parallel section. -4. `pr.airTrace.SetBase` / `SetExt` are not safe for concurrent map writes — - either lock around them, or buffer per goroutine and merge serially in - <1% time. -5. The downstream "commit by size group" block (`prover.go:474–531`) stays - serial — it groups across modules and already uses a parallel committer - internally. - -**Watch out for:** -- `pr.domainCache` is shared (`poly.WithDomainCache(&pr.domainCache)`). Check - `internal/poly/domain_cache.go` for thread-safety; if it isn't safe, either - give each goroutine its own cache or add an `sync.RWMutex` around it. This - is the single biggest correctness risk for this change. -- The per-goroutine workspace pool inside `dag.EvalWorkspace` - (`internal/dag/dag.go:1118-1140`) is currently per-call. If you reuse a - workspace across modules per goroutine you should also confirm it's reset - between calls. - -**Verification:** -- `go test ./...` must stay green (especially `prover/...` and integration). -- `/tmp/loom-bench` should show `prove`'s `par` jump from ~1.1x toward - GOMAXPROCS. Even a modest 8x would translate to ~3× Prove speedup. -- `go tool pprof -focus=ComputeAIRQuotients -top` should show `sync.(*WaitGroup).Go.func1` - cum time rise significantly. - ---- - -### #2 — Batch-invert in `DeepQuotientExt` (and parallelise the per-size loop) - -**Expected impact:** 4–8% Prove speedup from the batch invert alone; another -2–4% from parallelising sizes (limited by the small number of distinct sizes, -typically 3–4). - -**Evidence:** -- `internal/poly/ext.go:130` — `DeepQuotientExt` does, in a tight loop of N - iterations: - ```go - for j := 0; j < N; j++ { - var num ext.E4 - num.Sub(&v, &p[j]) - var den ext.E4 - den.Sub(&z, &ω[j]) - var inv ext.E4 - inv.Inverse(&den) // ← one full E4 inversion per element - out[j].Mul(&num, &inv) - } - ``` -- An E4 inversion (Frobenius + base-field inverse + multiplies) is ~10–30× - the cost of an E4 multiply. The denominators form a vector ideal for - Montgomery batch inversion: 1 inversion + 3(N-1) multiplies for N elements. -- `ext.BatchInvertE4` already exists and is used at - `internal/poly/iop_utils.go:375` and `:555` — same pattern. -- `prover.go:612` — `for i, N := range sizes { ... }` is the outer loop over - distinct module sizes. Each iteration owns its own `deepQuotient` map entry, - its own RS encoder, its own FRI level tree → trivially independent. -- `prover.go:629` — `C_s := make(poly.ExtPolynomial, N)` is allocated **per - shift, per size**. For N=2^17, ~20 shifts, ~3 sizes that's ~96 MiB of E4 - churn per Prove (and a contributor to the 12.9 GiB of `prove`-phase - allocations). - -**Proposed change:** - -1. In `DeepQuotientExt`: - - Build `dens[j] = z - ω[j]` into a scratch `[]ext.E4`. - - Call `ext.BatchInvertE4(dens)` once. - - Final loop becomes `out[j].Mul(&num, &dens[j])` — single E4 multiply per - iteration instead of an inversion. - - Stream-friendly: reuse a scratch buffer across calls via an optional - argument or a `sync.Pool`. - -2. In `ComputeDeepQuotient` (`prover.go:596`): - - Convert the `for i, N := range sizes` loop body into a closure dispatched - via `parallel.Execute(len(sizes), ...)`. - - `deepQuotients` (a `map[int]poly.ExtPolynomial`) is written from disjoint - keys per iteration → preallocate to `len(sizes)` and write under - `pr.mu`, **or** swap to `[]poly.ExtPolynomial` indexed by `i` and assemble - after the parallel section. - - The inner per-shift work (`prover.go:620–668`) also accumulates into a - per-size `deepQuotient` — that stays serial within a size, fine. - - `pr.Proof.DeepQuotientCommitment` and the FRI proving block (`prover.go:709–737`) - stay serial since `fri.Prove` is invoked once after all level trees are - built. - -3. Bonus, low-effort: reuse `C_s` across shifts. Currently allocated fresh at - `prover.go:629`; the only operation that writes it is `C_s[x].Add(&C_s[x], &term)` - inside the `k`-loop, so a zero-on-entry policy + a single allocation per - size eliminates the per-shift alloc. - -**Verification:** -- Existing prover tests cover this end-to-end; if they pass, the batch invert - is correct. -- Microbench (write a `BenchmarkDeepQuotientExt` in `internal/poly/ext_test.go` - parameterised on `N`) should show ~5–10× speedup on the function alone. -- `bench_profiles/cpu_prove.pprof` should show `ext.E4.Inverse` cum time drop - by ~95% and `DeepQuotientExt` flat time shrink to a fraction. - ---- - -### #3 — Stop running `bitReverseNaive` on E4 / E2 polynomials - -**Expected impact:** ~5.9% of Prove CPU, ~50 LOC of bit-reverse code, or -**zero work** if you can elide the bit-reverses entirely (see option B). - -**Evidence:** -- `utils.bitReverseNaive[E4-shaped]` shows up at **4.83s flat / 5.36s cum = - 5.86% of total Prove CPU**. The function (`gnark-crypto/utils/bitreverse.go:29`) - is literally a swap-pairs loop — no algorithm. -- `gnark-crypto/utils/bitreverse.go:20` dispatches: - ```go - if runtime.GOARCH == "arm64" || len(v) < (1<<21) || unsafe.Sizeof(v[0]) < 8 { - bitReverseNaive(v) - } else { - bitReverseCobra(v) - } - ``` - Our quotient FFTs land on `bigSize = eDeg·N`, typically 2^19–2^20 (N=2^17, - eDeg up to 4–8) — just under the 2^21 cutoff. So even though gnark-crypto - ships a cache-friendly Cobra variant and specialised `bitReverseCobraInPlace_9_*` - unrolls, we never reach them. -- Call sites: `internal/poly/compute_quotient.go:128, 131, 301, 304, 392, - 407, 423, 434` and `prover/prover.go:436, 443, 460, 467`. - -**Option A — Local fast bit-reverse for E4 (and base):** - -Add `internal/poly/bitreverse.go` with a Cobra-style or block-transpose -bit-reverse specialised for `ext.E4` (and likely `koalabear.Element`). -Route all loom call sites through it. - -The Cobra threshold is set conservatively in gnark-crypto for arrays with -non-trivial element size — for E4 (16 bytes) and N=2^19 the cache benefit is -real, the threshold could safely come down. Crib the algorithm from -`gnark-crypto/utils/bitreverse.go:128-160` and the `_9_21` family. - -**Option B — Eliminate the bit-reverses:** - -Several call sites are of the form "FFT (DIF) → BitReverse → … → BitReverse → -FFT (DIT)". When two BitReverses bracket a region with only pointwise -operations, both can be deleted. Inspect: -- `prover.go:435-437` then `:442-443` — `FFTInverseExt(DIF)` → BitReverse → - chunk → `FFTExt(DIF)` → BitReverse. Either keep everything in bit-reversed - order or use DIT/DIF pairing carefully. -- `compute_quotient.go:391-407` — same pattern around the coset FFT. - -If you can keep polynomials in bit-reversed order between `ComputeQuotient*` -and the per-chunk FFT, you may delete 4 BitReverse calls outright. This is -the loomwide win. - -**Verification:** -- Cross-check by running prover_test.go and integration_test — these polys - feed into Merkle leaves that the verifier re-checks, so any bit-order - mistake fails the proof. -- Profile: `bitReverseNaive[E4]` should disappear from `pprof -top`. - -**Risks:** -- BitReverse interacts with `fft.DIF` vs `fft.DIT` directionality; getting - the algebra right requires reading `gnark-crypto`'s FFT API. Don't ship - option B without exhaustive tests. - ---- - -## Secondary opportunities (medium confidence) - -These are visible in the profile but harder to size without prototyping. - -### #4 — Vectorise `dag.mulChildVectorIntoExt` and friends - -**Evidence:** `internal/dag/dag.go:1401` is a scalar Go loop: -```go -for j := range N { - dst[j].Mul(&dst[j], &src[j]) -} -``` -This is **15% of Prove CPU** (cumulative including the called E4.Mul). The same -file has analogous `addChildVectorToExt`, `subChildVectorFromExt`, -`copyChildVectorToExt`. - -gnark-crypto exposes `extensions.Vector.Mul`, `Vector.MulByElement` (already -shows up at 5.7% via FFT), `Vector.Add`, `Vector.Butterfly` with AVX-512 -implementations. Most of those entry points need a third argument shape -(`Vector.Mul(dst, a, b)`); the loom DAG site is `Mul(dst, dst, src)`. Either -adapt the call site to use a scratch buffer, or add an in-place -`Vector.MulSelf(dst, src)` to gnark-crypto. - -If a vector entry point cuts this 15% in half, that's ~7% Prove on its own, -**stackable with #1** (parallelising modules just means each module's inner -loop runs on one core; a SIMD inner loop multiplies the per-core throughput). - -### #5 — Reuse `EvalWorkspace` across modules (after #1) - -`internal/dag/dag.go:1118-1140` has a per-call `basePool` / `extPool`. When #1 -lands and modules run in parallel, give each worker goroutine a persistent -`EvalWorkspace` instead of creating one per `EvalOnAllEntriesMixedInto`. This -keeps the slices in the same goroutine's local cache and dodges the allocator -on every coset. - -The current `make([]ext.E4, N)` at `dag.go:1138` is allocated `n_nodes × n_cosets` -times per module. With degree ~5 and ~10 cosets per module, that's hundreds -of N-sized E4 slabs in the churn budget per Prove. - -### #6 — Parallelise `SampleEvaluations` over query positions - -`prover/prover.go:771-786`: NUM_QUERIES (~128) FRI queries, each opening every -committed tree at one position. Currently `for q, s := range pr.queryPositions` -is sequential. Each query owns its own `pr.Proof.PointSamplings[q]` slot — no -data races. `parallel.Execute(NQ, ...)`. - -This phase is ~6% of CPU; parallelising should push it to single-digit ms. - -### #7 — Parallelise `ExecuteSteps`' `GenCol` loop - -`prover.go:365`: `for _, m := range pr.program.Modules { ... gen.Gen(pr.t, &mCopy) }`. -Generators may write to disjoint trace columns, but they share `pr.t` (a -`trace.Trace`); concurrent map writes will race. Requires a per-module -`trace.Trace` produced under the goroutine, then merged. Speculative — confirm -no cross-module data dependencies in `GenCol` before attempting. - -### #8 — Poseidon2 `WriteExtBatch` element-by-element absorb - -`internal/hash/poseidon2_batch.go`'s `WriteExtBatch` shows up at **7.0s cum = -7.7% of Prove**. It calls `WriteElementBatch` four times per E4 (one per -coefficient). Each `WriteElementBatch` does an `absorbFullBlock` if the sponge -is at capacity. There's probably a vector store / unrolled writer hiding here -— inspect whether you can stage 4 elements at once instead of looping. - -Lower confidence: this kernel may already be well-tuned and the time is -fundamental work. But "WriteExt does 4× WriteElement" smells loose. - -### #9 — Memory churn: reduce `make([]ext.E4, N)` in inner loops - -`prover.go:629`, `prover.go:613`, `prover.go:672` each allocate a per-iteration -N-sized E4 slice. With N=2^17 (= 2 MiB per E4 slice), Prove allocates ~12.9 -GiB in total. GC is at 0.1% so the *runtime* cost is negligible — but the -peak heap (4.86 GiB) is driven by these large transient buffers, and Linux -page-fault cost on first-touch is non-zero. Pool / reuse them per Prove run. - -This is mostly a memory-footprint win (helpful for running larger workloads -without OOM), not a CPU win. - ---- - -## Things NOT worth optimising (per profile) - -- **GC tuning.** 0.1–0.6% of CPU. Already invisible. -- **Poseidon2 permutation kernel.** Already AVX-512 batch-16, near the SIMD - ceiling. -- **FFT kernels.** `kerDIFNP_512Ext`, `vectorButterfly_avx512` are gnark-crypto - AVX-512 code — leave alone. -- **`ComputeEvaluationsAtZeta`** (`prover.go:556`). Doesn't even show in top - 40 of pprof. Skip. -- **`merge-trace`**, **`compile`**. 0% and ~0.2% of total wall. - ---- - -## Suggested ordering for a coding-agent campaign - -1. **Land #3 Option A** first — small, isolated, low risk, ~6% Prove. Good warmup. -2. **Land #2 batch-invert** — small isolated change in `internal/poly/ext.go`, - ~5% Prove. Verify with a microbench. -3. **Land #1** (parallelise `ComputeAIRQuotients`) — biggest payoff but most - careful work (domain-cache thread-safety + per-goroutine workspace). After - this, re-run the bench and re-profile — the percentage breakdowns will - shift substantially. -4. **Land #6** while you're in the prover (`SampleEvaluations`). -5. Re-profile. If E4 inner loops are still the leaf cost, do **#4** - (vector E4 ops in `dag`). Otherwise pick the next-biggest item from the - refreshed profile. - -After 1–4, target should be Prove `par` ≥ 8x on 32 cores and wall time ≤ -15s on the default workload. - ---- - -## Open questions for the team - -- Is `poly.DomainCache` safe for concurrent reads from multiple goroutines? - This blocks #1. -- Are `board.Module.GenCol` generators required to be sequential by design, - or just by current implementation? This blocks #7. -- Why is gnark-crypto's `bitReverseCobra` threshold set at 2^21? Could it be - lowered upstream for E4-shaped slices (16 bytes/elem) where the cache - benefit kicks in earlier? From 703664d281aa708f8ff6d5d36278567b9973e113 Mon Sep 17 00:00:00 2001 From: Thomas Piellard Date: Wed, 27 May 2026 10:48:03 +0200 Subject: [PATCH 7/8] feat: evaluations at zeta in one go --- board/program.go | 1 + prover/prover.go | 48 +++++++++++++++++++++--------------------------- 2 files changed, 22 insertions(+), 27 deletions(-) diff --git a/board/program.go b/board/program.go index 7916b95..9c467a0 100644 --- a/board/program.go +++ b/board/program.go @@ -22,6 +22,7 @@ type Program struct { Steps [][]ProverStep } +// TODO group modules per size, fold the vanishing relations per modules of the same size -> less calls to ComputeQuotient func (pg *Program) SetSize(module string, size int) { _, ok := pg.Modules[module] if !ok { diff --git a/prover/prover.go b/prover/prover.go index eae4035..6fea1dc 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -459,7 +459,6 @@ func computeBaseAIRChunks(piBase map[string]poly.Polynomial, vrel *dag.DAG, N in } func (pr *proverRuntime) ComputeAIRQuotients() error { - chunkDomains := make(map[string]*fft.Domain) moduleNames := make([]string, 0, len(pr.program.Modules)) for name := range pr.program.Modules { @@ -511,12 +510,10 @@ func (pr *proverRuntime) ComputeAIRQuotients() error { for i, c := range baseChunks { name := constants.QuotientChunkName(moduleName, i) pr.airTrace.SetBase(name, c) - chunkDomains[name] = module.D } for i, c := range extChunks { name := constants.QuotientChunkName(moduleName, i) pr.airTrace.SetExt(name, c) - chunkDomains[name] = module.D } writeMu.Unlock() } @@ -592,30 +589,6 @@ func (pr *proverRuntime) ComputeAIRQuotients() error { pr.zeta = hash.OutputToExt(zeta) } - chunkNames := make([]string, 0, len(pr.airTrace.Base)+len(pr.airTrace.Ext)) - for name := range pr.airTrace.Base { - chunkNames = append(chunkNames, name) - } - for name := range pr.airTrace.Ext { - chunkNames = append(chunkNames, name) - } - sort.Strings(chunkNames) - - evals := make([]ext.E4, len(chunkNames)) - parallel.Execute(len(chunkNames), func(start, end int) { - for i := start; i < end; i++ { - name := chunkNames[i] - if p, ok := pr.airTrace.Base[name]; ok { - evals[i] = poly.EvaluateAtExt(p, chunkDomains[name], pr.zeta) - } else { - evals[i] = poly.ExtEvaluateAtExt(pr.airTrace.Ext[name], chunkDomains[name], pr.zeta) - } - } - }) - for i, name := range chunkNames { - pr.Proof.SetValueAtZetaExt(name, evals[i]) - } - return nil } @@ -644,6 +617,8 @@ func (pr *proverRuntime) ComputeEvaluationsAtZeta() error { errs := make([]error, len(moduleNames)) parallel.Execute(len(moduleNames), func(start, end int) { for idx := start; idx < end; idx++ { + + // trace polynomials module := pr.program.Modules[moduleNames[idx]] leaves := module.VanishingRelation.LeavesFull(config) local := make([]zetaEval, 0, len(leaves)) @@ -672,6 +647,24 @@ func (pr *proverRuntime) ComputeEvaluationsAtZeta() error { val := poly.EvaluateAtExt(p, module.D, evalPoint) local = append(local, zetaEval{key: leaf.String(), val: val}) } + + // air chunks + evalPoint := pr.zeta + for i := 0; ; i++ { + chunkName := constants.QuotientChunkName(moduleNames[idx], i) + if chunk, ok := pr.airTrace.Base[chunkName]; ok { + val := poly.EvaluateAtExt(chunk, module.D, evalPoint) + local = append(local, zetaEval{key: chunkName, val: val}) + continue + } + if chunk, ok := pr.airTrace.Ext[chunkName]; ok { + val := poly.ExtEvaluateAtExt(chunk, module.D, evalPoint) + local = append(local, zetaEval{key: chunkName, val: val}) + continue + } + break + } + results[idx] = local } }) @@ -685,6 +678,7 @@ func (pr *proverRuntime) ComputeEvaluationsAtZeta() error { pr.Proof.SetValueAtZetaExt(ev.key, ev.val) } } + return nil } From 58e62e267f4c9bfbf5ee69542a52d6534163c5ed Mon Sep 17 00:00:00 2001 From: Thomas Piellard Date: Wed, 27 May 2026 10:59:59 +0200 Subject: [PATCH 8/8] style: computeAIRChunks -> computeAIRQuotientChunks --- prover/prover.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/prover/prover.go b/prover/prover.go index 6fea1dc..39a1091 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -416,12 +416,12 @@ func (pr *proverRuntime) ExecuteSteps() error { return nil } -// computeExtAIRChunks computes the AIR quotient for a single ext-rooted module +// computeExtAIRQuotientChunks computes the AIR quotient for a single ext-rooted module // and returns its N-sized chunks in Lagrange form, ready to commit. Chunks // alias non-overlapping windows of the quotient backing array (no per-chunk // copy); after this returns the caller owns the chunks and the original // quotient is unreachable. -func computeExtAIRChunks(piBase map[string]poly.Polynomial, piExt map[string]poly.ExtPolynomial, vrel *dag.DAG, N int, D *fft.Domain, cache *poly.DomainCache) ([]poly.ExtPolynomial, error) { +func computeExtAIRQuotientChunks(piBase map[string]poly.Polynomial, piExt map[string]poly.ExtPolynomial, vrel *dag.DAG, N int, D *fft.Domain, cache *poly.DomainCache) ([]poly.ExtPolynomial, error) { quotient, err := poly.ComputeQuotientMixed(piBase, piExt, *vrel, N, poly.WithDomainCache(cache)) if err != nil { return nil, err @@ -440,8 +440,8 @@ func computeExtAIRChunks(piBase map[string]poly.Polynomial, piExt map[string]pol return chunks, nil } -// computeBaseAIRChunks is the base-field counterpart of computeExtAIRChunks. -func computeBaseAIRChunks(piBase map[string]poly.Polynomial, vrel *dag.DAG, N int, D *fft.Domain, cache *poly.DomainCache) ([]poly.Polynomial, error) { +// computeBaseAIRQuotientChunks is the base-field counterpart of computeExtAIRQuotientChunks. +func computeBaseAIRQuotientChunks(piBase map[string]poly.Polynomial, vrel *dag.DAG, N int, D *fft.Domain, cache *poly.DomainCache) ([]poly.Polynomial, error) { quotient, err := poly.ComputeQuotient(piBase, *vrel, N, poly.WithDomainCache(cache)) if err != nil { return nil, err @@ -493,9 +493,9 @@ func (pr *proverRuntime) ComputeAIRQuotients() error { err error ) if module.VanishingRelation.Root.Field == field.Ext { - extChunks, err = computeExtAIRChunks(pr.t.Base, pr.t.Ext, vrel, N, module.D, &pr.domainCache) + extChunks, err = computeExtAIRQuotientChunks(pr.t.Base, pr.t.Ext, vrel, N, module.D, &pr.domainCache) } else { - baseChunks, err = computeBaseAIRChunks(pr.t.Base, vrel, N, module.D, &pr.domainCache) + baseChunks, err = computeBaseAIRQuotientChunks(pr.t.Base, vrel, N, module.D, &pr.domainCache) } if err != nil { errMu.Lock()