From dcdada831899fbdbdf5d39cd07ad8f2667a77177 Mon Sep 17 00:00:00 2001 From: Gautam Botrel Date: Wed, 27 May 2026 21:31:30 +0000 Subject: [PATCH] perf: row-chunk per-node DAG vector evaluation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ComputeAIRQuotients (which is 41% of tall wall and 68% of wide wall on 100M-cell synth traces) drives a per-row sweep through every DAG node in the vanishing relation. The outer per-node walk is necessarily serial (later nodes consume earlier ones), but every per-node row sweep is embarrassingly parallel — disjoint slots in dst, read-only access to the child vectors. Wrap every per-row loop in evalBaseNodeOnAllEntries, evalExtNodeOnAllEntries, the lift/add/sub/mul child-vector helpers, fillBaseLeafVector and fillExtLeafVector with parallel.ExecuteWithThreshold(N, dagRowParallelThreshold=1<<12, …). For ext.Vector.{Add,Sub,Mul,MulByElement} (already SIMD-vectorised) the chunks just re-enter the same SIMD primitive on a slice. The DAG eval workspace's basePool/extPool allocators run only in the outer (serial) loop, so no concurrency protection is needed there. Bench on a 32-core EPYC 9R45, synth 100M cells, median of multiple runs: shape before after speedup tall (2^22 × 24) 5.55 s 4.44 s 1.25× wide (2^16 × 1536) 1.25 s 1.07 s 1.17× Phase breakdown — compute-air-quotients: tall 2.29 s → 1.35 s 1.70× wide 0.85 s → 0.66 s 1.29× Tests: go test ./... + -race ./internal/dag ./internal/poly ./prover — all clean. Existing prover end-to-end tests verify the parallel and serial paths agree (Prove → Verify). Co-Authored-By: Claude Opus 4.7 (1M context) --- internal/dag/dag.go | 166 ++++++++++++++++++++++++++++---------------- 1 file changed, 108 insertions(+), 58 deletions(-) diff --git a/internal/dag/dag.go b/internal/dag/dag.go index cf05302..0031c86 100644 --- a/internal/dag/dag.go +++ b/internal/dag/dag.go @@ -21,8 +21,14 @@ import ( ext "github.com/consensys/gnark-crypto/field/koalabear/extensions" "github.com/consensys/loom/expr" "github.com/consensys/loom/field" + "github.com/consensys/loom/internal/parallel" ) +// dagRowParallelThreshold is the smallest per-node row count at which fan-out +// across goroutines beats the scheduler overhead. Below this, evalBaseNode / +// evalExtNode iterate the row range serially. +const dagRowParallelThreshold = 1 << 12 + // NodeKind identifies the type of an expression DAG node. type NodeKind int @@ -1236,43 +1242,55 @@ func evalBaseNodeOnAllEntries(dst []koalabear.Element, n *DAGNode, baseVec [][]k copy(dst, baseVec[n.Children[0].Index]) for _, child := range n.Children[1:] { src := baseVec[child.Index] - for j := range N { - dst[j].Add(&dst[j], &src[j]) - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].Add(&dst[j], &src[j]) + } + }) } case KindSub: l, r := baseVec[n.Children[0].Index], baseVec[n.Children[1].Index] - for j := range N { - dst[j].Sub(&l[j], &r[j]) - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].Sub(&l[j], &r[j]) + } + }) case KindMul: copy(dst, baseVec[n.Children[0].Index]) for _, child := range n.Children[1:] { src := baseVec[child.Index] - for j := range N { - dst[j].Mul(&dst[j], &src[j]) - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].Mul(&dst[j], &src[j]) + } + }) } case KindPow: base := baseVec[n.Children[0].Index] tmp := alloc() copy(tmp, base) - for j := range N { - dst[j].SetOne() - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].SetOne() + } + }) exp := n.Exp for exp > 0 { if exp&1 == 1 { - for j := range N { - dst[j].Mul(&dst[j], &tmp[j]) - } - } - for j := range N { - tmp[j].Mul(&tmp[j], &tmp[j]) + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].Mul(&dst[j], &tmp[j]) + } + }) } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + tmp[j].Mul(&tmp[j], &tmp[j]) + } + }) exp >>= 1 } release(tmp) @@ -1287,22 +1305,28 @@ func evalBaseNodeOnAllEntries(dst []koalabear.Element, n *DAGNode, baseVec [][]k // regular columns using the same row-selection convention as EvalOnAllEntries. func fillBaseLeafVector(dst []koalabear.Element, n *DAGNode, PiBase [][]koalabear.Element, N int) { if n.IsConst { - for j := range N { - dst[j] = n.ConstVal - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j] = n.ConstVal + } + }) return } l := n.Leaf p := PiBase[l.Idx] if len(p) == 1 { - for j := range N { - dst[j] = p[0] - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j] = p[0] + } + }) } else if l.Type == expr.RotatedColumn { - for j := range N { - dst[j] = p[(j+N+l.Shift)%N] - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j] = p[(j+N+l.Shift)%N] + } + }) } else { copy(dst, p[:N]) } @@ -1335,19 +1359,25 @@ func evalExtNodeOnAllEntries(dst []ext.E4, n *DAGNode, baseVec [][]koalabear.Ele case KindPow: tmp := alloc() copyChildVectorToExt(tmp, n.Children[0], baseVec, extVec, N) - for j := range N { - dst[j].SetOne() - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].SetOne() + } + }) exp := n.Exp for exp > 0 { if exp&1 == 1 { - for j := range N { - dst[j].Mul(&dst[j], &tmp[j]) - } - } - for j := range N { - tmp[j].Mul(&tmp[j], &tmp[j]) + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].Mul(&dst[j], &tmp[j]) + } + }) } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + tmp[j].Mul(&tmp[j], &tmp[j]) + } + }) exp >>= 1 } release(tmp) @@ -1363,22 +1393,28 @@ func evalExtNodeOnAllEntries(dst []ext.E4, n *DAGNode, baseVec [][]koalabear.Ele func fillExtLeafVector(dst []ext.E4, n *DAGNode, PiExt [][]ext.E4, N int) { if n.IsConst { v := liftBaseToE4(n.ConstVal) - for j := range N { - dst[j] = v - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j] = v + } + }) return } l := n.Leaf p := PiExt[l.Idx] if len(p) == 1 { - for j := range N { - dst[j] = p[0] - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j] = p[0] + } + }) } else if l.Type == expr.RotatedColumn { - for j := range N { - dst[j] = p[(j+N+l.Shift)%N] - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j] = p[(j+N+l.Shift)%N] + } + }) } else { copy(dst, p[:N]) } @@ -1393,9 +1429,11 @@ func copyChildVectorToExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.El return } src := baseVec[child.Index] - for j := range N { - dst[j].Lift(&src[j]) - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].Lift(&src[j]) + } + }) } // addChildVectorToExt accumulates a child vector into an extension destination. @@ -1403,13 +1441,17 @@ func copyChildVectorToExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.El func addChildVectorToExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.Element, extVec [][]ext.E4, N int) { if child.Field == field.Ext { src := extVec[child.Index] - ext.Vector(dst).Add(ext.Vector(dst), ext.Vector(src)) + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + ext.Vector(dst[start:end]).Add(ext.Vector(dst[start:end]), ext.Vector(src[start:end])) + }) return } src := baseVec[child.Index] - for j := range N { - dst[j].B0.A0.Add(&dst[j].B0.A0, &src[j]) - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].B0.A0.Add(&dst[j].B0.A0, &src[j]) + } + }) } // subChildVectorFromExt subtracts a child vector from an extension destination. @@ -1417,13 +1459,17 @@ func addChildVectorToExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.Ele func subChildVectorFromExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.Element, extVec [][]ext.E4, N int) { if child.Field == field.Ext { src := extVec[child.Index] - ext.Vector(dst).Sub(ext.Vector(dst), ext.Vector(src)) + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + ext.Vector(dst[start:end]).Sub(ext.Vector(dst[start:end]), ext.Vector(src[start:end])) + }) return } src := baseVec[child.Index] - for j := range N { - dst[j].B0.A0.Sub(&dst[j].B0.A0, &src[j]) - } + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + for j := start; j < end; j++ { + dst[j].B0.A0.Sub(&dst[j].B0.A0, &src[j]) + } + }) } // mulChildVectorIntoExt multiplies an extension destination by a child vector. @@ -1431,11 +1477,15 @@ func subChildVectorFromExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.E func mulChildVectorIntoExt(dst []ext.E4, child *DAGNode, baseVec [][]koalabear.Element, extVec [][]ext.E4, N int) { if child.Field == field.Ext { src := extVec[child.Index] - ext.Vector(dst).Mul(ext.Vector(dst), ext.Vector(src)) + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + ext.Vector(dst[start:end]).Mul(ext.Vector(dst[start:end]), ext.Vector(src[start:end])) + }) return } src := baseVec[child.Index] - ext.Vector(dst).MulByElement(ext.Vector(dst), koalabear.Vector(src)) + parallel.ExecuteWithThreshold(N, dagRowParallelThreshold, func(start, end int) { + ext.Vector(dst[start:end]).MulByElement(ext.Vector(dst[start:end]), koalabear.Vector(src[start:end])) + }) } // EvalWithCache evaluates the DAG using the caller-supplied cache slice instead