From 3dfefbf77b480884ac19a713a69259e1c2b76150 Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Wed, 27 May 2026 20:28:11 +0000 Subject: [PATCH] perf: fan out ComputeEvaluationsAtZeta across columns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each per-column evaluation at ζ is an independent FFTInverse + Horner pass. The previous parallel.Execute only fanned out over modules — single-module programs (the common shape, including the synth wide benchmark) saw every column evaluated serially in one goroutine. Gather every (poly, evaluation point) pair across all modules into a single task list and parallel.Execute over it. Inner FFTs run with WithNbTasks(1) since the column-level fan-out already saturates the CPUs; the FFT option is forwarded through new variadic fft.Option args on poly.EvaluateAtExt and poly.ExtEvaluateAtExt. Drive-by: replace the O(shift) manual mul loop for the rotation twiddle with Element.Exp (binary exponentiation). Wide-trace benchmark (synth, log2_rows=12, repetitions=256, 32 cores): BenchmarkProveWide 232.8 ms → 178.0 ms (1.31×) prover/bench_evals_test.go is added so the wide-shape improvement remains visible in `go test -bench`. Co-Authored-By: Claude Opus 4.7 (1M context) --- internal/poly/ext.go | 10 +-- prover/bench_evals_test.go | 78 ++++++++++++++++++++++ prover/prover.go | 130 ++++++++++++++++++++----------------- 3 files changed, 155 insertions(+), 63 deletions(-) create mode 100644 prover/bench_evals_test.go diff --git a/internal/poly/ext.go b/internal/poly/ext.go index 2ac5887..f114157 100644 --- a/internal/poly/ext.go +++ b/internal/poly/ext.go @@ -82,7 +82,9 @@ func MulExt(P1, P2 ExtPolynomial) (ExtPolynomial, error) { // form over d, at the extension-field point zeta. Base coefficients are lifted // during Horner evaluation. // EvaluateAtExt assumes len(p) is a power of two; behaviour is undefined otherwise. -func EvaluateAtExt(p Polynomial, d *fft.Domain, zeta ext.E4) ext.E4 { +// Optional fftOpts are forwarded to the internal FFT (e.g. fft.WithNbTasks(1) +// when EvaluateAtExt is itself called from inside a parallel.Execute loop). +func EvaluateAtExt(p Polynomial, d *fft.Domain, zeta ext.E4, fftOpts ...fft.Option) ext.E4 { n := len(p) if n == 1 { return liftBaseToExt(p[0]) @@ -90,7 +92,7 @@ func EvaluateAtExt(p Polynomial, d *fft.Domain, zeta ext.E4) ext.E4 { _p := getBuf(n) copy(_p, p) nn := uint64(64 - bits.TrailingZeros64(uint64(n))) - d.FFTInverse(_p, fft.DIF) + d.FFTInverse(_p, fft.DIF, fftOpts...) var res ext.E4 for i := n - 1; i >= 0; i-- { @@ -106,7 +108,7 @@ func EvaluateAtExt(p Polynomial, d *fft.Domain, zeta ext.E4) ext.E4 { // ExtEvaluateAtExt evaluates an extension-field polynomial p, stored in // Lagrange normal form over d, at the extension-field point zeta. -func ExtEvaluateAtExt(p ExtPolynomial, d *fft.Domain, zeta ext.E4) ext.E4 { +func ExtEvaluateAtExt(p ExtPolynomial, d *fft.Domain, zeta ext.E4, fftOpts ...fft.Option) ext.E4 { n := len(p) if n == 1 { return p[0] @@ -114,7 +116,7 @@ func ExtEvaluateAtExt(p ExtPolynomial, d *fft.Domain, zeta ext.E4) ext.E4 { _p := getExtBuf(n) copy(_p, p) nn := uint64(64 - bits.TrailingZeros64(uint64(n))) - d.FFTInverseExt(_p, fft.DIF) + d.FFTInverseExt(_p, fft.DIF, fftOpts...) var res ext.E4 for i := n - 1; i >= 0; i-- { diff --git a/prover/bench_evals_test.go b/prover/bench_evals_test.go new file mode 100644 index 0000000..7433d66 --- /dev/null +++ b/prover/bench_evals_test.go @@ -0,0 +1,78 @@ +// Copyright Consensys Software Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on +// an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the +// specific language governing permissions and limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package prover + +import ( + "fmt" + "testing" + + "github.com/consensys/gnark-crypto/field/koalabear" + "github.com/consensys/loom/board" + "github.com/consensys/loom/expr" + "github.com/consensys/loom/setup" + "github.com/consensys/loom/trace" +) + +// BenchmarkProveWide exercises the wide-trace shape (single module, modest N, +// many columns) where evaluations-at-zeta is a sizeable share of total prove +// wall time. The R parameter scales width = 3*R and so the number of per-column +// evaluations at zeta — making this a useful proxy for the ComputeEvaluationsAtZeta +// fan-out. +func BenchmarkProveWide(b *testing.B) { + const ( + log2Rows = 12 + repetitions = 256 + ) + rows := 1 << log2Rows + + col := func(i int) string { return fmt.Sprintf("synth.c_%d", i) } + + builder := board.NewBuilder() + m := board.NewModule("synth") + m.N = rows + for k := 0; k < repetitions; k++ { + a := expr.Col(col(3 * k)) + bb := expr.Col(col(3*k + 1)) + c := expr.Col(col(3*k + 2)) + m.AssertZero(a.Mul(bb).Sub(c)) + } + builder.AddModule(m) + program, err := board.Compile(&builder) + if err != nil { + b.Fatalf("Compile: %v", err) + } + + t := trace.New(3 * repetitions) + for k := 0; k < repetitions; k++ { + a := make([]koalabear.Element, rows) + bs := make([]koalabear.Element, rows) + c := make([]koalabear.Element, rows) + for i := 0; i < rows; i++ { + a[i].SetUint64(uint64(i + 1 + k)) + bs[i].SetUint64(uint64(2*i + 3 + k)) + c[i].Mul(&a[i], &bs[i]) + } + t.SetBase(col(3*k), a) + t.SetBase(col(3*k+1), bs) + t.SetBase(col(3*k+2), c) + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := Prove(t, setup.ProvingKey{}, nil, program) + if err != nil { + b.Fatalf("Prove: %v", err) + } + } +} diff --git a/prover/prover.go b/prover/prover.go index c4e505c..b3df5a0 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -15,6 +15,7 @@ package prover import ( "fmt" + "math/big" "sort" "sync" @@ -594,6 +595,12 @@ func (pr *proverRuntime) ComputeAIRQuotients() error { // ComputeEvaluationsAtZeta computes the evaluations at zeta of every polynomial // appearing in every vanishing relation of every module. +// +// Each per-column evaluation is an independent FFTInverse + Horner pass, so we +// gather every (poly, evaluation point) pair across all modules into a single +// task list and fan it out across goroutines. Each inner FFT is capped to 1 +// goroutine (fft.WithNbTasks(1)) since the column-level fan-out already +// saturates the CPUs. func (pr *proverRuntime) ComputeEvaluationsAtZeta() error { config := expr.NewConfig( expr.WithoutLagrangeColumns(), @@ -602,9 +609,12 @@ func (pr *proverRuntime) ComputeEvaluationsAtZeta() error { expr.WithoutPublicColumns(), ) - type zetaEval struct { - key string - val ext.E4 + type task struct { + key string + basePoly poly.Polynomial // exactly one of basePoly / extPoly is non-nil + extPoly poly.ExtPolynomial // (single-element constant columns can take either rail) + D *fft.Domain + evalPoint ext.E4 } moduleNames := make([]string, 0, len(pr.program.Modules)) @@ -613,70 +623,72 @@ func (pr *proverRuntime) ComputeEvaluationsAtZeta() error { } 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) - } - - 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}) - } + // Compute leaves up front so we can size the task slice without re-walking + // the DAG, and so the parallel section has nothing to allocate. + moduleLeaves := make([][]*expr.Leaf, len(moduleNames)) + taskCap := 0 + for i, moduleName := range moduleNames { + leaves := pr.program.Modules[moduleName].VanishingRelation.LeavesFull(config) + moduleLeaves[i] = leaves + taskCap += len(leaves) + } + tasks := make([]task, 0, taskCap) + for i, moduleName := range moduleNames { + module := pr.program.Modules[moduleName] + leaves := moduleLeaves[i] - // air chunks + // trace polynomials + for _, leaf := range leaves { 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 + if leaf.Type == expr.RotatedColumn { + shift := ((leaf.Shift % module.N) + module.N) % module.N + var omegaPow koalabear.Element + omegaPow.Exp(module.D.Generator, big.NewInt(int64(shift))) + evalPoint.MulByElement(&evalPoint, &omegaPow) } - results[idx] = local + if p, ok := pr.t.Ext[leaf.Name]; ok { + tasks = append(tasks, task{key: leaf.String(), extPoly: p, D: module.D, evalPoint: evalPoint}) + continue + } + p, ok := pr.t.Base[leaf.Name] + if !ok { + return fmt.Errorf("ComputeEvaluationsAtZeta: column %q not found in trace", leaf.Name) + } + tasks = append(tasks, task{key: leaf.String(), basePoly: p, D: module.D, evalPoint: evalPoint}) } - }) - for _, err := range errs { - if err != nil { - return err + + // air chunks + for i := 0; ; i++ { + chunkName := constants.QuotientChunkName(moduleName, i) + if chunk, ok := pr.airTrace.Base[chunkName]; ok { + tasks = append(tasks, task{key: chunkName, basePoly: chunk, D: module.D, evalPoint: pr.zeta}) + continue + } + if chunk, ok := pr.airTrace.Ext[chunkName]; ok { + tasks = append(tasks, task{key: chunkName, extPoly: chunk, D: module.D, evalPoint: pr.zeta}) + continue + } + break } } - for _, local := range results { - for _, ev := range local { - pr.Proof.SetValueAtZetaExt(ev.key, ev.val) + + values := make([]ext.E4, len(tasks)) + // Inner FFTs run serial: the per-column fan-out already saturates the CPUs. + fftOpt := fft.WithNbTasks(1) + parallel.Execute(len(tasks), func(start, end int) { + for i := start; i < end; i++ { + t := tasks[i] + if t.extPoly != nil { + values[i] = poly.ExtEvaluateAtExt(t.extPoly, t.D, t.evalPoint, fftOpt) + } else { + values[i] = poly.EvaluateAtExt(t.basePoly, t.D, t.evalPoint, fftOpt) + } } + }) + + for i, t := range tasks { + pr.Proof.SetValueAtZetaExt(t.key, values[i]) } return nil