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