perf: share precomputed zeta-powers across ComputeEvaluationsAtZeta

internal/poly/ext.go

@@ -21,6 +21,7 @@ import (
 	"github.com/consensys/gnark-crypto/field/koalabear"
 	ext "github.com/consensys/gnark-crypto/field/koalabear/extensions"
 	"github.com/consensys/gnark-crypto/field/koalabear/fft"
+	"github.com/consensys/loom/internal/parallel"
 )
 
 // ExtPolynomial is a polynomial whose coefficients/evaluations live in the
@@ -129,6 +130,102 @@ func ExtEvaluateAtExt(p ExtPolynomial, d *fft.Domain, zeta ext.E4, fftOpts ...ff
 	return res
 }
 
+// BuildZPowBitReversed precomputes (zeta^bitRev(0), zeta^bitRev(1), …,
+// zeta^bitRev(n-1)) so that the SIMD InnerProduct[ByElement] can replace
+// the Horner loop in {Ext,}EvaluateAtExtWithZPow. The result is in
+// bit-reversed-canonical order to match the post-DIF FFTInverse layout
+// used by those functions.
+//
+// n must be a power of two. The work scales as O(n) ext-muls and is
+// chunked across goroutines using PowUint64 to seed each chunk
+// independently, so build cost amortises well when shared across tasks
+// that evaluate at the same zeta.
+func BuildZPowBitReversed(zeta ext.E4, n int) []ext.E4 {
+	if n <= 0 {
+		return nil
+	}
+	out := make([]ext.E4, n)
+	if n == 1 {
+		out[0].SetOne()
+		return out
+	}
+	nn := uint64(64 - bits.TrailingZeros64(uint64(n)))
+
+	// First sweep produces zeta^i in normal order using parallel chunks
+	// seeded with PowUint64Ext (binary exponentiation for ext.E4).
+	normal := make([]ext.E4, n)
+	parallel.ExecuteWithThreshold(n, zetaPowParallelThreshold, func(start, end int) {
+		acc := powUint64Ext(zeta, uint64(start))
+		for i := start; i < end; i++ {
+			normal[i].Set(&acc)
+			acc.Mul(&acc, &zeta)
+		}
+	})
+
+	// Permute into bit-reversed order.
+	parallel.ExecuteWithThreshold(n, zetaPowParallelThreshold, func(start, end int) {
+		for i := start; i < end; i++ {
+			iRev := bits.Reverse64(uint64(i)) >> nn
+			out[iRev] = normal[i]
+		}
+	})
+	return out
+}
+
+const zetaPowParallelThreshold = 1 << 12
+
+// powUint64Ext returns base^exp via binary exponentiation.
+func powUint64Ext(base ext.E4, exp uint64) ext.E4 {
+	var res ext.E4
+	res.SetOne()
+	b := base
+	for exp != 0 {
+		if exp&1 == 1 {
+			res.Mul(&res, &b)
+		}
+		b.Mul(&b, &b)
+		exp >>= 1
+	}
+	return res
+}
+
+// EvaluateAtExtWithZPow is EvaluateAtExt where the bit-reversed powers of
+// zeta have been precomputed by the caller (typically once per (n, zeta)
+// across many polynomials). Replaces the per-Horner ext-mul chain with a
+// single SIMD InnerProductByElement on the FFTInverse output.
+func EvaluateAtExtWithZPow(p Polynomial, d *fft.Domain, zPowBitRev []ext.E4, fftOpts ...fft.Option) ext.E4 {
+	n := len(p)
+	if n == 1 {
+		return liftBaseToExt(p[0])
+	}
+	if len(zPowBitRev) != n {
+		panic("EvaluateAtExtWithZPow: zPow length must equal len(p)")
+	}
+	_p := getBuf(n)
+	copy(_p, p)
+	d.FFTInverse(_p, fft.DIF, fftOpts...)
+	res := ext.Vector(zPowBitRev).InnerProductByElement(koalabear.Vector(_p))
+	putBuf(_p)
+	return res
+}
+
+// ExtEvaluateAtExtWithZPow is the ext-rail counterpart of EvaluateAtExtWithZPow.
+func ExtEvaluateAtExtWithZPow(p ExtPolynomial, d *fft.Domain, zPowBitRev []ext.E4, fftOpts ...fft.Option) ext.E4 {
+	n := len(p)
+	if n == 1 {
+		return p[0]
+	}
+	if len(zPowBitRev) != n {
+		panic("ExtEvaluateAtExtWithZPow: zPow length must equal len(p)")
+	}
+	_p := getExtBuf(n)
+	copy(_p, p)
+	d.FFTInverseExt(_p, fft.DIF, fftOpts...)
+	res := ext.Vector(zPowBitRev).InnerProduct(ext.Vector(_p))
+	putExtBuf(_p)
+	return res
+}
+
 // DeepQuotientExt computes q(X) = (v - p(X)) / (z - X) for an extension-field
 // polynomial p in Lagrange normal form over d. The domain points remain
 // base-field roots of unity and are lifted into E4 for the denominator.

internal/poly/ext_test.go

@@ -151,6 +151,68 @@ func TestExtEvaluateAtExt(t *testing.T) {
 	}
 }
 
+func TestEvaluateAtExtWithZPowMatchesHorner(t *testing.T) {
+	for _, n := range []int{2, 4, 8, 16, 64} {
+		t.Run("n="+itoa(n), func(t *testing.T) {
+			p := make(Polynomial, n)
+			for i := range p {
+				p[i].SetUint64(uint64(0xabc123 + i))
+			}
+			d := fft.NewDomain(uint64(n))
+			for _, zeta := range []ext.E4{
+				e4FromU64(2, 3, 5, 7),
+				e4FromU64(1, 0, 0, 0),
+				e4FromU64(0xdeadbeef, 0xcafef00d, 1, 42),
+			} {
+				want := EvaluateAtExt(p, d, zeta)
+				zPow := BuildZPowBitReversed(zeta, n)
+				got := EvaluateAtExtWithZPow(p, d, zPow)
+				if !got.Equal(&want) {
+					t.Fatalf("n=%d zeta=%s: WithZPow=%s Horner=%s", n, zeta.String(), got.String(), want.String())
+				}
+			}
+		})
+	}
+}
+
+func TestExtEvaluateAtExtWithZPowMatchesHorner(t *testing.T) {
+	for _, n := range []int{2, 4, 8, 16, 64} {
+		t.Run("n="+itoa(n), func(t *testing.T) {
+			p := make(ExtPolynomial, n)
+			for i := range p {
+				p[i] = e4FromU64(uint64(i+1), uint64(2*i+3), uint64(3*i+5), uint64(5*i+7))
+			}
+			d := fft.NewDomain(uint64(n))
+			for _, zeta := range []ext.E4{
+				e4FromU64(2, 3, 5, 7),
+				e4FromU64(1, 0, 0, 0),
+				e4FromU64(11, 22, 33, 44),
+			} {
+				want := ExtEvaluateAtExt(p, d, zeta)
+				zPow := BuildZPowBitReversed(zeta, n)
+				got := ExtEvaluateAtExtWithZPow(p, d, zPow)
+				if !got.Equal(&want) {
+					t.Fatalf("n=%d zeta=%s: WithZPow=%s Horner=%s", n, zeta.String(), got.String(), want.String())
+				}
+			}
+		})
+	}
+}
+
+func itoa(n int) string {
+	if n == 0 {
+		return "0"
+	}
+	var b [20]byte
+	i := len(b)
+	for n > 0 {
+		i--
+		b[i] = byte('0' + n%10)
+		n /= 10
+	}
+	return string(b[i:])
+}
+
 func TestDeepQuotientExt(t *testing.T) {
 	coeffs := ExtPolynomial{
 		e4FromU64(1, 2, 3, 4),

prover/prover.go

@@ -694,16 +694,52 @@ func (pr *proverRuntime) ComputeEvaluationsAtZeta() error {
 		}
 	}
 
+	// Precompute zeta-powers per distinct (n, evalPoint). Tasks evaluating the
+	// same poly size at the same point — which is the common case (no shifts,
+	// or one shift class) — share the bit-reversed power table and the per-
+	// task work drops to one SIMD InnerProduct{,ByElement}.
+	type zPowKey struct {
+		n    int
+		zeta ext.E4
+	}
+	zPowKeyOf := func(t task) zPowKey {
+		n := len(t.basePoly)
+		if t.extPoly != nil {
+			n = len(t.extPoly)
+		}
+		return zPowKey{n: n, zeta: t.evalPoint}
+	}
+	zPowCache := make(map[zPowKey][]ext.E4)
+	for _, t := range tasks {
+		k := zPowKeyOf(t)
+		if k.n <= 1 {
+			continue // EvaluateAtExt has a constant-poly fast path
+		}
+		if _, ok := zPowCache[k]; !ok {
+			zPowCache[k] = poly.BuildZPowBitReversed(k.zeta, k.n)
+		}
+	}
+
 	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]
+			k := zPowKeyOf(t)
+			zPow := zPowCache[k]
 			if t.extPoly != nil {
-				values[i] = poly.ExtEvaluateAtExt(t.extPoly, t.D, t.evalPoint, fftOpt)
+				if k.n <= 1 {
+					values[i] = poly.ExtEvaluateAtExt(t.extPoly, t.D, t.evalPoint, fftOpt)
+				} else {
+					values[i] = poly.ExtEvaluateAtExtWithZPow(t.extPoly, t.D, zPow, fftOpt)
+				}
 			} else {
-				values[i] = poly.EvaluateAtExt(t.basePoly, t.D, t.evalPoint, fftOpt)
+				if k.n <= 1 {
+					values[i] = poly.EvaluateAtExt(t.basePoly, t.D, t.evalPoint, fftOpt)
+				} else {
+					values[i] = poly.EvaluateAtExtWithZPow(t.basePoly, t.D, zPow, fftOpt)
+				}
 			}
 		}
 	})