Repository navigation
perf(ntt): proved parallel KoalaBear NTTs against Plonky3 - #395
Draft
mitschabaude-bot wants to merge 101 commits into
Draft
mitschabaude-bot wants to merge 101 commits into
mitschabaude-bot wants to merge 101 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
Sixteen workers on the eight-core host only add SMT siblings, which measures core sharing rather than scaling.
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.
mitschabaude-bot
changed the title
mitschabaude-bot
force-pushed
the
grant-closeout-followup
branch
from
October 7, 2026 11:06
539c933 to
6fa44c2
Compare
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.
…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.
An odd-size leaf now runs its radix-two layer first, with a sixteen-lane kernel (`pair16`), then radix-four passes down to the fused final layers and reversal (`lastRev`) that even leaves use. `runStagesDIFRadix4WithTwiddles_eq_top` proves the two stage orders equal. Leaves below 2^8 words run all passes and the scalar `decode`. With every leaf on one path, the sequential `stages` pipeline, the `leaf16` kernel, the separate `leafRev` pass and their refinement files are deleted. The native library root list now names `Packed.Columns` instead of the deleted `SliceTree`. Forward, ms, before → after: CPU 0 2^17 2.93 → 0.878, 2^19 13.2 → 3.92; CPUs 0-7 2^17 0.749 → 0.353, 2^19 2.96 → 1.18. Even sizes unchanged within noise.
Both plans' column transforms are now built from three task combinators in `CompPoly/Data/Tasks.lean`: `spawnAll` runs independent tasks, `spawnAfter` runs tasks once a set of tasks is done, and `join` collects results. `map_get_spawnAll` and `map_get_spawnAfter` state the results as plain maps, which replaces the task bookkeeping in each transform's proof. The fused kernels spill AVX registers with unaligned moves, and their speed depended on the stack alignment inherited from callers: builds with byte-identical kernels differed by 20% on one worker, and so did one binary under different environments. `CompPolyPackedNative` now compiles with `-mstackrealign`, which aligns the spill slots; timings are stable within 2% of the fast case.
The sixteen-lane kernels and their field-array models are now expanded by the macros of `Packed/Unrolled.lean` from one description per family: `step16`, `pair16`, `lastHalf`, `transposeStep` and `line` on packed words, `step16Field`, `leaf16Field` and `laneLine` on field arrays. The expansions are the terms that were written out before, so the refinement proofs are unchanged and the generated C is identical. About 1,700 lines of unrolled source are gone.
The worker count only chose between the column transform and one leaf, so both plans now take `parallel : Bool := true`; the runner passes it when it has more than one worker.
The packed kernels compile with -march=native, and the restored .lake/build may come from a runner with another CPU: the NTT bench died with SIGILL on objects Lake did not need to rebuild.
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.
Depends on #390. Review and merge #390 first. GitHub can't use its fork branch as the base, so this targets
mainfor now: review only this PR's changes.Summary
A fully kernel-proved, parallel KoalaBear NTT that is faster than Plonky3's parallel AVX2 NTT at a million points on one and eight workers and within 1.2× on four, plus
--suite nttto benchmark it against Plonky3.At$2^{12}$ Lean takes 0.34–0.59× Plonky3's time on every worker count; at $2^{16}$ , 0.75–0.81× on one worker, 0.92–0.97× on eight and 1.05–1.10× on four. Odd sizes cost the same per element as even ones.
There are two plans, both proved equal to the NTT specification for every input:
NTTFast.Packed.Plancomputes on packed 32-bit Montgomery words in aByteArraywith vectorized kernels.forwardPacked/inversePackedtake and return packed words, as in the table.NTTFast.NaturalPlanuses no extern at all: it computes on field arrays and exchanges packed words between tasks.forwardPacked/inversePackedhave the same packed API. It takes about 4× Plonky3's time on one or four workers and 2.2–2.4× on eight.Both plans run the same pipeline, described below.
New externs: core's proposed
ByteArrayword accessorsThe packed plan needs fast reads and writes of 32-bit words in a
ByteArray, which Lean doesn't have yet. leanprover/lean4#14053 proposes exactly these. This PR adds them inCompPoly.Data.ByteArray.Packas a backport:getUInt32LE!/setUInt32LE!,getUInt32LE/setUInt32LE,ugetUInt32LE/usetUInt32LE.tests/CompPolyTests/NTT/NativeStorage.leanchecks the C against them.memcpy, not four byte shifts. In code with many adjacent word accesses, clang vectorizes the byte shifts into shuffle sequences instead of combining them into word loads: 79 ms instead of 23 ms. A 4-bytememcpyis always a single load or store, with a byte swap on big-endian hosts.__builtin_assume(r->m_rc == 1)).Packed arrays
PackedRepr R(CompPoly.Data.ByteArray.Packed) is the single storage layer for fast field arrays. Every element takeswidthbytes;toBytes/ofBytesare the specification, and the fastread/writeat byte offsets are proved equal to it.PackedRepr.encode/decodeview aByteArrayas anArray R.UInt32and every 32-bit Montgomery field are instances on the word accessors above. Both NTTs store their intermediate results this way, and the column-parallel NTT is generic over the element width.How it's fast
Correctness
Packed.Plan.forwardPacked_correctandinversePacked_correct, andNaturalPlan.forwardPacked_eqandinversePacked_eq, state that the public APIs equalNTT.Forward.forwardSpecandNTT.Inverse.inverseSpecfor every input, parallel or not. There is nosorryand nonative_decide.Reproduce
python3 scripts/bench-fields.py --suite ntt --cpus 0 # one worker python3 scripts/bench-fields.py --suite ntt --cpus 0,1,2,3 python3 scripts/bench-fields.py --suite ntt --cpus 0,1,2,3,4,5,6,7COMPPOLY_NTT_IMPL=packed|externlessselects one Lean variant in the direct runner.Setup and validation
target-cpu=native.dc2fc3a, host load below 2.5 throughout. Median of five alternating rounds, each with 200 ms warmup and 50 samples.sorryor nonstandard axioms across 12,489 declarations.