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 new file mode 100644 index 0000000..acc424b --- /dev/null +++ b/bench/main.go @@ -0,0 +1,470 @@ +// 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" + "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" +) + +// 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() + + 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, *nbInstances) + for i := 0; i < *nbInstances; i++ { + t, sigma, size, err := gnarkplonk.GetIthPlonkTrace(n, i) + if err != nil { + fail("GetIthPlonkTrace[%d]: %v", i, err) + } + 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)), + } + sigmaGen := board.NewPermutationGen(sigma, gnarkplonk.Ith("plonk.S", i)) + if err := arguments.CopyConstraint(&builder, gnarkplonk.Ith("plonk", i), lro, sigmaGen); err != nil { + fail("CopyConstraint[%d]: %v", i, err) + } + } + phases = append(phases, tr.stop()) + + // ---- Phase 2: compile program ---------------------------------------- + tr = newTracker("compile", *sampleMs) + program, err := board.Compile(&builder) + if err != nil { + 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 { + 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.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) + } + 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 { + fail("create %s: %v", path, err) + } + return f +} + +func fail(format string, args ...any) { + fmt.Fprintf(os.Stderr, "bench: "+format+"\n", args...) + os.Exit(1) +} 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/internal/dag/dag.go b/internal/dag/dag.go index 6fa0e9b..cf05302 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. @@ -1367,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]) } } @@ -1384,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]) } } @@ -1401,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/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/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 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 } diff --git a/prover/prover.go b/prover/prover.go index 6bf088b..39a1091 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -27,9 +27,11 @@ import ( "github.com/consensys/loom/field" "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" + "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,113 @@ func (pr *proverRuntime) ExecuteSteps() error { return nil } +// 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 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 + } + // 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 +} + +// 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 + } + 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 = computeExtAIRQuotientChunks(pr.t.Base, pr.t.Ext, vrel, N, module.D, &pr.domainCache) + } else { + baseChunks, err = computeBaseAIRQuotientChunks(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) + } + for i, c := range extChunks { + name := constants.QuotientChunkName(moduleName, i) + pr.airTrace.SetExt(name, c) + } + 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 @@ -543,13 +589,6 @@ 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)) - } - for chunkName, chunkPoly := range pr.airTrace.Ext { - pr.Proof.SetValueAtZetaExt(chunkName, poly.ExtEvaluateAtExt(chunkPoly, chunkDomains[chunkName], pr.zeta)) - } - return nil } @@ -563,36 +602,111 @@ 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++ { + + // trace polynomials + 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 + 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}) } - p, ok := pr.t.Base[leaf.Name] - if !ok { - return fmt.Errorf("ComputeEvaluationsAtZeta: column %q not found in trace", leaf.Name) + + // 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 } - 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 @@ -616,6 +730,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 @@ -640,21 +755,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) @@ -681,15 +785,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)