From 6fb882a5641ed69a6766583e39538df94e2a789c Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Wed, 27 May 2026 20:40:50 +0000 Subject: [PATCH 1/2] perf: parallelize DEEP quotient assembly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the per-shift serial loop in ComputeDeepQuotient with a row-chunked parallel pass that fuses C_s materialisation, denominator inversion, and the DEEP quotient add into a single sweep: deepQuotient[x] += (v_s - Σ_k scales_k · cols_k[x]) / (z_s - ω^x) Three transformations make this fan-out possible: 1. The serial alpha-power chain (alphaAcc *= pr.alpha across columns and shifts) is unrolled up front into a per-column scales slice on each shift bundle, so per-row work has no cross-column dependency. 2. Constant (len-1) columns are folded into a single constContrib term per bundle so the per-row loop only touches real-width columns. 3. Denominators (z_s - ω^x) for the chunk are inverted via ext.BatchInvertE4 (one inversion + 3·chunkLen ext mul, instead of chunkLen ext inversions). The previous DeepQuotientExt helper is no longer reached from the prover and DeepQuotientExt itself becomes dead — kept as it's still exported. Drive-by: ω^shift now uses Element.Exp (binary exponentiation) instead of an O(shift) manual multiply loop. Numbers on the synth shapes (32 cores, KoalaBear + Poseidon2 + blowup 4 + 4 queries), median of 3 runs, applied on top of the step-1 and step-2 PRs: shape (cells) step-1+2 +step-3 speedup tall (2^20 rows × 24) 2.05 s 1.36 s 1.51× wide (2^14 rows × 1536) 0.288 s 0.285 s ~1.01× The wide shape was never DEEP-bound (8% of step-1+2 wall); the win shows up where DEEP assembly dominates. DEEP phase share on tall: 47% → 20%. vs Plonky3 (same shapes, same primitives): shape Plonky3 loom (step-1+2) loom (step-1+2+3) tall 1.56 s 2.05 s 1.36 s (1.15× faster than P3) wide 0.40 s 0.288 s 0.285 s (1.40× faster than P3) prover/bench_deep_quotient_test.go adds BenchmarkProveWide and BenchmarkProveTall so the wide- and tall-shape improvements stay visible in `go test -bench` without depending on the bench/synth command from PR #7. Co-Authored-By: Claude Opus 4.7 (1M context) --- prover/bench_deep_quotient_test.go | 92 +++++++++++ prover/prover.go | 247 +++++++++++++++++++---------- 2 files changed, 254 insertions(+), 85 deletions(-) create mode 100644 prover/bench_deep_quotient_test.go diff --git a/prover/bench_deep_quotient_test.go b/prover/bench_deep_quotient_test.go new file mode 100644 index 0000000..796c60e --- /dev/null +++ b/prover/bench_deep_quotient_test.go @@ -0,0 +1,92 @@ +// 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" +) + +// buildSynthProveInputs constructs a single-module synthetic AIR of the +// requested shape: log2_rows rows, repetitions × 3 columns, one deg-2 +// row-local constraint a*b - c = 0 per repetition. Used by the wide- and +// tall-shape micro benchmarks. +func buildSynthProveInputs(log2Rows, repetitions int) (board.Program, trace.Trace) { + 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)) + b := expr.Col(col(3*k + 1)) + c := expr.Col(col(3*k + 2)) + m.AssertZero(a.Mul(b).Sub(c)) + } + builder.AddModule(m) + program, err := board.Compile(&builder) + if err != nil { + panic(err) + } + + t := trace.New(3 * repetitions) + for k := 0; k < repetitions; k++ { + a := make([]koalabear.Element, rows) + b := make([]koalabear.Element, rows) + c := make([]koalabear.Element, rows) + for i := 0; i < rows; i++ { + a[i].SetUint64(uint64(i + 1 + k)) + b[i].SetUint64(uint64(2*i + 3 + k)) + c[i].Mul(&a[i], &b[i]) + } + t.SetBase(col(3*k), a) + t.SetBase(col(3*k+1), b) + t.SetBase(col(3*k+2), c) + } + return program, t +} + +// BenchmarkProveWide exercises the wide-trace shape (single module, modest N, +// many columns) where the DEEP quotient column accumulation dominates. +func BenchmarkProveWide(b *testing.B) { + program, t := buildSynthProveInputs(12, 256) + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := Prove(t, setup.ProvingKey{}, nil, program) + if err != nil { + b.Fatalf("Prove: %v", err) + } + } +} + +// BenchmarkProveTall exercises the tall-trace shape (single module, large N, +// few columns) — the regime where the per-row DEEP quotient assembly is the +// largest serial chunk before this PR's parallelisation. +func BenchmarkProveTall(b *testing.B) { + program, t := buildSynthProveInputs(18, 8) + 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..bae9508 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -15,6 +15,7 @@ package prover import ( "fmt" + "math/big" "sort" "sync" @@ -682,29 +683,26 @@ func (pr *proverRuntime) ComputeEvaluationsAtZeta() error { return nil } -func addScaledExtColumn(dst, col, scratch poly.ExtPolynomial, alpha *ext.E4) { - if len(col) == 1 { - var term ext.E4 - term.Mul(&col[0], alpha) - for i := range dst { - dst[i].Add(&dst[i], &term) - } - return - } - ext.Vector(scratch).ScalarMul(ext.Vector(col), alpha) - ext.Vector(dst).Add(ext.Vector(dst), ext.Vector(scratch)) -} - -func addScaledBaseColumn(dst poly.ExtPolynomial, col poly.Polynomial, alpha *ext.E4) { - if len(col) == 1 { - var term ext.E4 - term.MulByElement(alpha, &col[0]) - for i := range dst { - dst[i].Add(&dst[i], &term) - } - return - } - ext.Vector(dst).MulAccByElement(col, alpha) +// deepQuotientBundle aggregates everything needed to add one DEEP-quotient +// shift block's contribution to deepQuotient[*]: +// +// deepQuotient[x] += (vs - sum_k(scales_k * cols_k[x])) / (zs - omega^x) +// +// The serial alpha-power chain in the original code is unrolled into the +// scales slices below, so the per-row loop has no cross-column data dependency +// and can be chunked across goroutines. +type deepQuotientBundle struct { + zs ext.E4 // shifted evaluation point + vs ext.E4 // alpha-weighted sum of evaluations at zs + + // constContrib is the contribution from constant (len-1) columns, summed + // in advance so the per-row loop only iterates real-width columns. + constContrib ext.E4 + + extCols []poly.ExtPolynomial + extScales []ext.E4 + baseCols []poly.Polynomial + baseScales []ext.E4 } func (pr *proverRuntime) ComputeDeepQuotient() error { @@ -730,78 +728,47 @@ func (pr *proverRuntime) ComputeDeepQuotient() error { alphaAcc.SetOne() domainN := domainBySize[N] - scratch := make(poly.ExtPolynomial, N) + // ---- Phase 1: vanishing-relation columns, one bundle per shift ---- + bundles := make([]deepQuotientBundle, 0, len(dqLayout.Shifts[i])+1) for j, shift := range dqLayout.Shifts[i] { var omegaShift koalabear.Element - omegaShift.SetOne() - for k := 0; k < shift; k++ { - omegaShift.Mul(&omegaShift, &domainN.Generator) - } - z_s := pr.zeta - z_s.MulByElement(&z_s, &omegaShift) - - C_s := make(poly.ExtPolynomial, N) - var v_s ext.E4 - names := dqLayout.Names[i][j] - keys := dqLayout.Keys[i][j] - for k := range names { - evalAtZ, ok := pr.Proof.ValueAtZetaExt(keys[k]) - if !ok { - return fmt.Errorf("ComputeDeepQuotient: %q not found in ValuesAtZeta", keys[k]) - } - colExt, hasExt := pr.t.Ext[names[k]] - colBase, hasBase := pr.t.Base[names[k]] - if !hasExt && !hasBase { - return fmt.Errorf("ComputeDeepQuotient: column %q not found in trace", names[k]) - } - if hasExt { - addScaledExtColumn(C_s, colExt, scratch, &alphaAcc) - } else { - addScaledBaseColumn(C_s, colBase, &alphaAcc) - } - var term ext.E4 - term.Mul(&evalAtZ, &alphaAcc) - v_s.Add(&v_s, &term) - alphaAcc.Mul(&alphaAcc, &pr.alpha) - } - - DQ_s := poly.DeepQuotientExt(C_s, v_s, z_s, domainN) - for x := 0; x < N; x++ { - deepQuotient[x].Add(&deepQuotient[x], &DQ_s[x]) + omegaShift.Exp(domainN.Generator, big.NewInt(int64(shift))) + zs := pr.zeta + zs.MulByElement(&zs, &omegaShift) + + b, nextAlpha, err := pr.buildDeepQuotientBundle( + zs, + dqLayout.Names[i][j], + dqLayout.Keys[i][j], + pr.t.Base, pr.t.Ext, + alphaAcc, + ) + if err != nil { + return err } + bundles = append(bundles, b) + alphaAcc = nextAlpha } + // ---- Phase 2: AIR chunks at z=zeta (no shift) ---- if len(dqLayout.AIRChunks[i]) > 0 { - C_s := make(poly.ExtPolynomial, N) - var v_s ext.E4 - for _, chunkName := range dqLayout.AIRChunks[i] { - evalAtZ, ok := pr.Proof.ValueAtZetaExt(chunkName) - if !ok { - return fmt.Errorf("ComputeDeepQuotient: %q not found in ValuesAtZeta", chunkName) - } - chunkExt, hasExt := pr.airTrace.Ext[chunkName] - chunkBase, hasBase := pr.airTrace.Base[chunkName] - if !hasExt && !hasBase { - return fmt.Errorf("ComputeDeepQuotient: AIR chunk %q not found in trace", chunkName) - } - if hasExt { - addScaledExtColumn(C_s, chunkExt, scratch, &alphaAcc) - } else { - addScaledBaseColumn(C_s, chunkBase, &alphaAcc) - } - var term ext.E4 - term.Mul(&evalAtZ, &alphaAcc) - v_s.Add(&v_s, &term) - alphaAcc.Mul(&alphaAcc, &pr.alpha) - } - - DQ_air := poly.DeepQuotientExt(C_s, v_s, pr.zeta, domainN) - for x := 0; x < N; x++ { - deepQuotient[x].Add(&deepQuotient[x], &DQ_air[x]) + b, nextAlpha, err := pr.buildDeepQuotientBundle( + pr.zeta, + dqLayout.AIRChunks[i], + dqLayout.AIRChunks[i], // keys == names for AIR chunks + pr.airTrace.Base, pr.airTrace.Ext, + alphaAcc, + ) + if err != nil { + return err } + bundles = append(bundles, b) + alphaAcc = nextAlpha } + accumulateDeepQuotient(deepQuotient, bundles, domainN) + deepQuotients[N] = deepQuotient } @@ -836,6 +803,116 @@ func (pr *proverRuntime) ComputeDeepQuotient() error { return nil } +// buildDeepQuotientBundle collects column data + per-column alpha scale factors +// for one shift block and returns the bundle plus the next alpha accumulator. +// Constant (len-1) columns are folded into bundle.constContrib so the per-row +// loop only iterates real-width columns. +// +// names and keys are parallel: names[k] looks up the trace polynomial, +// keys[k] looks up the value at zeta in pr.Proof (for AIR chunks the two +// coincide). +func (pr *proverRuntime) buildDeepQuotientBundle( + zs ext.E4, + names, keys []string, + traceBase map[string]poly.Polynomial, + traceExt map[string]poly.ExtPolynomial, + alphaStart ext.E4, +) (deepQuotientBundle, ext.E4, error) { + b := deepQuotientBundle{zs: zs} + scale := alphaStart + for k, name := range names { + evalAtZ, ok := pr.Proof.ValueAtZetaExt(keys[k]) + if !ok { + return deepQuotientBundle{}, ext.E4{}, fmt.Errorf("ComputeDeepQuotient: %q not found in ValuesAtZeta", keys[k]) + } + colExt, hasExt := traceExt[name] + colBase, hasBase := traceBase[name] + if !hasExt && !hasBase { + return deepQuotientBundle{}, ext.E4{}, fmt.Errorf("ComputeDeepQuotient: column %q not found in trace", name) + } + + var term ext.E4 + term.Mul(&evalAtZ, &scale) + b.vs.Add(&b.vs, &term) + + switch { + case hasExt && len(colExt) == 1: + term.Mul(&colExt[0], &scale) + b.constContrib.Add(&b.constContrib, &term) + case hasExt: + b.extCols = append(b.extCols, colExt) + b.extScales = append(b.extScales, scale) + case len(colBase) == 1: + term.MulByElement(&scale, &colBase[0]) + b.constContrib.Add(&b.constContrib, &term) + default: + b.baseCols = append(b.baseCols, colBase) + b.baseScales = append(b.baseScales, scale) + } + + scale.Mul(&scale, &pr.alpha) + } + return b, scale, nil +} + +// accumulateDeepQuotient adds every bundle's DEEP-quotient contribution into +// deepQuotient using a single row-chunked parallel pass: each chunk computes +// (z_s - omega^x)^-1 in batch (BatchInvertE4), then sweeps every bundle to +// fold its (vs - sum_k scale_k * col_k[x]) numerator and add to deepQuotient[x]. +// One pass amortises C_s materialisation away — only a row-sized denominator +// buffer per chunk is allocated. +func accumulateDeepQuotient(deepQuotient poly.ExtPolynomial, bundles []deepQuotientBundle, domain *fft.Domain) { + N := len(deepQuotient) + if N == 0 || len(bundles) == 0 { + return + } + + parallel.Execute(N, func(start, end int) { + chunkLen := end - start + + // Compute denominators (z_s - omega^x) for every bundle, batch invert. + denoms := make([]ext.E4, chunkLen*len(bundles)) + var omegaX koalabear.Element + if start == 0 { + omegaX.SetOne() + } else { + omegaX.Exp(domain.Generator, big.NewInt(int64(start))) + } + for x := 0; x < chunkLen; x++ { + var omegaExt ext.E4 + omegaExt.Lift(&omegaX) + for b := range bundles { + denoms[b*chunkLen+x].Sub(&bundles[b].zs, &omegaExt) + } + omegaX.Mul(&omegaX, &domain.Generator) + } + invs := ext.BatchInvertE4(denoms) + + // Sweep bundles into deepQuotient row by row. + for b := range bundles { + bun := &bundles[b] + invRow := invs[b*chunkLen : (b+1)*chunkLen] + for x := start; x < end; x++ { + Cx := bun.constContrib + for k, col := range bun.extCols { + var term ext.E4 + term.Mul(&bun.extScales[k], &col[x]) + Cx.Add(&Cx, &term) + } + for k, col := range bun.baseCols { + var term ext.E4 + term.MulByElement(&bun.baseScales[k], &col[x]) + Cx.Add(&Cx, &term) + } + var num, dqx ext.E4 + num.Sub(&bun.vs, &Cx) + dqx.Mul(&num, &invRow[x-start]) + deepQuotient[x].Add(&deepQuotient[x], &dqx) + } + } + }) +} + // openWMerkleAt opens a WMerkleTree at the leaf index corresponding to FRI // query position `s`, reduced mod the tree's paired-leaf count (= // encoded_size/2 = RATE·N/2). From 7bbb828c439fb7d490f911e5521ceed9e8c1c335 Mon Sep 17 00:00:00 2001 From: Thomas Piellard Date: Thu, 28 May 2026 16:34:46 +0200 Subject: [PATCH 2/2] style: added comment --- prover/prover.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/prover/prover.go b/prover/prover.go index bae9508..ca3f2f5 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -772,6 +772,8 @@ func (pr *proverRuntime) ComputeDeepQuotient() error { deepQuotients[N] = deepQuotient } + // The DEEP quotients chunks are computed and grouped in decreasing size order, we can build the levels to call + // multi degree FRI levels := make([]fri.Level, len(sizes)) for li, N := range sizes { encoder := reedsolomon.NewEncoderWithDomainCache(uint64(constants.RATE)*uint64(N), &pr.domainCache)