NTT-based univariate multiplication

[erdkocak]

• Add an NTT domain carrying the transform size, primitive root, and normalization data
• Define direct forward and inverse NTT specs
• Implement iterative radix-2 forward and inverse transforms
• Compose them into an NTT-based multiplication pipeline
• Add executable tests for forward, inverse, and multiplication
• Add benchmarks for the fast multiplication path
• Prove forwardImpl = forwardSpec
• Prove inverseImpl = inverseSpec
• Prove fastMulImpl = p * q

[github-actions]

🤖 Gemini PR Summary

Core Components

• Domain structure manages radix-2 NTT parameters: primitive roots of unity, domain size constraints, and normalization factors.
• Formal specifications for forward and inverse transforms provided alongside iterative radix-2 implementations.
• Formalization of bit-reversal permutations and butterfly operations for Cooley-Tukey and Gentleman-Sande patterns.
• Fast multiplication pipeline utilizing pointwise NTT multiplication, including automatic domain size selection and polynomial truncation.
• Integration into the CompPoly library for univariate polynomial arithmetic.

Status of Proofs

• CRITICAL: This PR adds several sorry placeholders for fundamental correctness theorems.
• Proofs relating iterative implementations to high-level specifications (e.g., forwardImpl = forwardSpec and inverseImpl = inverseSpec) are incomplete.
• The final correctness proof relating the NTT-based multiplication pipeline (fastMulImpl) to standard polynomial multiplication is marked with sorry.

Testing and Benchmarking

• Executable tests using #guard verify NTT implementations against concrete cases in the KoalaBear field.
• Benchmarking utilities measure the performance of NTT-based multiplication against standard methods.
• Automated script identifies the efficiency crossover point where the NTT-based approach outperforms naive multiplication.

---

Statistics

Metric	Count
📝 Files Changed	11
✅ Lines Added	680
❌ Lines Removed	0

---

Lean Declarations

✏️ **Added:** 41 declaration(s)

• theorem fastMulSpec_eq_mul (D : Domain R) (p q : CPolynomial.Raw R) in CompPoly/Univariate/NTT/FastMul.lean
• def testBits32 : Fin (KoalaBear.twoAdicity + 1) in tests/CompPolyTests/Univariate/NTT/Common.lean
• def testDomain32 : Domain KoalaBear.Field where in tests/CompPolyTests/Univariate/NTT/Common.lean
• def bestLogN (requiredLen : Nat) : Nat in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• theorem forwardImpl_correct (D : Domain R) (p : CPolynomial.Raw R) : in CompPoly/Univariate/NTT/Forward.lean
• def q : CPolynomial.Raw KoalaBear.Field in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• def testLogN32 : Nat in tests/CompPolyTests/Univariate/NTT/Common.lean
• def requiredLength (p q : CPolynomial.Raw R) : Nat in CompPoly/Univariate/NTT/Domain.lean
• def bitRevPermute (D : Domain R) (a : Array R) : Array R in CompPoly/Univariate/NTT/Inverse.lean
• def benchSizes : Array Nat in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• def bestDomainForLength? (requiredLen : Nat) : Option (Domain KoalaBear.Field) in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• def inverseSpec (D : Domain R) (v : Array R) : Array R in CompPoly/Univariate/NTT/Inverse.lean
• def testDomain64 : Domain KoalaBear.Field where in tests/CompPolyTests/Univariate/NTT/Common.lean
• def avgMsString (totalMs reps : Nat) : String in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• def timeRepeated {α : Type} (reps : Nat) (f : Unit → α) : IO (Nat × α) in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• def normalize (D : Domain R) (a : Array R) : Array R in CompPoly/Univariate/NTT/Inverse.lean
• def mkPoly (n seed : Nat) : CPolynomial.Raw KoalaBear.Field in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• def testLogN : Nat in tests/CompPolyTests/Univariate/NTT/Common.lean
• def testLogN64 : Nat in tests/CompPolyTests/Univariate/NTT/Common.lean
• def bitRevPermute (D : Domain R) (a : Array R) : Array R in CompPoly/Univariate/NTT/Forward.lean
• def butterflyStage (D : Domain R) (stage : Nat) (a : Array R) : Array R in CompPoly/Univariate/NTT/Forward.lean
• def truncate (m : Nat) (p : CPolynomial.Raw R) : CPolynomial.Raw R in CompPoly/Univariate/NTT/Domain.lean
• def p : CPolynomial.Raw KoalaBear.Field in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• def bitRevNat : Nat → Nat → Nat in CompPoly/Univariate/NTT/Forward.lean
• theorem fastMulSpec_coeff (D : Domain R) (p q : CPolynomial.Raw R) (i : Nat) : in CompPoly/Univariate/NTT/FastMul.lean
• def fits (D : Domain R) (p q : CPolynomial.Raw R) : Prop in CompPoly/Univariate/NTT/Domain.lean
• def repeatsFor (n : Nat) : Nat in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• def runStages (D : Domain R) (a : Array R) : Array R in CompPoly/Univariate/NTT/Inverse.lean
• theorem inverseImpl_correct (D : Domain R) (v : Array R) : in CompPoly/Univariate/NTT/Inverse.lean
• def runStages (D : Domain R) (a : Array R) : Array R in CompPoly/Univariate/NTT/Forward.lean
• def checkSize : Nat in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• def testBits : Fin (KoalaBear.twoAdicity + 1) in tests/CompPolyTests/Univariate/NTT/Common.lean
• def testDomain : Domain KoalaBear.Field where in tests/CompPolyTests/Univariate/NTT/Common.lean
• def inttAt (D : Domain R) (v : Array R) (k : D.Idx) : R in CompPoly/Univariate/NTT/Inverse.lean
• def testBits64 : Fin (KoalaBear.twoAdicity + 1) in tests/CompPolyTests/Univariate/NTT/Common.lean
• theorem fastMulImpl_correct (D : Domain R) (p q : CPolynomial.Raw R) : in CompPoly/Univariate/NTT/FastMul.lean
• def inverseImpl (D : Domain R) (v : Array R) : CPolynomial.Raw R in CompPoly/Univariate/NTT/Inverse.lean
• abbrev Idx (D : Domain R) in CompPoly/Univariate/NTT/Domain.lean
• def butterflyStage (D : Domain R) (stage : Nat) (a : Array R) : Array R in CompPoly/Univariate/NTT/Inverse.lean
• def speedupString (nttMs rawMs : Nat) : String in tests/CompPolyTests/Univariate/NTT/Benchmark.lean
• theorem fastMulImpl_eq_mul (D : Domain R) (p q : CPolynomial.Raw R) in CompPoly/Univariate/NTT/FastMul.lean

---

sorry Tracking

❌ **Added:** 5 `sorry`(s)

• theorem fastMulSpec_eq_mul (D : Domain R) (p q : CPolynomial.Raw R) in CompPoly/Univariate/NTT/FastMul.lean (L63)
• theorem forwardImpl_correct (D : Domain R) (p : CPolynomial.Raw R) : in CompPoly/Univariate/NTT/Forward.lean (L79)
• theorem fastMulSpec_coeff (D : Domain R) (p q : CPolynomial.Raw R) (i : Nat) : in CompPoly/Univariate/NTT/FastMul.lean (L59)
• theorem fastMulImpl_eq_mul (D : Domain R) (p q : CPolynomial.Raw R) in CompPoly/Univariate/NTT/FastMul.lean (L67)
• theorem inverseImpl_correct (D : Domain R) (v : Array R) : in CompPoly/Univariate/NTT/Inverse.lean (L81)

---

🎨 **Style Guide Adherence**

There are more than 20 violations in this diff. They are grouped by rule below:

• Documentation Standards

  • Rule: "Every definition and major theorem should have a docstring." (39 violations)
    • Domain.lean (line 56): def requiredLength (p q : CPolynomial.Raw R) : Nat :=
    • Forward.lean (line 33): def bitRevNat : Nat → Nat → Nat
    • Benchmark.lean (line 27): def bestLogN (requiredLen : Nat) : Nat :=
  • Rule: "Each file should start with a /-! ... -/ block containing a title, summary, notation, and references." (8 violations: All new files lack Notation and References sections)
    • Domain.lean (line 10)
    • FastMul.lean (line 11)
    • Forward.lean (line 8)

• Syntax and Formatting

  • Rule: "Avoid parentheses where possible. Use <| (pipe left) and |> (pipe right) to reduce nesting." (11 violations)
    • Domain.lean (line 27): natCast_ne_zero : (((2 ^ logN : Nat) : R) ≠ 0)
    • Forward.lean (line 27): ... D.omega ^ ((k : Nat) * (j : Nat))
    • Benchmark.lean (line 69): s!"{(Float.ofNat totalMs) / (Float.ofNat reps)}"
  • Rule: "Prefer fun x ↦ ... over λ x, ...." (10 violations: The guide specifies the ↦ delimiter)
    • FastMul.lean (line 26): Array.ofFn (fun i : D.Idx => a.getD i.1 0 * b.getD i.1 0)
    • Forward.lean (line 31): Array.ofFn (fun k : D.Idx => nttAt D a k)
    • Benchmark.lean (line 99): (fun _ => FastMul.fastMulImpl benchDomain p q)
  • Rule: "Indentation: Use 2 spaces for indentation." (5 violations: Unnecessary indentation of module docstring content)
    • Benchmark.lean (line 9): # Univariate Multiplication Benchmark
    • Common.lean (line 9): # Univariate NTT Test Helpers
    • FastMul.lean (line 9 - tests): # Univariate NTT FastMul Tests

---

📄 **Per-File Summaries**

• CompPoly.lean: This update expands the CompPoly library by importing new modules related to univariate Lagrange interpolation and the Number Theoretic Transform (NTT), including its application for fast multiplication. No new theorems or proofs are implemented directly in this file, and no sorry placeholders are introduced.
• CompPoly/Univariate/NTT/Domain.lean: This file introduces the Domain structure and associated definitions to represent radix-2 Number Theoretic Transform (NTT) parameters, including roots of unity and domain size constraints. It also provides helper functions for polynomial truncation and determining the required domain size for convolutions. No sorry or admit placeholders are included.
• CompPoly/Univariate/NTT/FastMul.lean: This new file defines the specification and implementation pipelines for fast polynomial multiplication using Number Theoretic Transforms (NTT), including the pointwise multiplication of transformed arrays. It introduces the fastMulSpec and fastMulImpl definitions and proves the implementation's correctness relative to the specification; however, the final proofs relating these functions to standard polynomial multiplication are currently marked with sorry.
• CompPoly/Univariate/NTT/Inverse.lean: This file introduces definitions and an implementation for the Inverse Number Theoretic Transform (INTT), including a direct formula specification and an iterative radix-2 butterfly algorithm. It provides several helper functions for bit-reversal and normalization alongside a correctness theorem relating the implementation to the specification, which currently contains a sorry placeholder.
• tests/CompPolyTests.lean: This update expands the univariate polynomial test suite by integrating new test modules for Number Theoretic Transform (NTT) operations, specifically covering forward transforms, inverse transforms, and fast multiplication. No new theorems, definitions, or sorry placeholders are introduced in this file.
• tests/CompPolyTests/Univariate/NTT/Common.lean: This new file defines shared helper constants and Domain instances for the KoalaBear field to support univariate NTT testing. It introduces concrete definitions and proofs for domain properties across multiple sizes, including primitive root verifications, without any sorry placeholders.
• tests/CompPolyTests/Univariate/NTT/FastMul.lean: This file introduces a suite of executable tests for the Number Theoretic Transform (NTT) fast multiplication implementation. It uses #guard statements to verify that FastMul.fastMulImpl correctly matches standard polynomial multiplication across various concrete cases in the KoalaBear field. No sorry or admit placeholders are included.
• tests/CompPolyTests/Univariate/NTT/Inverse.lean: This new test file introduces concrete executable checks for the univariate inverse Number Theoretic Transform (NTT). It uses #guard commands to verify that the inverse NTT implementation matches its specification across several test cases, and it contains no sorry or admit placeholders.
• tests/CompPolyTests/Univariate/NTT/Forward.lean: This new test file introduces concrete executable checks for the forward Number Theoretic Transform (NTT) implementation. It uses #guard commands to verify consistency between the implementation and its specification across several polynomial examples, containing no sorry placeholders.
• CompPoly/Univariate/NTT/Forward.lean: This file defines the forward Number Theoretic Transform (NTT) by providing both a direct mathematical specification and an iterative radix-2 implementation. It introduces necessary components such as bit-reversal permutations and butterfly stages, and includes a placeholder theorem forwardImpl_correct (marked with sorry) for the eventual proof of correctness.
• tests/CompPolyTests/Univariate/NTT/Benchmark.lean: This file introduces a benchmarking suite to compare the performance of NTT-based polynomial multiplication against standard "raw" multiplication using the KoalaBear field. It defines several utility functions for performance measurement and NTT domain selection, concluding with an executable script that identifies the efficiency crossover point across various operand sizes.

---

Last updated: 2026-03-21 19:17 UTC.

[dhsorens]

at first glance, the code is clean and well-organized and spec looks to be right. Just a few comments for now

[dhsorens]

it looks like bitRevPermute, butterflyStage, and runStages are nearly identical between Forward.lean and Inverse.lean. to avoid code duplication, it's probably a good idea to factor these into a shared definition (parametric over omega/omegaInv) in a shared file. you might want to put all the shared operations in there as well (e.g. including bitRevNat)

[dhsorens]

it might be worth trying to weaken the field assumption, if possible. we can always generalize results once they're proved though so don't worry too much about this right away

[dhsorens]

note that in Nat, 0 - 1 = 0 which can lead to undesirable behavior. worth checking that this underflow won't affect the rest