perf: fan out ComputeEvaluationsAtZeta across columns

internal/poly/ext.go

@@ -82,15 +82,17 @@ 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])
 	}
 	_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,15 +108,15 @@ 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]
 	}
 	_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-- {

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)
+		}
+	}
+}

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