Repository navigation
perf(poly): parallel one-point evaluation against arkworks - #398
Draft
mitschabaude-bot wants to merge 107 commits into
Draft
mitschabaude-bot wants to merge 107 commits into
mitschabaude-bot wants to merge 107 commits into
Conversation
…ocks Fold the 64x64 product's middle terms without carry bookkeeping so clang emits one widening multiply for squarings too; `mul x x` inlined into `square` used to compile to a four-imul chain. Run the Fermat inverse and the pow ladder on unreduced words (`mulLazy`, `squareNLazy`, `invLazy`, `powLazy`) with a single canonicalization at the end, and make `squareN` tail recursive. Correctness rests on word-level cast lemmas (`foldUInt128Lazy_cast`, `mulLazy_cast`, `squareNLazy_cast`, `invLazy_cast`, `powLazy_cast`). A/B against master on the same core: inv 311 -> 167 ns, pow 99 -> 74 ns, mul and add unchanged.
`conditionalSubtract` was a plain def, so every KoalaBear and BabyBear `add` paid an out-of-line call. A/B: add latency 844 -> 606 ps, throughput 876 -> 166 ps.
…nreduced words `univariate-dense-goldilocks` now runs the native-word carrier beside ZMod like the other fast fields, with rows named in the shared `-fast` scheme. Add tests that unreduced `ofUInt64` words compare and multiply as their residues.
Compute the canonical sum as `x - (p - y)`: the subtraction borrows exactly when `x + y < p`, so one borrow-selected correction by `2^32 - 1` yields the reduced result, three dependent instructions instead of a carry check plus a separate compare against `p`. `addRaw_toNat` carries the bound and the cast lemma; the carry-based reducer and its lemmas are removed. A/B against master on the same core: add latency 1021 -> 616 ps, throughput 381 -> 230 ps; mul, inv and pow unchanged.
# Conflicts: # bench/CompPolyBench/Fields/Arith.lean
… the Goldilocks fold Two changes to the 128-bit fold behind every fast Goldilocks multiply. The borrow of the low word against the top limb has probability about 2^-32, but a Lean `if` compiled to a select on the critical path; moving that arm into a `@[noinline]` function forces a predicted branch. The middle term `hi_lo * (2^32 - 1)` compiled to a three-cycle `imul` inside loops (clang holds the constant in a register); it is now `(hi <<< 32) - hi_lo`, two one-cycle operations. `shiftLeft32_sub_low` and `foldTail_cast` carry the proofs; the fold lemma cases on the borrow. A/B against master on the same core: mul latency 3077 -> 2485 ps, throughput 815 -> 739 ps, inv 314 -> 136 ns, pow 101 -> 64 ns; add unchanged.
One-sentence docstrings, `2 ^ k` literals instead of decimals, a `_toNat` bridge for the product's high word, implicit bounds on the `addRaw` lemmas, camelCase limb names, and README rows for the Goldilocks files.
# Conflicts: # CompPoly/Fields/README.md
A generated kernel runs the last four DIF layers of eight leaf blocks at once, one block per vector lane, and writes the leaf in natural order. It replaces the scalar leaf16 kernel and the separate leafRev pass for even leaves of at least 2^8 words (2^16 leaf: 0.63 -> 0.48 ms). The slice tree takes the leaf computation as a parameter; NativeLastCorrectness proves the fused leaf equal to leafRev after stages, and runPacked/runFields with their DFT theorems move there. 2^20 packed I/O: 14.8 -> 12.4 ms on one worker, 4.10 -> 3.47 ms on four.
pairGo and pairInputGo now check the whole batch range once and run a USize loop whose proof carries the remaining range, instead of checking and converting indices on every sixteen-word batch. Packed-io NTT at one million points: 5-7% faster on one and four workers.
inner16 now checks the range of all whole sixteen-cell batches once and runs inner16Loop over USize indices before the scalar tail. The proof uses new size lemmas for overwriting stores and step16. Packed-io NTT at one million points: about 5% faster on one and four workers.
The one-point evaluation work (EvalFast, its lazy field kernels, the matched Rust evaluator and the poly-eval benchmark suite) moves to the poly-eval branch, to return against stock reference implementations in its own PR. This PR keeps the proved parallel NTTs and their Plonky3 comparison.
This reverts commit 6fa44c2.
Restores the parallel one-point evaluation work on its own branch and replaces the hand-written Rust evaluators with ark-poly's DensePolynomial::evaluate (parallel feature) in a Rayon pool of the selected size; one thread gives its plain Horner path. KoalaBear and Goldilocks use arkworks' generic Montgomery fields, since arkworks ships neither. The binary tower rows are dropped, as arkworks cannot represent the field.
With as many CPUs as workers, the caller was still running when the leaves were spawned, so one woken worker usually queued behind another and its leaf started only when the first finished. The caller now walks the tree's lower edge, spawning the upper subtrees, and evaluates the lowest leaf itself. Million-coefficient KoalaBear on four workers: 1.94 to 0.95 ms.
mitschabaude-bot
force-pushed
the
poly-eval-ark
branch
from
October 7, 2026 11:06
c3935d0 to
528bc15
Compare
With the 50 ms warmup, idle worker cores were often still at reduced clocks when timing began, so eight-worker evalFast measured 0.50 or 0.85 ms depending on the run's calibration. The poly-eval suite now uses the large budget (200 ms warmup, 50 samples) in Lean and Rust.
EvalKernel gains minLeaf, the smallest range worth a task (8192 coefficients by default, 1024 for BN254), and evalFast caps its depth so that no leaf is smaller. Small polynomials are evaluated on the calling thread with the lazy leaf kernel.
Contributor
🤖 PR Summary
Failed to generate AI summary. Please check the per-file summaries and statistics below. Statistics
Lean Declarations ✏️ Removed: 35 declaration(s)
✏️ Added: 944 declaration(s)
…and 794 more not listed. ✏️ Affected: 21 declaration(s) (line number changed)
Coverage Notes
Partially Analyzed Files
Last updated: 2026-10-13 03:05 UTC. |
Idle worker cores take a few hundred milliseconds of load before the frequency governor raises their clocks; the 50 ms warmup often ended before that. The NTT suite now uses the large budget (200 ms warmup, 50 samples) in Lean and Rust; the Rust harness gains Budget presets matching Lean's.
# Conflicts: # bench/rust/src/ntt.rs
…rray word accessors The packed NTT now reads and writes words with `ByteArray.ugetUInt32LE` and `ByteArray.usetUInt32LE` from lean4#14053, backported in `CompPoly.Data.ByteArray.Pack` with the same names, signatures, Lean definitions and lemmas. `Native.readRaw` and the `Native.storeWords` extern are gone: `storeWords` is a Lean batch of word stores, appends grow the buffer with `ByteArray.copySlice`, and checked reads are `readAt`. `storeWords_eq` and the read lemmas prove them equal to the previous storage semantics, so the refinement proofs carry over. The backport's C differs from the PR's: loads and stores use a fixed-size `memcpy`, and a store states that the array is still unshared, so a run of stores keeps one sharing check. With the PR's byte-shift C the million-point transform took 79 ms instead of 23 ms; without the assumption, 23 ms instead of 13. Range checks run on machine words; `storeWords` reads the append offset before the growth consumes the buffer, which would otherwise copy it on every batch.
The binary split tree, the sliced pair tree, the transpose and the reversed leaves now store each 16-word batch at a cursor into a buffer pre-sized with zeros, instead of growing the output by appending. Zero buffers are filled by appending copies of a separate 4 KiB zero block into a reserved destination, so the destination stays exclusive.
The transpose extends its output by one zero block right before each tile overwrites it, so the zeros are still in cache. encode overwrites a zero buffer word by word instead of appending one word at a time; this restores the small transforms that end in an encode.
storeWords loses its append mode, which only push, push16 and the old generate used. generate now overwrites a zero buffer word by word, and encode is generate over the field array. Removed: push, push16, zeroBytes, the natural-index write16, Storage.wordAt and lemmas no longer referenced. The native smoke test keeps only the accessor checks: storeWords and read are Lean definitions proved against Storage.pack.
PackedRepr R stores every element in width bytes, with toBytes/ofBytes as the specification and fast read/write at byte offsets proved equal to it; PackedRepr.encode/decode view a ByteArray as an Array R. UInt32 and every Native32 FastField are one-word instances on the backported ByteArray word accessors. The packed NTT's Storage.pack and packFields become PackedRepr.encode at UInt32 and at the field, and the executable encoder is Native.encodeFields. The column-parallel NTT drops Word32Repr and the byte-by-byte ByteWords layer for PackedRepr.getAt/setAt with element-count invariants, so it is generic over the element width. Zero buffers are ByteArray.zeros. Data/Bytes/Words.lean and Packed/Storage.lean are gone.
CompPoly is extern-free by design; only the packed NTT kernels opt into the backported ByteArray word accessors. The one-word PackedRepr instance read and wrote through those accessors, which put them under the extern-free column-parallel NTT too. It now loads and stores the four bytes with ByteArray.uget/uset, proved equal to the accessors' Lean definitions, and the column-parallel NTT runs at its speed before the PackedRepr refactor.
Plan.forward/inverse took field arrays through a separate input path: field-array split kernels, their models and refinement proofs, and a field-input task tree. Consumers hold packed words (PackedRepr.encode, or Native.encodeFields and Native.unpack for fast conversion), so the plan now exposes only forwardPacked/inversePacked. The packed decode path shared by small and custom-depth transforms is Native.runDecoded. The benchmark's field-array variant goes with the API; packed-io remains.
NaturalPlan.forwardPacked/inversePacked take and return packed elements (PackedRepr), like the packed plan; forwardPacked_eq and inversePacked_eq prove forwardPacked (encode a) = encode (forward a) and the same for the inverse. Even sizes from 2^16 run Columns.transformPacked: column tasks decode their own columns from the packed buffer, and sixteen tasks each interleave a contiguous range of rows into packed words, joined by one copy. The proofs reuse the array transform: a generic simulation lemma over the interleave loop relates packed stores to array stores and offset chunks to slices of the whole interleave. Other sizes unpack with Columns.decodeAll, run the array transform and pack with Columns.encodeAll. PackedRepr gains replicate (filled by doubling copies) and setIf, a guarded store whose machine-size bound is proved once per buffer. The benchmark's variants are now packed and externless, both with packed words in and out.
The column transform now handles odd sizes: a leaf finishes with the radix-two layer (lastPair), which commutes with cutting the array into leaves. So every size from 2^15 with more than one worker takes the column path, and smaller transforms or a single worker run the sequential transform as one leaf. The segment path for odd sizes (Parallel.chunksTask) and the inverse pass lists it needed are gone. Odd sizes on several workers get 2-4x faster (2^19 on 8 workers: 17.7 ms to 4.2 ms); other shapes are unchanged. The NTT suite validates 2^15 and 2^17 against Plonky3.
The packed plan now runs the extern-free plan's pipeline on packed words: from 2^15 elements and with more than one worker, column tasks gather their columns and run the top two radix-four passes with `inner16`, sixteen leaf tasks finish with `leafNatural`, and `naturalLeaves` interleaves; otherwise one leaf. `Packed/Columns.lean` holds the pipeline, `Packed/ColumnsCorrectness.lean` its proof, which reuses the extern-free column and leaf lemmas. The split tree, slice tree, tiled decoder, `Plan.defaultDepth` and the in-stage normalization they needed are deleted. The spec files that remain are merged into `Packed/DifSpec.lean`. Both plans use sixteen column tasks at every worker count, which keeps a task's rows and twiddles in L2 with four workers. Forward, ms, previous packed plan → this: CPUs 0-7 2^12 0.079 → 0.020, 2^16 0.410 → 0.246, 2^20 2.53 → 2.42; CPUs 0-3 2^18 0.687 → 0.672, 2^20 3.54 → 3.23; CPU 0 2^20 10.8 → 7.69.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Depends on #395, which depends on #390. Review and merge those first. This branch starts at #395's tip. GitHub cannot use a fork branch as an upstream base, so this draft targets
mainand its diff includes both earlier PRs. Review only the new changes.Verified multicore evaluation of one polynomial at one point, benchmarked against an existing fast implementation: arkworks' stock
DensePolynomial::evaluate(ark-poly0.5.0 with itsparallelfeature), and against CompPoly's existing sequentialevalHorner.CPolynomial.evalFast x p logWorkerssplits contiguous coefficient ranges without copying, evaluates the leaves concurrently with lazy-reduction Horner kernels (KoalaBear, Goldilocks, BN254 scalar) and joins with square-and-multiply powers.evalFast_eq_evalHornerandevalFast_eq_evalprove it correct over any semiring.EvalKernel.minLeaf(8,192 coefficients by default, 1,024 for BN254), so small polynomials stay on the calling thread instead of paying for tasks.evalFastmeasure 0.50 or 0.85 ms depending on the run.--suite poly-evalruns LeanevalHorner, LeanevalFast, and ark-poly on one thread and on the selected worker count, from identical fixtures; all four results must agree before timing.Fp64via#[derive(MontConfig)]); those rows compare against stock arkworks code, not Plonky3-grade field arithmetic. BN254 usesark_bn254::Fr. Binary-tower rows are gone, since arkworks cannot represent the field.Results
One million coefficients, eight workers on the eight physical cores of a Ryzen 7 3700X (no SMT siblings). Milliseconds per evaluation; median of five paired run medians, each after 200 ms of warmup. Lean / ark-poly > 1 means arkworks is faster.
evalHorner(1 core)evalFast(8 workers)evalFastis 8–10× faster than sequentialevalHornerand beats ark-poly on KoalaBear and Goldilocks. On BN254 it trails by 1.2×, all of it field arithmetic: Lean's sequential Horner is already 1.6× slower than arkworks' BN254 multiplication.Measured source:
7b8d30d(the working tree differed only in an uncommitted scratch benchmark). Lean 4.34.0; Rust 1.95.0-nightly with-C target-cpu=native. Shared host, load about 0.5–2.5.Reproduce