From 390bf165084f6441b577f6730269784df1340088 Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Wed, 27 May 2026 21:35:03 +0000 Subject: [PATCH 1/2] perf: share precomputed zeta-powers across ComputeEvaluationsAtZeta tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Horner loop in EvaluateAtExt / ExtEvaluateAtExt runs sequentially inside each task: n ext-multiplies + n ext-adds + a bits.Reverse64 per row. Profiling showed this loop at 10–16% flat CPU on the synth bench (8.5s of CPU on tall) — second only to the Merkle-node permutation. Replace it with one SIMD-accelerated InnerProduct: precompute the bit-reversed powers of zeta once (shared across every task that evaluates an n-row polynomial at the same point), then the per-task inner work becomes: d.FFTInverse(_p, fft.DIF, ...) // unchanged res = ext.Vector(zPowBitRev).InnerProductByElement(_p) Both Vector.InnerProduct (ext × ext) and Vector.InnerProductByElement (ext × base) have AVX-512 paths in gnark-crypto. Powers are computed in parallel via chunked binary-exponentiation seeding (powUint64Ext) and permuted into bit-reversed order to match the FFTInverse DIF output, so EvaluateAtExtWithZPow needs no extra BitReverse pass. Numbers, 32-core EPYC 9R45, synth (100M cells), median of 3: shape before after speedup tall (2^22 × 24) 5.55 s 4.82 s 1.15× wide (2^16 × 1536) 1.25 s 1.20 s 1.04× Phase breakdown — evaluations-at-zeta: tall 450 ms → 149 ms 3.0× wide 57 ms → 7 ms 8.1× Tests: go test ./... + -race ./internal/poly ./prover. Existing prover end-to-end tests exercise Verify on every Prove; an off-by-one in the power-table order or a wrong-rail inner-product would surface. Co-Authored-By: Claude Opus 4.7 (1M context) --- internal/poly/ext.go | 97 ++++++++++++++++++++++++++++++++++++++++++++ prover/prover.go | 40 +++++++++++++++++- 2 files changed, 135 insertions(+), 2 deletions(-) diff --git a/internal/poly/ext.go b/internal/poly/ext.go index f114157..58d9f08 100644 --- a/internal/poly/ext.go +++ b/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. diff --git a/prover/prover.go b/prover/prover.go index aef7fcd..70f5141 100644 --- a/prover/prover.go +++ b/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) + } } } }) From 9d8a728f01c326201c58b04903290a6aec3a3fe4 Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Wed, 27 May 2026 22:18:25 +0000 Subject: [PATCH 2/2] review: cover EvaluateAtExtWithZPow with unit tests PR #13 review feedback (Copilot): the new bit-reversed power-table evaluators only had end-to-end coverage through the prover. Add internal/poly tests that build zPow and cross-check {Ext,}EvaluateAtExtWithZPow against the existing Horner evaluators across multiple sizes and zeta values (incl. identity / mixed coords). --- internal/poly/ext_test.go | 62 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/internal/poly/ext_test.go b/internal/poly/ext_test.go index 826fbd7..1b9a155 100644 --- a/internal/poly/ext_test.go +++ b/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),