perf: parallelize LDE / coset-FFT / Merkle pipeline

bench/synth/main.go

@@ -0,0 +1,177 @@
+// Command synth proves a synthetic AIR with configurable trace shape
+// (rows × cols) and reports per-phase wall time, prove/verify totals and
+// proof size. It is the loom counterpart of plonky3's prove_prime_field_31
+// example: both stacks build a deg-2 row-local constraint over a tall or
+// wide trace and use matching FRI parameters, so the two can be compared
+// apple-to-apple.
+//
+// Per row group of 3 columns (a, b, c), one constraint is enforced:
+//
+//	a * b - c = 0
+//
+// Run examples:
+//
+//	go run ./bench/synth -log2-rows 20 -repetitions 8     # tall
+//	go run ./bench/synth -log2-rows 14 -repetitions 512   # wide
+package main
+
+import (
+	"bytes"
+	"encoding/gob"
+	"flag"
+	"fmt"
+	"os"
+	"runtime"
+	"sort"
+	"sync"
+	"time"
+
+	"github.com/consensys/gnark-crypto/field/koalabear"
+	"github.com/consensys/loom"
+	"github.com/consensys/loom/board"
+	"github.com/consensys/loom/expr"
+	"github.com/consensys/loom/prover"
+	"github.com/consensys/loom/setup"
+	"github.com/consensys/loom/trace"
+	"github.com/consensys/loom/verifier"
+)
+
+var (
+	log2Rows    = flag.Int("log2-rows", 20, "log2 of trace height (rows)")
+	repetitions = flag.Int("repetitions", 16, "number of (a,b,c) tuples per row; trace width = 3 * repetitions")
+	hashName    = flag.String("hash", "poseidon2", "hash backend: poseidon2 | sha256")
+	gomaxprocs  = flag.Int("gomaxprocs", 0, "override GOMAXPROCS (0 = leave default)")
+)
+
+const moduleName = "synth"
+
+func colName(i int) string { return fmt.Sprintf("%s.c_%d", moduleName, i) }
+
+func main() {
+	flag.Parse()
+
+	if *gomaxprocs > 0 {
+		runtime.GOMAXPROCS(*gomaxprocs)
+	}
+	procs := runtime.GOMAXPROCS(0)
+	rows := 1 << *log2Rows
+	width := 3 * *repetitions
+
+	fmt.Printf("loom synth   rows=2^%d (=%d)  width=%d (3 cols × %d reps)  cells=%d  hash=%s  GOMAXPROCS=%d  NumCPU=%d\n\n",
+		*log2Rows, rows, width, *repetitions, rows*width, *hashName, procs, runtime.NumCPU())
+
+	// Build module + compile program.
+	builder := board.NewBuilder()
+	m := board.NewModule(moduleName)
+	m.N = rows
+	for k := 0; k < *repetitions; k++ {
+		a := expr.Col(colName(3 * k))
+		b := expr.Col(colName(3*k + 1))
+		c := expr.Col(colName(3*k + 2))
+		m.AssertZero(a.Mul(b).Sub(c)) // a*b - c = 0
+	}
+	builder.AddModule(m)
+	program, err := board.Compile(&builder)
+	if err != nil {
+		fail("Compile: %v", err)
+	}
+
+	// Synthesize trace (column-major). Values are derived from cheap
+	// deterministic arithmetic so trace generation isn't the bottleneck.
+	t := trace.New(width)
+	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(colName(3*k), a)
+		t.SetBase(colName(3*k+1), b)
+		t.SetBase(colName(3*k+2), c)
+	}
+
+	hashBackend := resolveHashBackend(*hashName)
+
+	// Prove with phase timings collected via the public WithPhaseCallback option.
+	type phase struct {
+		name string
+		d    time.Duration
+	}
+	var (
+		mu     sync.Mutex
+		phases []phase
+	)
+	opts := []prover.Option{
+		prover.WithHashBackend(hashBackend),
+		prover.WithPhaseCallback(func(name string, d time.Duration) {
+			mu.Lock()
+			phases = append(phases, phase{name, d})
+			mu.Unlock()
+		}),
+	}
+
+	t0 := time.Now()
+	prf, err := prover.Prove(t, setup.ProvingKey{}, nil, program, opts...)
+	if err != nil {
+		fail("Prove: %v", err)
+	}
+	proveWall := time.Since(t0)
+
+	// Serialize for size, then verify.
+	var buf bytes.Buffer
+	if err := gob.NewEncoder(&buf).Encode(&prf); err != nil {
+		fail("encode proof: %v", err)
+	}
+
+	t0 = time.Now()
+	if err := verifier.Verify(nil, setup.VerificationKey{}, program, prf,
+		verifier.WithHashBackend(hashBackend)); err != nil {
+		fail("Verify: %v", err)
+	}
+	verifyWall := time.Since(t0)
+
+	// Report.
+	fmt.Println("prove-phase breakdown:")
+	for _, p := range phases {
+		share := 100 * float64(p.d) / float64(proveWall)
+		fmt.Printf("  %-26s %s  %5.1f%%\n", p.name, fmtDur(p.d), share)
+	}
+	sort.Slice(phases, func(i, j int) bool { return phases[i].d > phases[j].d })
+
+	fmt.Printf("\nprove wall : %s\n", fmtDur(proveWall))
+	fmt.Printf("verify wall: %s\n", fmtDur(verifyWall))
+	fmt.Printf("proof size : %d B (gob)\n", buf.Len())
+	fmt.Printf("proof      : %d commitments, %d FRI levels, %d query samplings\n",
+		len(prf.Commitments), len(prf.DeepQuotientCommitment), len(prf.PointSamplings))
+}
+
+func resolveHashBackend(name string) loom.HashBackend {
+	switch name {
+	case "poseidon2":
+		return loom.Poseidon2HashBackend()
+	case "sha256":
+		return loom.SHA256HashBackend()
+	default:
+		fail("unknown hash backend %q (want poseidon2 | sha256)", name)
+		return loom.HashBackend{}
+	}
+}
+
+func fmtDur(d time.Duration) string {
+	switch {
+	case d >= time.Second:
+		return fmt.Sprintf("%.2fs", d.Seconds())
+	case d >= time.Millisecond:
+		return fmt.Sprintf("%.1fms", float64(d.Microseconds())/1000)
+	default:
+		return d.String()
+	}
+}
+
+func fail(format string, args ...any) {
+	fmt.Fprintf(os.Stderr, "synth: "+format+"\n", args...)
+	os.Exit(1)
+}

internal/commitment/commitment.go

@@ -202,18 +202,23 @@ func (rs *RSCommit) Commit(basePolys []poly.Polynomial, extPolys []poly.ExtPolyn
 		domainCache = &poly.DomainCache{}
 	}
 
-	// 1- encode every polynomial on its rail
+	// 1- encode every polynomial on its rail. Each Encode is independent
+	//    (disjoint input/output slices, shared read-only domain).
 	encodedBase := make([]poly.Polynomial, len(basePolys))
-	for i, pol := range basePolys {
-		n := len(pol)
-		encodedBase[i] = rs.Encoder.Encode(pol, domainCache.Get(uint64(n)))
-	}
+	parallel.Execute(len(basePolys), func(start, end int) {
+		for i := start; i < end; i++ {
+			pol := basePolys[i]
+			encodedBase[i] = rs.Encoder.Encode(pol, domainCache.Get(uint64(len(pol))))
+		}
+	})
 
 	encodedExt := make([]poly.ExtPolynomial, len(extPolys))
-	for i, pol := range extPolys {
-		n := len(pol)
-		encodedExt[i] = rs.Encoder.EncodeExt(pol, domainCache.Get(uint64(n)))
-	}
+	parallel.Execute(len(extPolys), func(start, end int) {
+		for i := start; i < end; i++ {
+			pol := extPolys[i]
+			encodedExt[i] = rs.Encoder.EncodeExt(pol, domainCache.Get(uint64(len(pol))))
+		}
+	})
 
 	// 2- build the merkle tree, with rs.N/2 leafs
 	// the i-th leaf is base pairs followed by extension pairs.
@@ -235,13 +240,7 @@ func (rs *RSCommit) Commit(basePolys []poly.Polynomial, extPolys []poly.ExtPolyn
 		Ext:        encodedExt,
 		PairOffset: halfN,
 	}
-	if batchHasher, ok := rs.LeafHasher.(BatchLeafHasher); ok {
-		hashLeavesBatchParallel(batchHasher, leaves, src)
-	} else {
-		parallel.Execute(halfN, func(start, end int) {
-			hashLeavesScalar(rs.LeafHasher, leaves[start:end], src, start)
-		})
-	}
+	HashLeavesParallel(rs.LeafHasher, leaves, src)
 
 	if err := tree.Build(leaves); err != nil {
 		return WMerkleTree{}, err
@@ -250,6 +249,19 @@ func (rs *RSCommit) Commit(basePolys []poly.Polynomial, extPolys []poly.ExtPolyn
 	return wTree, nil
 }
 
+// HashLeavesParallel hashes len(dst) paired leaves from src into dst, using
+// the batched leaf hasher when available (rate-16 Poseidon2 sponge) and
+// fanning the work out across goroutines.
+func HashLeavesParallel(lh LeafHasher, dst []hash.Digest, src LeafSource) {
+	if batchHasher, ok := lh.(BatchLeafHasher); ok {
+		hashLeavesBatchParallel(batchHasher, dst, src)
+		return
+	}
+	parallel.Execute(len(dst), func(start, end int) {
+		hashLeavesScalar(lh, dst[start:end], src, start)
+	})
+}
+
 func hashLeavesBatchParallel(lh BatchLeafHasher, dst []hash.Digest, src LeafSource) {
 	batchSize := lh.BatchSize()
 	if batchSize <= 0 {

internal/fri/fri.go

@@ -26,9 +26,15 @@ import (
 	fiatshamir "github.com/consensys/loom/internal/fiat-shamir"
 	"github.com/consensys/loom/internal/hash"
 	"github.com/consensys/loom/internal/merkle"
+	"github.com/consensys/loom/internal/parallel"
+	"github.com/consensys/loom/internal/poly"
 	"github.com/consensys/loom/internal/reedsolomon"
 )
 
+// foldParallelThreshold is the smallest half-layer size at which fan-out
+// across goroutines beats the precomputed-xInv seeding overhead per chunk.
+const foldParallelThreshold = 1 << 12
+
 // Params holds the FRI configuration and precomputed per-level data.
 // Build once with NewParams; reuse across many Prove/Verify calls.
 type Params struct {
@@ -855,18 +861,18 @@ func log2(n int) int {
 }
 
 // buildTreeBase builds a Merkle tree of Nⱼ/2 leaves where
-// leaf k = LeafHasher(encode(layer[k]) || encode(layer[k + Nⱼ/2])).
+// leaf k = LeafHasher(layer[k] || layer[k + Nⱼ/2]).
 func buildTreeBase(layer []koalabear.Element, lh commitment.LeafHasher, nh commitment.NodeHasher) (*merkle.Tree, error) {
 	half := len(layer) / 2
 	tree, err := merkle.New(half, nh)
 	if err != nil {
 		return nil, err
 	}
 	leaves := make([]hash.Digest, half)
-	for k := 0; k < half; k++ {
-		pair := []commitment.PairBase{{layer[k], layer[k+half]}}
-		leaves[k] = lh.HashLeaf(pair, nil)
-	}
+	commitment.HashLeavesParallel(lh, leaves, commitment.LeafSource{
+		Base:       []poly.Polynomial{layer},
+		PairOffset: half,
+	})
 	return tree, tree.Build(leaves)
 }
 
@@ -878,67 +884,65 @@ func buildTreeExt(layer []ext.E4, lh commitment.LeafHasher, nh commitment.NodeHa
 		return nil, err
 	}
 	leaves := make([]hash.Digest, half)
-	for k := 0; k < half; k++ {
-		pair := []commitment.PairExt{{layer[k], layer[k+half]}}
-		leaves[k] = lh.HashLeaf(nil, pair)
-	}
+	commitment.HashLeavesParallel(lh, leaves, commitment.LeafSource{
+		Ext:        []poly.ExtPolynomial{layer},
+		PairOffset: half,
+	})
 	return tree, tree.Build(leaves)
 }
 
 // foldLayerBase folds a base-field layer of size Nⱼ into a layer of size Nⱼ/2.
+//
+// The naive loop carries a serial dependency on xInv = GeneratorInv^i; each
+// parallel chunk seeds xInv with GeneratorInv^chunkStart so chunks run
+// independently.
 func foldLayerBase(layer []koalabear.Element, alpha koalabear.Element, domain *fft.Domain, invTwo koalabear.Element) []koalabear.Element {
 	half := len(layer) / 2
 	next := make([]koalabear.Element, half)
+	parallel.ExecuteWithThreshold(half, foldParallelThreshold, func(start, end int) {
+		xInv := poly.PowUint64(domain.GeneratorInv, uint64(start))
+		var sum, diff koalabear.Element
+		for i := start; i < end; i++ {
+			p, q := layer[i], layer[i+half]
 
-	var xInv, sum, diff koalabear.Element
-	xInv.SetOne()
-
-	for i := 0; i < half; i++ {
-		p, q := layer[i], layer[i+half]
-
-		sum.Add(&p, &q)
-		sum.Mul(&sum, &invTwo)
-
-		diff.Sub(&p, &q)
-		diff.Mul(&diff, &invTwo)
-		diff.Mul(&diff, &xInv)
-		diff.Mul(&diff, &alpha)
-
-		next[i].Add(&sum, &diff)
+			sum.Add(&p, &q)
+			sum.Mul(&sum, &invTwo)
 
-		xInv.Mul(&xInv, &domain.GeneratorInv)
-	}
+			diff.Sub(&p, &q)
+			diff.Mul(&diff, &invTwo)
+			diff.Mul(&diff, &xInv)
+			diff.Mul(&diff, &alpha)
 
+			next[i].Add(&sum, &diff)
+			xInv.Mul(&xInv, &domain.GeneratorInv)
+		}
+	})
 	return next
 }
 
-// foldLayerExt folds an extension-field layer of size Nⱼ into a layer of size
-// Nⱼ/2. Domain factors remain base-field elements and are multiplied with
-// MulByElement.
+// foldLayerExt is the extension-field counterpart of foldLayerBase; domain
+// factors stay in the base field and are multiplied via MulByElement.
 func foldLayerExt(layer []ext.E4, alpha ext.E4, domain *fft.Domain, invTwo koalabear.Element) []ext.E4 {
 	half := len(layer) / 2
 	next := make([]ext.E4, half)
+	parallel.ExecuteWithThreshold(half, foldParallelThreshold, func(start, end int) {
+		xInv := poly.PowUint64(domain.GeneratorInv, uint64(start))
+		for i := start; i < end; i++ {
+			p, q := layer[i], layer[i+half]
 
-	var xInv koalabear.Element
-	xInv.SetOne()
-
-	for i := 0; i < half; i++ {
-		p, q := layer[i], layer[i+half]
+			var sum, diff ext.E4
+			sum.Add(&p, &q)
+			sum.MulByElement(&sum, &invTwo)
 
-		var sum, diff ext.E4
-		sum.Add(&p, &q)
-		sum.MulByElement(&sum, &invTwo)
-
-		diff.Sub(&p, &q)
-		diff.MulByElement(&diff, &invTwo)
-		diff.MulByElement(&diff, &xInv)
-		diff.Mul(&diff, &alpha)
-
-		next[i].Add(&sum, &diff)
-
-		xInv.Mul(&xInv, &domain.GeneratorInv)
-	}
+			diff.Sub(&p, &q)
+			diff.MulByElement(&diff, &invTwo)
+			diff.MulByElement(&diff, &xInv)
+			diff.Mul(&diff, &alpha)
 
+			next[i].Add(&sum, &diff)
+			xInv.Mul(&xInv, &domain.GeneratorInv)
+		}
+	})
 	return next
 }