Repository navigation
feat(fields): fast KoalaBear quartic extension and coset interpolation - #399
Draft
mitschabaude-bot wants to merge 111 commits into
Draft
mitschabaude-bot wants to merge 111 commits into
mitschabaude-bot wants to merge 111 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
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.
KoalaBear.Fast.Ext4 is KoalaBear[X]/(X^4 - 3) over Montgomery coefficients, proved isomorphic to the spec KoalaBear.Ext4 operation by operation. Multiplication reduces each output coefficient once (Native32.mulAdd4); inversion uses the norm to the base field, with one base-field inversion. The byte codec matches the spec field's, and Plonky3 known answers are tested. Multiplication 28 ns and inversion 161 ns, against 7.3 us and 913 us for the spec field.
aeval_eq_cosetBarycentric proves Plonky3's interpolate_coset formula: a polynomial of degree below n evaluates at an extension-field point z outside the coset s·H as (zⁿ - sⁿ)/(n sⁿ) · Σ xᵢ/(z - xᵢ) · p(xᵢ). The coset's nodal polynomial is Xⁿ - sⁿ, so each barycentric weight is xᵢ/(n sⁿ). List.batchInv is Montgomery's trick, proved correct for inputs without zeros and transferable along maps preserving *, ⁻¹, 1.
KoalaBear.Fast.interpolateCoset evaluates each column of a row-major matrix of coset values at an extension-field point, like Plonky3's interpolate_coset: coset nodes by prefix products, one batch inversion for the weights x_i/(z - x_i), one dot product per column and a final scale by (zⁿ - sⁿ)/(n sⁿ). toSpec_interpolateCoset proves every output equal to aeval z of the column's polynomial.
Adds --suite interpolate: CompPolyInterpolateBench times the proved reference KoalaBear.Fast.interpolateCoset, and the Rust crate calls p3_interpolation::interpolate_coset 0.4.2 on the same fixtures (root, shift, two alternating extension points, row-major matrix). Full output digests agree before timing. The sequential reference is 3-58x slower than parallel Plonky3 on eight workers.
mitschabaude-bot
changed the title
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: 1152 declaration(s)
…and 1002 more not listed. ✏️ Affected: 21 declaration(s) (line number changed)
Coverage Notes
Partially Analyzed Files
Last updated: 2026-10-13 02:52 UTC. |
…the reference `interpolateCosetPacked` takes Plonky3's row-major matrix of Montgomery words in a `ByteArray`: - 1/(z - x) = (b0 + t b1 + t^2 b2 + t^3) / N(t) for t = z0 - x, so each row's weight is four base-field values and only the base-field norms need inverting; - blocks of 256 rows: a forward pass stores nodes, t^2, norms and prefix products, one inversion per block, and a backward pass overwrites each row with its weights; - column sums accumulate unreduced 64-bit products with one conditional subtraction each and one Montgomery reduction per block, four columns sharing each row's weights; - row ranges run on parallel tasks, the caller taking the lowest. `interpolateCosetPacked_eq` proves the result equal to `interpolateCoset` on the decoded words for every point off the coset; the native smoke test compares the two on 38 shapes.
Interleaved runs over 2^12..2^20 rows and 1..16 columns: 512 rows beat 256 by up to 8% and 1024 by up to 20%; only 2^16 x 16 prefers 128. Eight tasks stay ahead of sixteen, which oversubscribe the cores by 1.4x.
Fixtures hold canonical values; the packed path reads Montgomery words, so the runner packs the decoded matrix before timing, as Plonky3 builds its matrix from the same values.
…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.
Port the packed coset interpolation from the removed readRaw extern to the backported ByteArray.getUInt32LE! accessor.
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.
# Conflicts: # bench/README.md # tests/CompPolyNativeSmoke.lean
interpolateCosetPacked is Lean over the backported word accessors and is proved equal to the reference by interpolateCosetPacked_eq; the accessors' inline C keeps its compiled checks. The removed checks spawned tasks, whose 1 GiB stack reservations exceeded the smoke test's 4 GiB address-space limit.
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.
Move the packed interpolation proofs to PackedRepr.encode.
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. 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.A fast, fully proved degree-4 extension of KoalaBear,
KoalaBear[X] / (X^4 - 3): the same field as the specKoalaBear.Ext4and Plonky3'sBinomialExtensionField<KoalaBear, 4>. On top of it, a proved coset interpolation at an extension point, benchmarked against Plonky3's stockinterpolate_coset.KoalaBear.Fast.Ext4holds four Montgomery coefficients.toSpecmaps it into the spec field;toSpec_add,toSpec_mul,toSpec_inv,toSpec_divandtoSpec_ofBaseprove agreement operation by operation, andspecEquivis the ring isomorphism.3into the left factors, so each output coefficient is a sum of four products with a single Montgomery reduction. The newMontgomery.Native32.mulAdd4does this for every 32-bit Montgomery field: four products stay below2 · p · 2^32, so one conditional subtraction brings them intoreduce's range.σ : X ↦ -X,a · σ(a) = g₀ + g₁ X²anda⁻¹ = σ(a) (g₀ - g₁ X²) / (g₀² - 3 g₁²), one base-field inversion. The norm of a nonzero element is nonzero becauseσis an injective ring endomorphism of the quotient field and3is a quadratic non-residue modp.toBytes_toSpec), with round trip throughByteCodec.Results
The existing
fields-extension-koalabear-ext4-muland-invbenchmark groups now time both types on the same operands; the harness checks that both produce the same checksum. Per call, including the harness's operand lookup and the result allocation:KoalaBear.Ext4KoalaBear.Fast.Ext4The tests check known answers computed with Plonky3 (
p3-field0.4.2) for a product and two inverses, field laws, agreement with the spec field, and the codec.Coset interpolation
KoalaBear.Fast.interpolateCosetis Plonky3'sinterpolate_cosetfor KoalaBear and its quartic extension: given a row-major2^logN × widthmatrix whose columns hold the values of polynomials on the cosets · ⟨ω⟩, it evaluates every column's polynomial at a pointzof the extension, by(zⁿ - sⁿ) / (n sⁿ) · ∑ᵢ xᵢ / (z - xᵢ) · p(xᵢ).NTT.Domain.aeval_eq_cosetBarycentricproves that formula for any field extensionE / F: the coset's nodal polynomial isXⁿ - sⁿ, so each barycentric weight isxᵢ / (n sⁿ).List.batchInvis Montgomery's batch inversion, proved correct for inputs without zeros and transferable along maps preserving*,⁻¹and1, so it runs on the fast extension, which has noFieldinstance.toSpec_interpolateCosetproves every output equal toaeval z pⱼfor the column polynomialspⱼ.--suite interpolateruns it againstp3_interpolation::interpolate_coset0.4.2 on the same fixtures; full output digests agree before timing, from single rows to 2²⁰ × 16.KoalaBear.Fast.interpolateCosetPackedis the fast version. It takes the matrix as Montgomery words in aByteArray, the layout of Plonky3'sRowMajorMatrix, andinterpolateCosetPacked_eqproves it equal to the reference for every point off the coset.t = z₀ - x,(z - x)⁻¹ = (b₀ + t b₁ + t² b₂ + t³) / N(t), with extension constantsbₖthat depend onzonly andN(t)the norm ofz - x. So each row's weight is four base-field valuesx tᵏ / N, only base-field norms need inverting, and a column sum is∑ₖ Sₖ bₖfor four base-field sums. That is about twelve base-field multiplications per row.t², norm and prefix product, one inversion per block inverts the product, and a backward pass overwrites each row with its weights.Montgomery.Native32.LazyAcc) and one Montgomery reduction per block. Four columns share each row's weights.It reads and writes words with the packed NTT's backported
ByteArrayword accessors (lean4#14053) and adds no extern of its own. The benchmark suite also checks its full outputs against the reference and Plonky3 before timing.Five rounds of
--suite interpolateon eight workers at130c36a, milliseconds per call: