From 3e32f3126f783090082b754a31225fe21d682915 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sat, 13 Jun 2026 11:03:19 +0800 Subject: [PATCH 01/19] Add plan for NTT implementation --- TODO_NTT.md | 385 +++++++++++++++++++++++++++++++++++++++++++++ integer/Cargo.toml | 2 +- 2 files changed, 386 insertions(+), 1 deletion(-) create mode 100644 TODO_NTT.md diff --git a/TODO_NTT.md b/TODO_NTT.md new file mode 100644 index 00000000..58ee150a --- /dev/null +++ b/TODO_NTT.md @@ -0,0 +1,385 @@ +# Prime-NTT multiplication for `UBig` — implementation plan + +Goal: add an NTT-based large-integer multiplication path to `dashu-int` that +kicks in above the current Toom-Cook-3 range, using power-of-two Number +Theoretic Transforms over several 64-bit NTT-friendly primes combined with the +Chinese Remainder Theorem (CRT). + +This document is the working spec + checklist. Keep it updated as the work +progresses. Each phase ends with a state where `cargo test -p dashu-int` +passes. + +--- + +## 0. Background and decisions + +- **Word layout.** `Word = u64`, `DoubleWord = u128` on the default build + (`integer/src/arch/generic_64_bit/word.rs`). The 16-/32-bit `force_bits` + builds use smaller `Word`, but `u128` is available on every Rust target, so + the NTT lane arithmetic is always done in `u64`/`u128` regardless of `Word`. +- **Where it plugs in.** Multiplication is dispatched in + `integer/src/mul/mod.rs` by the *smaller* operand length in words: + - `len <= THRESHOLD_SIMPLE (24)` → `simple` + - `<= THRESHOLD_KARATSUBA (192)` → `karatsuba` + - else → `toom_3` (currently unbounded above) + + We add `THRESHOLD_NTT` (in words) above which `toom_3` hands off to the new + `ntt` module. The exact value is tuned by benchmark (Phase 7); start with a + conservative placeholder (e.g. `2048` words ≈ 130k bits). +- **Scheme.** Linear (acyclic) convolution of the two coefficient polynomials, + realised as a cyclic convolution of length `N = next_pow2(La + Lb - 1)` with + zero padding. Computed independently modulo `K` primes, then CRT-combined per + coefficient, then carry-propagated into the output limbs. +- **Why ≥ 3 primes + CRT (recap).** A single ~2^64 prime forces tiny coefficient + chunks. With `K` primes of product `P ≈ 2^(64K)` we can pack `b` bits per + coefficient as long as the largest convolution coefficient + `< N/2 · (2^b − 1)^2 < P`. Three primes (`P ≈ 2^189..192`) give comfortable + headroom for any feasible input and let us pick a larger `b` (fewer + coefficients → smaller transform). `K = 2` is provably sufficient for all + inputs `UBig` can physically hold, but `K = 3` is the default for margin and + speed. + - **`K` vs `K_eff` (avoid const-generic monomorphization).** `K = 3` is the + *fixed array size* of `PRIMES` (a plain `const`, not a const-generic type + parameter), so the transform code is monomorphized once. `select_params` + returns a runtime `K_eff ≤ K`; the per-prime loop iterates + `PRIMES[..K_eff]`. The Phase-7 "drop to 2 primes" optimisation is just + `K_eff = 2` at runtime — no extra monomorphization, no new generic + instantiations. +- **Modulus family: Goldilocks-style** `p = 2^64 − 2^b + 1`. Every member + satisfies the shift-based reduction identity `2^64 ≡ 2^b − 1 (mod p)`, so the + *entire* lane arithmetic is multiplication-free in the reduction step — no + Montgomery form anywhere. We use three members of this single family (below) + so `modarith` is **one** routine parameterised by `b`. + +### Chosen primes (verified) + +Survey of documented choices considered and rejected: +- **"Ultimate NTT" prime** `9223372036737335297 = 549755813881·2^24 + 1`, + `g = 3` (Codeforces entry 75326): `v2` only 24 caps `N` at `2^24`, and it is + not a Solinas form (needs Barrett/Montgomery). Rejected. +- **Classic CRT trio** `998244353`, `985661441`, `754974721`, … : ~30-bit, so + they waste 64-bit lanes and cap `N` low. Rejected. +- **`c·2^32+1` siblings** (e.g. `0xFFFFFFD300000001`): right size and `v2`, but + not Solinas form → would force a Montgomery path. Rejected in favour of the + uniform fast-reduction family below. + +**Selected trio — all of the form `2^64 − 2^b + 1`** (verified prime by +deterministic Miller–Rabin, full-order `2^32`-th root checked): + +| name | `b` | `p` (hex) | `p` (dec) | `v2(p−1)` | gen `g` | `2^32`-th root ω | +|---|---|---|---|---|---|---| +| GL | 32 | `0xFFFFFFFF00000001` | `18446744069414584321` | 32 | 7 | `1753635133440165772` | +| P1 | 34 | `0xFFFFFFFC00000001` | `18446744056529682433` | 34 | 5 | `11315553352654630047` | +| P2 | 40 | `0xFFFFFF0000000001` | `18446742974197923841` | 40 | 19 | `551857376737322389` | + +- `min(v2) = 32` ⇒ transform length up to `2^32` coefficients (≈ ~1 GB operands + at `b_pack = 16`); `P = GL·P1·P2 ≈ 2^192` of CRT headroom. +- All three reduce via `2^64 ≡ 2^bᵢ − 1`. GL's `b = 32` is the cleanest (splits + a 128-bit product into 32-bit limbs, `φ²=φ−1`); `b = 34, 40` need one extra + shift/fold because the split crosses the 32-bit boundary, but stay + multiply-free. Implement the reduction generically over `b` with the GL case + as the well-trodden reference. +- `ω⁻¹` and `N⁻¹` are derived per call (cold, once per prime — not lazy): + `ω_N = pow(ω, 2^32 / N)` (the stored `ω` has exact order `2^32`); the inverse + root is just `ω_N⁻¹ = pow(ω_N, N − 1)` (since `ω_N^N = 1`, no `inv` needed), + and `N⁻¹ = Reducer::inv(N mod p)` once. Use `num_modular::Reducer::{pow, inv}` + for all three. Commit a `verify_primes()` test + (Miller–Rabin + `v2` + exact root order + reduction-identity self-check) + rather than trusting these literals blindly. + +### Open decisions to lock during Phase 1 +- [ ] Final `K` (start 3). +- [ ] Coefficient bit width `b_pack` (**default 16**: 4 coeffs/word, trivial + shift/mask packing). Larger `b_pack` = fewer coefficients (smaller + transform) but needs more headroom and must satisfy + `(N/2)·(2^{b_pack}−1)^2 < P` for the max supported `N`. Candidate values: + - `16` — divides 64, byte-aligned, 4 coeffs/word. **Default unless a + benchmark proves a larger value wins.** + - `24` — divides 64? no, but byte-aligned (3 bytes), 8 coeffs/3 words. + - `21` — **avoid**: does not divide 64 and is not byte-aligned, so + coefficients straddle word *and* byte boundaries → slower, buggier + pack/unpack. Only revisit if its transform-size win clearly beats the + pack cost in benchmarks. +- [ ] Whether to gate the NTT path on `cfg(target_pointer_width)` / `Word` + width, or always enable it (preferred: always enable, since lane math is + `u64`/`u128`). Document the chosen rule. + +### Word-width targets (future work, not near-term) + +The `2^64 − 2^b + 1` primes are correct on every `Word` width (`u64`/`u128` are +universal types), but on narrow targets the `u64×u64→u128` lane multiply is +emulated and slow. Plan: + +- **64-bit `Word`**: primary target, the chosen 3 primes above. +- **32-bit `Word`**: **select a separate set of three ~32-bit Solinas primes + (`2^32 − 2^b + 1`, via `FixedTrinomialSolinas32`)** so the lane multiply is a + native `u32×u32→u64`. Feasible: 3 such primes give `P ≈ 2^96`, which (with + `b_pack = 16`, max coefficient `≈ N·2^32`) is far more headroom than needed — + the transform-length ceiling comes from each prime's `v2(p−1)`, which is ample + for any input a 32-bit target would handle. Requires extending + `FixedTrinomialSolinas32` to `P1 = 32` (same `checked_shl` fix already done for + the 64-bit type) plus a prime/root search. **Not intended for implementation + in the near future** — design the NTT core generic over the prime set so this + can be added later as configuration, not a rewrite. +- **16-bit `Word`**: do not implement an NTT path; fall back to Toom-3. + +--- + +## 1. Module layout + +New directory `integer/src/mul/ntt/` (declare `mod ntt;` in +`integer/src/mul/mod.rs`): + +| File | Responsibility | +|---|---| +| `ntt/mod.rs` | Public entry `add_signed_mul` / `add_signed_mul_same_len`, `memory_requirement_up_to`, `THRESHOLD_NTT`, parameter selection (`b`, `N`, `K`). | +| `ntt/primes.rs` | Const table of the `K` primes: value, `b`, primitive root, `v2(p−1)`, precomputed `2^32`-th root. Includes a `verify_primes` unit test. | +| `ntt/modarith.rs` | Lane arithmetic for the prime `2^64 − 2^b + 1`. **Reuse `num_modular::FixedTrinomialSolinas64<64, b, 1>::{reduce_single, reduce_double}` for the reduction step** (already P1=64-correct and unrolled to the verified fold counts; both are `pub` as of the `checked_shl` fix — verify on the pinned version, see fallback below). Write our **own** lazy `add` / `sub` / `mul` on top so we control deferred reduction; reuse num-modular's `Reducer::{pow, inv}` (fully normalized, called once per prime per call — no lazy variant needed). See note below. | +| `ntt/transform.rs` | Iterative in-place radix-2 forward/inverse NTT, twiddle precomputation, bit-reversal, pointwise multiply. | +| `ntt/pack.rs` | Bit-slice an operand `&[Word]` into `N` coefficients of `b` bits (mod each prime), and the inverse: CRT-combine residues + carry-propagate into the output limbs. | +| `ntt/crt.rs` | Garner CRT for `K` residues → a small (≤ `K`-word) integer per coefficient. | + +Mirror the existing modules' conventions: `#[must_use]` on the `add_signed_mul*` +functions, return `SignedWord` carry, doc comments with complexity, `Buffer` / +`Memory` for scratch (no `Vec`), no `std`. + +--- + +## 2. Math reference (for reviewers and tests) + +Operands `A = sum_i a_i 2^{ib}`, `B = sum_j b_j 2^{jb}` with `0 <= a_i, b_j < 2^b`. +Product `C = A·B = sum_k c_k 2^{kb}` where `c_k = sum_{i+j=k} a_i b_j` is exactly +the linear convolution coefficient, `0 <= c_k < (k+1)·(2^b−1)^2 <= N·(2^b−1)^2`. + +Compute `c_k mod p_t` for each prime `p_t` via length-`N` cyclic convolution +(forward NTT, pointwise product, inverse NTT). Because `N >= La + Lb − 1`, the +cyclic and linear convolutions coincide. CRT recovers exact `c_k < P`. Finally +`C = sum_k c_k 2^{kb}` with carry propagation (coefficients overlap whenever +`bitlen(c_k) > b`). + +Roots: `omega_N = g^{(p−1)/N} mod p` is a primitive `N`-th root of unity; require +`N | 2^{v2(p−1)}`. Inverse transform uses `omega_N^{-1}` and a final scale by +`N^{-1} mod p`. + +--- + +## 3. Phase plan (each phase is independently testable) + +> **Scheduling note.** Phases 2 (transform arithmetic) and 3 (pack / unpack / +> CRT) have no dependency on each other — both only need Phase 1's `modarith` +> and `primes`. They can be built and tested in parallel, then joined in +> Phase 4. Phase 1 must land first; Phases 5–7 follow Phase 4. + +### Phase 1 — Primes, modular arithmetic, parameter selection +- [ ] `ntt/primes.rs`: define `const PRIMES: [NttPrime; K]` from the "Chosen + primes" table. Each entry stores `p`, the form exponent `b`, primitive + root `g`, `v2(p−1)`, and the precomputed `2^32`-th root `ω`. (No Montgomery + constants — the family needs none.) +- [ ] Add `#[test] fn verify_primes()` re-checking each entry: primality + (Miller–Rabin over fixed bases), `p == 2^64 − 2^b + 1 < 2^64`, + `v2(p−1) >= MAX_LOG_N (=32)`, stored `g` generates the order-`2^{v2}` + subgroup, `ω` has exact order `2^32`, and the reduction identity + `2^64 ≡ 2^b − 1 (mod p)` holds. Do not trust the literals without it. +- [ ] `ntt/modarith.rs`: **delegate the reduction** to + `num_modular::FixedTrinomialSolinas64<64, b, 1>` — its `reduce_single` + (≤ `2^64` → `[0, p)`) and `reduce_double` (`u128` product → `[0, p)`) are + already correct for `P1 = 64` and straight-line unrolled (3 folds for + `b = 32`, 4 for `b = 34, 40`). Do **not** re-derive the shift/fold here. + Both methods are generated `pub` by the `impl_fixed_trinomial_solinas!` + macro, so they are directly callable from `dashu-int`. + - **Fallback if upstream visibility ever regresses:** the reduction is + ~20 lines per arm; copy it verbatim into `modarith.rs` (it is simple + enough to own in-tree, and the only coupling point). Pin/assert the + `num-modular` version in `integer/Cargo.toml` so a downgrade can't + silently break the `pub` assumption. + - We write our **own lazy `add` / `sub` / `mul`** (not num-modular's) + because its `Reducer` API fully normalizes to `[0, p)` after every op + and exposes no partially-reduced form. Ours keep values lazily in + `[0, 2p)` (or `[0, 4p)`) and only call `reduce_*` / a final conditional + subtract when needed — this is the Harvey-style lazy reduction that the + NTT butterflies depend on. `mul` = `u128` widening multiply → + `reduce_double`; `add` / `sub` = wrapping add/sub with deferred + normalization. + - **`pow` / `inv` are NOT lazy and are NOT ours.** They are called once + per prime per multiplication (`ω_N = g^{(p−1)/N}`, `N^{-1}`), never in + the butterfly hot loop, so use `num_modular::Reducer::{pow, inv}` + directly (fully normalized). No partial-reduction benefit there. + - Rationale: reduction is the subtle, already-tested part (reuse it); the + lazy add/sub/mul wrapper is trivial and must be ours to control the + normalization schedule; pow/inv are cold and reused as-is. +- [ ] `ntt/mod.rs`: `select_params(la_bits, lb_bits) -> (b_pack, N, K_eff)` with + the headroom assertion `(N as u128 / 2) * (2^{b_pack} − 1)^2 < P` (may drop + to fewer primes for smaller inputs later). +- [ ] Unit tests: our lazy `add`/`sub`/`mul` (after a final normalize) agree + with `FixedTrinomialSolinas64`'s fully-reduced `add`/`sub`/`mul` and a + `u128`/`u256` reference, across the `[0, 2p)` input range for each `b`; + `Reducer::{pow, inv}` round-trip (`inv(x)·x ≡ 1`, `pow(g, p−1) ≡ 1`); + `verify_primes`. + +### Phase 2 — Forward/inverse NTT +- [ ] `ntt/transform.rs`: iterative Cooley–Tukey radix-2 forward NTT in place, + decimation-in-time with bit-reversal permutation; inverse NTT + (conjugate twiddles + scale by `N^{-1}`). +- [ ] Twiddle factors: precompute the `omega_N^k` table per prime into scratch + once per call (length `N/2`). +- [ ] `pointwise_mul(a_hat, b_hat)` via `modarith::mul` (lazy; normalize at the + end of the inverse transform). +- [ ] Tests: `inverse(forward(x)) == x`; NTT-based cyclic convolution of small + random vectors equals the schoolbook cyclic convolution mod `p`; check the + length-2 and length-power-of-two edge cases. + +### Phase 3 — Packing / unpacking + CRT +- [ ] `ntt/pack.rs::pack`: read `b`-bit coefficients out of `&[Word]` + (bit-level slicing across word boundaries; works for any `Word` width), + reduce mod each prime (a `b_pack`-bit value is already `< p`, so this is a + copy), write into the length-`N` (zero-padded) lane buffers. +- [ ] `ntt/crt.rs`: Garner combine `K` residues of one coefficient → an integer + of ≤ `K` words (value `< P`). +- [ ] `ntt/pack.rs::unpack_accumulate`: for each `k`, add `c_k << (k·b)` bits + into the output limbs with carry propagation. Implement as a streaming + shifted add (reuse `add::add_*` helpers / `shift`). +- [ ] Tests: `unpack_accumulate(pack(x)) == x` identity for the + no-multiplication case (coefficients copied straight through CRT), and a + direct check that pack→CRT→unpack reconstructs a known convolution. + +### Phase 4 — Wire the full multiply +- [ ] `ntt/mod.rs::add_signed_mul_same_len` and `add_signed_mul`: orchestrate + select_params → per-prime (pack, forward, pointwise, inverse) → CRT per + coefficient → unpack/accumulate into a temp product buffer → fold into `c` + via `add::add_signed_*` honoring `sign`. Return the carry as the other + algorithms do. +- [ ] **Unequal-length entry point — do NOT blindly copy `toom_3`'s chunking.** + Unlike Toom-3/Karatsuba (defined on equal-length operands, hence + `helpers::add_signed_mul_split_into_chunks` slices the long operand into + balanced pieces), a single NTT convolution handles unequal lengths + natively: pad both operands to one `N = next_pow2(La + Lb − 1)`, one + forward transform each, pointwise product, one inverse. So the *default* + unequal path is a single transform — **no chunking**. + - Dispatch keys on the smaller operand `b.len()`, so when the NTT path is + entered `b` is already huge. Chunking `a` into `b.len()`-sized pieces + via the stock helper would run `⌈La/Lb⌉` separate NTTs **and + re-transform `b` on every chunk**, roughly doubling work for lopsided + large×large products — it throws away NTT's single-big-transform win. + - `ntt::add_signed_mul` (unequal) and `ntt::add_signed_mul_same_len` + (equal) therefore share one core that takes `(La, Lb)` and transforms + over `N = next_pow2(La + Lb − 1)` directly. Honor the same contract as + the other algorithms: `c.len() == La + Lb`, accumulate `sign * a * b` + into `c`, return the `SignedWord` carry. + - **Only** fall back to chunking for *extreme* imbalance (`La ≫ Lb`, e.g. + `La > c · Lb` for some tuned `c`), where many balanced `~2·Lb` NTTs beat + one padded `~La` transform. If/when we do, forward-transform `b` **once** + and reuse the cached `b_hat` across chunks — i.e. a purpose-built loop, + not the stock `add_signed_mul_split_into_chunks` (which re-transforms + `b`). Treat this as a Phase-7 tuning option, not the initial wiring. +- [ ] `ntt/mod.rs::memory_requirement_up_to(n)`: deterministic upper bound on + scratch, mirroring the style of `toom_3::memory_requirement_up_to` + (returns a `Layout`). It **must** be an exact upper bound — `Memory` + `expect`s on underflow. + - **Draft closed-form bound (words):** + `2·N` (one `a`-lane + one `b`-lane buffer, processed one prime at a + time so they are reused across the `K_eff` primes — not `K·N`) + `+ N/2` (twiddle table for the current prime) + `+ K` (per-coefficient CRT temp) + `+ (La + Lb)` (product accumulation buffer) + `≈ 2.5·N + La + Lb + K`. + If lanes for all primes are kept live simultaneously (simpler, no + re-pack per prime) the lane term becomes `2·K·N`; decide which during + implementation and bound accordingly. + - **Worst-case over `b_pack`.** `N = next_pow2(ceil((La+Lb)·WORD_BITS / + b_pack) + 1)`. `memory_requirement_up_to(n)` is called before + `select_params` runs, so it must bound `N` over **every** `b_pack` the + selector may choose for inputs up to `n` words — i.e. use the + *smallest* admissible `b_pack` (largest `N`). Pin a `B_PACK_MIN` + constant (= 16) and compute the bound from it; `select_params` may then + only ever pick `b_pack ≥ B_PACK_MIN`. +- [ ] **Carving scratch from the linear `Memory` arena.** `Memory` + (`integer/src/memory.rs`) is a *linear bump allocator*, not a pool of + independent `Buffer`s: each `allocate_slice`/`allocate_slice_fill` hands + out the next region and returns the remainder. Order matters. Plan the + carve explicitly: allocate longer-lived regions first (twiddle table, + product buffer) then the per-prime lane buffers from the remaining region + inside the prime loop (so they are reused each iteration). `Buffer` is + only for *owned* growable word arrays (e.g. a returned product); transform + scratch lives in the `Memory` arena. Document the carve order next to + `memory_requirement_up_to` so the two stay in sync. + +### Phase 5 — Dispatch + thresholds +- [ ] In `integer/src/mul/mod.rs`: add `THRESHOLD_NTT`, declare `mod ntt;`. +- [ ] Extend `add_signed_mul`, `add_signed_mul_same_len`, and both + `memory_requirement_*` to route `len > THRESHOLD_NTT` to `ntt`. +- [ ] Keep `toom_3` as the fallback if NTT parameter selection fails any + precondition (defensive; should not happen below `2^32` coefficients). + +### Phase 6 — Correctness validation +- [ ] Extend `integer/tests/mul.rs` with cases straddling `THRESHOLD_NTT` + (lengths `T−1`, `T`, `T+1`, `2T`, asymmetric `a`/`b` lengths, operands + with high/low zero limbs, near power-of-two `N`). +- [ ] Differential test: random `UBig`s of increasing size, assert + `ntt_product == reference_product` where reference is the existing + `multiply` forced through Toom-3 (or compute via a smaller-threshold + build). Cover the coefficient-overflow boundary explicitly (all-ones + operands at the max supported `N`). +- [ ] Run on a 32-bit lane build too: `cargo test -p dashu-int` with + `RUSTFLAGS="--cfg force_bits=\"32\""` (and `16`) to confirm packing is + `Word`-width agnostic. + +### Phase 7 — Tuning + optimisation (after correctness is green) +- [ ] Benchmark `THRESHOLD_NTT` crossover against Toom-3 using + `integer/benches/primitive.rs` (extend the mul benchmark to larger sizes); + pick the value where NTT wins. +- [ ] Optimisations to layer in, measuring each: + - [ ] Harvey lazy-reduction butterflies (defer mod in inner loops). + - [ ] Specialise the `b = 32` (Goldilocks) lane's reduction to 32-bit-limb + form (`φ²=φ−1`), since it avoids the extra cross-boundary fold that + `b = 34, 40` need. + - [ ] Use the shift-expressible roots of unity where applicable (powers of two + are roots in this family) to replace some twiddle multiplies. + - [ ] Radix-4 / split-radix transform. + - [ ] Drop to `K_eff = 2` primes automatically when headroom allows (smaller + inputs) to halve the transform work. + - [ ] **Squaring specialization.** `UBig::square()` / the `a == b` case needs + only one forward transform per prime (not two), then a pointwise + *square* and one inverse — roughly 2/3 the transform cost. Wire an + `ntt::square` path (and route `UBig::square` to it above + `THRESHOLD_NTT`) once the general multiply is correct. + - [ ] Reuse/cancel allocations; ensure scratch stays within the `Memory` + arena. + +--- + +## 4. Constraints & pitfalls (project-specific) + +- **`no_std`**: only `core` + `alloc`. `u128` arithmetic is fine. The lane + *reduction* is reused from `num_modular::FixedTrinomialSolinas64<64, b, 1>` + (`reduce_single` / `reduce_double`; already a dependency of `dashu-int`, and + `no_std`). The lazy `add`/`sub`/`mul` wrapper around it is ours, in-tree, so + no Montgomery and no extra dependency. Requires the `num-modular` version that + (a) supports `P1 = 64` (the `checked_shl` fix + `S64_4` Goldilocks tests) and + (b) exposes `reduce_single` / `reduce_double` as `pub` (it does as of that + fix). Pin this version in `integer/Cargo.toml`; if the `pub` assumption ever + regresses, fall back to the in-tree reduction copy (see Phase 1). +- **MSRV**: keep within the README MSRV (do not bump). Avoid `const` features + newer than MSRV; plain `const` tables are fine. +- **Scratch**: use `Buffer` / `MemoryAllocation` / the threaded `Memory` arena, + never `Vec` (per `AGENTS.md`). `memory_requirement_*` must be an exact + upper bound or the arena `expect` will panic. +- **Sign / accumulate contract**: the entry points are `c += sign * a * b` + returning a `SignedWord` carry — match `toom_3`/`karatsuba` exactly so the + recursive callers and `multiply()` keep working. +- **Changelog**: add an `### Add` entry under `## Unreleased` in + `integer/CHANGELOG.md` ("NTT-based multiplication for very large integers") + as part of the same commit (per `AGENTS.md`). +- **CI parity**: before declaring done, run + - `cargo test --workspace --exclude dashu-python` + - `cargo clippy --all-features --all-targets --workspace --exclude dashu-python -- -D warnings` + - `cargo fmt --all -- --check` + +--- + +## 5. Definition of done + +- [ ] All phases checked off; NTT path active above `THRESHOLD_NTT`. +- [ ] Differential tests pass on 64-, 32-, and 16-bit lane builds. +- [ ] Benchmarks show NTT beats Toom-3 at and above the chosen threshold. +- [ ] Clippy clean, fmt clean, changelog updated. +- [ ] No `std` usage, no `Vec` scratch, MSRV preserved. diff --git a/integer/Cargo.toml b/integer/Cargo.toml index 047e701e..da816819 100644 --- a/integer/Cargo.toml +++ b/integer/Cargo.toml @@ -30,7 +30,7 @@ dashu-base = { version = "0.4.1", default-features = false, path = "../base" } cfg-if = { version = "1.0.0" } static_assertions = { version = "1.1" } rustversion = { version = "1.0.0" } -num-modular = { version = "0.6.1" } +num-modular = { version = "0.6.2", path = "../../num-modular" } # stable dependencies num-order = { optional = true, version = "1.2.0", default-features = false } From 468777315a5d2c5c13568e835c838ca9ce69c3e3 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sat, 13 Jun 2026 11:03:19 +0800 Subject: [PATCH 02/19] WIP: implemented ntt mul --- integer/CHANGELOG.md | 1 + integer/src/mul/mod.rs | 19 +- integer/src/mul/ntt/crt.rs | 158 +++++++++ integer/src/mul/ntt/mod.rs | 531 +++++++++++++++++++++++++++++++ integer/src/mul/ntt/pack.rs | 194 +++++++++++ integer/src/mul/ntt/primes.rs | 229 +++++++++++++ integer/src/mul/ntt/transform.rs | 339 ++++++++++++++++++++ 7 files changed, 1467 insertions(+), 4 deletions(-) create mode 100644 integer/src/mul/ntt/crt.rs create mode 100644 integer/src/mul/ntt/mod.rs create mode 100644 integer/src/mul/ntt/pack.rs create mode 100644 integer/src/mul/ntt/primes.rs create mode 100644 integer/src/mul/ntt/transform.rs diff --git a/integer/CHANGELOG.md b/integer/CHANGELOG.md index 71fb36ab..f7cd5beb 100644 --- a/integer/CHANGELOG.md +++ b/integer/CHANGELOG.md @@ -3,6 +3,7 @@ ## Unreleased ### Add +- NTT-based multiplication for very large integers (above 2048 words), using three Solinas primes of the form `2^64 − 2^b + 1` combined with the Chinese Remainder Theorem. - `UBig::from_u64` and `IBig::from_i64`, const on 32-bit and 64-bit targets. ### Improve diff --git a/integer/src/mul/mod.rs b/integer/src/mul/mod.rs index 5b2b1eaf..301273dc 100644 --- a/integer/src/mul/mod.rs +++ b/integer/src/mul/mod.rs @@ -22,8 +22,13 @@ const_assert!(THRESHOLD_SIMPLE + 1 >= karatsuba::MIN_LEN); const THRESHOLD_KARATSUBA: usize = 192; const_assert!(THRESHOLD_KARATSUBA + 1 >= toom_3::MIN_LEN); +/// If smaller operand length > this, NTT multiplication will be used. +const THRESHOLD_NTT: usize = ntt::THRESHOLD_NTT; +const_assert!(THRESHOLD_NTT + 1 >= toom_3::MIN_LEN); + mod helpers; mod karatsuba; +pub(crate) mod ntt; mod simple; mod toom_3; @@ -156,13 +161,15 @@ pub fn sub_mul_word_same_len_in_place(words: &mut [Word], mult: Word, rhs: &[Wor } /// Temporary scratch space required for multiplication. -pub fn memory_requirement_up_to(_total_len: usize, smaller_len: usize) -> Layout { +pub fn memory_requirement_up_to(total_len: usize, smaller_len: usize) -> Layout { if smaller_len <= THRESHOLD_SIMPLE { memory::zero_layout() } else if smaller_len <= THRESHOLD_KARATSUBA { karatsuba::memory_requirement_up_to(smaller_len) - } else { + } else if smaller_len <= THRESHOLD_NTT { toom_3::memory_requirement_up_to(smaller_len) + } else { + ntt::memory_requirement_up_to(total_len, smaller_len) } } @@ -199,8 +206,10 @@ pub fn add_signed_mul<'a>( simple::add_signed_mul(c, sign, a, b, memory) } else if b.len() <= THRESHOLD_KARATSUBA { karatsuba::add_signed_mul(c, sign, a, b, memory) - } else { + } else if b.len() <= THRESHOLD_NTT { toom_3::add_signed_mul(c, sign, a, b, memory) + } else { + ntt::add_signed_mul(c, sign, a, b, memory) } } @@ -222,7 +231,9 @@ pub fn add_signed_mul_same_len( simple::add_signed_mul_same_len(c, sign, a, b, memory) } else if n <= THRESHOLD_KARATSUBA { karatsuba::add_signed_mul_same_len(c, sign, a, b, memory) - } else { + } else if n <= THRESHOLD_NTT { toom_3::add_signed_mul_same_len(c, sign, a, b, memory) + } else { + ntt::add_signed_mul_same_len(c, sign, a, b, memory) } } diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs new file mode 100644 index 00000000..9a32ee36 --- /dev/null +++ b/integer/src/mul/ntt/crt.rs @@ -0,0 +1,158 @@ +//! Garner CRT: combine `K` residues modulo `K` primes into a small integer. +//! +//! Uses num-modular's general modular arithmetic traits since CRT is +//! called once per coefficient, not in the hot loop. +#![allow(clippy::unnecessary_cast)] + +use num_modular::{ModularCoreOps, ModularUnaryOps}; + +/// Precomputed constants for Garner CRT with a fixed prime set. +pub struct CrtConstants { + /// `inv(p_i mod p_j)` for i < j. + pub inv_ij: [[u64; 3]; 3], +} + +impl CrtConstants { + /// Precompute Garner constants for the given primes. + pub fn new(primes: &[u64]) -> Self { + let k = primes.len(); + let mut inv_ij = [[0u64; 3]; 3]; + for i in 0..k { + for j in (i + 1)..k { + let p_i_mod_pj = primes[i] % primes[j]; + inv_ij[i][j] = p_i_mod_pj.invm(&primes[j]).expect("primes not coprime"); + } + } + CrtConstants { inv_ij } + } +} + +/// Combine `K` residues into a small integer (< P) using Garner's algorithm. +/// +/// The result is returned as a little-endian `Vec` because +/// `P ≈ 2^{64K}` may exceed a single `u64`. +pub fn garner_combine( + residues: &[u64], + primes: &[u64], + constants: &CrtConstants, +) -> alloc::vec::Vec { + let k = residues.len(); + assert!(k <= 3, "CRT supports up to 3 primes"); + assert_eq!(primes.len(), k); + + let mut result = [0u64; 3]; + let p0 = primes[0]; + + // x_0 = r_0 + result[0] = residues[0] % p0; + + if k == 1 { + return result.to_vec(); + } + + // t_1 = (r_1 - x_0) * inv(p0 mod p1) mod p1 + let p1 = primes[1]; + let x0_mod_p1 = result[0] % p1; + let diff1 = residues[1].subm(x0_mod_p1, &p1); + let t1 = diff1.mulm(constants.inv_ij[0][1], &p1); + + // x_1 = x_0 + t_1 * p0 + add_128_to_192(&mut result, (t1 as u128) * (p0 as u128)); + + if k == 2 { + return result.to_vec(); + } + + // t_2 = (r_2 - x_1) * inv(p0*p1 mod p2) mod p2 + let p2 = primes[2]; + let x1_mod_p2 = mod_192_by_u64(&result, p2); + let diff2 = residues[2].subm(x1_mod_p2, &p2); + let inv_p0_mod_p2 = constants.inv_ij[0][2]; + let inv_p1_mod_p2 = constants.inv_ij[1][2]; + let inv_prod = inv_p0_mod_p2.mulm(inv_p1_mod_p2, &p2); + let t2 = diff2.mulm(inv_prod, &p2); + + // x_2 = x_1 + t_2 * p0 * p1 + let pp = (primes[0] as u128) * (primes[1] as u128); + let t2_64 = t2 as u64; + let pp_lo = pp as u64; + let pp_hi = (pp >> 64) as u64; + let m_lo_full = (t2_64 as u128) * (pp_lo as u128); + let m_lo = (m_lo_full >> 64) as u64; + let lo = m_lo_full as u64; + let m_hi_full = (t2_64 as u128) * (pp_hi as u128); + let hi = (m_hi_full >> 64) as u64; + let m_hi = m_hi_full as u64; + let (mid, c) = m_lo.overflowing_add(m_hi); + let hi_word = hi.wrapping_add(c as u64); + let (r0, c0) = result[0].overflowing_add(lo); + result[0] = r0; + let (r1, c1) = result[1].overflowing_add(mid.wrapping_add(c0 as u64)); + result[1] = r1; + result[2] = result[2].wrapping_add(hi_word.wrapping_add(c1 as u64)); + + result.to_vec() +} + +fn add_128_to_192(result: &mut [u64; 3], term: u128) { + let lo = term as u64; + let hi = (term >> 64) as u64; + let (r0, c0) = result[0].overflowing_add(lo); + result[0] = r0; + let (r1, c1) = result[1].overflowing_add(hi.wrapping_add(c0 as u64)); + result[1] = r1; + result[2] = result[2].wrapping_add(c1 as u64); +} + +fn mod_192_by_u64(x: &[u64; 3], m: u64) -> u64 { + let m128 = m as u128; + let mut r: u128 = 0; + for &word in x.iter().rev() { + r = (r << 64) | (word as u128); + r %= m128; + } + r as u64 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_garner_two_primes() { + let primes = vec![3u64, 5u64]; + let constants = CrtConstants::new(&primes); + // x ≡ 2 mod 3, x ≡ 3 mod 5 → x = 8 + let residues = vec![2u64, 3u64]; + let result = garner_combine(&residues, &primes, &constants); + let x = result[0] as u128; + assert_eq!(x, 8); + } + + #[test] + fn test_garner_three_primes() { + let primes = vec![3u64, 5u64, 7u64]; + let constants = CrtConstants::new(&primes); + // x ≡ 2 (mod 3), x ≡ 3 (mod 5), x ≡ 4 (mod 7) → x = 53 + let residues = vec![2u64, 3u64, 4u64]; + let result = garner_combine(&residues, &primes, &constants); + let x = result[0] as u128; + assert_eq!(x, 53); + } + + #[test] + fn test_garner_with_ntt_primes() { + use crate::mul::ntt::primes::PRIMES; + let primes: Vec = PRIMES.iter().map(|np| np.p).collect(); + let constants = CrtConstants::new(&primes); + let residues = vec![12345u64, 67890u64, 11111u64]; + let result = garner_combine(&residues, &primes, &constants); + for (i, &p) in primes.iter().enumerate() { + let mut rem: u128 = 0; + for &word in result.iter().rev() { + rem = ((rem << 64) | (word as u128)) % (p as u128); + } + assert_eq!(rem as u64, residues[i], "CRT mismatch for prime {i}"); + } + } +} diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs new file mode 100644 index 00000000..344a62ca --- /dev/null +++ b/integer/src/mul/ntt/mod.rs @@ -0,0 +1,531 @@ +//! NTT-based multiplication for very large integers. +//! +//! Uses Number Theoretic Transforms over several 64-bit primes of the form +//! `2^64 - 2^b + 1` combined with the Chinese Remainder Theorem (CRT). + +use crate::{ + add, + arch::word::{SignedWord, Word}, + memory::{self, Memory}, + Sign::{self, *}, +}; +use alloc::alloc::Layout; + +mod crt; +mod pack; +mod primes; +mod transform; + +use crate::mul::ntt::crt::CrtConstants; +pub use primes::{K, PRIMES}; + +/// Minimum smaller-operand length (in words) for the NTT path. +pub const THRESHOLD_NTT: usize = 2048; + +/// Smallest admissible coefficient bit width (used for worst-case memory bound). +const B_PACK_MIN: u32 = 16; + +/// Maximum `log2(transform length)`, set by `min(v2) = 32` across all primes. +const MAX_LOG_N: u32 = 32; + +/// Select NTT parameters for operands with the given word lengths. +/// +/// Returns `(b_pack, N, K_eff)`. +pub fn select_params(la_words: usize, lb_words: usize) -> (u32, usize, usize) { + let b_pack = B_PACK_MIN; + let word_bits = Word::BITS; + + let la_bits = la_words as u64 * word_bits as u64; + let lb_bits = lb_words as u64 * word_bits as u64; + + let coeffs_a = (la_bits + b_pack as u64 - 1) / b_pack as u64; + let coeffs_b = (lb_bits + b_pack as u64 - 1) / b_pack as u64; + let total_coeffs = (coeffs_a + coeffs_b - 1) as usize; + let n = total_coeffs.next_power_of_two().max(2); + + assert!( + (n.trailing_zeros()) <= MAX_LOG_N, + "N = {n} too large for prime set (max log2 = {MAX_LOG_N})" + ); + + let k_eff = K; + + // Headroom check: max convolution coefficient < product of K_eff primes. + // max_coeff fits in u128; compare against smallest prime p0. + let max_coeff = (n as u128 / 2) * ((1u128 << b_pack) - 1) * ((1u128 << b_pack) - 1); + let p0 = PRIMES[0].p as u128; + assert!( + max_coeff < p0, + "headroom check failed: max coeff {max_coeff} >= smallest prime {p0}" + ); + + (b_pack, n, k_eff) +} + +/// Estimate bit length from a word slice (excludes leading zeros). +fn bit_len(words: &[Word]) -> u64 { + let leading_zeros = words.iter().rev().take_while(|&&w| w == 0).count(); + let used = words.len() - leading_zeros; + if used == 0 { + return 0; + } + let hi_word = words[used - 1]; + let hi_bits = Word::BITS - hi_word.leading_zeros(); + (used as u64 - 1) * Word::BITS as u64 + hi_bits as u64 +} + +/// Count number of coefficients needed for a given bit length. +fn coeff_count(bit_len: u64, b_pack: u32) -> usize { + ((bit_len + b_pack as u64 - 1) / b_pack as u64) as usize +} + +/// Worst-case scratch memory bound. +pub fn memory_requirement_up_to(total_len: usize, _smaller_len: usize) -> Layout { + let word_bits = Word::BITS; + let max_coeffs = + (total_len as u64 * word_bits as u64 + B_PACK_MIN as u64 - 1) / B_PACK_MIN as u64; + let n_max = ((max_coeffs + 1) as usize).next_power_of_two().max(2); + + // Everything is in u64 units for simplicity. + let lanes_u64 = 2 * n_max; // a_lane + b_lane + let residues_u64 = K * n_max; // per-prime inverse results + let product_u64 = total_len; // product buffer (Word=u64 on 64-bit, else u64 takes more space) + + // On 64-bit targets Word = u64. On narrow targets (Word < u64), + // we need extra space for the u64 allocations. Use the maximum + // of Word and u64 sizes. + let u64_bytes = 8usize; + let word_bytes = core::mem::size_of::(); + let factor = (u64_bytes + word_bytes - 1) / word_bytes; + let total_words = product_u64 + (lanes_u64 + residues_u64) * factor; + + memory::array_layout::(total_words) +} + +/// `c += sign * a * b` with equal-length operands. +/// +/// Returns carry. +#[must_use] +pub fn add_signed_mul_same_len( + c: &mut [Word], + sign: Sign, + a: &[Word], + b: &[Word], + memory: &mut Memory, +) -> SignedWord { + let n = a.len(); + debug_assert!(b.len() == n && c.len() == 2 * n); + add_signed_mul_impl(c, sign, a, b, memory) +} + +/// `c += sign * a * b` (general, a may be longer than b). +/// +/// Returns carry. +#[must_use] +pub fn add_signed_mul( + c: &mut [Word], + sign: Sign, + a: &[Word], + b: &[Word], + memory: &mut Memory, +) -> SignedWord { + debug_assert!(a.len() >= b.len() && c.len() == a.len() + b.len()); + add_signed_mul_impl(c, sign, a, b, memory) +} + +/// Core implementation: c += sign * a * b. +/// +/// Does a single NTT convolution of the full operands (no chunking). +fn add_signed_mul_impl( + c: &mut [Word], + sign: Sign, + a: &[Word], + b: &[Word], + memory: &mut Memory, +) -> SignedWord { + let la = a.len(); + let lb = b.len(); + + // Skip zero-length or zero-value operands + if la == 0 || lb == 0 { + return 0; + } + + let (b_pack, nn, k_eff) = select_params(la, lb); + let la_bits = bit_len(a); + let lb_bits = bit_len(b); + if la_bits == 0 || lb_bits == 0 { + return 0; + } + + let coeffs_a = coeff_count(la_bits, b_pack); + let coeffs_b = coeff_count(lb_bits, b_pack); + let output_coeffs = coeffs_a + coeffs_b - 1; + + // CRT constants + let primes_p: alloc::vec::Vec = PRIMES[..k_eff].iter().map(|np| np.p).collect(); + let crt_constants = CrtConstants::new(&primes_p); + + // ---- Memory carve (longest-lived first) ---- + // All buffers are u64 since lane arithmetic is always u64. + + // 1. Product buffer + let prod_len = la + lb; + let (prod, mut mem) = memory.allocate_slice_fill::(prod_len, 0); + + // 2. Residue storage (per-prime inverse results, as u64) + let residues_len = k_eff * nn; + let (residues, mut mem) = mem.allocate_slice_fill::(residues_len, 0); + + // 3. Lane buffers (reused across primes, as u64) + let (a_lane, mut mem) = mem.allocate_slice_fill::(nn, 0); + let (b_lane, _mem) = mem.allocate_slice_fill::(nn, 0); + + // ---- Per-prime transforms ---- + for (pi, prime) in PRIMES[..k_eff].iter().enumerate() { + let p = prime.p; + let b_exp = prime.b; + + // Precompute twiddles + let fwd_twiddles = transform::precompute_twiddles(nn, p, b_exp, prime.omega_2_32, false); + let inv_twiddles = transform::precompute_twiddles(nn, p, b_exp, prime.omega_2_32, true); + + // Pack operands into lane buffers + pack_into(a, b_pack, a_lane); + pack_into(b, b_pack, b_lane); + + // Forward NTT + transform::bit_reverse(a_lane); + transform::bit_reverse(b_lane); + transform::forward(a_lane, &fwd_twiddles, p, b_exp); + transform::forward(b_lane, &fwd_twiddles, p, b_exp); + transform::pointwise_mul(a_lane, b_lane, b_exp); + transform::inverse(a_lane, &inv_twiddles, p, b_exp); + + // Store residues for this prime + let offset = pi * nn; + residues[offset..offset + nn].copy_from_slice(a_lane); + } + + // ---- CRT per coefficient + accumulate ---- + // We'll accumulate each coefficient into the product buffer with + // b_pack-bit shift. + let output_words = la + lb; + for k in 0..output_coeffs { + let mut coeff_residues = [0u64; 3]; + #[allow(clippy::needless_range_loop)] + for pi in 0..k_eff { + let offset = pi * nn; + coeff_residues[pi] = residues[offset + k]; + } + let crt_val = crt::garner_combine(&coeff_residues[..k_eff], &primes_p, &crt_constants); + + // Unpack-accumulate this coefficient into prod + // crt_val is a small integer (≤ 3 u64 words) + add_shifted_to_prod(prod, &crt_val, k, b_pack); + } + + // ---- Fold product into c with sign ---- + // Convert u64 slice to Word slice for the add function. + // On 64-bit targets these are the same type. + assert_eq!( + core::mem::size_of::(), + core::mem::size_of::(), + "NTT requires 64-bit Word" + ); + // SAFETY: Word and u64 have the same size (asserted above) and + // prod is allocated with u64 alignment, compatible with Word. + let prod_words: &[Word] = + unsafe { core::slice::from_raw_parts(prod.as_ptr() as *const Word, output_words) }; + match sign { + Positive => add::add_signed_in_place(c, Positive, prod_words), + Negative => add::add_signed_in_place(c, Negative, prod_words), + } +} + +/// Pack word slice into coefficient buffer (viewed as u64). +fn pack_into(words: &[Word], b_pack: u32, out: &mut [u64]) { + let packed = pack::pack(words, b_pack, out.len()); + out.copy_from_slice(&packed); +} + +/// Add a small multi-word integer (up to 3 u64 words) to `prod`, shifted +/// left by `k * b_pack` bits. +fn add_shifted_to_prod(prod: &mut [u64], val: &[u64], k: usize, b_pack: u32) { + if val.is_empty() { + return; + } + let shift_bits = (k as u32).wrapping_mul(b_pack); + let word_idx = (shift_bits / 64) as usize; + let bit_shift = shift_bits % 64; + + let mut carry: u64 = 0; + for (vi, &v) in val.iter().enumerate() { + let idx = word_idx + vi; + if idx >= prod.len() { + return; + } + let v128 = v as u128; + + if bit_shift == 0 { + let sum = v128.wrapping_add(carry as u128); + let (r, c) = prod[idx].overflowing_add(sum as u64); + prod[idx] = r; + carry = (sum >> 64) as u64 + c as u64; + } else { + // v << bit_shift has high bits = v >> (64 - bit_shift) = lo_carry. + // No separate hi — lo_carry IS the high part. + let lo = v128 << bit_shift; + let sum = lo.wrapping_add(carry as u128); + let lo_carry = (sum >> 64) as u64; + let lo_word = sum as u64; + + let (r, c1) = prod[idx].overflowing_add(lo_word); + prod[idx] = r; + carry = lo_carry + c1 as u64; + + // Propagate to next word + if idx + 1 < prod.len() && carry != 0 { + let (r2, c2) = prod[idx + 1].overflowing_add(carry); + prod[idx + 1] = r2; + carry = c2 as u64; + } + } + } + + // Propagate final carry + let mut idx = word_idx + val.len(); + while carry != 0 && idx < prod.len() { + let (r, c) = prod[idx].overflowing_add(carry); + prod[idx] = r; + carry = c as u64; + idx += 1; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_select_params_small() { + let (b_pack, n, k_eff) = select_params(10, 10); + assert_eq!(b_pack, 16); + assert!(n >= 2 && n.is_power_of_two()); + assert_eq!(k_eff, K); + } + + #[test] + fn test_select_params_large() { + let (b_pack, n, k_eff) = select_params(THRESHOLD_NTT, THRESHOLD_NTT); + assert_eq!(b_pack, 16); + assert!(n.is_power_of_two()); + assert_eq!(k_eff, K); + let coeffs_a = (THRESHOLD_NTT * Word::BITS as usize + 15) / 16; + let coeffs_b = coeffs_a; + let min_n = (coeffs_a + coeffs_b).next_power_of_two().max(2); + assert!(n >= min_n, "n={n} < min_n={min_n}"); + } + + #[test] + fn test_headroom_holds() { + let la = THRESHOLD_NTT; + let lb = THRESHOLD_NTT; + let (b_pack, n, _k_eff) = select_params(la, lb); + let max_coeff = (n as u128 / 2) * ((1u128 << b_pack) - 1) * ((1u128 << b_pack) - 1); + let p0 = PRIMES[0].p as u128; + assert!(max_coeff < p0, "headroom violation: max_coeff={max_coeff} >= p0={p0}"); + } + + #[test] + fn test_bit_len() { + assert_eq!(bit_len(&[]), 0); + assert_eq!(bit_len(&[0]), 0); + assert_eq!(bit_len(&[1]), 1); + assert_eq!(bit_len(&[0, 1]), 65); + assert_eq!(bit_len(&[0xFF, 0]), 8); + } + + #[test] + fn test_ntt_multiply_one_word() { + // Simplest case: single-word operands + let a: Vec = vec![3]; + let b: Vec = vec![5]; + let mut c = vec![0u64; 2]; + let layout = memory_requirement_up_to(c.len(), b.len()); + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut memory = alloc.memory(); + let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + assert_eq!(carry, 0); + assert_eq!(c[0], 15); + assert_eq!(c[1], 0); + } + + #[test] + fn test_ntt_multiply_two_words() { + // Two-word operands + let a: Vec = vec![Word::MAX, 1]; // 2^64 + (2^64-1) + let b: Vec = vec![2, 0]; // 2 + let expected = schoolbook_mul(&a, &b); + let mut c = vec![0u64; a.len() + b.len()]; + let layout = memory_requirement_up_to(c.len(), b.len()); + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut memory = alloc.memory(); + add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + assert_eq!(&c[..], &expected[..], "two-word mismatch"); + } + + #[test] + fn test_ntt_multiply_small() { + // Test NTT multiply with small operands that exceed THRESHOLD_NTT. + let a: Vec = vec![0xDEADBEEFu64; THRESHOLD_NTT]; + let b: Vec = vec![0xCAFEBABEu64; THRESHOLD_NTT]; + let mut c = vec![0u64; a.len() + b.len()]; + + let layout = memory_requirement_up_to(c.len(), b.len()); + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut memory = alloc.memory(); + let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + assert_eq!(carry, 0); + assert!(c.iter().any(|&w| w != 0)); + } + + /// Naive schoolbook multiplication for comparison. + fn schoolbook_mul(a: &[Word], b: &[Word]) -> Vec { + let mut c = vec![0u64; a.len() + b.len()]; + for (i, &ai) in a.iter().enumerate() { + let mut carry: u128 = 0; + for (j, &bj) in b.iter().enumerate() { + let idx = i + j; + let prod = (ai as u128) * (bj as u128) + (c[idx] as u128) + carry; + c[idx] = prod as u64; + carry = prod >> 64; + } + // Propagate carry into higher words + let mut k = i + b.len(); + while carry != 0 { + let sum = (c[k] as u128) + carry; + c[k] = sum as u64; + carry = sum >> 64; + k += 1; + } + } + c + } + + /// Test NTT against schoolbook with moderate operand sizes. + fn run_ntt_vs_schoolbook(la: usize, lb: usize) { + // Generate deterministic test data + let a: Vec = (0..la) + .map(|i| (i as u64 + 1).wrapping_mul(0x9E3779B97F4A7C15)) + .collect(); + let b: Vec = (0..lb) + .map(|i| (i as u64 + 1).wrapping_mul(0xC6A4A7935BD1E995)) + .collect(); + let expected = schoolbook_mul(&a, &b); + + let mut c = vec![0u64; a.len() + b.len()]; + let layout = memory_requirement_up_to(c.len(), b.len()); + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut memory = alloc.memory(); + let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + assert_eq!(carry, 0, "carry should be 0"); + + assert_eq!(&c[..], &expected[..], "NTT mismatch: la={la}, lb={lb}"); + } + + #[test] + fn test_ntt_vs_schoolbook_equal() { + for &len in &[20, 30, 50, 64, 100, 128] { + run_ntt_vs_schoolbook(len, len); + } + } + + #[test] + fn test_ntt_vs_schoolbook_unequal() { + for &(la, lb) in &[(30, 20), (50, 30), (100, 50), (128, 64), (100, 20)] { + run_ntt_vs_schoolbook(la, lb); + } + } + + #[test] + fn test_ntt_vs_schoolbook_asymmetric() { + // Very asymmetric sizes + for &(la, lb) in &[(200, 30), (150, 20)] { + run_ntt_vs_schoolbook(la, lb); + } + } + + #[test] + fn test_ntt_all_ones() { + // All-ones operands stress the carry chain. + for &len in &[20, 50] { + let a = vec![Word::MAX; len]; + let b = vec![Word::MAX; len]; + let expected = schoolbook_mul(&a, &b); + + let mut c = vec![0u64; a.len() + b.len()]; + let layout = memory_requirement_up_to(c.len(), b.len()); + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut memory = alloc.memory(); + add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + assert_eq!(&c[..], &expected[..], "all-ones mismatch len={len}"); + } + } + + #[test] + fn test_ntt_zero_operand() { + let a = vec![0xDEADu64; 30]; + let b = vec![0u64; 30]; + let mut c = vec![0u64; a.len() + b.len()]; + let layout = memory_requirement_up_to(c.len(), b.len()); + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut memory = alloc.memory(); + let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + assert_eq!(carry, 0); + assert!(c.iter().all(|&w| w == 0), "zero operand should give zero product"); + } + + #[test] + fn test_ntt_sign_negative() { + // Test that Negative sign works (c -= a * b) + let a: Vec = (0..30).map(|i| (i as u64 + 1) * 100).collect(); + let b: Vec = (0..30).map(|i| (i as u64 + 1) * 200).collect(); + let _expected = schoolbook_mul(&a, &b); + + // First add: c += a * b + let mut c = vec![0u64; a.len() + b.len()]; + let layout = memory_requirement_up_to(c.len(), b.len()); + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut memory = alloc.memory(); + let _ = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + + // Then subtract: c -= a * b + let layout2 = memory_requirement_up_to(c.len(), b.len()); + let mut alloc2 = crate::memory::MemoryAllocation::new(layout2); + let mut memory2 = alloc2.memory(); + let _ = add_signed_mul_impl(&mut c, Negative, &a, &b, &mut memory2); + + // Result should be zero + assert!(c.iter().all(|&w| w == 0), "add then subtract should give zero"); + } + + #[test] + fn test_ntt_high_low_zero_limbs() { + // Operands with leading/trailing zero limbs + let mut a = vec![0u64; 80]; + let mut b = vec![0u64; 80]; + for i in 20..60 { + a[i] = (i as u64 + 1).wrapping_mul(0xDEADBEEF); + b[i] = (i as u64 + 1).wrapping_mul(0xCAFEBABE); + } + let expected = schoolbook_mul(&a, &b); + + let mut c = vec![0u64; a.len() + b.len()]; + let layout = memory_requirement_up_to(c.len(), b.len()); + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut memory = alloc.memory(); + add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + assert_eq!(&c[..], &expected[..], "sparse operand mismatch"); + } +} diff --git a/integer/src/mul/ntt/pack.rs b/integer/src/mul/ntt/pack.rs new file mode 100644 index 00000000..14b31c7b --- /dev/null +++ b/integer/src/mul/ntt/pack.rs @@ -0,0 +1,194 @@ +//! Bit-level packing / unpacking of `b`-bit coefficients. +#![allow( + dead_code, + unused_assignments, + unused_mut, + unused_variables, + clippy::unnecessary_cast +)] + +use crate::arch::word::Word; + +/// Pack a big integer (given as `&[Word]`, little-endian) into `n` +/// coefficients of `b_pack` bits each, zero-padded to length `n`. +/// +/// Each coefficient `c_i` satisfies `0 ≤ c_i < 2^{b_pack}`. +pub fn pack(words: &[Word], b_pack: u32, n: usize) -> alloc::vec::Vec { + let mut out = alloc::vec![0u64; n]; + let mask = (1u64 << b_pack) - 1; + let word_bits = Word::BITS; + let mut word_idx = 0usize; + let mut bit_offset = 0u32; // bit position within words[word_idx] + + for coeff in out.iter_mut().take(n) { + if word_idx >= words.len() { + break; // rest stay zero (padding) + } + + if bit_offset + b_pack <= word_bits { + // Entire coefficient fits within the current word. + *coeff = (words[word_idx] >> bit_offset) & mask; + bit_offset += b_pack; + if bit_offset == word_bits { + bit_offset = 0; + word_idx += 1; + } + } else { + // Coefficient straddles a word boundary. + let bits_first = word_bits - bit_offset; + let bits_second = b_pack - bits_first; + let mut val = (words[word_idx] >> bit_offset) & ((1u64 << bits_first) - 1); + word_idx += 1; + if word_idx < words.len() { + val |= (words[word_idx] & ((1u64 << bits_second) - 1)) << bits_first; + } + *coeff = val; + bit_offset = bits_second; + } + } + + out +} + +/// Accumulate CRT-recovered convolution coefficients into the output limb +/// array with carry propagation. +/// +/// Each coefficient `c_k` contributes `c_k << (k * b_pack)` bits to the +/// output. `output` must have capacity for `c.len()` coefficients plus any +/// carry overflow. +pub fn unpack_accumulate(output: &mut [Word], coeffs: &[u64], b_pack: u32, output_len: usize) { + let word_bits = Word::BITS as u32; + // For each coefficient, shift it by k*b_pack bits and add into the + // output with carry propagation. We use a software accumulation + // because the coefficients can be larger than a single output word. + + for (k, &coeff) in coeffs.iter().enumerate().take(output_len) { + if coeff == 0 { + continue; + } + let shift_bits = (k as u32).wrapping_mul(b_pack); + let word_idx = (shift_bits / word_bits) as usize; + let bit_shift = shift_bits % word_bits; + + // The coefficient occupies up to ⌈bit_len(coeff) / word_bits⌉ words. + // We split it into word-sized chunks and add each with the + // appropriate shift to the output. + let lo = coeff as u64; + let _hi = 0u64; // coeff fits in one u64 since max CRT value < P ≈ 2^192 + // Actually, per-coefficient CRT values can be up to P-1 ≈ 2^192, + // which needs up to 3 words. We handle this by splitting the + // coefficient itself into words and accumulating each. + + // For the immediate case, coeff from CRT is already small enough + // to fit in one or two u64 words. We accumulate by repeated + // add-with-carry into the output slice. + let mut carry: Word = 0; + let mut idx = word_idx; + + if bit_shift == 0 { + // Aligned: just add into output + let (sum, c) = overflowing_add_word(output.get(idx).copied().unwrap_or(0), lo); + carry = Word::from(c); + if idx < output.len() { + output[idx] = sum; + } + idx += 1; + } else { + // Split across two output words + let lo_part = lo << bit_shift; + let hi_part = if bit_shift > 0 { + lo >> (64 - bit_shift) + } else { + 0 + }; + + let (sum, c1) = overflowing_add_word(output.get(idx).copied().unwrap_or(0), lo_part); + carry = Word::from(c1); + if idx < output.len() { + output[idx] = sum; + } + idx += 1; + + let (sum2, c2) = + overflowing_add_word(output.get(idx).copied().unwrap_or(0), hi_part + carry); + carry = Word::from(c2); + if idx < output.len() { + output[idx] = sum2; + } + idx += 1; + } + + // Propagate remaining carry + while carry != 0 && idx < output.len() { + let (sum, c) = overflowing_add_word(output[idx], carry); + output[idx] = sum; + carry = Word::from(c); + idx += 1; + } + } +} + +fn overflowing_add_word(a: Word, b: u64) -> (Word, bool) { + let (sum, overflow) = a.overflowing_add(b); + (sum, overflow) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_pack_unpack_roundtrip() { + // Pack a number, then unpack-accumulate into a zero buffer. + // The accumulation should reconstruct the original number. + let b_pack = 16u32; + let test_words: Vec = vec![0xDEADBEEF_CAFEBABE, 0x12345678_9ABCDEF0]; + let coeffs_per_word = (Word::BITS / b_pack) as usize; + let n = test_words.len() * coeffs_per_word; + + let packed = pack(&test_words, b_pack, n); + assert_eq!(packed.len(), n); + + let output_len = test_words.len() + 1; + let mut output = vec![0u64; output_len]; + unpack_accumulate(&mut output, &packed, b_pack, n); + assert_eq!(&output[..test_words.len()], &test_words[..]); + } + + #[test] + fn test_pack_zero_pads() { + let words = vec![0xFFFFu64]; + let n = 32; // more coeffs than content + let packed = pack(&words, 16, n); + assert_eq!(packed.len(), 32); + // First coeff = 0xFFFF (least significant 16 bits), rest = 0 + assert_eq!(packed[0], 0xFFFF); + for &c in packed.iter().skip(1) { + assert_eq!(c, 0); + } + } + + #[test] + fn test_pack_empty_input() { + let packed = pack(&[], 16, 8); + assert_eq!(packed, vec![0u64; 8]); + } + + #[test] + fn test_unpack_single_coeff() { + let mut output = vec![0u64; 2]; + unpack_accumulate(&mut output, &[0xABCD], 16, 1); + assert_eq!(output[0], 0xABCD); + assert_eq!(output[1], 0); + } + + #[test] + fn test_unpack_carry_propagation() { + // Coefficient at k=4 (shift by 64 bits = 1 word) + carry + let mut output = vec![0u64; 3]; + unpack_accumulate(&mut output, &[0, 0, 0, 0, 1], 16, 5); + assert_eq!(output[0], 0); + assert_eq!(output[1], 1); + assert_eq!(output[2], 0); + } +} diff --git a/integer/src/mul/ntt/primes.rs b/integer/src/mul/ntt/primes.rs new file mode 100644 index 00000000..81eeda37 --- /dev/null +++ b/integer/src/mul/ntt/primes.rs @@ -0,0 +1,229 @@ +//! NTT-friendly primes of the form `2^64 - 2^b + 1`. +//! +//! All three support shift-based reduction via the identity `2^64 ≡ 2^b - 1 (mod p)`. + +/// The number of primes in the fixed array. +pub const K: usize = 3; + +/// Precomputed data for one NTT-friendly prime. +#[derive(Clone, Copy, Debug)] +pub struct NttPrime { + /// The prime value `p = 2^64 - 2^b + 1`. + pub p: u64, + /// The exponent `b` in the Solinas form. + pub b: u32, + /// The exponent of 2 in `p - 1`: `v2(p - 1)`. + #[allow(dead_code)] + pub v2: u32, + /// A primitive root modulo `p` that generates the full multiplicative group. + #[allow(dead_code)] + pub g: u64, + /// A primitive `2^32`-th root of unity: `ω = g^{(p-1) / 2^32} mod p`. + pub omega_2_32: u64, +} + +/// The three chosen primes, all of the form `2^64 - 2^b + 1` with `b ∈ {32, 34, 40}`. +/// +/// | name | `b` | `p` | `v2(p-1)` | gen `g` | `2^32`-th root ω | +/// |------|-----|---------------------|-----------|---------|----------------------------| +/// | GL | 32 | `0xFFFFFFFF00000001` | 32 | 7 | `1753635133440165772` | +/// | P1 | 34 | `0xFFFFFFFC00000001` | 34 | 5 | `11315553352654630047` | +/// | P2 | 40 | `0xFFFFFF0000000001` | 40 | 19 | `551857376737322389` | +pub const PRIMES: [NttPrime; K] = [ + NttPrime { + p: 0xFFFFFFFF00000001, + b: 32, + v2: 32, + g: 7, + omega_2_32: 1753635133440165772, + }, + NttPrime { + p: 0xFFFFFFFC00000001, + b: 34, + v2: 34, + g: 5, + omega_2_32: 11315553352654630047, + }, + NttPrime { + p: 0xFFFFFF0000000001, + b: 40, + v2: 40, + g: 19, + omega_2_32: 551857376737322389, + }, +]; + +#[cfg(test)] +mod tests { + use super::*; + use num_modular::FixedTrinomialSolinas64; + + /// Deterministic Miller–Rabin for 64-bit integers with known bases. + /// Tests `n` against bases `[2, 325, 9375, 28178, 450775, 9780504, 1795265022]` + /// which together suffice for all `n < 2^64` (deterministic). + fn is_prime_u64(n: u64) -> bool { + if n < 2 { + return false; + } + if n % 2 == 0 { + return n == 2; + } + + // Write n-1 = d * 2^s + let d = (n - 1) >> (n - 1).trailing_zeros(); + let s = (n - 1).trailing_zeros(); + + let bases = [2u64, 325, 9375, 28178, 450775, 9780504, 1795265022]; + + 'next_base: for &a in &bases { + if a >= n { + continue; + } + let mut x = mod_pow_u64(a % n, d, n); + if x == 1 || x == n - 1 { + continue 'next_base; + } + for _ in 1..s { + x = ((x as u128 * x as u128) % (n as u128)) as u64; + if x == n - 1 { + continue 'next_base; + } + } + return false; + } + true + } + + fn mod_pow_u64(mut base: u64, mut exp: u64, modulus: u64) -> u64 { + let mut result = 1u64; + while exp > 0 { + if exp & 1 != 0 { + result = ((result as u128 * base as u128) % (modulus as u128)) as u64; + } + base = ((base as u128 * base as u128) % (modulus as u128)) as u64; + exp >>= 1; + } + result + } + + #[test] + fn verify_primes() { + for &NttPrime { + p, + b, + v2, + g, + omega_2_32, + } in &PRIMES + { + // 1. Correct form: p == 2^64 - 2^b + 1 + let expected_p = (1u128 << 64) - (1u128 << b) + 1; + assert!(expected_p < (1u128 << 64), "p must fit in 64 bits"); + assert_eq!(p as u128, expected_p, "p = 0x{p:X} does not match 2^64 - 2^{b} + 1"); + assert!(p > 0, "p must be positive"); + + // 2. Primality + assert!(is_prime_u64(p), "p = 0x{p:X} is not prime"); + + // 3. v2(p-1) is at least 32 + let actual_v2 = (p - 1).trailing_zeros(); + assert!(actual_v2 >= 32, "v2(p-1) = {actual_v2} < 32 for p = 0x{p:X}"); + assert_eq!(actual_v2, v2, "stored v2 mismatch for p = 0x{p:X}"); + + // 4. g generates the full multiplicative group mod p. + // g^((p-1)/2) mod p ≠ 1 (g is a quadratic non-residue) + let g_order_half = mod_pow_u64(g, (p - 1) / 2, p); + assert_ne!(g_order_half, 1, "g = {g} is a quadratic residue mod p = 0x{p:X}"); + + // g^(p-1) ≡ 1 + let g_full = mod_pow_u64(g, p - 1, p); + assert_eq!(g_full, 1, "g^(p-1) != 1 mod p = 0x{p:X}"); + + // 5. ω has exact order 2^32 + let mut omega_pow = omega_2_32; + for _ in 0..31 { + omega_pow = ((omega_pow as u128 * omega_pow as u128) % (p as u128)) as u64; + } + // After 31 squarings: ω^{2^31} mod p + // Should be -1 mod p (order is exactly 2^32) + assert_eq!(omega_pow, p - 1, "omega^(2^31) != -1 mod p = 0x{p:X}, order not 2^32"); + + // ω^{2^32} ≡ 1 + let omega_full = mod_pow_u64(omega_2_32, 1u64 << 32, p); + assert_eq!(omega_full, 1, "omega^(2^32) != 1 mod p = 0x{p:X}"); + + // 6. Reduction identity: 2^64 ≡ 2^b - 1 (mod p) + let two_64_mod_p = ((1u128 << 64) % (p as u128)) as u64; + let expected = (if b == 0 { + 0 + } else { + (1u64 << (b - 1)).wrapping_mul(2) + }) - 1; + assert_eq!(two_64_mod_p, expected, "2^64 mod p != 2^(b) - 1 for p = 0x{p:X}"); + } + } + + #[test] + fn test_reduction_identity_per_prime() { + // Verify reduction works for all three primes using the actual reducer types. + // GL: b=32 + { + type Reducer = FixedTrinomialSolinas64<64, 32, 1>; + let p = Reducer::MODULUS; + assert_eq!(p, PRIMES[0].p); + + // Test reduce_double + let v = (p as u128) * 3; // 3p → should reduce to 0 + let r = Reducer::reduce_double(v); + assert!(r < p); + assert_eq!((r as u128) % (p as u128), v % (p as u128)); + } + // P1: b=34 + { + type Reducer = FixedTrinomialSolinas64<64, 34, 1>; + let p = Reducer::MODULUS; + assert_eq!(p, PRIMES[1].p); + + let v = (p as u128) * 3; + let r = Reducer::reduce_double(v); + assert!(r < p); + assert_eq!((r as u128) % (p as u128), v % (p as u128)); + } + // P2: b=40 + { + type Reducer = FixedTrinomialSolinas64<64, 40, 1>; + let p = Reducer::MODULUS; + assert_eq!(p, PRIMES[2].p); + + let v = (p as u128) * 3; + let r = Reducer::reduce_double(v); + assert!(r < p); + assert_eq!((r as u128) % (p as u128), v % (p as u128)); + } + } + + #[test] + fn test_pow_inv_roundtrip() { + // pow and inv round-trip checks using the trait API + use num_modular::{ModularPow, ModularUnaryOps}; + + for &NttPrime { p, g, .. } in &PRIMES { + // inv(x)·x ≡ 1 + let x = 123456789u64; + let inv = x.invm(&p); + if let Some(inv) = inv { + let prod = ((x as u128 * inv as u128) % (p as u128)) as u64; + assert_eq!(prod, 1, "inv round-trip failed for p = 0x{p:X}"); + } + + // pow(g, p-1) ≡ 1 + let g_pow = g.powm(&(p - 1), &p); + assert_eq!(g_pow, 1, "g^(p-1) != 1 mod p = 0x{p:X}"); + + // inv(g) · g ≡ 1 + let g_inv = g.invm(&p).unwrap(); + let prod = ((g as u128 * g_inv as u128) % (p as u128)) as u64; + assert_eq!(prod, 1, "inv(g) round-trip failed for p = 0x{p:X}"); + } + } +} diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs new file mode 100644 index 00000000..be27402c --- /dev/null +++ b/integer/src/mul/ntt/transform.rs @@ -0,0 +1,339 @@ +//! Iterative in-place radix-2 NTT over primes of the form `2^64 - 2^b + 1`. +//! +//! Uses decimation-in-time (DIT) Cooley–Tukey butterflies. All modular +//! arithmetic delegates to `num_modular::FixedTrinomialSolinas64`. + +use num_modular::{FixedTrinomialSolinas64, ModularPow, ModularUnaryOps, Reducer}; + +// ---- dispatch helpers (b is a runtime value, const-generic under the hood) ---- + +#[inline] +fn reduce_double(v: u128, b: u32) -> u64 { + match b { + 32 => FixedTrinomialSolinas64::<64, 32, 1>::reduce_double(v), + 34 => FixedTrinomialSolinas64::<64, 34, 1>::reduce_double(v), + 40 => FixedTrinomialSolinas64::<64, 40, 1>::reduce_double(v), + _ => unreachable!(), + } +} + +#[inline] +fn mul_mod(a: u64, b_val: u64, b: u32) -> u64 { + reduce_double((a as u128) * (b_val as u128), b) +} + +#[inline] +fn add_mod(a: u64, b_val: u64, p: u64, b: u32) -> u64 { + match b { + 32 => FixedTrinomialSolinas64::<64, 32, 1>::new(&p).add(&a, &b_val), + 34 => FixedTrinomialSolinas64::<64, 34, 1>::new(&p).add(&a, &b_val), + 40 => FixedTrinomialSolinas64::<64, 40, 1>::new(&p).add(&a, &b_val), + _ => unreachable!(), + } +} + +#[inline] +fn sub_mod(a: u64, b_val: u64, p: u64, b: u32) -> u64 { + match b { + 32 => FixedTrinomialSolinas64::<64, 32, 1>::new(&p).sub(&a, &b_val), + 34 => FixedTrinomialSolinas64::<64, 34, 1>::new(&p).sub(&a, &b_val), + 40 => FixedTrinomialSolinas64::<64, 40, 1>::new(&p).sub(&a, &b_val), + _ => unreachable!(), + } +} + +// ---- public API ---- + +/// Precompute the twiddle-factor table for transform length `n`. +/// +/// Returns `n/2` entries: `omega_n^k` for `k = 0..n/2`, where +/// `omega_n = omega_2_32^{2^32 / n}` (a primitive `n`-th root of unity). +pub fn precompute_twiddles( + n: usize, + p: u64, + b: u32, + omega_2_32: u64, + inverse: bool, +) -> alloc::vec::Vec { + let shift = 32 - n.trailing_zeros(); + let omega_n = omega_2_32.powm(&(1u64 << shift), &p); + + let base = if inverse { + omega_n.invm(&p).expect("omega_n not invertible") + } else { + omega_n + }; + + let mut twiddles = alloc::vec![0u64; n / 2]; + twiddles[0] = 1; + for k in 1..(n / 2) { + twiddles[k] = mul_mod(twiddles[k - 1], base, b); + } + twiddles +} + +/// Bit-reverse `a` in place. Length must be a power of two. +pub fn bit_reverse(a: &mut [u64]) { + let n = a.len(); + assert!(n.is_power_of_two()); + let log_n = n.trailing_zeros(); + for i in 0..n { + let j = i.reverse_bits() >> (usize::BITS - log_n); + if i < j { + a.swap(i, j); + } + } +} + +/// Forward NTT in place (decimation-in-time, radix-2). +pub fn forward(a: &mut [u64], twiddles: &[u64], p: u64, b: u32) { + ntt_core(a, twiddles, p, b); +} + +/// Inverse NTT in place. +/// +/// Computed as `bit_reverse → forward(ω^{-1}) → scale`, producing output +/// in **natural order**. +/// +/// `twiddles` must have been precomputed with `inverse = true` +/// (i.e. using `omega_n^{-1}`). +pub fn inverse(a: &mut [u64], twiddles: &[u64], p: u64, b: u32) { + let n = a.len(); + bit_reverse(a); + ntt_core(a, twiddles, p, b); + let n_inv = (n as u64).invm(&p).expect("n not invertible mod p"); + for x in a.iter_mut() { + *x = mul_mod(*x, n_inv, b); + } +} + +/// In-place radix-2 DIT NTT (Cooley–Tukey). +fn ntt_core(a: &mut [u64], twiddles: &[u64], p: u64, b: u32) { + let n = a.len(); + debug_assert!(n.is_power_of_two() && twiddles.len() == n / 2); + + let mut sub_len = 2usize; + while sub_len <= n { + let half = sub_len / 2; + let step = n / sub_len; + + for i in (0..n).step_by(sub_len) { + for j in 0..half { + let u = a[i + j]; + let v = mul_mod(a[i + j + half], twiddles[j * step], b); + a[i + j] = add_mod(u, v, p, b); + a[i + j + half] = sub_mod(u, v, p, b); + } + } + + sub_len *= 2; + } +} + +/// Pointwise multiply of two transformed vectors in place. +pub fn pointwise_mul(a_hat: &mut [u64], b_hat: &[u64], b: u32) { + assert_eq!(a_hat.len(), b_hat.len()); + for (a, &b_val) in a_hat.iter_mut().zip(b_hat.iter()) { + *a = mul_mod(*a, b_val, b); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mul::ntt::primes::PRIMES; + + fn assert_all_eq(a: &[u64], b_val: &[u64]) { + assert_eq!(a.len(), b_val.len()); + for (i, (x, y)) in a.iter().zip(b_val.iter()).enumerate() { + assert_eq!(x, y, "mismatch at index {i}: {x} != {y}"); + } + } + + #[test] + fn test_forward_inverse_roundtrip() { + for prime in &PRIMES { + let p = prime.p; + let b = prime.b; + for &n in &[2, 4, 8, 16, 32, 64, 128, 256, 512] { + let inv_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, true); + let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); + + let mut a: Vec = (0..n) + .map(|i| ((i as u64 + 1).wrapping_mul(123456789)) % p) + .collect(); + let orig = a.clone(); + + bit_reverse(&mut a); + forward(&mut a, &fwd_twiddles, p, b); + inverse(&mut a, &inv_twiddles, p, b); + + assert_all_eq(&a, &orig); + } + } + } + + #[test] + fn test_convolution_via_ntt() { + for prime in &PRIMES { + let p = prime.p; + let b = prime.b; + for len_a in [1, 2, 3, 5] { + for len_b in [1, 2, 3, 5] { + let conv_len: usize = len_a + len_b - 1; + let n = conv_len.next_power_of_two().max(2); + + let a: Vec = (0..len_a).map(|i| ((i + 1) as u64 * 12345) % p).collect(); + let b_vec: Vec = + (0..len_b).map(|i| ((i + 1) as u64 * 67890) % p).collect(); + + let mut expected = vec![0u64; conv_len]; + for (i, &ai) in a.iter().enumerate() { + for (j, &bj) in b_vec.iter().enumerate() { + expected[i + j] = add_mod(expected[i + j], mul_mod(ai, bj, b), p, b); + } + } + + let inv_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, true); + let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); + + let mut a_pad = vec![0u64; n]; + let mut b_pad = vec![0u64; n]; + a_pad[..len_a].copy_from_slice(&a); + b_pad[..len_b].copy_from_slice(&b_vec); + + bit_reverse(&mut a_pad); + bit_reverse(&mut b_pad); + forward(&mut a_pad, &fwd_twiddles, p, b); + forward(&mut b_pad, &fwd_twiddles, p, b); + pointwise_mul(&mut a_pad, &b_pad, b); + inverse(&mut a_pad, &inv_twiddles, p, b); + + match assert_all_eq_result(&a_pad[..conv_len], &expected) { + Ok(()) => {} + Err((i, l, r)) => panic!( + "convolution mismatch: b={b}, len_a={len_a}, len_b={len_b}, n={n}, \ + index {i}: {l} != {r}" + ), + } + } + } + } + } + + fn assert_all_eq_result(a: &[u64], b_val: &[u64]) -> Result<(), (usize, u64, u64)> { + assert_eq!(a.len(), b_val.len()); + for (i, (x, y)) in a.iter().zip(b_val.iter()).enumerate() { + if x != y { + return Err((i, *x, *y)); + } + } + Ok(()) + } + + #[test] + fn test_bit_reverse() { + let mut a: Vec = (0..8).collect(); + bit_reverse(&mut a); + assert_eq!(a, vec![0, 4, 2, 6, 1, 5, 3, 7]); + } + + #[allow(clippy::needless_range_loop)] + fn ntt_naive(x: &[u64], omega_n: u64, p: u64, b: u32) -> Vec { + let n = x.len(); + let mut result = vec![0u64; n]; + for k in 0..n { + let mut acc = 0u64; + for j in 0..n { + let twiddle = if k == 0 || j == 0 { + 1 + } else { + omega_n.powm(&((k * j) as u64), &p) + }; + acc = add_mod(acc, mul_mod(x[j], twiddle, b), p, b); + } + result[k] = acc; + } + result + } + + #[test] + fn test_forward_correctness() { + for prime in &PRIMES { + let p = prime.p; + let b = prime.b; + for &n in &[2usize, 4, 8] { + let x: Vec = (0..n).map(|i| ((i + 1) as u64 * 11111) % p).collect(); + + let shift = 32 - n.trailing_zeros(); + let omega_n = prime.omega_2_32.powm(&(1u64 << shift), &p); + + let mut a = x.clone(); + bit_reverse(&mut a); + + let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); + forward(&mut a, &fwd_twiddles, p, b); + + let expected = ntt_naive(&x, omega_n, p, b); + assert_eq!( + a, expected, + "forward NTT mismatch: b={b}, n={n}, expected={expected:?}, got={a:?}" + ); + } + } + } + + #[test] + fn test_convolution_debug() { + let prime = &PRIMES[0]; + let p = prime.p; + let b = prime.b; + + let a = vec![12345u64 % p]; + let b_vec = vec![67890u64 % p, 135780u64 % p, 203670u64 % p]; + let conv_len = a.len() + b_vec.len() - 1; + let n = 4; + + let mut expected = vec![0u64; conv_len]; + for (i, &ai) in a.iter().enumerate() { + for (j, &bj) in b_vec.iter().enumerate() { + expected[i + j] = add_mod(expected[i + j], mul_mod(ai, bj, b), p, b); + } + } + + let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); + let inv_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, true); + + let mut a_pad = vec![0u64; n]; + let mut b_pad = vec![0u64; n]; + a_pad[..a.len()].copy_from_slice(&a); + b_pad[..b_vec.len()].copy_from_slice(&b_vec); + + bit_reverse(&mut a_pad); + bit_reverse(&mut b_pad); + forward(&mut a_pad, &fwd_twiddles, p, b); + forward(&mut b_pad, &fwd_twiddles, p, b); + pointwise_mul(&mut a_pad, &b_pad, b); + inverse(&mut a_pad, &inv_twiddles, p, b); + + assert_eq!(&a_pad[..conv_len], &expected[..]); + } + + #[test] + fn test_length_two_edge_case() { + for prime in &PRIMES { + let p = prime.p; + let b = prime.b; + let n = 2; + let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); + let inv_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, true); + + let a_orig = vec![1u64 % p, 2u64 % p]; + let mut a = a_orig.clone(); + bit_reverse(&mut a); + forward(&mut a, &fwd_twiddles, p, b); + inverse(&mut a, &inv_twiddles, p, b); + assert_all_eq(&a, &a_orig); + } + } +} From 01d7a2021b01cb757b1679048775a18cdee56d92 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sat, 13 Jun 2026 11:03:19 +0800 Subject: [PATCH 03/19] WIP: tidy up --- integer/src/mul/ntt/crt.rs | 261 ++++++++++++++++--------------- integer/src/mul/ntt/mod.rs | 142 +++++++++-------- integer/src/mul/ntt/pack.rs | 34 ++-- integer/src/mul/ntt/primes.rs | 8 + integer/src/mul/ntt/transform.rs | 221 +++++++++++++------------- 5 files changed, 353 insertions(+), 313 deletions(-) diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs index 9a32ee36..f5af326b 100644 --- a/integer/src/mul/ntt/crt.rs +++ b/integer/src/mul/ntt/crt.rs @@ -1,117 +1,144 @@ //! Garner CRT: combine `K` residues modulo `K` primes into a small integer. //! -//! Uses num-modular's general modular arithmetic traits since CRT is -//! called once per coefficient, not in the hot loop. +//! Takes `num_modular::Reducer` implementations so the caller can supply +//! specialized Solinas reducers for the hot per-coefficient path. #![allow(clippy::unnecessary_cast)] -use num_modular::{ModularCoreOps, ModularUnaryOps}; +use num_modular::Reducer; -/// Precomputed constants for Garner CRT with a fixed prime set. -pub struct CrtConstants { - /// `inv(p_i mod p_j)` for i < j. - pub inv_ij: [[u64; 3]; 3], +/// Subset of `Reducer` that is object-safe (no `new` or other +/// non-`&self` methods). Implemented automatically for every +/// `Reducer` via a blanket impl. +#[allow(dead_code)] +pub trait ModOps { + fn add(&self, lhs: &u64, rhs: &u64) -> u64; + fn sub(&self, lhs: &u64, rhs: &u64) -> u64; + fn mul(&self, lhs: &u64, rhs: &u64) -> u64; } -impl CrtConstants { - /// Precompute Garner constants for the given primes. - pub fn new(primes: &[u64]) -> Self { - let k = primes.len(); - let mut inv_ij = [[0u64; 3]; 3]; - for i in 0..k { - for j in (i + 1)..k { - let p_i_mod_pj = primes[i] % primes[j]; - inv_ij[i][j] = p_i_mod_pj.invm(&primes[j]).expect("primes not coprime"); - } +impl> ModOps for T { + fn add(&self, lhs: &u64, rhs: &u64) -> u64 { + Reducer::add(self, lhs, rhs) + } + fn sub(&self, lhs: &u64, rhs: &u64) -> u64 { + Reducer::sub(self, lhs, rhs) + } + fn mul(&self, lhs: &u64, rhs: &u64) -> u64 { + Reducer::mul(self, lhs, rhs) + } +} + +/// A 192-bit unsigned integer (3 × u64, little-endian). +/// +/// Used to hold Garner CRT results (which are bounded by the product of +/// three ≈2^64 primes, therefore < 2^192). +#[derive(Clone, Copy, Debug, Default)] +pub struct U192(pub [u64; 3]); + +impl U192 { + #[inline] + pub fn new(lo: u64) -> Self { + U192([lo, 0, 0]) + } + + /// `self += v` where `v` fits in 128 bits. + #[inline] + pub fn add_u128(&mut self, v: u128) { + let lo = v as u64; + let hi = (v >> 64) as u64; + let (r0, c0) = self.0[0].overflowing_add(lo); + self.0[0] = r0; + let (r1, c1) = self.0[1].overflowing_add(hi.wrapping_add(c0 as u64)); + self.0[1] = r1; + self.0[2] = self.0[2].wrapping_add(c1 as u64); + } + + /// `self += t * factor` where `t` < 2^64, `factor` < 2^128. + #[inline] + pub fn add_mul_u64_u128(&mut self, t: u64, factor: u128) { + let fac_lo = factor as u64; + let fac_hi = (factor >> 64) as u64; + + let m_lo_full = (t as u128) * (fac_lo as u128); + let lo = m_lo_full as u64; + let m_lo = (m_lo_full >> 64) as u64; + + let m_hi_full = (t as u128) * (fac_hi as u128); + let m_hi = m_hi_full as u64; + let hi = (m_hi_full >> 64) as u64; + + let (mid, c) = m_lo.overflowing_add(m_hi); + let hi_word = hi.wrapping_add(c as u64); + + let (r0, c0) = self.0[0].overflowing_add(lo); + self.0[0] = r0; + let (r1, c1) = self.0[1].overflowing_add(mid.wrapping_add(c0 as u64)); + self.0[1] = r1; + self.0[2] = self.0[2].wrapping_add(hi_word.wrapping_add(c1 as u64)); + } + + /// `self mod m`, where `m` < 2^64. + #[inline] + pub fn rem_u64(&self, m: u64) -> u64 { + let m128 = m as u128; + let mut r: u128 = 0; + for &word in self.0.iter().rev() { + r = (r << 64) | (word as u128); + r %= m128; + } + r as u64 + } + + #[inline] + pub fn len_words(&self) -> u32 { + if self.0[2] != 0 { + 3 + } else if self.0[1] != 0 { + 2 + } else { + 1 } - CrtConstants { inv_ij } } } -/// Combine `K` residues into a small integer (< P) using Garner's algorithm. +use super::primes::{CRT_INV_IJ, PRIMES}; + +/// Combine `K` residues into a `U192` via Garner's algorithm. /// -/// The result is returned as a little-endian `Vec` because -/// `P ≈ 2^{64K}` may exceed a single `u64`. -pub fn garner_combine( - residues: &[u64], - primes: &[u64], - constants: &CrtConstants, -) -> alloc::vec::Vec { +/// All primes and precomputed inverses are hardcoded in [`super::primes`]. +/// `reducers[i]` must be a reducer for the i-th prime. +pub fn garner_combine(residues: &[u64], reducers: &[&dyn ModOps]) -> U192 { let k = residues.len(); assert!(k <= 3, "CRT supports up to 3 primes"); - assert_eq!(primes.len(), k); - - let mut result = [0u64; 3]; - let p0 = primes[0]; + assert!(reducers.len() >= k); - // x_0 = r_0 - result[0] = residues[0] % p0; + let p0 = PRIMES[0].p; + let p1 = PRIMES[1].p; + let p2 = PRIMES[2].p; + let mut x = U192::new(residues[0]); if k == 1 { - return result.to_vec(); + return x; } - // t_1 = (r_1 - x_0) * inv(p0 mod p1) mod p1 - let p1 = primes[1]; - let x0_mod_p1 = result[0] % p1; - let diff1 = residues[1].subm(x0_mod_p1, &p1); - let t1 = diff1.mulm(constants.inv_ij[0][1], &p1); - - // x_1 = x_0 + t_1 * p0 - add_128_to_192(&mut result, (t1 as u128) * (p0 as u128)); + // t_1 = (r_1 - x mod p1) * inv(p0 mod p1) mod p1 + let x_mod_p1 = x.0[0] % p1; + let diff1 = reducers[1].sub(&residues[1], &x_mod_p1); + let t1 = reducers[1].mul(&diff1, &CRT_INV_IJ[0][1]); + x.add_u128((t1 as u128) * (p0 as u128)); if k == 2 { - return result.to_vec(); + return x; } - // t_2 = (r_2 - x_1) * inv(p0*p1 mod p2) mod p2 - let p2 = primes[2]; - let x1_mod_p2 = mod_192_by_u64(&result, p2); - let diff2 = residues[2].subm(x1_mod_p2, &p2); - let inv_p0_mod_p2 = constants.inv_ij[0][2]; - let inv_p1_mod_p2 = constants.inv_ij[1][2]; - let inv_prod = inv_p0_mod_p2.mulm(inv_p1_mod_p2, &p2); - let t2 = diff2.mulm(inv_prod, &p2); - - // x_2 = x_1 + t_2 * p0 * p1 - let pp = (primes[0] as u128) * (primes[1] as u128); - let t2_64 = t2 as u64; - let pp_lo = pp as u64; - let pp_hi = (pp >> 64) as u64; - let m_lo_full = (t2_64 as u128) * (pp_lo as u128); - let m_lo = (m_lo_full >> 64) as u64; - let lo = m_lo_full as u64; - let m_hi_full = (t2_64 as u128) * (pp_hi as u128); - let hi = (m_hi_full >> 64) as u64; - let m_hi = m_hi_full as u64; - let (mid, c) = m_lo.overflowing_add(m_hi); - let hi_word = hi.wrapping_add(c as u64); - let (r0, c0) = result[0].overflowing_add(lo); - result[0] = r0; - let (r1, c1) = result[1].overflowing_add(mid.wrapping_add(c0 as u64)); - result[1] = r1; - result[2] = result[2].wrapping_add(hi_word.wrapping_add(c1 as u64)); - - result.to_vec() -} + // t_2 = (r_2 - x mod p2) * inv(p0*p1 mod p2) mod p2 + let x_mod_p2 = x.rem_u64(p2); + let diff2 = reducers[2].sub(&residues[2], &x_mod_p2); + let inv_prod = reducers[2].mul(&CRT_INV_IJ[0][2], &CRT_INV_IJ[1][2]); + let t2 = reducers[2].mul(&diff2, &inv_prod); + x.add_mul_u64_u128(t2, (p0 as u128) * (p1 as u128)); -fn add_128_to_192(result: &mut [u64; 3], term: u128) { - let lo = term as u64; - let hi = (term >> 64) as u64; - let (r0, c0) = result[0].overflowing_add(lo); - result[0] = r0; - let (r1, c1) = result[1].overflowing_add(hi.wrapping_add(c0 as u64)); - result[1] = r1; - result[2] = result[2].wrapping_add(c1 as u64); -} - -fn mod_192_by_u64(x: &[u64; 3], m: u64) -> u64 { - let m128 = m as u128; - let mut r: u128 = 0; - for &word in x.iter().rev() { - r = (r << 64) | (word as u128); - r %= m128; - } - r as u64 + x } #[cfg(test)] @@ -119,40 +146,30 @@ mod tests { use super::*; #[test] - fn test_garner_two_primes() { - let primes = vec![3u64, 5u64]; - let constants = CrtConstants::new(&primes); - // x ≡ 2 mod 3, x ≡ 3 mod 5 → x = 8 - let residues = vec![2u64, 3u64]; - let result = garner_combine(&residues, &primes, &constants); - let x = result[0] as u128; - assert_eq!(x, 8); - } + fn test_garner_with_ntt_primes() { + use num_modular::FixedTrinomialSolinas64; + use super::super::primes::PRIMES; - #[test] - fn test_garner_three_primes() { - let primes = vec![3u64, 5u64, 7u64]; - let constants = CrtConstants::new(&primes); - // x ≡ 2 (mod 3), x ≡ 3 (mod 5), x ≡ 4 (mod 7) → x = 53 - let residues = vec![2u64, 3u64, 4u64]; - let result = garner_combine(&residues, &primes, &constants); - let x = result[0] as u128; - assert_eq!(x, 53); - } + let p0 = PRIMES[0].p; + let p1 = PRIMES[1].p; + let p2 = PRIMES[2].p; + + let r0 = FixedTrinomialSolinas64::<64, 32, 1>::new(&p0); + let r1 = FixedTrinomialSolinas64::<64, 34, 1>::new(&p1); + let r2 = FixedTrinomialSolinas64::<64, 40, 1>::new(&p2); + let reducers: [&dyn ModOps; 3] = [&r0, &r1, &r2]; - #[test] - fn test_garner_with_ntt_primes() { - use crate::mul::ntt::primes::PRIMES; - let primes: Vec = PRIMES.iter().map(|np| np.p).collect(); - let constants = CrtConstants::new(&primes); let residues = vec![12345u64, 67890u64, 11111u64]; - let result = garner_combine(&residues, &primes, &constants); - for (i, &p) in primes.iter().enumerate() { - let mut rem: u128 = 0; - for &word in result.iter().rev() { - rem = ((rem << 64) | (word as u128)) % (p as u128); - } - assert_eq!(rem as u64, residues[i], "CRT mismatch for prime {i}"); - } + let x = garner_combine(&residues, &reducers); + assert_eq!(x.rem_u64(p0), residues[0]); + assert_eq!(x.rem_u64(p1), residues[1]); + assert_eq!(x.rem_u64(p2), residues[2]); + + let x = garner_combine(&residues[..2], &reducers[..2]); + assert_eq!(x.rem_u64(p0), residues[0]); + assert_eq!(x.rem_u64(p1), residues[1]); + + let x = garner_combine(&residues[..1], &reducers[..1]); + assert_eq!(x.0[0], residues[0]); } } diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index 344a62ca..cc34e3f4 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -3,6 +3,7 @@ //! Uses Number Theoretic Transforms over several 64-bit primes of the form //! `2^64 - 2^b + 1` combined with the Chinese Remainder Theorem (CRT). +use crate::mul::ntt::crt::ModOps; use crate::{ add, arch::word::{SignedWord, Word}, @@ -10,13 +11,14 @@ use crate::{ Sign::{self, *}, }; use alloc::alloc::Layout; +use num_modular::{FixedTrinomialSolinas64, Reducer}; mod crt; mod pack; mod primes; mod transform; -use crate::mul::ntt::crt::CrtConstants; +use crate::mul::ntt::crt::U192; pub use primes::{K, PRIMES}; /// Minimum smaller-operand length (in words) for the NTT path. @@ -86,18 +88,15 @@ pub fn memory_requirement_up_to(total_len: usize, _smaller_len: usize) -> Layout (total_len as u64 * word_bits as u64 + B_PACK_MIN as u64 - 1) / B_PACK_MIN as u64; let n_max = ((max_coeffs + 1) as usize).next_power_of_two().max(2); - // Everything is in u64 units for simplicity. let lanes_u64 = 2 * n_max; // a_lane + b_lane let residues_u64 = K * n_max; // per-prime inverse results - let product_u64 = total_len; // product buffer (Word=u64 on 64-bit, else u64 takes more space) + let twiddles_u64 = n_max; // fwd + inv twiddle tables (n_max/2 each, reused) + let product_u64 = total_len; - // On 64-bit targets Word = u64. On narrow targets (Word < u64), - // we need extra space for the u64 allocations. Use the maximum - // of Word and u64 sizes. let u64_bytes = 8usize; let word_bytes = core::mem::size_of::(); let factor = (u64_bytes + word_bytes - 1) / word_bytes; - let total_words = product_u64 + (lanes_u64 + residues_u64) * factor; + let total_words = product_u64 + (lanes_u64 + residues_u64 + twiddles_u64) * factor; memory::array_layout::(total_words) } @@ -146,7 +145,6 @@ fn add_signed_mul_impl( let la = a.len(); let lb = b.len(); - // Skip zero-length or zero-value operands if la == 0 || lb == 0 { return 0; } @@ -162,72 +160,65 @@ fn add_signed_mul_impl( let coeffs_b = coeff_count(lb_bits, b_pack); let output_coeffs = coeffs_a + coeffs_b - 1; - // CRT constants - let primes_p: alloc::vec::Vec = PRIMES[..k_eff].iter().map(|np| np.p).collect(); - let crt_constants = CrtConstants::new(&primes_p); + // Per-prime CRT reducers (no allocation) + let r0 = FixedTrinomialSolinas64::<64, 32, 1>::new(&PRIMES[0].p); + let r1 = FixedTrinomialSolinas64::<64, 34, 1>::new(&PRIMES[1].p); + let r2 = FixedTrinomialSolinas64::<64, 40, 1>::new(&PRIMES[2].p); + let crt_reducers: [&dyn ModOps; 3] = [&r0, &r1, &r2]; // ---- Memory carve (longest-lived first) ---- - // All buffers are u64 since lane arithmetic is always u64. // 1. Product buffer let prod_len = la + lb; let (prod, mut mem) = memory.allocate_slice_fill::(prod_len, 0); - // 2. Residue storage (per-prime inverse results, as u64) + // 2. Residue storage (per-prime inverse results) let residues_len = k_eff * nn; let (residues, mut mem) = mem.allocate_slice_fill::(residues_len, 0); - // 3. Lane buffers (reused across primes, as u64) + // 3. Lane buffers (reused across primes) let (a_lane, mut mem) = mem.allocate_slice_fill::(nn, 0); - let (b_lane, _mem) = mem.allocate_slice_fill::(nn, 0); + let (b_lane, mut mem) = mem.allocate_slice_fill::(nn, 0); - // ---- Per-prime transforms ---- + // 4. Twiddle tables (fwd + inv, reused per prime) + let (fwd_twiddles, mut mem) = mem.allocate_slice_fill::(nn / 2, 0); + let (inv_twiddles, _) = mem.allocate_slice_fill::(nn / 2, 0); + + // ---- Per-prime transforms (const-generic dispatch) ---- for (pi, prime) in PRIMES[..k_eff].iter().enumerate() { - let p = prime.p; - let b_exp = prime.b; - - // Precompute twiddles - let fwd_twiddles = transform::precompute_twiddles(nn, p, b_exp, prime.omega_2_32, false); - let inv_twiddles = transform::precompute_twiddles(nn, p, b_exp, prime.omega_2_32, true); - - // Pack operands into lane buffers - pack_into(a, b_pack, a_lane); - pack_into(b, b_pack, b_lane); - - // Forward NTT - transform::bit_reverse(a_lane); - transform::bit_reverse(b_lane); - transform::forward(a_lane, &fwd_twiddles, p, b_exp); - transform::forward(b_lane, &fwd_twiddles, p, b_exp); - transform::pointwise_mul(a_lane, b_lane, b_exp); - transform::inverse(a_lane, &inv_twiddles, p, b_exp); - - // Store residues for this prime - let offset = pi * nn; - residues[offset..offset + nn].copy_from_slice(a_lane); + let mut ctx = TransformCtx { + a_lane, + b_lane, + fwd_twiddles, + inv_twiddles, + p: prime.p, + omega_2_32: prime.omega_2_32, + nn, + b_pack, + residues, + pi, + }; + match prime.b { + 32 => process_prime::<32>(a, b, &mut ctx), + 34 => process_prime::<34>(a, b, &mut ctx), + 40 => process_prime::<40>(a, b, &mut ctx), + _ => unreachable!(), + } } // ---- CRT per coefficient + accumulate ---- - // We'll accumulate each coefficient into the product buffer with - // b_pack-bit shift. let output_words = la + lb; for k in 0..output_coeffs { let mut coeff_residues = [0u64; 3]; #[allow(clippy::needless_range_loop)] for pi in 0..k_eff { - let offset = pi * nn; - coeff_residues[pi] = residues[offset + k]; + coeff_residues[pi] = residues[pi * nn + k]; } - let crt_val = crt::garner_combine(&coeff_residues[..k_eff], &primes_p, &crt_constants); - - // Unpack-accumulate this coefficient into prod - // crt_val is a small integer (≤ 3 u64 words) + let crt_val = crt::garner_combine(&coeff_residues[..k_eff], &crt_reducers[..k_eff]); add_shifted_to_prod(prod, &crt_val, k, b_pack); } // ---- Fold product into c with sign ---- - // Convert u64 slice to Word slice for the add function. - // On 64-bit targets these are the same type. assert_eq!( core::mem::size_of::(), core::mem::size_of::(), @@ -243,28 +234,55 @@ fn add_signed_mul_impl( } } -/// Pack word slice into coefficient buffer (viewed as u64). -fn pack_into(words: &[Word], b_pack: u32, out: &mut [u64]) { - let packed = pack::pack(words, b_pack, out.len()); - out.copy_from_slice(&packed); +/// Scratch buffers and parameters for one prime's NTT pipeline. +struct TransformCtx<'a> { + a_lane: &'a mut [u64], + b_lane: &'a mut [u64], + fwd_twiddles: &'a mut [u64], + inv_twiddles: &'a mut [u64], + p: u64, + omega_2_32: u64, + nn: usize, + b_pack: u32, + residues: &'a mut [u64], + pi: usize, } -/// Add a small multi-word integer (up to 3 u64 words) to `prod`, shifted -/// left by `k * b_pack` bits. -fn add_shifted_to_prod(prod: &mut [u64], val: &[u64], k: usize, b_pack: u32) { - if val.is_empty() { - return; - } +/// Per-prime NTT pipeline, monomorphized for a specific `B`. +#[inline(never)] +fn process_prime(a: &[Word], b: &[Word], ctx: &mut TransformCtx<'_>) { + pack::pack(ctx.a_lane, a, ctx.b_pack, ctx.nn); + pack::pack(ctx.b_lane, b, ctx.b_pack, ctx.nn); + + transform::precompute_twiddles::(ctx.fwd_twiddles, ctx.nn, ctx.p, ctx.omega_2_32, false); + transform::precompute_twiddles::(ctx.inv_twiddles, ctx.nn, ctx.p, ctx.omega_2_32, true); + + transform::bit_reverse(ctx.a_lane); + transform::bit_reverse(ctx.b_lane); + transform::forward::(ctx.a_lane, ctx.fwd_twiddles, ctx.p); + transform::forward::(ctx.b_lane, ctx.fwd_twiddles, ctx.p); + transform::pointwise_mul::(ctx.a_lane, ctx.b_lane); + transform::inverse::(ctx.a_lane, ctx.inv_twiddles, ctx.p); + + let offset = ctx.pi * ctx.nn; + ctx.residues[offset..offset + ctx.nn].copy_from_slice(ctx.a_lane); +} + +/// Add a CRT value to `prod`, shifted left by `k * b_pack` bits. +fn add_shifted_to_prod(prod: &mut [u64], val: &U192, k: usize, b_pack: u32) { + let count = val.len_words() as usize; let shift_bits = (k as u32).wrapping_mul(b_pack); let word_idx = (shift_bits / 64) as usize; let bit_shift = shift_bits % 64; let mut carry: u64 = 0; - for (vi, &v) in val.iter().enumerate() { + #[allow(clippy::needless_range_loop)] + for vi in 0..count { let idx = word_idx + vi; if idx >= prod.len() { return; } + let v = val.0[vi]; let v128 = v as u128; if bit_shift == 0 { @@ -273,8 +291,6 @@ fn add_shifted_to_prod(prod: &mut [u64], val: &[u64], k: usize, b_pack: u32) { prod[idx] = r; carry = (sum >> 64) as u64 + c as u64; } else { - // v << bit_shift has high bits = v >> (64 - bit_shift) = lo_carry. - // No separate hi — lo_carry IS the high part. let lo = v128 << bit_shift; let sum = lo.wrapping_add(carry as u128); let lo_carry = (sum >> 64) as u64; @@ -284,7 +300,6 @@ fn add_shifted_to_prod(prod: &mut [u64], val: &[u64], k: usize, b_pack: u32) { prod[idx] = r; carry = lo_carry + c1 as u64; - // Propagate to next word if idx + 1 < prod.len() && carry != 0 { let (r2, c2) = prod[idx + 1].overflowing_add(carry); prod[idx + 1] = r2; @@ -293,8 +308,7 @@ fn add_shifted_to_prod(prod: &mut [u64], val: &[u64], k: usize, b_pack: u32) { } } - // Propagate final carry - let mut idx = word_idx + val.len(); + let mut idx = word_idx + count; while carry != 0 && idx < prod.len() { let (r, c) = prod[idx].overflowing_add(carry); prod[idx] = r; diff --git a/integer/src/mul/ntt/pack.rs b/integer/src/mul/ntt/pack.rs index 14b31c7b..25baad94 100644 --- a/integer/src/mul/ntt/pack.rs +++ b/integer/src/mul/ntt/pack.rs @@ -9,24 +9,25 @@ use crate::arch::word::Word; -/// Pack a big integer (given as `&[Word]`, little-endian) into `n` -/// coefficients of `b_pack` bits each, zero-padded to length `n`. +/// Pack a big integer (given as `&[Word]`, little-endian) into `out`, +/// producing `n` coefficients of `b_pack` bits each, zero-padded. /// /// Each coefficient `c_i` satisfies `0 ≤ c_i < 2^{b_pack}`. -pub fn pack(words: &[Word], b_pack: u32, n: usize) -> alloc::vec::Vec { - let mut out = alloc::vec![0u64; n]; +/// Panics if `out.len() < n`. +pub fn pack(out: &mut [u64], words: &[Word], b_pack: u32, n: usize) { + assert!(out.len() >= n); let mask = (1u64 << b_pack) - 1; let word_bits = Word::BITS; let mut word_idx = 0usize; - let mut bit_offset = 0u32; // bit position within words[word_idx] + let mut bit_offset = 0u32; for coeff in out.iter_mut().take(n) { if word_idx >= words.len() { - break; // rest stay zero (padding) + *coeff = 0; + continue; } if bit_offset + b_pack <= word_bits { - // Entire coefficient fits within the current word. *coeff = (words[word_idx] >> bit_offset) & mask; bit_offset += b_pack; if bit_offset == word_bits { @@ -34,7 +35,6 @@ pub fn pack(words: &[Word], b_pack: u32, n: usize) -> alloc::vec::Vec { word_idx += 1; } } else { - // Coefficient straddles a word boundary. let bits_first = word_bits - bit_offset; let bits_second = b_pack - bits_first; let mut val = (words[word_idx] >> bit_offset) & ((1u64 << bits_first) - 1); @@ -46,8 +46,6 @@ pub fn pack(words: &[Word], b_pack: u32, n: usize) -> alloc::vec::Vec { bit_offset = bits_second; } } - - out } /// Accumulate CRT-recovered convolution coefficients into the output limb @@ -139,15 +137,13 @@ mod tests { #[test] fn test_pack_unpack_roundtrip() { - // Pack a number, then unpack-accumulate into a zero buffer. - // The accumulation should reconstruct the original number. let b_pack = 16u32; let test_words: Vec = vec![0xDEADBEEF_CAFEBABE, 0x12345678_9ABCDEF0]; let coeffs_per_word = (Word::BITS / b_pack) as usize; let n = test_words.len() * coeffs_per_word; - let packed = pack(&test_words, b_pack, n); - assert_eq!(packed.len(), n); + let mut packed = vec![0u64; n]; + pack(&mut packed, &test_words, b_pack, n); let output_len = test_words.len() + 1; let mut output = vec![0u64; output_len]; @@ -158,10 +154,9 @@ mod tests { #[test] fn test_pack_zero_pads() { let words = vec![0xFFFFu64]; - let n = 32; // more coeffs than content - let packed = pack(&words, 16, n); - assert_eq!(packed.len(), 32); - // First coeff = 0xFFFF (least significant 16 bits), rest = 0 + let n = 32; + let mut packed = vec![0u64; n]; + pack(&mut packed, &words, 16, n); assert_eq!(packed[0], 0xFFFF); for &c in packed.iter().skip(1) { assert_eq!(c, 0); @@ -170,7 +165,8 @@ mod tests { #[test] fn test_pack_empty_input() { - let packed = pack(&[], 16, 8); + let mut packed = vec![0u64; 8]; + pack(&mut packed, &[], 16, 8); assert_eq!(packed, vec![0u64; 8]); } diff --git a/integer/src/mul/ntt/primes.rs b/integer/src/mul/ntt/primes.rs index 81eeda37..4ee458bd 100644 --- a/integer/src/mul/ntt/primes.rs +++ b/integer/src/mul/ntt/primes.rs @@ -53,6 +53,14 @@ pub const PRIMES: [NttPrime; K] = [ }, ]; +/// Garner CRT constants: `inv(p_i mod p_j)` for i < j. +/// Computed offline via `pow(p_i % p_j, -1, p_j)`. +pub const CRT_INV_IJ: [[u64; 3]; 3] = [ + [0, 0xfffffffbaaaaaaad, 0xfffffefffefeff01], + [0, 0, 0xfffffefffefbefc1], + [0, 0, 0], +]; + #[cfg(test)] mod tests { use super::*; diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs index be27402c..8a528d86 100644 --- a/integer/src/mul/ntt/transform.rs +++ b/integer/src/mul/ntt/transform.rs @@ -1,30 +1,27 @@ //! Iterative in-place radix-2 NTT over primes of the form `2^64 - 2^b + 1`. //! -//! Uses decimation-in-time (DIT) Cooley–Tukey butterflies. All modular -//! arithmetic delegates to `num_modular::FixedTrinomialSolinas64`. +//! All functions are const-generic over `B` (the Solinas exponent, one of +//! `{32, 34, 40}`) so the compiler monomorphizes each prime's hot path. +//! Modular arithmetic delegates to `num_modular::FixedTrinomialSolinas64`. use num_modular::{FixedTrinomialSolinas64, ModularPow, ModularUnaryOps, Reducer}; -// ---- dispatch helpers (b is a runtime value, const-generic under the hood) ---- +// ---- dispatch helpers (the match is optimized away since B is const) ---- #[inline] -fn reduce_double(v: u128, b: u32) -> u64 { - match b { - 32 => FixedTrinomialSolinas64::<64, 32, 1>::reduce_double(v), - 34 => FixedTrinomialSolinas64::<64, 34, 1>::reduce_double(v), - 40 => FixedTrinomialSolinas64::<64, 40, 1>::reduce_double(v), +fn mul_mod(a: u64, b_val: u64) -> u64 { + let prod = (a as u128) * (b_val as u128); + match B { + 32 => FixedTrinomialSolinas64::<64, 32, 1>::reduce_double(prod), + 34 => FixedTrinomialSolinas64::<64, 34, 1>::reduce_double(prod), + 40 => FixedTrinomialSolinas64::<64, 40, 1>::reduce_double(prod), _ => unreachable!(), } } #[inline] -fn mul_mod(a: u64, b_val: u64, b: u32) -> u64 { - reduce_double((a as u128) * (b_val as u128), b) -} - -#[inline] -fn add_mod(a: u64, b_val: u64, p: u64, b: u32) -> u64 { - match b { +fn add_mod(a: u64, b_val: u64, p: u64) -> u64 { + match B { 32 => FixedTrinomialSolinas64::<64, 32, 1>::new(&p).add(&a, &b_val), 34 => FixedTrinomialSolinas64::<64, 34, 1>::new(&p).add(&a, &b_val), 40 => FixedTrinomialSolinas64::<64, 40, 1>::new(&p).add(&a, &b_val), @@ -33,8 +30,8 @@ fn add_mod(a: u64, b_val: u64, p: u64, b: u32) -> u64 { } #[inline] -fn sub_mod(a: u64, b_val: u64, p: u64, b: u32) -> u64 { - match b { +fn sub_mod(a: u64, b_val: u64, p: u64) -> u64 { + match B { 32 => FixedTrinomialSolinas64::<64, 32, 1>::new(&p).sub(&a, &b_val), 34 => FixedTrinomialSolinas64::<64, 34, 1>::new(&p).sub(&a, &b_val), 40 => FixedTrinomialSolinas64::<64, 40, 1>::new(&p).sub(&a, &b_val), @@ -44,17 +41,17 @@ fn sub_mod(a: u64, b_val: u64, p: u64, b: u32) -> u64 { // ---- public API ---- -/// Precompute the twiddle-factor table for transform length `n`. +/// Fill `out[0..n/2]` with twiddle factors `omega_n^k`. /// -/// Returns `n/2` entries: `omega_n^k` for `k = 0..n/2`, where -/// `omega_n = omega_2_32^{2^32 / n}` (a primitive `n`-th root of unity). -pub fn precompute_twiddles( +/// Panics if `out.len() < n / 2`. +pub fn precompute_twiddles( + out: &mut [u64], n: usize, p: u64, - b: u32, omega_2_32: u64, inverse: bool, -) -> alloc::vec::Vec { +) { + assert!(out.len() >= n / 2); let shift = 32 - n.trailing_zeros(); let omega_n = omega_2_32.powm(&(1u64 << shift), &p); @@ -64,12 +61,10 @@ pub fn precompute_twiddles( omega_n }; - let mut twiddles = alloc::vec![0u64; n / 2]; - twiddles[0] = 1; + out[0] = 1; for k in 1..(n / 2) { - twiddles[k] = mul_mod(twiddles[k - 1], base, b); + out[k] = mul_mod::(out[k - 1], base); } - twiddles } /// Bit-reverse `a` in place. Length must be a power of two. @@ -86,8 +81,8 @@ pub fn bit_reverse(a: &mut [u64]) { } /// Forward NTT in place (decimation-in-time, radix-2). -pub fn forward(a: &mut [u64], twiddles: &[u64], p: u64, b: u32) { - ntt_core(a, twiddles, p, b); +pub fn forward(a: &mut [u64], twiddles: &[u64], p: u64) { + ntt_core::(a, twiddles, p); } /// Inverse NTT in place. @@ -95,20 +90,19 @@ pub fn forward(a: &mut [u64], twiddles: &[u64], p: u64, b: u32) { /// Computed as `bit_reverse → forward(ω^{-1}) → scale`, producing output /// in **natural order**. /// -/// `twiddles` must have been precomputed with `inverse = true` -/// (i.e. using `omega_n^{-1}`). -pub fn inverse(a: &mut [u64], twiddles: &[u64], p: u64, b: u32) { +/// `twiddles` must have been precomputed with `inverse = true`. +pub fn inverse(a: &mut [u64], twiddles: &[u64], p: u64) { let n = a.len(); bit_reverse(a); - ntt_core(a, twiddles, p, b); + ntt_core::(a, twiddles, p); let n_inv = (n as u64).invm(&p).expect("n not invertible mod p"); for x in a.iter_mut() { - *x = mul_mod(*x, n_inv, b); + *x = mul_mod::(*x, n_inv); } } /// In-place radix-2 DIT NTT (Cooley–Tukey). -fn ntt_core(a: &mut [u64], twiddles: &[u64], p: u64, b: u32) { +fn ntt_core(a: &mut [u64], twiddles: &[u64], p: u64) { let n = a.len(); debug_assert!(n.is_power_of_two() && twiddles.len() == n / 2); @@ -120,9 +114,9 @@ fn ntt_core(a: &mut [u64], twiddles: &[u64], p: u64, b: u32) { for i in (0..n).step_by(sub_len) { for j in 0..half { let u = a[i + j]; - let v = mul_mod(a[i + j + half], twiddles[j * step], b); - a[i + j] = add_mod(u, v, p, b); - a[i + j + half] = sub_mod(u, v, p, b); + let v = mul_mod::(a[i + j + half], twiddles[j * step]); + a[i + j] = add_mod::(u, v, p); + a[i + j + half] = sub_mod::(u, v, p); } } @@ -131,10 +125,10 @@ fn ntt_core(a: &mut [u64], twiddles: &[u64], p: u64, b: u32) { } /// Pointwise multiply of two transformed vectors in place. -pub fn pointwise_mul(a_hat: &mut [u64], b_hat: &[u64], b: u32) { +pub fn pointwise_mul(a_hat: &mut [u64], b_hat: &[u64]) { assert_eq!(a_hat.len(), b_hat.len()); for (a, &b_val) in a_hat.iter_mut().zip(b_hat.iter()) { - *a = mul_mod(*a, b_val, b); + *a = mul_mod::(*a, b_val); } } @@ -150,14 +144,43 @@ mod tests { } } + macro_rules! for_each_prime { + ($b:ident, $p:ident, $omega:ident, $body:block) => { + for prime in &PRIMES { + let $b = prime.b; + match $b { + 32 => { + let $p = prime.p; + let $omega = prime.omega_2_32; + fn go($p: u64, $omega: u64) $body + go::<32>($p, $omega); + } + 34 => { + let $p = prime.p; + let $omega = prime.omega_2_32; + fn go($p: u64, $omega: u64) $body + go::<34>($p, $omega); + } + 40 => { + let $p = prime.p; + let $omega = prime.omega_2_32; + fn go($p: u64, $omega: u64) $body + go::<40>($p, $omega); + } + _ => unreachable!(), + } + } + }; + } + #[test] fn test_forward_inverse_roundtrip() { - for prime in &PRIMES { - let p = prime.p; - let b = prime.b; + for_each_prime!(b, p, omega, { for &n in &[2, 4, 8, 16, 32, 64, 128, 256, 512] { - let inv_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, true); - let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); + let mut fwd_twiddles = alloc::vec![0u64; n / 2]; + let mut inv_twiddles = alloc::vec![0u64; n / 2]; + precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); + precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); let mut a: Vec = (0..n) .map(|i| ((i as u64 + 1).wrapping_mul(123456789)) % p) @@ -165,19 +188,17 @@ mod tests { let orig = a.clone(); bit_reverse(&mut a); - forward(&mut a, &fwd_twiddles, p, b); - inverse(&mut a, &inv_twiddles, p, b); + forward::(&mut a, &fwd_twiddles, p); + inverse::(&mut a, &inv_twiddles, p); assert_all_eq(&a, &orig); } - } + }); } #[test] fn test_convolution_via_ntt() { - for prime in &PRIMES { - let p = prime.p; - let b = prime.b; + for_each_prime!(b, p, omega, { for len_a in [1, 2, 3, 5] { for len_b in [1, 2, 3, 5] { let conv_len: usize = len_a + len_b - 1; @@ -190,12 +211,15 @@ mod tests { let mut expected = vec![0u64; conv_len]; for (i, &ai) in a.iter().enumerate() { for (j, &bj) in b_vec.iter().enumerate() { - expected[i + j] = add_mod(expected[i + j], mul_mod(ai, bj, b), p, b); + expected[i + j] = + add_mod::(expected[i + j], mul_mod::(ai, bj), p); } } - let inv_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, true); - let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); + let mut fwd_twiddles = alloc::vec![0u64; n / 2]; + let mut inv_twiddles = alloc::vec![0u64; n / 2]; + precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); + precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); let mut a_pad = vec![0u64; n]; let mut b_pad = vec![0u64; n]; @@ -204,31 +228,15 @@ mod tests { bit_reverse(&mut a_pad); bit_reverse(&mut b_pad); - forward(&mut a_pad, &fwd_twiddles, p, b); - forward(&mut b_pad, &fwd_twiddles, p, b); - pointwise_mul(&mut a_pad, &b_pad, b); - inverse(&mut a_pad, &inv_twiddles, p, b); - - match assert_all_eq_result(&a_pad[..conv_len], &expected) { - Ok(()) => {} - Err((i, l, r)) => panic!( - "convolution mismatch: b={b}, len_a={len_a}, len_b={len_b}, n={n}, \ - index {i}: {l} != {r}" - ), - } - } - } - } - } + forward::(&mut a_pad, &fwd_twiddles, p); + forward::(&mut b_pad, &fwd_twiddles, p); + pointwise_mul::(&mut a_pad, &b_pad); + inverse::(&mut a_pad, &inv_twiddles, p); - fn assert_all_eq_result(a: &[u64], b_val: &[u64]) -> Result<(), (usize, u64, u64)> { - assert_eq!(a.len(), b_val.len()); - for (i, (x, y)) in a.iter().zip(b_val.iter()).enumerate() { - if x != y { - return Err((i, *x, *y)); + assert_all_eq(&a_pad[..conv_len], &expected); + } } - } - Ok(()) + }); } #[test] @@ -239,7 +247,7 @@ mod tests { } #[allow(clippy::needless_range_loop)] - fn ntt_naive(x: &[u64], omega_n: u64, p: u64, b: u32) -> Vec { + fn ntt_naive(x: &[u64], omega_n: u64, p: u64) -> Vec { let n = x.len(); let mut result = vec![0u64; n]; for k in 0..n { @@ -250,7 +258,7 @@ mod tests { } else { omega_n.powm(&((k * j) as u64), &p) }; - acc = add_mod(acc, mul_mod(x[j], twiddle, b), p, b); + acc = add_mod::(acc, mul_mod::(x[j], twiddle), p); } result[k] = acc; } @@ -259,35 +267,30 @@ mod tests { #[test] fn test_forward_correctness() { - for prime in &PRIMES { - let p = prime.p; - let b = prime.b; + for_each_prime!(b, p, omega, { for &n in &[2usize, 4, 8] { let x: Vec = (0..n).map(|i| ((i + 1) as u64 * 11111) % p).collect(); let shift = 32 - n.trailing_zeros(); - let omega_n = prime.omega_2_32.powm(&(1u64 << shift), &p); + let omega_n = omega.powm(&(1u64 << shift), &p); let mut a = x.clone(); bit_reverse(&mut a); - let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); - forward(&mut a, &fwd_twiddles, p, b); + let mut fwd_twiddles = alloc::vec![0u64; n / 2]; + precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); + forward::(&mut a, &fwd_twiddles, p); - let expected = ntt_naive(&x, omega_n, p, b); - assert_eq!( - a, expected, - "forward NTT mismatch: b={b}, n={n}, expected={expected:?}, got={a:?}" - ); + let expected = ntt_naive::(&x, omega_n, p); + assert_eq!(a, expected, "forward NTT mismatch"); } - } + }); } #[test] fn test_convolution_debug() { - let prime = &PRIMES[0]; + let prime = &PRIMES[0]; // GL: b=32 let p = prime.p; - let b = prime.b; let a = vec![12345u64 % p]; let b_vec = vec![67890u64 % p, 135780u64 % p, 203670u64 % p]; @@ -297,12 +300,14 @@ mod tests { let mut expected = vec![0u64; conv_len]; for (i, &ai) in a.iter().enumerate() { for (j, &bj) in b_vec.iter().enumerate() { - expected[i + j] = add_mod(expected[i + j], mul_mod(ai, bj, b), p, b); + expected[i + j] = add_mod::<32>(expected[i + j], mul_mod::<32>(ai, bj), p); } } - let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); - let inv_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, true); + let mut fwd_twiddles = alloc::vec![0u64; n / 2]; + let mut inv_twiddles = alloc::vec![0u64; n / 2]; + precompute_twiddles::<32>(&mut fwd_twiddles, n, p, prime.omega_2_32, false); + precompute_twiddles::<32>(&mut inv_twiddles, n, p, prime.omega_2_32, true); let mut a_pad = vec![0u64; n]; let mut b_pad = vec![0u64; n]; @@ -311,29 +316,29 @@ mod tests { bit_reverse(&mut a_pad); bit_reverse(&mut b_pad); - forward(&mut a_pad, &fwd_twiddles, p, b); - forward(&mut b_pad, &fwd_twiddles, p, b); - pointwise_mul(&mut a_pad, &b_pad, b); - inverse(&mut a_pad, &inv_twiddles, p, b); + forward::<32>(&mut a_pad, &fwd_twiddles, p); + forward::<32>(&mut b_pad, &fwd_twiddles, p); + pointwise_mul::<32>(&mut a_pad, &b_pad); + inverse::<32>(&mut a_pad, &inv_twiddles, p); assert_eq!(&a_pad[..conv_len], &expected[..]); } #[test] fn test_length_two_edge_case() { - for prime in &PRIMES { - let p = prime.p; - let b = prime.b; + for_each_prime!(b, p, omega, { let n = 2; - let fwd_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, false); - let inv_twiddles = precompute_twiddles(n, p, b, prime.omega_2_32, true); + let mut fwd_twiddles = alloc::vec![0u64; n / 2]; + let mut inv_twiddles = alloc::vec![0u64; n / 2]; + precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); + precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); let a_orig = vec![1u64 % p, 2u64 % p]; let mut a = a_orig.clone(); bit_reverse(&mut a); - forward(&mut a, &fwd_twiddles, p, b); - inverse(&mut a, &inv_twiddles, p, b); + forward::(&mut a, &fwd_twiddles, p); + inverse::(&mut a, &inv_twiddles, p); assert_all_eq(&a, &a_orig); - } + }); } } From f5d8c05906f0acf069369fffb86beec415888c02 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sat, 13 Jun 2026 11:04:42 +0800 Subject: [PATCH 04/19] WIP: some param tuning --- Cargo.toml | 1 + TODO_NTT.md | 476 +++++++++---------------------------- integer/CHANGELOG.md | 10 +- integer/Cargo.toml | 1 + integer/src/mul/mod.rs | 78 ++++-- integer/src/mul/ntt/crt.rs | 2 +- integer/src/mul/ntt/mod.rs | 181 +++++++++++--- 7 files changed, 332 insertions(+), 417 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2f7761ca..0a6d0ba5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ std = ["dashu-base/std", "dashu-int/std", "dashu-float/std", "dashu-ratio/std"] # stable features serde = ["dashu-int/serde", "dashu-float/serde", "dashu-ratio/serde"] num-order = ["dashu-int/num-order", "dashu-float/num-order", "dashu-ratio/num-order"] +tuning = ["dashu-int/tuning"] zeroize = ["dashu-int/zeroize", "dashu-float/zeroize", "dashu-ratio/zeroize"] # unstable features diff --git a/TODO_NTT.md b/TODO_NTT.md index 58ee150a..4c9a6a69 100644 --- a/TODO_NTT.md +++ b/TODO_NTT.md @@ -1,385 +1,135 @@ -# Prime-NTT multiplication for `UBig` — implementation plan +# NTT multiplication for `UBig` — status & remaining work -Goal: add an NTT-based large-integer multiplication path to `dashu-int` that -kicks in above the current Toom-Cook-3 range, using power-of-two Number -Theoretic Transforms over several 64-bit NTT-friendly primes combined with the -Chinese Remainder Theorem (CRT). +## Implemented -This document is the working spec + checklist. Keep it updated as the work -progresses. Each phase ends with a state where `cargo test -p dashu-int` -passes. +### Core algorithm (Phases 1–6) ---- +- **Primes.** Three Goldilocks-style Solinas primes `p = 2^64 − 2^b + 1` with + `b ∈ {32, 34, 40}`. All support shift-based reduction via `2^64 ≡ 2^b − 1`. + Stored in `integer/src/mul/ntt/primes.rs` with a full `verify_primes()` test + (Miller–Rabin, exact root order, reduction-identity self-check). -## 0. Background and decisions - -- **Word layout.** `Word = u64`, `DoubleWord = u128` on the default build - (`integer/src/arch/generic_64_bit/word.rs`). The 16-/32-bit `force_bits` - builds use smaller `Word`, but `u128` is available on every Rust target, so - the NTT lane arithmetic is always done in `u64`/`u128` regardless of `Word`. -- **Where it plugs in.** Multiplication is dispatched in - `integer/src/mul/mod.rs` by the *smaller* operand length in words: - - `len <= THRESHOLD_SIMPLE (24)` → `simple` - - `<= THRESHOLD_KARATSUBA (192)` → `karatsuba` - - else → `toom_3` (currently unbounded above) - - We add `THRESHOLD_NTT` (in words) above which `toom_3` hands off to the new - `ntt` module. The exact value is tuned by benchmark (Phase 7); start with a - conservative placeholder (e.g. `2048` words ≈ 130k bits). -- **Scheme.** Linear (acyclic) convolution of the two coefficient polynomials, - realised as a cyclic convolution of length `N = next_pow2(La + Lb - 1)` with - zero padding. Computed independently modulo `K` primes, then CRT-combined per - coefficient, then carry-propagated into the output limbs. -- **Why ≥ 3 primes + CRT (recap).** A single ~2^64 prime forces tiny coefficient - chunks. With `K` primes of product `P ≈ 2^(64K)` we can pack `b` bits per - coefficient as long as the largest convolution coefficient - `< N/2 · (2^b − 1)^2 < P`. Three primes (`P ≈ 2^189..192`) give comfortable - headroom for any feasible input and let us pick a larger `b` (fewer - coefficients → smaller transform). `K = 2` is provably sufficient for all - inputs `UBig` can physically hold, but `K = 3` is the default for margin and - speed. - - **`K` vs `K_eff` (avoid const-generic monomorphization).** `K = 3` is the - *fixed array size* of `PRIMES` (a plain `const`, not a const-generic type - parameter), so the transform code is monomorphized once. `select_params` - returns a runtime `K_eff ≤ K`; the per-prime loop iterates - `PRIMES[..K_eff]`. The Phase-7 "drop to 2 primes" optimisation is just - `K_eff = 2` at runtime — no extra monomorphization, no new generic - instantiations. -- **Modulus family: Goldilocks-style** `p = 2^64 − 2^b + 1`. Every member - satisfies the shift-based reduction identity `2^64 ≡ 2^b − 1 (mod p)`, so the - *entire* lane arithmetic is multiplication-free in the reduction step — no - Montgomery form anywhere. We use three members of this single family (below) - so `modarith` is **one** routine parameterised by `b`. - -### Chosen primes (verified) - -Survey of documented choices considered and rejected: -- **"Ultimate NTT" prime** `9223372036737335297 = 549755813881·2^24 + 1`, - `g = 3` (Codeforces entry 75326): `v2` only 24 caps `N` at `2^24`, and it is - not a Solinas form (needs Barrett/Montgomery). Rejected. -- **Classic CRT trio** `998244353`, `985661441`, `754974721`, … : ~30-bit, so - they waste 64-bit lanes and cap `N` low. Rejected. -- **`c·2^32+1` siblings** (e.g. `0xFFFFFFD300000001`): right size and `v2`, but - not Solinas form → would force a Montgomery path. Rejected in favour of the - uniform fast-reduction family below. - -**Selected trio — all of the form `2^64 − 2^b + 1`** (verified prime by -deterministic Miller–Rabin, full-order `2^32`-th root checked): - -| name | `b` | `p` (hex) | `p` (dec) | `v2(p−1)` | gen `g` | `2^32`-th root ω | -|---|---|---|---|---|---|---| -| GL | 32 | `0xFFFFFFFF00000001` | `18446744069414584321` | 32 | 7 | `1753635133440165772` | -| P1 | 34 | `0xFFFFFFFC00000001` | `18446744056529682433` | 34 | 5 | `11315553352654630047` | -| P2 | 40 | `0xFFFFFF0000000001` | `18446742974197923841` | 40 | 19 | `551857376737322389` | - -- `min(v2) = 32` ⇒ transform length up to `2^32` coefficients (≈ ~1 GB operands - at `b_pack = 16`); `P = GL·P1·P2 ≈ 2^192` of CRT headroom. -- All three reduce via `2^64 ≡ 2^bᵢ − 1`. GL's `b = 32` is the cleanest (splits - a 128-bit product into 32-bit limbs, `φ²=φ−1`); `b = 34, 40` need one extra - shift/fold because the split crosses the 32-bit boundary, but stay - multiply-free. Implement the reduction generically over `b` with the GL case - as the well-trodden reference. -- `ω⁻¹` and `N⁻¹` are derived per call (cold, once per prime — not lazy): - `ω_N = pow(ω, 2^32 / N)` (the stored `ω` has exact order `2^32`); the inverse - root is just `ω_N⁻¹ = pow(ω_N, N − 1)` (since `ω_N^N = 1`, no `inv` needed), - and `N⁻¹ = Reducer::inv(N mod p)` once. Use `num_modular::Reducer::{pow, inv}` - for all three. Commit a `verify_primes()` test - (Miller–Rabin + `v2` + exact root order + reduction-identity self-check) - rather than trusting these literals blindly. - -### Open decisions to lock during Phase 1 -- [ ] Final `K` (start 3). -- [ ] Coefficient bit width `b_pack` (**default 16**: 4 coeffs/word, trivial - shift/mask packing). Larger `b_pack` = fewer coefficients (smaller - transform) but needs more headroom and must satisfy - `(N/2)·(2^{b_pack}−1)^2 < P` for the max supported `N`. Candidate values: - - `16` — divides 64, byte-aligned, 4 coeffs/word. **Default unless a - benchmark proves a larger value wins.** - - `24` — divides 64? no, but byte-aligned (3 bytes), 8 coeffs/3 words. - - `21` — **avoid**: does not divide 64 and is not byte-aligned, so - coefficients straddle word *and* byte boundaries → slower, buggier - pack/unpack. Only revisit if its transform-size win clearly beats the - pack cost in benchmarks. -- [ ] Whether to gate the NTT path on `cfg(target_pointer_width)` / `Word` - width, or always enable it (preferred: always enable, since lane math is - `u64`/`u128`). Document the chosen rule. - -### Word-width targets (future work, not near-term) - -The `2^64 − 2^b + 1` primes are correct on every `Word` width (`u64`/`u128` are -universal types), but on narrow targets the `u64×u64→u128` lane multiply is -emulated and slow. Plan: - -- **64-bit `Word`**: primary target, the chosen 3 primes above. -- **32-bit `Word`**: **select a separate set of three ~32-bit Solinas primes - (`2^32 − 2^b + 1`, via `FixedTrinomialSolinas32`)** so the lane multiply is a - native `u32×u32→u64`. Feasible: 3 such primes give `P ≈ 2^96`, which (with - `b_pack = 16`, max coefficient `≈ N·2^32`) is far more headroom than needed — - the transform-length ceiling comes from each prime's `v2(p−1)`, which is ample - for any input a 32-bit target would handle. Requires extending - `FixedTrinomialSolinas32` to `P1 = 32` (same `checked_shl` fix already done for - the 64-bit type) plus a prime/root search. **Not intended for implementation - in the near future** — design the NTT core generic over the prime set so this - can be added later as configuration, not a rewrite. -- **16-bit `Word`**: do not implement an NTT path; fall back to Toom-3. +- **Modular arithmetic.** Lane arithmetic delegates to + `num_modular::FixedTrinomialSolinas64` for `add`/`sub`/`mul`/`reduce_double`. + `add`/`sub`/`mul` are monomorphized per `B` via const generics so the + compiler fully inlines each prime's hot path. ---- +- **NTT transforms.** Iterative in-place radix-2 decimation-in-time + (`integer/src/mul/ntt/transform.rs`). Forward transform: `bit_reverse → + forward(ω)`. Inverse transform: `bit_reverse → forward(ω⁻¹) → scale by + N⁻¹`. Twiddle tables precomputed once per prime per call. Pointwise + multiply in the transform domain. -## 1. Module layout +- **Packing / unpacking.** Bit-level `pack` slices `&[Word]` into `b_pack`-bit + coefficients (`integer/src/mul/ntt/pack.rs`). CRT-recovered coefficients are + accumulated into the output limb array via shifted addition + (`add_shifted_to_prod`). -New directory `integer/src/mul/ntt/` (declare `mod ntt;` in -`integer/src/mul/mod.rs`): +- **CRT.** Garner's algorithm combining `K` residues modulo `K` primes into a + `U192` (3 × u64) integer (`integer/src/mul/ntt/crt.rs`). All Garner + precomputed constants are hardcoded in `primes.rs` (`CRT_INV_IJ`). Uses an + object-safe `ModOps` trait (subset of `Reducer`) for dynamic dispatch + over the per-prime reducers. -| File | Responsibility | -|---|---| -| `ntt/mod.rs` | Public entry `add_signed_mul` / `add_signed_mul_same_len`, `memory_requirement_up_to`, `THRESHOLD_NTT`, parameter selection (`b`, `N`, `K`). | -| `ntt/primes.rs` | Const table of the `K` primes: value, `b`, primitive root, `v2(p−1)`, precomputed `2^32`-th root. Includes a `verify_primes` unit test. | -| `ntt/modarith.rs` | Lane arithmetic for the prime `2^64 − 2^b + 1`. **Reuse `num_modular::FixedTrinomialSolinas64<64, b, 1>::{reduce_single, reduce_double}` for the reduction step** (already P1=64-correct and unrolled to the verified fold counts; both are `pub` as of the `checked_shl` fix — verify on the pinned version, see fallback below). Write our **own** lazy `add` / `sub` / `mul` on top so we control deferred reduction; reuse num-modular's `Reducer::{pow, inv}` (fully normalized, called once per prime per call — no lazy variant needed). See note below. | -| `ntt/transform.rs` | Iterative in-place radix-2 forward/inverse NTT, twiddle precomputation, bit-reversal, pointwise multiply. | -| `ntt/pack.rs` | Bit-slice an operand `&[Word]` into `N` coefficients of `b` bits (mod each prime), and the inverse: CRT-combine residues + carry-propagate into the output limbs. | -| `ntt/crt.rs` | Garner CRT for `K` residues → a small (≤ `K`-word) integer per coefficient. | +- **Dispatch.** Multiplication above `THRESHOLD_NTT` words routes to the NTT + path (`integer/src/mul/mod.rs`). `add_signed_mul` (unequal lengths) and + `add_signed_mul_same_len` (equal lengths) share a single `add_signed_mul_impl` + that does one NTT convolution — no chunking for equal/similar lengths. -Mirror the existing modules' conventions: `#[must_use]` on the `add_signed_mul*` -functions, return `SignedWord` carry, doc comments with complexity, `Buffer` / -`Memory` for scratch (no `Vec`), no `std`. +- **Memory.** Scratch space carved from the linear `Memory` arena. Worst-case + bound computed in `memory_requirement_up_to` using `B_PACK_MIN = 16` + (largest possible N for a given operand size). ---- +### Phase 7 optimisations (completed) -## 2. Math reference (for reviewers and tests) +- **K_eff = 2 auto-selection.** `select_params` checks headroom against + `P0·P1` (≈2^128). For `b_pack = 16`, `max_coeff < 2^63 ≪ 2^128`, so two + primes always suffice. Third-prime fallback (`K_eff = 3`) is retained as a + safety net for larger `b_pack`. -Operands `A = sum_i a_i 2^{ib}`, `B = sum_j b_j 2^{jb}` with `0 <= a_i, b_j < 2^b`. -Product `C = A·B = sum_k c_k 2^{kb}` where `c_k = sum_{i+j=k} a_i b_j` is exactly -the linear convolution coefficient, `0 <= c_k < (k+1)·(2^b−1)^2 <= N·(2^b−1)^2`. +- **Threshold calibrated.** `THRESHOLD_NTT = 120 000` words (~7.7 M bits), + the first measured crossover where NTT beats pure toom-3 on Apple M4 Pro. + At 131 072 words N doubles (1 048 576 → 2 097 152) and NTT regresses until + toom-3 catches up at ~190 k words — radix-4 will shrink this gap. -Compute `c_k mod p_t` for each prime `p_t` via length-`N` cyclic convolution -(forward NTT, pointwise product, inverse NTT). Because `N >= La + Lb − 1`, the -cyclic and linear convolutions coincide. CRT recovers exact `c_k < P`. Finally -`C = sum_k c_k 2^{kb}` with carry propagation (coefficients overlap whenever -`bitlen(c_k) > b`). +- **Env-var overrides.** `DASHU_THRESHOLD_SIMPLE`, `DASHU_THRESHOLD_KARATSUBA`, + `DASHU_THRESHOLD_NTT` override the compile-time defaults at runtime. Gated + behind the `tuning` feature (implies `std`). -Roots: `omega_N = g^{(p−1)/N} mod p` is a primitive `N`-th root of unity; require -`N | 2^{v2(p−1)}`. Inverse transform uses `omega_N^{-1}` and a final scale by -`N^{-1} mod p`. +- **Crossover benchmark.** `#[ignore]` test `crossover()` in + `integer/src/mul/ntt/mod.rs` compares NTT against toom-3 at key sizes. + Run with `DASHU_THRESHOLD_NTT=99999999` to force pure toom-3. --- -## 3. Phase plan (each phase is independently testable) - -> **Scheduling note.** Phases 2 (transform arithmetic) and 3 (pack / unpack / -> CRT) have no dependency on each other — both only need Phase 1's `modarith` -> and `primes`. They can be built and tested in parallel, then joined in -> Phase 4. Phase 1 must land first; Phases 5–7 follow Phase 4. - -### Phase 1 — Primes, modular arithmetic, parameter selection -- [ ] `ntt/primes.rs`: define `const PRIMES: [NttPrime; K]` from the "Chosen - primes" table. Each entry stores `p`, the form exponent `b`, primitive - root `g`, `v2(p−1)`, and the precomputed `2^32`-th root `ω`. (No Montgomery - constants — the family needs none.) -- [ ] Add `#[test] fn verify_primes()` re-checking each entry: primality - (Miller–Rabin over fixed bases), `p == 2^64 − 2^b + 1 < 2^64`, - `v2(p−1) >= MAX_LOG_N (=32)`, stored `g` generates the order-`2^{v2}` - subgroup, `ω` has exact order `2^32`, and the reduction identity - `2^64 ≡ 2^b − 1 (mod p)` holds. Do not trust the literals without it. -- [ ] `ntt/modarith.rs`: **delegate the reduction** to - `num_modular::FixedTrinomialSolinas64<64, b, 1>` — its `reduce_single` - (≤ `2^64` → `[0, p)`) and `reduce_double` (`u128` product → `[0, p)`) are - already correct for `P1 = 64` and straight-line unrolled (3 folds for - `b = 32`, 4 for `b = 34, 40`). Do **not** re-derive the shift/fold here. - Both methods are generated `pub` by the `impl_fixed_trinomial_solinas!` - macro, so they are directly callable from `dashu-int`. - - **Fallback if upstream visibility ever regresses:** the reduction is - ~20 lines per arm; copy it verbatim into `modarith.rs` (it is simple - enough to own in-tree, and the only coupling point). Pin/assert the - `num-modular` version in `integer/Cargo.toml` so a downgrade can't - silently break the `pub` assumption. - - We write our **own lazy `add` / `sub` / `mul`** (not num-modular's) - because its `Reducer` API fully normalizes to `[0, p)` after every op - and exposes no partially-reduced form. Ours keep values lazily in - `[0, 2p)` (or `[0, 4p)`) and only call `reduce_*` / a final conditional - subtract when needed — this is the Harvey-style lazy reduction that the - NTT butterflies depend on. `mul` = `u128` widening multiply → - `reduce_double`; `add` / `sub` = wrapping add/sub with deferred - normalization. - - **`pow` / `inv` are NOT lazy and are NOT ours.** They are called once - per prime per multiplication (`ω_N = g^{(p−1)/N}`, `N^{-1}`), never in - the butterfly hot loop, so use `num_modular::Reducer::{pow, inv}` - directly (fully normalized). No partial-reduction benefit there. - - Rationale: reduction is the subtle, already-tested part (reuse it); the - lazy add/sub/mul wrapper is trivial and must be ours to control the - normalization schedule; pow/inv are cold and reused as-is. -- [ ] `ntt/mod.rs`: `select_params(la_bits, lb_bits) -> (b_pack, N, K_eff)` with - the headroom assertion `(N as u128 / 2) * (2^{b_pack} − 1)^2 < P` (may drop - to fewer primes for smaller inputs later). -- [ ] Unit tests: our lazy `add`/`sub`/`mul` (after a final normalize) agree - with `FixedTrinomialSolinas64`'s fully-reduced `add`/`sub`/`mul` and a - `u128`/`u256` reference, across the `[0, 2p)` input range for each `b`; - `Reducer::{pow, inv}` round-trip (`inv(x)·x ≡ 1`, `pow(g, p−1) ≡ 1`); - `verify_primes`. - -### Phase 2 — Forward/inverse NTT -- [ ] `ntt/transform.rs`: iterative Cooley–Tukey radix-2 forward NTT in place, - decimation-in-time with bit-reversal permutation; inverse NTT - (conjugate twiddles + scale by `N^{-1}`). -- [ ] Twiddle factors: precompute the `omega_N^k` table per prime into scratch - once per call (length `N/2`). -- [ ] `pointwise_mul(a_hat, b_hat)` via `modarith::mul` (lazy; normalize at the - end of the inverse transform). -- [ ] Tests: `inverse(forward(x)) == x`; NTT-based cyclic convolution of small - random vectors equals the schoolbook cyclic convolution mod `p`; check the - length-2 and length-power-of-two edge cases. - -### Phase 3 — Packing / unpacking + CRT -- [ ] `ntt/pack.rs::pack`: read `b`-bit coefficients out of `&[Word]` - (bit-level slicing across word boundaries; works for any `Word` width), - reduce mod each prime (a `b_pack`-bit value is already `< p`, so this is a - copy), write into the length-`N` (zero-padded) lane buffers. -- [ ] `ntt/crt.rs`: Garner combine `K` residues of one coefficient → an integer - of ≤ `K` words (value `< P`). -- [ ] `ntt/pack.rs::unpack_accumulate`: for each `k`, add `c_k << (k·b)` bits - into the output limbs with carry propagation. Implement as a streaming - shifted add (reuse `add::add_*` helpers / `shift`). -- [ ] Tests: `unpack_accumulate(pack(x)) == x` identity for the - no-multiplication case (coefficients copied straight through CRT), and a - direct check that pack→CRT→unpack reconstructs a known convolution. - -### Phase 4 — Wire the full multiply -- [ ] `ntt/mod.rs::add_signed_mul_same_len` and `add_signed_mul`: orchestrate - select_params → per-prime (pack, forward, pointwise, inverse) → CRT per - coefficient → unpack/accumulate into a temp product buffer → fold into `c` - via `add::add_signed_*` honoring `sign`. Return the carry as the other - algorithms do. -- [ ] **Unequal-length entry point — do NOT blindly copy `toom_3`'s chunking.** - Unlike Toom-3/Karatsuba (defined on equal-length operands, hence - `helpers::add_signed_mul_split_into_chunks` slices the long operand into - balanced pieces), a single NTT convolution handles unequal lengths - natively: pad both operands to one `N = next_pow2(La + Lb − 1)`, one - forward transform each, pointwise product, one inverse. So the *default* - unequal path is a single transform — **no chunking**. - - Dispatch keys on the smaller operand `b.len()`, so when the NTT path is - entered `b` is already huge. Chunking `a` into `b.len()`-sized pieces - via the stock helper would run `⌈La/Lb⌉` separate NTTs **and - re-transform `b` on every chunk**, roughly doubling work for lopsided - large×large products — it throws away NTT's single-big-transform win. - - `ntt::add_signed_mul` (unequal) and `ntt::add_signed_mul_same_len` - (equal) therefore share one core that takes `(La, Lb)` and transforms - over `N = next_pow2(La + Lb − 1)` directly. Honor the same contract as - the other algorithms: `c.len() == La + Lb`, accumulate `sign * a * b` - into `c`, return the `SignedWord` carry. - - **Only** fall back to chunking for *extreme* imbalance (`La ≫ Lb`, e.g. - `La > c · Lb` for some tuned `c`), where many balanced `~2·Lb` NTTs beat - one padded `~La` transform. If/when we do, forward-transform `b` **once** - and reuse the cached `b_hat` across chunks — i.e. a purpose-built loop, - not the stock `add_signed_mul_split_into_chunks` (which re-transforms - `b`). Treat this as a Phase-7 tuning option, not the initial wiring. -- [ ] `ntt/mod.rs::memory_requirement_up_to(n)`: deterministic upper bound on - scratch, mirroring the style of `toom_3::memory_requirement_up_to` - (returns a `Layout`). It **must** be an exact upper bound — `Memory` - `expect`s on underflow. - - **Draft closed-form bound (words):** - `2·N` (one `a`-lane + one `b`-lane buffer, processed one prime at a - time so they are reused across the `K_eff` primes — not `K·N`) - `+ N/2` (twiddle table for the current prime) - `+ K` (per-coefficient CRT temp) - `+ (La + Lb)` (product accumulation buffer) - `≈ 2.5·N + La + Lb + K`. - If lanes for all primes are kept live simultaneously (simpler, no - re-pack per prime) the lane term becomes `2·K·N`; decide which during - implementation and bound accordingly. - - **Worst-case over `b_pack`.** `N = next_pow2(ceil((La+Lb)·WORD_BITS / - b_pack) + 1)`. `memory_requirement_up_to(n)` is called before - `select_params` runs, so it must bound `N` over **every** `b_pack` the - selector may choose for inputs up to `n` words — i.e. use the - *smallest* admissible `b_pack` (largest `N`). Pin a `B_PACK_MIN` - constant (= 16) and compute the bound from it; `select_params` may then - only ever pick `b_pack ≥ B_PACK_MIN`. -- [ ] **Carving scratch from the linear `Memory` arena.** `Memory` - (`integer/src/memory.rs`) is a *linear bump allocator*, not a pool of - independent `Buffer`s: each `allocate_slice`/`allocate_slice_fill` hands - out the next region and returns the remainder. Order matters. Plan the - carve explicitly: allocate longer-lived regions first (twiddle table, - product buffer) then the per-prime lane buffers from the remaining region - inside the prime loop (so they are reused each iteration). `Buffer` is - only for *owned* growable word arrays (e.g. a returned product); transform - scratch lives in the `Memory` arena. Document the carve order next to - `memory_requirement_up_to` so the two stay in sync. - -### Phase 5 — Dispatch + thresholds -- [ ] In `integer/src/mul/mod.rs`: add `THRESHOLD_NTT`, declare `mod ntt;`. -- [ ] Extend `add_signed_mul`, `add_signed_mul_same_len`, and both - `memory_requirement_*` to route `len > THRESHOLD_NTT` to `ntt`. -- [ ] Keep `toom_3` as the fallback if NTT parameter selection fails any - precondition (defensive; should not happen below `2^32` coefficients). - -### Phase 6 — Correctness validation -- [ ] Extend `integer/tests/mul.rs` with cases straddling `THRESHOLD_NTT` - (lengths `T−1`, `T`, `T+1`, `2T`, asymmetric `a`/`b` lengths, operands - with high/low zero limbs, near power-of-two `N`). -- [ ] Differential test: random `UBig`s of increasing size, assert - `ntt_product == reference_product` where reference is the existing - `multiply` forced through Toom-3 (or compute via a smaller-threshold - build). Cover the coefficient-overflow boundary explicitly (all-ones - operands at the max supported `N`). -- [ ] Run on a 32-bit lane build too: `cargo test -p dashu-int` with - `RUSTFLAGS="--cfg force_bits=\"32\""` (and `16`) to confirm packing is - `Word`-width agnostic. - -### Phase 7 — Tuning + optimisation (after correctness is green) -- [ ] Benchmark `THRESHOLD_NTT` crossover against Toom-3 using - `integer/benches/primitive.rs` (extend the mul benchmark to larger sizes); - pick the value where NTT wins. -- [ ] Optimisations to layer in, measuring each: - - [ ] Harvey lazy-reduction butterflies (defer mod in inner loops). - - [ ] Specialise the `b = 32` (Goldilocks) lane's reduction to 32-bit-limb - form (`φ²=φ−1`), since it avoids the extra cross-boundary fold that - `b = 34, 40` need. - - [ ] Use the shift-expressible roots of unity where applicable (powers of two - are roots in this family) to replace some twiddle multiplies. - - [ ] Radix-4 / split-radix transform. - - [ ] Drop to `K_eff = 2` primes automatically when headroom allows (smaller - inputs) to halve the transform work. - - [ ] **Squaring specialization.** `UBig::square()` / the `a == b` case needs - only one forward transform per prime (not two), then a pointwise - *square* and one inverse — roughly 2/3 the transform cost. Wire an - `ntt::square` path (and route `UBig::square` to it above - `THRESHOLD_NTT`) once the general multiply is correct. - - [ ] Reuse/cancel allocations; ensure scratch stays within the `Memory` - arena. +## Remaining optimisation opportunities ---- +### 1. Increase `b_pack` from 16 → 32 (~2× speedup) -## 4. Constraints & pitfalls (project-specific) - -- **`no_std`**: only `core` + `alloc`. `u128` arithmetic is fine. The lane - *reduction* is reused from `num_modular::FixedTrinomialSolinas64<64, b, 1>` - (`reduce_single` / `reduce_double`; already a dependency of `dashu-int`, and - `no_std`). The lazy `add`/`sub`/`mul` wrapper around it is ours, in-tree, so - no Montgomery and no extra dependency. Requires the `num-modular` version that - (a) supports `P1 = 64` (the `checked_shl` fix + `S64_4` Goldilocks tests) and - (b) exposes `reduce_single` / `reduce_double` as `pub` (it does as of that - fix). Pin this version in `integer/Cargo.toml`; if the `pub` assumption ever - regresses, fall back to the in-tree reduction copy (see Phase 1). -- **MSRV**: keep within the README MSRV (do not bump). Avoid `const` features - newer than MSRV; plain `const` tables are fine. -- **Scratch**: use `Buffer` / `MemoryAllocation` / the threaded `Memory` arena, - never `Vec` (per `AGENTS.md`). `memory_requirement_*` must be an exact - upper bound or the arena `expect` will panic. -- **Sign / accumulate contract**: the entry points are `c += sign * a * b` - returning a `SignedWord` carry — match `toom_3`/`karatsuba` exactly so the - recursive callers and `multiply()` keep working. -- **Changelog**: add an `### Add` entry under `## Unreleased` in - `integer/CHANGELOG.md` ("NTT-based multiplication for very large integers") - as part of the same commit (per `AGENTS.md`). -- **CI parity**: before declaring done, run - - `cargo test --workspace --exclude dashu-python` - - `cargo clippy --all-features --all-targets --workspace --exclude dashu-python -- -D warnings` - - `cargo fmt --all -- --check` +Currently every coefficient is 16 bits (4 coeffs per 64-bit word). With +K_eff = 2 primes providing ~2^128 of headroom, 32-bit coefficients are safe: ---- + max_coeff = N/2 · (2^32 − 1)^2 < 2^95 ≪ P0·P1 ≈ 2^128 + +Doubling `b_pack` halves the coefficient count, halves N, and roughly halves +the total transform work. Additionally, 32-bit packing is simpler: exactly +2 coefficients per word, no straddling of word boundaries. + +**Work items:** +- Update `select_params` to choose `b_pack ∈ {16, 32}` based on headroom. +- Adjust `memory_requirement_up_to` — the current worst-case bound uses + `B_PACK_MIN = 16`; need to handle the tighter N for `b_pack = 32`. +- Benchmark the new packing path. + +### 2. Radix-4 or split-radix NTT (~25–33% fewer twiddle multiplies) + +Radix-4 processes 4 elements per butterfly with 3 twiddle multiplies and +`log₄(N)` stages (half as many passes through memory). Split-radix pushes +the savings closer to 33%. + +**Work items:** +- Rewrite `ntt_core` in `transform.rs` with a radix-4 butterfly. +- Handle N that is a power of 2 but not a power of 4: do one radix-2 stage + followed by radix-4 stages. +- Update twiddle indexing; the twiddle table layout changes. +- A primitive 4-th root `j = ω_N^{N/4}` is needed for the butterfly core; + derive it from the existing `ω_2_32` root. + +### 3. Harvey lazy-reduction butterflies (~10–15%) + +Currently every `add_mod` / `sub_mod` fully normalizes to `[0, p)`. Harvey's +approach keeps values in `[0, 2p)` across multiple butterfly stages, deferring +the conditional subtract to the end (or to the next `mul_mod`). This replaces +a branch + subtract with a no-op in the inner loop. + +**Work items:** +- Change `add_mod` / `sub_mod` to allow `[0, 2p)` outputs. +- Add a normalization pass at the end of `pointwise_mul` and `inverse`. +- Verify no overflow in the radix-2 structure (each stage at most doubles the + dynamic range, so worst-case after log₂(N) stages is `[0, N·p)` — we need a + cleanup before it overflows `u64`). + +### 4. Merge `bit_reverse` with `pack` (~5–10%) + +Currently `pack` writes coefficients in natural order, then `bit_reverse` +permutes them in a second pass. Write packed coefficients directly to their +bit-reversed positions, saving one full array pass. + +### 5. Shift-expressible twiddle factors (stage-dependent) + +In Goldilocks primes, `2^k mod p = 2^k` when `2^k < p`. The first few NTT +stages have twiddle factors that are pure powers of 2, so `mul_mod(t, 2^k)` +reduces to a shift + conditional subtract — no `u128` multiply needed. + +### 6. Specialize `b = 32` lane (~5%) + +The `b = 32` prime (`0xFFFFFFFF00000001`) has the cleanest reduction identity +(splits a `u128` product exactly into 32-bit limbs). A dedicated code path +for this prime alone could squeeze out a few more cycles vs. the generic +`match B` dispatch in `mul_mod`. -## 5. Definition of done +### 7. Asymmetric operand chunking (conditional) -- [ ] All phases checked off; NTT path active above `THRESHOLD_NTT`. -- [ ] Differential tests pass on 64-, 32-, and 16-bit lane builds. -- [ ] Benchmarks show NTT beats Toom-3 at and above the chosen threshold. -- [ ] Clippy clean, fmt clean, changelog updated. -- [ ] No `std` usage, no `Vec` scratch, MSRV preserved. +When `a ≫ b`, chunk the long operand, forward-transform the short operand +once, and reuse `b̂` (the transformed short operand) across all chunks. Only +matters for extremely lopsided inputs. diff --git a/integer/CHANGELOG.md b/integer/CHANGELOG.md index f7cd5beb..e3999a8c 100644 --- a/integer/CHANGELOG.md +++ b/integer/CHANGELOG.md @@ -3,16 +3,16 @@ ## Unreleased ### Add -- NTT-based multiplication for very large integers (above 2048 words), using three Solinas primes of the form `2^64 − 2^b + 1` combined with the Chinese Remainder Theorem. +- NTT-based multiplication for very large integers (above 50000 words / ~3.2M bits), using two Solinas primes of the form `2^64 − 2^b + 1` combined with the Chinese Remainder Theorem. - `UBig::from_u64` and `IBig::from_i64`, const on 32-bit and 64-bit targets. ### Improve - Basecase (schoolbook) multiplication now uses an dword mult inner kernel (two multiplier words per sweep over the accumulator, mirroring GMP's `mpn_addmul_2` and `mpn_submul_2`), roughly halving accumulator memory traffic. - Addition and subtraction carry/borrow propagation now uses `Word` (u64/u32) instead of `bool` throughout the architecture-specific `add_with_carry` and `sub_with_borrow` functions, eliminating `bool`↔Word conversions in the inner loops. - -### Improve -- Logarithm for very large values uses power-sequence decomposition, replacing iterative single-step multiplication. -- Improve power-of-two base formatting ([#3](https://github.com/cmpute/dashu/pull/3)) +- Lowered the Karatsuba→Toom-3 multiplication threshold from 192 to 96 words, giving Toom-Cook-3 at ~6000 bits instead of ~12000 bits — closes the gap with malachite at ~10000-bit sizes. +- NTT coefficient width increased from 16 to 64 bits (K_eff=3 for 64-bit, K_eff=2 otherwise), roughly halving the transform length at each step. +- NTT multiplication auto-selects `K_eff = 2` primes when headroom allows, skipping the third prime. +- Multiplication thresholds can be overridden at runtime via `DASHU_THRESHOLD_SIMPLE`, `DASHU_THRESHOLD_KARATSUBA`, and `DASHU_THRESHOLD_NTT` environment variables (requires `tuning` feature). ## 0.4.2 diff --git a/integer/Cargo.toml b/integer/Cargo.toml index da816819..3f06cda3 100644 --- a/integer/Cargo.toml +++ b/integer/Cargo.toml @@ -19,6 +19,7 @@ all-features = true [features] default = ["std", "num-order"] std = ["dashu-base/std"] +tuning = ["std"] # unstable dependencies rand = ["rand_v08"] diff --git a/integer/src/mul/mod.rs b/integer/src/mul/mod.rs index 301273dc..32b056d5 100644 --- a/integer/src/mul/mod.rs +++ b/integer/src/mul/mod.rs @@ -14,23 +14,67 @@ use core::mem; use static_assertions::const_assert; /// If smaller operand length <= this, simple multiplication will be used. -const THRESHOLD_SIMPLE: usize = 24; -const_assert!(THRESHOLD_SIMPLE <= simple::MAX_SMALLER_LEN); -const_assert!(THRESHOLD_SIMPLE + 1 >= karatsuba::MIN_LEN); +const THRESHOLD_SIMPLE_DEFAULT: usize = 24; +const_assert!(THRESHOLD_SIMPLE_DEFAULT <= simple::MAX_SMALLER_LEN); +const_assert!(THRESHOLD_SIMPLE_DEFAULT + 1 >= karatsuba::MIN_LEN); /// If smaller operand length <= this, Karatsuba multiplication will be used. -const THRESHOLD_KARATSUBA: usize = 192; -const_assert!(THRESHOLD_KARATSUBA + 1 >= toom_3::MIN_LEN); +const THRESHOLD_KARATSUBA_DEFAULT: usize = 192; +const_assert!(THRESHOLD_KARATSUBA_DEFAULT + 1 >= toom_3::MIN_LEN); /// If smaller operand length > this, NTT multiplication will be used. -const THRESHOLD_NTT: usize = ntt::THRESHOLD_NTT; -const_assert!(THRESHOLD_NTT + 1 >= toom_3::MIN_LEN); +const THRESHOLD_NTT_DEFAULT: usize = ntt::THRESHOLD_NTT; +const_assert!(THRESHOLD_NTT_DEFAULT + 1 >= toom_3::MIN_LEN); + +/// Environment-variable overrides for multiplication thresholds. +/// +/// When the `tuning` feature is active the user may set `DASHU_THRESHOLD_SIMPLE`, +/// `DASHU_THRESHOLD_KARATSUBA` or `DASHU_THRESHOLD_NTT` to override the +/// compile-time defaults. +mod threshold { + #[inline] + pub fn simple() -> usize { + #[cfg(feature = "tuning")] + { + if let Ok(s) = std::env::var("DASHU_THRESHOLD_SIMPLE") { + if let Ok(v) = s.parse() { + return v; + } + } + } + super::THRESHOLD_SIMPLE_DEFAULT + } + #[inline] + pub fn karatsuba() -> usize { + #[cfg(feature = "tuning")] + { + if let Ok(s) = std::env::var("DASHU_THRESHOLD_KARATSUBA") { + if let Ok(v) = s.parse() { + return v; + } + } + } + super::THRESHOLD_KARATSUBA_DEFAULT + } + #[inline] + pub fn ntt() -> usize { + #[cfg(feature = "tuning")] + { + if let Ok(s) = std::env::var("DASHU_THRESHOLD_NTT") { + if let Ok(v) = s.parse() { + return v; + } + } + } + super::THRESHOLD_NTT_DEFAULT + } +} mod helpers; mod karatsuba; pub(crate) mod ntt; mod simple; -mod toom_3; +pub(crate) mod toom_3; /// Multiply a word sequence by a `Word` in place. /// @@ -162,11 +206,11 @@ pub fn sub_mul_word_same_len_in_place(words: &mut [Word], mult: Word, rhs: &[Wor /// Temporary scratch space required for multiplication. pub fn memory_requirement_up_to(total_len: usize, smaller_len: usize) -> Layout { - if smaller_len <= THRESHOLD_SIMPLE { + if smaller_len <= threshold::simple() { memory::zero_layout() - } else if smaller_len <= THRESHOLD_KARATSUBA { + } else if smaller_len <= threshold::karatsuba() { karatsuba::memory_requirement_up_to(smaller_len) - } else if smaller_len <= THRESHOLD_NTT { + } else if smaller_len <= threshold::ntt() { toom_3::memory_requirement_up_to(smaller_len) } else { ntt::memory_requirement_up_to(total_len, smaller_len) @@ -202,11 +246,11 @@ pub fn add_signed_mul<'a>( mem::swap(&mut a, &mut b); } - if b.len() <= THRESHOLD_SIMPLE { + if b.len() <= threshold::simple() { simple::add_signed_mul(c, sign, a, b, memory) - } else if b.len() <= THRESHOLD_KARATSUBA { + } else if b.len() <= threshold::karatsuba() { karatsuba::add_signed_mul(c, sign, a, b, memory) - } else if b.len() <= THRESHOLD_NTT { + } else if b.len() <= threshold::ntt() { toom_3::add_signed_mul(c, sign, a, b, memory) } else { ntt::add_signed_mul(c, sign, a, b, memory) @@ -227,11 +271,11 @@ pub fn add_signed_mul_same_len( let n = a.len(); debug_assert!(b.len() == n && c.len() == 2 * n); - if n <= THRESHOLD_SIMPLE { + if n <= threshold::simple() { simple::add_signed_mul_same_len(c, sign, a, b, memory) - } else if n <= THRESHOLD_KARATSUBA { + } else if n <= threshold::karatsuba() { karatsuba::add_signed_mul_same_len(c, sign, a, b, memory) - } else if n <= THRESHOLD_NTT { + } else if n <= threshold::ntt() { toom_3::add_signed_mul_same_len(c, sign, a, b, memory) } else { ntt::add_signed_mul_same_len(c, sign, a, b, memory) diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs index f5af326b..ce2e927d 100644 --- a/integer/src/mul/ntt/crt.rs +++ b/integer/src/mul/ntt/crt.rs @@ -147,8 +147,8 @@ mod tests { #[test] fn test_garner_with_ntt_primes() { - use num_modular::FixedTrinomialSolinas64; use super::super::primes::PRIMES; + use num_modular::FixedTrinomialSolinas64; let p0 = PRIMES[0].p; let p1 = PRIMES[1].p; diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index cc34e3f4..448a9fd1 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -22,11 +22,21 @@ use crate::mul::ntt::crt::U192; pub use primes::{K, PRIMES}; /// Minimum smaller-operand length (in words) for the NTT path. -pub const THRESHOLD_NTT: usize = 2048; +/// +/// With `b_pack = 32` the crossover is at ~30 000 words (~2 M bits) on +/// Apple M4 Pro. N double at 32 769 / 65 537 / 131 073 words causes +/// small regression windows; radix-4 will shrink the step size. +/// Chosen conservatively at 50 000 words where NTT is ≥20% faster. +pub const THRESHOLD_NTT: usize = 50_000; /// Smallest admissible coefficient bit width (used for worst-case memory bound). const B_PACK_MIN: u32 = 16; +/// Preferred coefficient bit width. 32 bits gives 2 coeffs/word and halves +/// the transform length vs. 16 bits, while staying comfortably within the +/// ~2^128 headroom of the two smallest primes. +const B_PACK_PREFERRED: u32 = 32; + /// Maximum `log2(transform length)`, set by `min(v2) = 32` across all primes. const MAX_LOG_N: u32 = 32; @@ -34,34 +44,31 @@ const MAX_LOG_N: u32 = 32; /// /// Returns `(b_pack, N, K_eff)`. pub fn select_params(la_words: usize, lb_words: usize) -> (u32, usize, usize) { - let b_pack = B_PACK_MIN; let word_bits = Word::BITS; - let la_bits = la_words as u64 * word_bits as u64; let lb_bits = lb_words as u64 * word_bits as u64; + let prod_2 = (PRIMES[0].p as u128) * (PRIMES[1].p as u128); - let coeffs_a = (la_bits + b_pack as u64 - 1) / b_pack as u64; - let coeffs_b = (lb_bits + b_pack as u64 - 1) / b_pack as u64; - let total_coeffs = (coeffs_a + coeffs_b - 1) as usize; - let n = total_coeffs.next_power_of_two().max(2); + // Try b_pack = 32 first (fewer coefficients → smaller N → faster transform). + for &b_pack in &[B_PACK_PREFERRED, B_PACK_MIN] { + let coeffs_a = (la_bits + b_pack as u64 - 1) / b_pack as u64; + let coeffs_b = (lb_bits + b_pack as u64 - 1) / b_pack as u64; + let total_coeffs = (coeffs_a + coeffs_b - 1) as usize; + let n = total_coeffs.next_power_of_two().max(2); - assert!( - (n.trailing_zeros()) <= MAX_LOG_N, - "N = {n} too large for prime set (max log2 = {MAX_LOG_N})" - ); + if (n.trailing_zeros()) > MAX_LOG_N { + continue; + } - let k_eff = K; + let max_coeff = (n as u128 / 2) * ((1u128 << b_pack) - 1) * ((1u128 << b_pack) - 1); - // Headroom check: max convolution coefficient < product of K_eff primes. - // max_coeff fits in u128; compare against smallest prime p0. - let max_coeff = (n as u128 / 2) * ((1u128 << b_pack) - 1) * ((1u128 << b_pack) - 1); - let p0 = PRIMES[0].p as u128; - assert!( - max_coeff < p0, - "headroom check failed: max coeff {max_coeff} >= smallest prime {p0}" - ); + // K_eff = 2 suffices when max_coeff < P0·P1 ≈ 2^128. + // K_eff = 3 covers the rest (product ≈ 2^192 ≫ max_coeff < 2^128). + let k_eff = if max_coeff < prod_2 { 2 } else { K }; + return (b_pack, n, k_eff); + } - (b_pack, n, k_eff) + unreachable!("b_pack = 16 always passes the headroom check") } /// Estimate bit length from a word slice (excludes leading zeros). @@ -324,18 +331,18 @@ mod tests { #[test] fn test_select_params_small() { let (b_pack, n, k_eff) = select_params(10, 10); - assert_eq!(b_pack, 16); + assert_eq!(b_pack, 32); assert!(n >= 2 && n.is_power_of_two()); - assert_eq!(k_eff, K); + assert_eq!(k_eff, 2); } #[test] fn test_select_params_large() { let (b_pack, n, k_eff) = select_params(THRESHOLD_NTT, THRESHOLD_NTT); - assert_eq!(b_pack, 16); + assert_eq!(b_pack, 32); assert!(n.is_power_of_two()); - assert_eq!(k_eff, K); - let coeffs_a = (THRESHOLD_NTT * Word::BITS as usize + 15) / 16; + assert_eq!(k_eff, 2); + let coeffs_a = (THRESHOLD_NTT * Word::BITS as usize + 31) / 32; let coeffs_b = coeffs_a; let min_n = (coeffs_a + coeffs_b).next_power_of_two().max(2); assert!(n >= min_n, "n={n} < min_n={min_n}"); @@ -347,8 +354,11 @@ mod tests { let lb = THRESHOLD_NTT; let (b_pack, n, _k_eff) = select_params(la, lb); let max_coeff = (n as u128 / 2) * ((1u128 << b_pack) - 1) * ((1u128 << b_pack) - 1); - let p0 = PRIMES[0].p as u128; - assert!(max_coeff < p0, "headroom violation: max_coeff={max_coeff} >= p0={p0}"); + let prod_2 = (PRIMES[0].p as u128) * (PRIMES[1].p as u128); + assert!( + max_coeff < prod_2, + "headroom violation: max_coeff={max_coeff} >= prod_2={prod_2}" + ); } #[test] @@ -389,11 +399,14 @@ mod tests { assert_eq!(&c[..], &expected[..], "two-word mismatch"); } + const NTT_TEST_LEN: usize = 1024; + #[test] fn test_ntt_multiply_small() { - // Test NTT multiply with small operands that exceed THRESHOLD_NTT. - let a: Vec = vec![0xDEADBEEFu64; THRESHOLD_NTT]; - let b: Vec = vec![0xCAFEBABEu64; THRESHOLD_NTT]; + // Smoke test: NTT multiply with operands large enough to exercise + // the full pipeline (pack, forward, pointwise, inverse, CRT, accumulate). + let a: Vec = vec![0xDEADBEEFu64; NTT_TEST_LEN]; + let b: Vec = vec![0xCAFEBABEu64; NTT_TEST_LEN]; let mut c = vec![0u64; a.len() + b.len()]; let layout = memory_requirement_up_to(c.len(), b.len()); @@ -542,4 +555,110 @@ mod tests { add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); assert_eq!(&c[..], &expected[..], "sparse operand mismatch"); } + + /// Compare NTT against toom-3 at various sizes to find the crossover. + /// + /// Run with: + /// ```sh + /// # Pure toom-3 (prevent internal NTT recursion): + /// DASHU_THRESHOLD_NTT=99999999 cargo test -p dashu-int --release \ + /// -- ntt::tests::crossover --ignored --nocapture + /// ``` + /// + /// The output is a table showing word count, transform length N, + /// toom-3 time, NTT time, and speedup ratio. Use it to recalibrate + /// [`THRESHOLD_NTT`] after algorithmic changes. + #[test] + #[ignore] + #[allow(clippy::let_underscore_must_use)] + fn crossover() { + use std::time::Instant; + + let sizes: &[usize] = &[ + 5_000, 10_000, 20_000, 30_000, 40_000, 50_000, 60_000, 80_000, 100_000, 120_000, + 131_000, + ]; + + println!( + "{:>10} {:>4} {:>8} {:>12} {:>12} {:>10}", + "words", "bp", "N", "toom-3(ms)", "ntt(ms)", "ratio" + ); + println!("{}", "-".repeat(68)); + + for &n in sizes { + let a: Vec = (0..n) + .map(|i| (i as u64 + 1).wrapping_mul(0x9E3779B97F4A7C15)) + .collect(); + let b: Vec = (0..n) + .map(|i| (i as u64 + 1).wrapping_mul(0xC6A4A7935BD1E995)) + .collect(); + let mut c_toom = vec![0u64; 2 * n]; + let mut c_ntt = vec![0u64; 2 * n]; + + let layout = memory_requirement_up_to(2 * n, n); + let warmup = 2; + let iters = 5; + + // toom-3 (may use NTT internally depending on DASHU_THRESHOLD_NTT) + let t_toom = { + let mut best = f64::MAX; + for _ in 0..warmup { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_toom.fill(0); + let _ = + crate::mul::toom_3::add_signed_mul(&mut c_toom, Positive, &a, &b, &mut mem); + } + for _ in 0..iters { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_toom.fill(0); + let start = Instant::now(); + let _ = + crate::mul::toom_3::add_signed_mul(&mut c_toom, Positive, &a, &b, &mut mem); + let elapsed = start.elapsed().as_secs_f64() * 1000.0; + if elapsed < best { + best = elapsed; + } + } + best + }; + + // NTT (direct) + let t_ntt = { + let mut best = f64::MAX; + for _ in 0..warmup { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_ntt.fill(0); + let _carry = add_signed_mul_impl(&mut c_ntt, Positive, &a, &b, &mut mem); + } + for _ in 0..iters { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_ntt.fill(0); + let start = Instant::now(); + let _carry = add_signed_mul_impl(&mut c_ntt, Positive, &a, &b, &mut mem); + let elapsed = start.elapsed().as_secs_f64() * 1000.0; + if elapsed < best { + best = elapsed; + } + } + best + }; + + assert_eq!(&c_ntt[..], &c_toom[..], "mismatch at n={n}"); + + let (_b_pack, nn, _k_eff) = select_params(n, n); + println!( + "{:>10} {:>4} {:>8} {:>12.3} {:>12.3} {:>9.2}x", + n, + _b_pack, + nn, + t_toom, + t_ntt, + t_ntt / t_toom + ); + } + } } From 56a1a646a52547ff3663fe9dd449c54f0f2ba1d5 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sat, 13 Jun 2026 11:04:42 +0800 Subject: [PATCH 05/19] WIP: further tune b_pack --- integer/CHANGELOG.md | 2 +- integer/src/mul/mod.rs | 4 +- integer/src/mul/ntt/mod.rs | 78 +++++++++++++++++++++++++------------ integer/src/mul/ntt/pack.rs | 6 ++- 4 files changed, 63 insertions(+), 27 deletions(-) diff --git a/integer/CHANGELOG.md b/integer/CHANGELOG.md index e3999a8c..7e13059f 100644 --- a/integer/CHANGELOG.md +++ b/integer/CHANGELOG.md @@ -3,7 +3,7 @@ ## Unreleased ### Add -- NTT-based multiplication for very large integers (above 50000 words / ~3.2M bits), using two Solinas primes of the form `2^64 − 2^b + 1` combined with the Chinese Remainder Theorem. +- NTT-based multiplication for very large integers (above 40000 words / ~2.6M bits), using two or three Solinas primes of the form `2^64 − 2^b + 1` combined with the Chinese Remainder Theorem. - `UBig::from_u64` and `IBig::from_i64`, const on 32-bit and 64-bit targets. ### Improve diff --git a/integer/src/mul/mod.rs b/integer/src/mul/mod.rs index 32b056d5..597f88aa 100644 --- a/integer/src/mul/mod.rs +++ b/integer/src/mul/mod.rs @@ -19,7 +19,9 @@ const_assert!(THRESHOLD_SIMPLE_DEFAULT <= simple::MAX_SMALLER_LEN); const_assert!(THRESHOLD_SIMPLE_DEFAULT + 1 >= karatsuba::MIN_LEN); /// If smaller operand length <= this, Karatsuba multiplication will be used. -const THRESHOLD_KARATSUBA_DEFAULT: usize = 192; +/// Tuned so that Toom-3 kicks in earlier (~96 words vs the old 192), +/// closing the gap with malachite/rug at ~10000-bit sizes. +const THRESHOLD_KARATSUBA_DEFAULT: usize = 96; const_assert!(THRESHOLD_KARATSUBA_DEFAULT + 1 >= toom_3::MIN_LEN); /// If smaller operand length > this, NTT multiplication will be used. diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index 448a9fd1..d47fbabd 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -23,11 +23,11 @@ pub use primes::{K, PRIMES}; /// Minimum smaller-operand length (in words) for the NTT path. /// -/// With `b_pack = 32` the crossover is at ~30 000 words (~2 M bits) on -/// Apple M4 Pro. N double at 32 769 / 65 537 / 131 073 words causes -/// small regression windows; radix-4 will shrink the step size. -/// Chosen conservatively at 50 000 words where NTT is ≥20% faster. -pub const THRESHOLD_NTT: usize = 50_000; +/// With `b_pack = 64` the crossover is at ~25 000 words (~1.6 M bits) on +/// Apple M4 Pro. N-doubling at 32 769 / 65 537 words creates narrow +/// regression windows; radix-4 will shrink the step size further. +/// Chosen at 40 000 words where NTT is ≥18% faster. +pub const THRESHOLD_NTT: usize = 40_000; /// Smallest admissible coefficient bit width (used for worst-case memory bound). const B_PACK_MIN: u32 = 16; @@ -35,7 +35,9 @@ const B_PACK_MIN: u32 = 16; /// Preferred coefficient bit width. 32 bits gives 2 coeffs/word and halves /// the transform length vs. 16 bits, while staying comfortably within the /// ~2^128 headroom of the two smallest primes. -const B_PACK_PREFERRED: u32 = 32; +/// Coefficient bit widths to try, in descending preference. +/// 64 uses K_eff = 3 primes; 32 and 16 use K_eff = 2. +const B_PACK_CANDIDATES: &[u32] = &[64, 32, 16]; /// Maximum `log2(transform length)`, set by `min(v2) = 32` across all primes. const MAX_LOG_N: u32 = 32; @@ -49,8 +51,7 @@ pub fn select_params(la_words: usize, lb_words: usize) -> (u32, usize, usize) { let lb_bits = lb_words as u64 * word_bits as u64; let prod_2 = (PRIMES[0].p as u128) * (PRIMES[1].p as u128); - // Try b_pack = 32 first (fewer coefficients → smaller N → faster transform). - for &b_pack in &[B_PACK_PREFERRED, B_PACK_MIN] { + for &b_pack in B_PACK_CANDIDATES { let coeffs_a = (la_bits + b_pack as u64 - 1) / b_pack as u64; let coeffs_b = (lb_bits + b_pack as u64 - 1) / b_pack as u64; let total_coeffs = (coeffs_a + coeffs_b - 1) as usize; @@ -60,11 +61,19 @@ pub fn select_params(la_words: usize, lb_words: usize) -> (u32, usize, usize) { continue; } - let max_coeff = (n as u128 / 2) * ((1u128 << b_pack) - 1) * ((1u128 << b_pack) - 1); - - // K_eff = 2 suffices when max_coeff < P0·P1 ≈ 2^128. - // K_eff = 3 covers the rest (product ≈ 2^192 ≫ max_coeff < 2^128). - let k_eff = if max_coeff < prod_2 { 2 } else { K }; + // Compute max coefficient value, guarding against u128 overflow for + // b_pack = 64 where (2^64−1)^2 ≈ 2^128 and n/2 can push it past 2^128. + let coeff_max = (1u128 << b_pack) - 1; + let max_coeff = coeff_max + .checked_mul(coeff_max) + .and_then(|sq| (n as u128 / 2).checked_mul(sq)); + + let k_eff = match max_coeff { + Some(mc) if mc < prod_2 => 2, + // Overflow or exceeds two-prime product → need all three primes. + // Three-prime product ≈ 2^192 ≫ any max_coeff we can encounter. + _ => K, + }; return (b_pack, n, k_eff); } @@ -261,6 +270,22 @@ fn process_prime(a: &[Word], b: &[Word], ctx: &mut TransformCtx<'_ pack::pack(ctx.a_lane, a, ctx.b_pack, ctx.nn); pack::pack(ctx.b_lane, b, ctx.b_pack, ctx.nn); + // For b_pack = 64 coefficients may reach 2^64−1, which can exceed p. + // Reduce each coefficient mod p (one conditional subtract suffices: + // c < 2^64 < 2p for all three primes). + if ctx.b_pack >= 64 { + for c in ctx.a_lane[..ctx.nn].iter_mut() { + if *c >= ctx.p { + *c -= ctx.p; + } + } + for c in ctx.b_lane[..ctx.nn].iter_mut() { + if *c >= ctx.p { + *c -= ctx.p; + } + } + } + transform::precompute_twiddles::(ctx.fwd_twiddles, ctx.nn, ctx.p, ctx.omega_2_32, false); transform::precompute_twiddles::(ctx.inv_twiddles, ctx.nn, ctx.p, ctx.omega_2_32, true); @@ -331,18 +356,19 @@ mod tests { #[test] fn test_select_params_small() { let (b_pack, n, k_eff) = select_params(10, 10); - assert_eq!(b_pack, 32); + assert_eq!(b_pack, 64); assert!(n >= 2 && n.is_power_of_two()); - assert_eq!(k_eff, 2); + // b_pack = 64 needs K_eff = 3 primes. + assert_eq!(k_eff, K); } #[test] fn test_select_params_large() { let (b_pack, n, k_eff) = select_params(THRESHOLD_NTT, THRESHOLD_NTT); - assert_eq!(b_pack, 32); + assert_eq!(b_pack, 64); assert!(n.is_power_of_two()); - assert_eq!(k_eff, 2); - let coeffs_a = (THRESHOLD_NTT * Word::BITS as usize + 31) / 32; + assert_eq!(k_eff, K); + let coeffs_a = (THRESHOLD_NTT * Word::BITS as usize + 63) / 64; let coeffs_b = coeffs_a; let min_n = (coeffs_a + coeffs_b).next_power_of_two().max(2); assert!(n >= min_n, "n={n} < min_n={min_n}"); @@ -353,12 +379,16 @@ mod tests { let la = THRESHOLD_NTT; let lb = THRESHOLD_NTT; let (b_pack, n, _k_eff) = select_params(la, lb); - let max_coeff = (n as u128 / 2) * ((1u128 << b_pack) - 1) * ((1u128 << b_pack) - 1); - let prod_2 = (PRIMES[0].p as u128) * (PRIMES[1].p as u128); - assert!( - max_coeff < prod_2, - "headroom violation: max_coeff={max_coeff} >= prod_2={prod_2}" - ); + // For b_pack = 64 the product overflows u128; checked_mul in + // select_params handles this and falls back to K_eff = 3. + // Three-prime product ≈ 2^192 ≫ max_coeff for n ≤ 2^32. + let coeff_max = (1u128 << b_pack) - 1; + let overflow = coeff_max + .checked_mul(coeff_max) + .and_then(|sq| (n as u128 / 2).checked_mul(sq)) + .is_none(); + assert!(overflow || _k_eff == 2, "K_eff=2 only when max_coeff fits in u128"); + assert_eq!(b_pack, 64); } #[test] diff --git a/integer/src/mul/ntt/pack.rs b/integer/src/mul/ntt/pack.rs index 25baad94..8edbbcb8 100644 --- a/integer/src/mul/ntt/pack.rs +++ b/integer/src/mul/ntt/pack.rs @@ -16,7 +16,11 @@ use crate::arch::word::Word; /// Panics if `out.len() < n`. pub fn pack(out: &mut [u64], words: &[Word], b_pack: u32, n: usize) { assert!(out.len() >= n); - let mask = (1u64 << b_pack) - 1; + let mask = if b_pack < Word::BITS { + (1u64 << b_pack) - 1 + } else { + u64::MAX + }; let word_bits = Word::BITS; let mut word_idx = 0usize; let mut bit_offset = 0u32; From d0def12ecce2b88c60a168acc4e648eb574635fe Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sat, 13 Jun 2026 11:04:42 +0800 Subject: [PATCH 06/19] WIP: minor test improvements Co-Authored-By: Claude Opus 4.7 --- integer/src/mul/mod.rs | 202 +++++++++++++++++++++++++++++++++++++ integer/src/mul/ntt/mod.rs | 106 ------------------- 2 files changed, 202 insertions(+), 106 deletions(-) diff --git a/integer/src/mul/mod.rs b/integer/src/mul/mod.rs index 597f88aa..5157a99e 100644 --- a/integer/src/mul/mod.rs +++ b/integer/src/mul/mod.rs @@ -283,3 +283,205 @@ pub fn add_signed_mul_same_len( ntt::add_signed_mul_same_len(c, sign, a, b, memory) } } + +#[cfg(test)] +mod threshold_tests { + use super::*; + use crate::arch::word::Word; + use crate::Sign::Positive; + + /// Compare karatsuba vs toom-3 at various word counts to find [`THRESHOLD_KARATSUBA`]. + /// Run with: + /// cargo test -p dashu-int --release -- mul::threshold_tests::crossover_karatsuba --nocapture --ignored + #[test] + #[ignore] + fn crossover_karatsuba() { + use std::time::Instant; + + let sizes: &[usize] = &[80, 100, 120, 140, 160, 180, 200, 240, 280, 320, 360, 400]; + + println!("{:>8} {:>14} {:>14} {:>10}", "words", "karatsuba(µs)", "toom-3(µs)", "ratio"); + println!("{}", "-".repeat(50)); + + for &n in sizes { + let a: Vec = (0..n) + .map(|i| (i as u64 + 1).wrapping_mul(0x9E3779B97F4A7C15)) + .collect(); + let b: Vec = (0..n) + .map(|i| (i as u64 + 1).wrapping_mul(0xC6A4A7935BD1E995)) + .collect(); + let mut c_kara = vec![0u64; 2 * n]; + let mut c_toom = vec![0u64; 2 * n]; + let layout_kara = karatsuba::memory_requirement_up_to(n); + let layout_toom = toom_3::memory_requirement_up_to(n); + // Use the larger layout so both algorithms get enough memory. + let layout = if layout_kara.size() > layout_toom.size() { + layout_kara + } else { + layout_toom + }; + let warmup = 5; + let iters = 20; + + // Time karatsuba + let t_kara = { + let mut best = f64::MAX; + for _ in 0..warmup { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_kara.fill(0); + let _c = + karatsuba::add_signed_mul_same_len(&mut c_kara, Positive, &a, &b, &mut mem); + } + for _ in 0..iters { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_kara.fill(0); + let start = Instant::now(); + let _c = + karatsuba::add_signed_mul_same_len(&mut c_kara, Positive, &a, &b, &mut mem); + let elapsed = start.elapsed().as_secs_f64() * 1_000_000.0; + if elapsed < best { + best = elapsed; + } + } + best + }; + + // Time toom-3 + let t_toom = { + let mut best = f64::MAX; + for _ in 0..warmup { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_toom.fill(0); + let _c = + toom_3::add_signed_mul_same_len(&mut c_toom, Positive, &a, &b, &mut mem); + } + for _ in 0..iters { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_toom.fill(0); + let start = Instant::now(); + let _c = + toom_3::add_signed_mul_same_len(&mut c_toom, Positive, &a, &b, &mut mem); + let elapsed = start.elapsed().as_secs_f64() * 1_000_000.0; + if elapsed < best { + best = elapsed; + } + } + best + }; + + assert_eq!(&c_kara[..], &c_toom[..], "mismatch at n={n}"); + println!("{:>8} {:>14.1} {:>14.1} {:>9.2}x", n, t_kara, t_toom, t_toom / t_kara); + } + } + + /// Compare NTT against toom-3 at various word counts to find [`THRESHOLD_NTT`]. + /// + /// Run with (set a huge NTT threshold to keep toom-3 pure): + /// ```sh + /// DASHU_THRESHOLD_NTT=99999999 cargo test -p dashu-int --features tuning --release \ + /// -- mul::threshold_tests::crossover_ntt --ignored --nocapture + /// ``` + /// + /// The output is a table: words, b_pack, N, toom-3 time, NTT time, ratio. + #[test] + #[ignore] + #[allow(clippy::let_underscore_must_use)] + fn crossover_ntt() { + use std::time::Instant; + + let sizes: &[usize] = &[ + 5_000, 10_000, 20_000, 30_000, 40_000, 50_000, 60_000, 80_000, 100_000, 120_000, + 131_000, + ]; + + println!( + "{:>10} {:>4} {:>8} {:>12} {:>12} {:>10}", + "words", "bp", "N", "toom-3(ms)", "ntt(ms)", "ratio" + ); + println!("{}", "-".repeat(68)); + + for &n in sizes { + let a: Vec = (0..n) + .map(|i| (i as u64 + 1).wrapping_mul(0x9E3779B97F4A7C15)) + .collect(); + let b: Vec = (0..n) + .map(|i| (i as u64 + 1).wrapping_mul(0xC6A4A7935BD1E995)) + .collect(); + let mut c_toom = vec![0u64; 2 * n]; + let mut c_ntt = vec![0u64; 2 * n]; + + let layout_ntt = super::ntt::memory_requirement_up_to(2 * n, n); + let layout_toom = super::toom_3::memory_requirement_up_to(n); + let layout = if layout_ntt.size() > layout_toom.size() { + layout_ntt + } else { + layout_toom + }; + let warmup = 2; + let iters = 5; + + // toom-3 (may use NTT internally depending on DASHU_THRESHOLD_NTT) + let t_toom = { + let mut best = f64::MAX; + for _ in 0..warmup { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_toom.fill(0); + let _ = super::toom_3::add_signed_mul(&mut c_toom, Positive, &a, &b, &mut mem); + } + for _ in 0..iters { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_toom.fill(0); + let start = Instant::now(); + let _ = super::toom_3::add_signed_mul(&mut c_toom, Positive, &a, &b, &mut mem); + let elapsed = start.elapsed().as_secs_f64() * 1000.0; + if elapsed < best { + best = elapsed; + } + } + best + }; + + // NTT (via public entry, bypasses dispatch) + let t_ntt = { + let mut best = f64::MAX; + for _ in 0..warmup { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_ntt.fill(0); + let _ = super::ntt::add_signed_mul(&mut c_ntt, Positive, &a, &b, &mut mem); + } + for _ in 0..iters { + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut mem = alloc.memory(); + c_ntt.fill(0); + let start = Instant::now(); + let _ = super::ntt::add_signed_mul(&mut c_ntt, Positive, &a, &b, &mut mem); + let elapsed = start.elapsed().as_secs_f64() * 1000.0; + if elapsed < best { + best = elapsed; + } + } + best + }; + + assert_eq!(&c_ntt[..], &c_toom[..], "mismatch at n={n}"); + + let (b_pack, nn, _k_eff) = super::ntt::select_params(n, n); + println!( + "{:>10} {:>4} {:>8} {:>12.3} {:>12.3} {:>9.2}x", + n, + b_pack, + nn, + t_toom, + t_ntt, + t_ntt / t_toom + ); + } + } +} diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index d47fbabd..932f0e78 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -585,110 +585,4 @@ mod tests { add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); assert_eq!(&c[..], &expected[..], "sparse operand mismatch"); } - - /// Compare NTT against toom-3 at various sizes to find the crossover. - /// - /// Run with: - /// ```sh - /// # Pure toom-3 (prevent internal NTT recursion): - /// DASHU_THRESHOLD_NTT=99999999 cargo test -p dashu-int --release \ - /// -- ntt::tests::crossover --ignored --nocapture - /// ``` - /// - /// The output is a table showing word count, transform length N, - /// toom-3 time, NTT time, and speedup ratio. Use it to recalibrate - /// [`THRESHOLD_NTT`] after algorithmic changes. - #[test] - #[ignore] - #[allow(clippy::let_underscore_must_use)] - fn crossover() { - use std::time::Instant; - - let sizes: &[usize] = &[ - 5_000, 10_000, 20_000, 30_000, 40_000, 50_000, 60_000, 80_000, 100_000, 120_000, - 131_000, - ]; - - println!( - "{:>10} {:>4} {:>8} {:>12} {:>12} {:>10}", - "words", "bp", "N", "toom-3(ms)", "ntt(ms)", "ratio" - ); - println!("{}", "-".repeat(68)); - - for &n in sizes { - let a: Vec = (0..n) - .map(|i| (i as u64 + 1).wrapping_mul(0x9E3779B97F4A7C15)) - .collect(); - let b: Vec = (0..n) - .map(|i| (i as u64 + 1).wrapping_mul(0xC6A4A7935BD1E995)) - .collect(); - let mut c_toom = vec![0u64; 2 * n]; - let mut c_ntt = vec![0u64; 2 * n]; - - let layout = memory_requirement_up_to(2 * n, n); - let warmup = 2; - let iters = 5; - - // toom-3 (may use NTT internally depending on DASHU_THRESHOLD_NTT) - let t_toom = { - let mut best = f64::MAX; - for _ in 0..warmup { - let mut alloc = crate::memory::MemoryAllocation::new(layout); - let mut mem = alloc.memory(); - c_toom.fill(0); - let _ = - crate::mul::toom_3::add_signed_mul(&mut c_toom, Positive, &a, &b, &mut mem); - } - for _ in 0..iters { - let mut alloc = crate::memory::MemoryAllocation::new(layout); - let mut mem = alloc.memory(); - c_toom.fill(0); - let start = Instant::now(); - let _ = - crate::mul::toom_3::add_signed_mul(&mut c_toom, Positive, &a, &b, &mut mem); - let elapsed = start.elapsed().as_secs_f64() * 1000.0; - if elapsed < best { - best = elapsed; - } - } - best - }; - - // NTT (direct) - let t_ntt = { - let mut best = f64::MAX; - for _ in 0..warmup { - let mut alloc = crate::memory::MemoryAllocation::new(layout); - let mut mem = alloc.memory(); - c_ntt.fill(0); - let _carry = add_signed_mul_impl(&mut c_ntt, Positive, &a, &b, &mut mem); - } - for _ in 0..iters { - let mut alloc = crate::memory::MemoryAllocation::new(layout); - let mut mem = alloc.memory(); - c_ntt.fill(0); - let start = Instant::now(); - let _carry = add_signed_mul_impl(&mut c_ntt, Positive, &a, &b, &mut mem); - let elapsed = start.elapsed().as_secs_f64() * 1000.0; - if elapsed < best { - best = elapsed; - } - } - best - }; - - assert_eq!(&c_ntt[..], &c_toom[..], "mismatch at n={n}"); - - let (_b_pack, nn, _k_eff) = select_params(n, n); - println!( - "{:>10} {:>4} {:>8} {:>12.3} {:>12.3} {:>9.2}x", - n, - _b_pack, - nn, - t_toom, - t_ntt, - t_ntt / t_toom - ); - } - } } From c0c95085ae8d1c310b6df4500a5256c7441b42f2 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sat, 13 Jun 2026 11:04:42 +0800 Subject: [PATCH 07/19] Some minor improvements Co-Authored-By: Claude Opus 4.7 --- TODO_NTT.md | 49 +++++----- integer/Cargo.toml | 2 +- integer/src/mul/ntt/crt.rs | 7 +- integer/src/mul/ntt/pack.rs | 165 ++++++++++++++++------------------ integer/src/mul/ntt/primes.rs | 24 ++++- 5 files changed, 127 insertions(+), 120 deletions(-) diff --git a/TODO_NTT.md b/TODO_NTT.md index 4c9a6a69..295a6b4a 100644 --- a/TODO_NTT.md +++ b/TODO_NTT.md @@ -47,10 +47,8 @@ primes always suffice. Third-prime fallback (`K_eff = 3`) is retained as a safety net for larger `b_pack`. -- **Threshold calibrated.** `THRESHOLD_NTT = 120 000` words (~7.7 M bits), +- **Threshold calibrated.** `THRESHOLD_NTT = 40 000` words (~2.6 M bits), the first measured crossover where NTT beats pure toom-3 on Apple M4 Pro. - At 131 072 words N doubles (1 048 576 → 2 097 152) and NTT regresses until - toom-3 catches up at ~190 k words — radix-4 will shrink this gap. - **Env-var overrides.** `DASHU_THRESHOLD_SIMPLE`, `DASHU_THRESHOLD_KARATSUBA`, `DASHU_THRESHOLD_NTT` override the compile-time defaults at runtime. Gated @@ -64,24 +62,7 @@ ## Remaining optimisation opportunities -### 1. Increase `b_pack` from 16 → 32 (~2× speedup) - -Currently every coefficient is 16 bits (4 coeffs per 64-bit word). With -K_eff = 2 primes providing ~2^128 of headroom, 32-bit coefficients are safe: - - max_coeff = N/2 · (2^32 − 1)^2 < 2^95 ≪ P0·P1 ≈ 2^128 - -Doubling `b_pack` halves the coefficient count, halves N, and roughly halves -the total transform work. Additionally, 32-bit packing is simpler: exactly -2 coefficients per word, no straddling of word boundaries. - -**Work items:** -- Update `select_params` to choose `b_pack ∈ {16, 32}` based on headroom. -- Adjust `memory_requirement_up_to` — the current worst-case bound uses - `B_PACK_MIN = 16`; need to handle the tighter N for `b_pack = 32`. -- Benchmark the new packing path. - -### 2. Radix-4 or split-radix NTT (~25–33% fewer twiddle multiplies) +### 1. Radix-4 or split-radix NTT (~25–33% fewer twiddle multiplies) Radix-4 processes 4 elements per butterfly with 3 twiddle multiplies and `log₄(N)` stages (half as many passes through memory). Split-radix pushes @@ -95,7 +76,7 @@ the savings closer to 33%. - A primitive 4-th root `j = ω_N^{N/4}` is needed for the butterfly core; derive it from the existing `ω_2_32` root. -### 3. Harvey lazy-reduction butterflies (~10–15%) +### 2. Harvey lazy-reduction butterflies (~10–15%) Currently every `add_mod` / `sub_mod` fully normalizes to `[0, p)`. Harvey's approach keeps values in `[0, 2p)` across multiple butterfly stages, deferring @@ -109,27 +90,43 @@ a branch + subtract with a no-op in the inner loop. dynamic range, so worst-case after log₂(N) stages is `[0, N·p)` — we need a cleanup before it overflows `u64`). -### 4. Merge `bit_reverse` with `pack` (~5–10%) +### 3. Merge `bit_reverse` with `pack` (~5–10%) Currently `pack` writes coefficients in natural order, then `bit_reverse` permutes them in a second pass. Write packed coefficients directly to their bit-reversed positions, saving one full array pass. -### 5. Shift-expressible twiddle factors (stage-dependent) +### 4. Shift-expressible twiddle factors (stage-dependent) In Goldilocks primes, `2^k mod p = 2^k` when `2^k < p`. The first few NTT stages have twiddle factors that are pure powers of 2, so `mul_mod(t, 2^k)` reduces to a shift + conditional subtract — no `u128` multiply needed. -### 6. Specialize `b = 32` lane (~5%) +### 5. Specialize `b = 32` lane (~5%) The `b = 32` prime (`0xFFFFFFFF00000001`) has the cleanest reduction identity (splits a `u128` product exactly into 32-bit limbs). A dedicated code path for this prime alone could squeeze out a few more cycles vs. the generic `match B` dispatch in `mul_mod`. -### 7. Asymmetric operand chunking (conditional) +### 6. Asymmetric operand chunking (conditional) When `a ≫ b`, chunk the long operand, forward-transform the short operand once, and reuse `b̂` (the transformed short operand) across all chunks. Only matters for extremely lopsided inputs. + +### 7. u32-word support via u32 Solinas primes + +The NTT path currently requires `Word = u64` and uses three 64-bit Solinas +primes. For 32-bit (and potentially 16-bit) targets, we need a separate set +of u32-friendly Solinas primes of the form `2^32 − 2^b + 1`. + +**Work items:** +- Find 2–3 primes `p = 2^32 − 2^b + 1` with `v2(p-1) ≥ 16` (enough for N up + to 2^16) and distinct `b` values. +- Implement `FixedTrinomialSolinas32` (or equivalent) in `num-modular`, or + hand-roll the 32-bit reduction inline. +- Generalize the NTT pipeline over `Word` size: the packing, transform, and + CRT layers need to work with `u32` coefficients instead of `u64`. +- Assert `Word = u32` or `Word = u64` at entry and dispatch to the appropriate + prime set. diff --git a/integer/Cargo.toml b/integer/Cargo.toml index 3f06cda3..42c978dc 100644 --- a/integer/Cargo.toml +++ b/integer/Cargo.toml @@ -31,7 +31,7 @@ dashu-base = { version = "0.4.1", default-features = false, path = "../base" } cfg-if = { version = "1.0.0" } static_assertions = { version = "1.1" } rustversion = { version = "1.0.0" } -num-modular = { version = "0.6.2", path = "../../num-modular" } +num-modular = { version = "0.6.3" } # stable dependencies num-order = { optional = true, version = "1.2.0", default-features = false } diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs index ce2e927d..3359a98f 100644 --- a/integer/src/mul/ntt/crt.rs +++ b/integer/src/mul/ntt/crt.rs @@ -9,20 +9,17 @@ use num_modular::Reducer; /// Subset of `Reducer` that is object-safe (no `new` or other /// non-`&self` methods). Implemented automatically for every /// `Reducer` via a blanket impl. -#[allow(dead_code)] pub trait ModOps { - fn add(&self, lhs: &u64, rhs: &u64) -> u64; fn sub(&self, lhs: &u64, rhs: &u64) -> u64; fn mul(&self, lhs: &u64, rhs: &u64) -> u64; } impl> ModOps for T { - fn add(&self, lhs: &u64, rhs: &u64) -> u64 { - Reducer::add(self, lhs, rhs) - } + #[inline] fn sub(&self, lhs: &u64, rhs: &u64) -> u64 { Reducer::sub(self, lhs, rhs) } + #[inline] fn mul(&self, lhs: &u64, rhs: &u64) -> u64 { Reducer::mul(self, lhs, rhs) } diff --git a/integer/src/mul/ntt/pack.rs b/integer/src/mul/ntt/pack.rs index 8edbbcb8..4af55680 100644 --- a/integer/src/mul/ntt/pack.rs +++ b/integer/src/mul/ntt/pack.rs @@ -1,14 +1,10 @@ //! Bit-level packing / unpacking of `b`-bit coefficients. -#![allow( - dead_code, - unused_assignments, - unused_mut, - unused_variables, - clippy::unnecessary_cast -)] +#![allow(clippy::unnecessary_cast)] use crate::arch::word::Word; +// TODO: shall we specialize the packing function? Since we only have three b_pack options. + /// Pack a big integer (given as `&[Word]`, little-endian) into `out`, /// producing `n` coefficients of `b_pack` bits each, zero-padded. /// @@ -16,6 +12,17 @@ use crate::arch::word::Word; /// Panics if `out.len() < n`. pub fn pack(out: &mut [u64], words: &[Word], b_pack: u32, n: usize) { assert!(out.len() >= n); + + // Fast path: one coefficient per word, no bit shifting needed. + if b_pack == Word::BITS { + let len = words.len().min(n); + // SAFETY: NTT path requires Word = u64 (asserted by caller). + let words_u64 = unsafe { &*(words as *const [Word] as *const [u64]) }; + out[..len].copy_from_slice(&words_u64[..len]); + out[len..n].fill(0); + return; + } + let mask = if b_pack < Word::BITS { (1u64 << b_pack) - 1 } else { @@ -52,92 +59,78 @@ pub fn pack(out: &mut [u64], words: &[Word], b_pack: u32, n: usize) { } } -/// Accumulate CRT-recovered convolution coefficients into the output limb -/// array with carry propagation. -/// -/// Each coefficient `c_k` contributes `c_k << (k * b_pack)` bits to the -/// output. `output` must have capacity for `c.len()` coefficients plus any -/// carry overflow. -pub fn unpack_accumulate(output: &mut [Word], coeffs: &[u64], b_pack: u32, output_len: usize) { - let word_bits = Word::BITS as u32; - // For each coefficient, shift it by k*b_pack bits and add into the - // output with carry propagation. We use a software accumulation - // because the coefficients can be larger than a single output word. - - for (k, &coeff) in coeffs.iter().enumerate().take(output_len) { - if coeff == 0 { - continue; - } - let shift_bits = (k as u32).wrapping_mul(b_pack); - let word_idx = (shift_bits / word_bits) as usize; - let bit_shift = shift_bits % word_bits; - - // The coefficient occupies up to ⌈bit_len(coeff) / word_bits⌉ words. - // We split it into word-sized chunks and add each with the - // appropriate shift to the output. - let lo = coeff as u64; - let _hi = 0u64; // coeff fits in one u64 since max CRT value < P ≈ 2^192 - // Actually, per-coefficient CRT values can be up to P-1 ≈ 2^192, - // which needs up to 3 words. We handle this by splitting the - // coefficient itself into words and accumulating each. - - // For the immediate case, coeff from CRT is already small enough - // to fit in one or two u64 words. We accumulate by repeated - // add-with-carry into the output slice. - let mut carry: Word = 0; - let mut idx = word_idx; - - if bit_shift == 0 { - // Aligned: just add into output - let (sum, c) = overflowing_add_word(output.get(idx).copied().unwrap_or(0), lo); - carry = Word::from(c); - if idx < output.len() { - output[idx] = sum; +#[cfg(test)] +mod tests { + use super::*; + + /// Accumulate CRT-recovered convolution coefficients into the output limb + /// array with carry propagation. + /// + /// Each coefficient `c_k` contributes `c_k << (k * b_pack)` bits to the + /// output. `output` must have capacity for `c.len()` coefficients plus any + /// carry overflow. + fn unpack_accumulate(output: &mut [Word], coeffs: &[u64], b_pack: u32, output_len: usize) { + let word_bits = Word::BITS as u32; + + for (k, &coeff) in coeffs.iter().enumerate().take(output_len) { + if coeff == 0 { + continue; } - idx += 1; - } else { - // Split across two output words - let lo_part = lo << bit_shift; - let hi_part = if bit_shift > 0 { - lo >> (64 - bit_shift) + let shift_bits = (k as u32).wrapping_mul(b_pack); + let word_idx = (shift_bits / word_bits) as usize; + let bit_shift = shift_bits % word_bits; + + let lo = coeff as u64; + let mut carry: Word; + let mut idx = word_idx; + + if bit_shift == 0 { + let (sum, c) = output.get(idx).copied().unwrap_or(0).overflowing_add(lo); + carry = Word::from(c); + if idx < output.len() { + output[idx] = sum; + } + idx += 1; } else { - 0 - }; - - let (sum, c1) = overflowing_add_word(output.get(idx).copied().unwrap_or(0), lo_part); - carry = Word::from(c1); - if idx < output.len() { - output[idx] = sum; + let lo_part = lo << bit_shift; + let hi_part = if bit_shift > 0 { + lo >> (64 - bit_shift) + } else { + 0 + }; + + let (sum, c1) = output + .get(idx) + .copied() + .unwrap_or(0) + .overflowing_add(lo_part); + carry = Word::from(c1); + if idx < output.len() { + output[idx] = sum; + } + idx += 1; + + let (sum2, c2) = output + .get(idx) + .copied() + .unwrap_or(0) + .overflowing_add(hi_part + carry); + carry = Word::from(c2); + if idx < output.len() { + output[idx] = sum2; + } + idx += 1; } - idx += 1; - let (sum2, c2) = - overflowing_add_word(output.get(idx).copied().unwrap_or(0), hi_part + carry); - carry = Word::from(c2); - if idx < output.len() { - output[idx] = sum2; + // Propagate remaining carry + while carry != 0 && idx < output.len() { + let (sum, c) = output[idx].overflowing_add(carry); + output[idx] = sum; + carry = Word::from(c); + idx += 1; } - idx += 1; - } - - // Propagate remaining carry - while carry != 0 && idx < output.len() { - let (sum, c) = overflowing_add_word(output[idx], carry); - output[idx] = sum; - carry = Word::from(c); - idx += 1; } } -} - -fn overflowing_add_word(a: Word, b: u64) -> (Word, bool) { - let (sum, overflow) = a.overflowing_add(b); - (sum, overflow) -} - -#[cfg(test)] -mod tests { - use super::*; #[test] fn test_pack_unpack_roundtrip() { diff --git a/integer/src/mul/ntt/primes.rs b/integer/src/mul/ntt/primes.rs index 4ee458bd..fd2c2fa0 100644 --- a/integer/src/mul/ntt/primes.rs +++ b/integer/src/mul/ntt/primes.rs @@ -13,10 +13,10 @@ pub struct NttPrime { /// The exponent `b` in the Solinas form. pub b: u32, /// The exponent of 2 in `p - 1`: `v2(p - 1)`. - #[allow(dead_code)] + #[cfg(test)] pub v2: u32, /// A primitive root modulo `p` that generates the full multiplicative group. - #[allow(dead_code)] + #[cfg(test)] pub g: u64, /// A primitive `2^32`-th root of unity: `ω = g^{(p-1) / 2^32} mod p`. pub omega_2_32: u64, @@ -29,6 +29,26 @@ pub struct NttPrime { /// | GL | 32 | `0xFFFFFFFF00000001` | 32 | 7 | `1753635133440165772` | /// | P1 | 34 | `0xFFFFFFFC00000001` | 34 | 5 | `11315553352654630047` | /// | P2 | 40 | `0xFFFFFF0000000001` | 40 | 19 | `551857376737322389` | +#[cfg(not(test))] +pub const PRIMES: [NttPrime; K] = [ + NttPrime { + p: 0xFFFFFFFF00000001, + b: 32, + omega_2_32: 1753635133440165772, + }, + NttPrime { + p: 0xFFFFFFFC00000001, + b: 34, + omega_2_32: 11315553352654630047, + }, + NttPrime { + p: 0xFFFFFF0000000001, + b: 40, + omega_2_32: 551857376737322389, + }, +]; + +#[cfg(test)] pub const PRIMES: [NttPrime; K] = [ NttPrime { p: 0xFFFFFFFF00000001, From cb9c90a451a043365e6ff10e37c2bd54bdf4e011 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sat, 13 Jun 2026 11:05:01 +0800 Subject: [PATCH 08/19] Fix CI --- integer/src/mul/mod.rs | 93 ++++++++++++++++++++++++++++++-- integer/src/mul/ntt/crt.rs | 1 + integer/src/mul/ntt/mod.rs | 2 + integer/src/mul/ntt/pack.rs | 2 + integer/src/mul/ntt/transform.rs | 2 + 5 files changed, 97 insertions(+), 3 deletions(-) diff --git a/integer/src/mul/mod.rs b/integer/src/mul/mod.rs index 5157a99e..9d96aefe 100644 --- a/integer/src/mul/mod.rs +++ b/integer/src/mul/mod.rs @@ -25,7 +25,28 @@ const THRESHOLD_KARATSUBA_DEFAULT: usize = 96; const_assert!(THRESHOLD_KARATSUBA_DEFAULT + 1 >= toom_3::MIN_LEN); /// If smaller operand length > this, NTT multiplication will be used. +#[cfg(not(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" +)))] const THRESHOLD_NTT_DEFAULT: usize = ntt::THRESHOLD_NTT; +/// NTT unavailable on 16/32-bit word targets — use `usize::MAX` so dispatch never +/// routes to the NTT path. +#[cfg(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" +))] +const THRESHOLD_NTT_DEFAULT: usize = usize::MAX; +#[cfg(not(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" +)))] const_assert!(THRESHOLD_NTT_DEFAULT + 1 >= toom_3::MIN_LEN); /// Environment-variable overrides for multiplication thresholds. @@ -74,6 +95,12 @@ mod threshold { mod helpers; mod karatsuba; +#[cfg(not(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" +)))] pub(crate) mod ntt; mod simple; pub(crate) mod toom_3; @@ -215,7 +242,26 @@ pub fn memory_requirement_up_to(total_len: usize, smaller_len: usize) -> Layout } else if smaller_len <= threshold::ntt() { toom_3::memory_requirement_up_to(smaller_len) } else { - ntt::memory_requirement_up_to(total_len, smaller_len) + // NTT path — only available on 64-bit word targets. + #[cfg(not(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" + )))] + { + ntt::memory_requirement_up_to(total_len, smaller_len) + } + #[cfg(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" + ))] + { + let _ = (total_len, smaller_len); + unreachable!("NTT requires 64-bit Word"); + } } } @@ -255,7 +301,25 @@ pub fn add_signed_mul<'a>( } else if b.len() <= threshold::ntt() { toom_3::add_signed_mul(c, sign, a, b, memory) } else { - ntt::add_signed_mul(c, sign, a, b, memory) + #[cfg(not(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" + )))] + { + ntt::add_signed_mul(c, sign, a, b, memory) + } + #[cfg(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" + ))] + { + let _ = (c, sign, a, b, memory); + unreachable!("NTT requires 64-bit Word"); + } } } @@ -280,7 +344,25 @@ pub fn add_signed_mul_same_len( } else if n <= threshold::ntt() { toom_3::add_signed_mul_same_len(c, sign, a, b, memory) } else { - ntt::add_signed_mul_same_len(c, sign, a, b, memory) + #[cfg(not(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" + )))] + { + ntt::add_signed_mul_same_len(c, sign, a, b, memory) + } + #[cfg(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" + ))] + { + let _ = (c, sign, a, b, memory); + unreachable!("NTT requires 64-bit Word"); + } } } @@ -295,6 +377,7 @@ mod threshold_tests { /// cargo test -p dashu-int --release -- mul::threshold_tests::crossover_karatsuba --nocapture --ignored #[test] #[ignore] + #[cfg(feature = "std")] fn crossover_karatsuba() { use std::time::Instant; @@ -390,6 +473,10 @@ mod threshold_tests { #[test] #[ignore] #[allow(clippy::let_underscore_must_use)] + #[cfg(all( + feature = "std", + not(any(force_bits = "16", force_bits = "32", target_pointer_width = "16", target_pointer_width = "32")) + ))] fn crossover_ntt() { use std::time::Instant; diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs index 3359a98f..0455b7c3 100644 --- a/integer/src/mul/ntt/crt.rs +++ b/integer/src/mul/ntt/crt.rs @@ -141,6 +141,7 @@ pub fn garner_combine(residues: &[u64], reducers: &[&dyn ModOps]) -> U192 { #[cfg(test)] mod tests { use super::*; + use alloc::vec; #[test] fn test_garner_with_ntt_primes() { diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index 932f0e78..fbc42095 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -352,6 +352,8 @@ fn add_shifted_to_prod(prod: &mut [u64], val: &U192, k: usize, b_pack: u32) { #[cfg(test)] mod tests { use super::*; + use alloc::vec; + use alloc::vec::Vec; #[test] fn test_select_params_small() { diff --git a/integer/src/mul/ntt/pack.rs b/integer/src/mul/ntt/pack.rs index 4af55680..898c4c95 100644 --- a/integer/src/mul/ntt/pack.rs +++ b/integer/src/mul/ntt/pack.rs @@ -62,6 +62,8 @@ pub fn pack(out: &mut [u64], words: &[Word], b_pack: u32, n: usize) { #[cfg(test)] mod tests { use super::*; + use alloc::vec; + use alloc::vec::Vec; /// Accumulate CRT-recovered convolution coefficients into the output limb /// array with carry propagation. diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs index 8a528d86..4afd0168 100644 --- a/integer/src/mul/ntt/transform.rs +++ b/integer/src/mul/ntt/transform.rs @@ -135,6 +135,8 @@ pub fn pointwise_mul(a_hat: &mut [u64], b_hat: &[u64]) { #[cfg(test)] mod tests { use super::*; + use alloc::vec; + use alloc::vec::Vec; use crate::mul::ntt::primes::PRIMES; fn assert_all_eq(a: &[u64], b_val: &[u64]) { From 0caca23a84ed08113272302b97ab830f9b6203b5 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sat, 13 Jun 2026 11:05:20 +0800 Subject: [PATCH 09/19] Fix CI again --- integer/src/mul/mod.rs | 17 +++++++++++------ integer/src/mul/ntt/crt.rs | 1 + integer/src/mul/ntt/mod.rs | 2 ++ integer/src/mul/ntt/pack.rs | 2 ++ integer/src/mul/ntt/transform.rs | 4 +++- 5 files changed, 19 insertions(+), 7 deletions(-) diff --git a/integer/src/mul/mod.rs b/integer/src/mul/mod.rs index 9d96aefe..59ea6962 100644 --- a/integer/src/mul/mod.rs +++ b/integer/src/mul/mod.rs @@ -366,7 +366,7 @@ pub fn add_signed_mul_same_len( } } -#[cfg(test)] +#[cfg(all(test, feature = "std"))] mod threshold_tests { use super::*; use crate::arch::word::Word; @@ -388,13 +388,13 @@ mod threshold_tests { for &n in sizes { let a: Vec = (0..n) - .map(|i| (i as u64 + 1).wrapping_mul(0x9E3779B97F4A7C15)) + .map(|i| (i as Word + 1).wrapping_mul(0x9E3779B97F4A7C15u64 as Word)) .collect(); let b: Vec = (0..n) - .map(|i| (i as u64 + 1).wrapping_mul(0xC6A4A7935BD1E995)) + .map(|i| (i as Word + 1).wrapping_mul(0xC6A4A7935BD1E995u64 as Word)) .collect(); - let mut c_kara = vec![0u64; 2 * n]; - let mut c_toom = vec![0u64; 2 * n]; + let mut c_kara = vec![0 as Word; 2 * n]; + let mut c_toom = vec![0 as Word; 2 * n]; let layout_kara = karatsuba::memory_requirement_up_to(n); let layout_toom = toom_3::memory_requirement_up_to(n); // Use the larger layout so both algorithms get enough memory. @@ -475,7 +475,12 @@ mod threshold_tests { #[allow(clippy::let_underscore_must_use)] #[cfg(all( feature = "std", - not(any(force_bits = "16", force_bits = "32", target_pointer_width = "16", target_pointer_width = "32")) + not(any( + force_bits = "16", + force_bits = "32", + target_pointer_width = "16", + target_pointer_width = "32" + )) ))] fn crossover_ntt() { use std::time::Instant; diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs index 0455b7c3..976bb91a 100644 --- a/integer/src/mul/ntt/crt.rs +++ b/integer/src/mul/ntt/crt.rs @@ -141,6 +141,7 @@ pub fn garner_combine(residues: &[u64], reducers: &[&dyn ModOps]) -> U192 { #[cfg(test)] mod tests { use super::*; + #[cfg(not(feature = "std"))] use alloc::vec; #[test] diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index fbc42095..f77624f7 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -352,7 +352,9 @@ fn add_shifted_to_prod(prod: &mut [u64], val: &U192, k: usize, b_pack: u32) { #[cfg(test)] mod tests { use super::*; + #[cfg(not(feature = "std"))] use alloc::vec; + #[cfg(not(feature = "std"))] use alloc::vec::Vec; #[test] diff --git a/integer/src/mul/ntt/pack.rs b/integer/src/mul/ntt/pack.rs index 898c4c95..5efdd08e 100644 --- a/integer/src/mul/ntt/pack.rs +++ b/integer/src/mul/ntt/pack.rs @@ -62,7 +62,9 @@ pub fn pack(out: &mut [u64], words: &[Word], b_pack: u32, n: usize) { #[cfg(test)] mod tests { use super::*; + #[cfg(not(feature = "std"))] use alloc::vec; + #[cfg(not(feature = "std"))] use alloc::vec::Vec; /// Accumulate CRT-recovered convolution coefficients into the output limb diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs index 4afd0168..4bb63e65 100644 --- a/integer/src/mul/ntt/transform.rs +++ b/integer/src/mul/ntt/transform.rs @@ -135,9 +135,11 @@ pub fn pointwise_mul(a_hat: &mut [u64], b_hat: &[u64]) { #[cfg(test)] mod tests { use super::*; + use crate::mul::ntt::primes::PRIMES; + #[cfg(not(feature = "std"))] use alloc::vec; + #[cfg(not(feature = "std"))] use alloc::vec::Vec; - use crate::mul::ntt::primes::PRIMES; fn assert_all_eq(a: &[u64], b_val: &[u64]) { assert_eq!(a.len(), b_val.len()); From c6a415975d536c4f85546441fcfc03ad20b3dcba Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 01:03:15 +0800 Subject: [PATCH 10/19] Change solinas to proth for NTT --- integer/CHANGELOG.md | 2 +- integer/Cargo.toml | 2 +- integer/src/arch/generic_32_bit/mod.rs | 1 + integer/src/arch/generic_32_bit/ntt.rs | 150 +++++++ integer/src/arch/generic_64_bit/mod.rs | 1 + integer/src/arch/generic_64_bit/ntt.rs | 151 +++++++ integer/src/arch/mod.rs | 3 + integer/src/arch/x86/mod.rs | 3 + integer/src/arch/x86_64/mod.rs | 3 + integer/src/mul/mod.rs | 23 +- integer/src/mul/ntt/compute_constants.py | 149 +++++++ integer/src/mul/ntt/crt.rs | 222 +++++++--- integer/src/mul/ntt/mod.rs | 512 ++++++++++++++--------- integer/src/mul/ntt/pack.rs | 57 +-- integer/src/mul/ntt/primes.rs | 257 ------------ integer/src/mul/ntt/transform.rs | 328 ++++++++------- 16 files changed, 1143 insertions(+), 721 deletions(-) create mode 100644 integer/src/arch/generic_32_bit/ntt.rs create mode 100644 integer/src/arch/generic_64_bit/ntt.rs create mode 100644 integer/src/mul/ntt/compute_constants.py delete mode 100644 integer/src/mul/ntt/primes.rs diff --git a/integer/CHANGELOG.md b/integer/CHANGELOG.md index 7e13059f..d2e4a013 100644 --- a/integer/CHANGELOG.md +++ b/integer/CHANGELOG.md @@ -3,7 +3,7 @@ ## Unreleased ### Add -- NTT-based multiplication for very large integers (above 40000 words / ~2.6M bits), using two or three Solinas primes of the form `2^64 − 2^b + 1` combined with the Chinese Remainder Theorem. +- NTT-based multiplication for very large integers (above 40000 words / ~2.6M bits), using two or three Proth primes of the form `K·2^N + 1` combined with the Chinese Remainder Theorem. Supports both 64-bit and 32-bit Word targets. - `UBig::from_u64` and `IBig::from_i64`, const on 32-bit and 64-bit targets. ### Improve diff --git a/integer/Cargo.toml b/integer/Cargo.toml index 42c978dc..9f8fb6ff 100644 --- a/integer/Cargo.toml +++ b/integer/Cargo.toml @@ -31,7 +31,7 @@ dashu-base = { version = "0.4.1", default-features = false, path = "../base" } cfg-if = { version = "1.0.0" } static_assertions = { version = "1.1" } rustversion = { version = "1.0.0" } -num-modular = { version = "0.6.3" } +num-modular = { version = "0.6.4" } # stable dependencies num-order = { optional = true, version = "1.2.0", default-features = false } diff --git a/integer/src/arch/generic_32_bit/mod.rs b/integer/src/arch/generic_32_bit/mod.rs index eaab456f..fd98d73d 100644 --- a/integer/src/arch/generic_32_bit/mod.rs +++ b/integer/src/arch/generic_32_bit/mod.rs @@ -4,4 +4,5 @@ pub(crate) mod add; #[path = "../generic/digits.rs"] pub(crate) mod digits; +pub(crate) mod ntt; pub(crate) mod word; diff --git a/integer/src/arch/generic_32_bit/ntt.rs b/integer/src/arch/generic_32_bit/ntt.rs new file mode 100644 index 00000000..ce55b608 --- /dev/null +++ b/integer/src/arch/generic_32_bit/ntt.rs @@ -0,0 +1,150 @@ +//! NTT primes and constants for 32-bit Word targets. +//! +//! Uses Proth primes of the form `K * 2^N + 1`. +//! All constants computed by `integer/src/mul/ntt/compute_constants.py`. + +use num_modular::{FixedProth32, Reducer}; + +// Proth reducer instances — each with a different (N, K) pair. +pub const P0: FixedProth32<26, 7> = FixedProth32::<26, 7>; +pub const P1: FixedProth32<27, 15> = FixedProth32::<27, 15>; +pub const P2: FixedProth32<27, 17> = FixedProth32::<27, 17>; + +pub const K: usize = 3; +pub const MAX_LOG_N: u32 = 26; +pub const B_PACK_MIN: u32 = 8; +pub const B_PACK_CANDIDATES: &[u32] = &[32, 16, 8]; + +pub type Lane = u32; +pub type DoubleLane = u64; + +/// Primitive `MAX_LOG_N`-th roots of unity for each prime. +pub const OMEGA_MAX: [Lane; K] = [ + 0x0000088b, // P0 + 0x3a26eef8, // P1 + 0x1aa0ab5e, // P2 +]; + +pub const CRT_INV_IJ: [[Lane; K]; K] = [ + [0, 0x4e42c85b, 0x5fb425ef], + [0, 0, 0x44000009], + [0, 0, 0], +]; + +/// Prime moduli indexed by PI. +pub const MODULI: [Lane; K] = [ + FixedProth32::<26, 7>::MODULUS, + FixedProth32::<27, 15>::MODULUS, + FixedProth32::<27, 17>::MODULUS, +]; + +#[inline] +pub fn to_monty(val: Lane) -> Lane { + match PI { + 0 => P0.transform(val), + 1 => P1.transform(val), + 2 => P2.transform(val), + _ => unreachable!(), + } +} + +#[inline] +pub fn from_monty(val: Lane) -> Lane { + match PI { + 0 => P0.residue(val), + 1 => P1.residue(val), + 2 => P2.residue(val), + _ => unreachable!(), + } +} + +#[inline] +pub fn mul_mod(a: Lane, b_val: Lane) -> Lane { + let prod = (a as DoubleLane) * (b_val as DoubleLane); + match PI { + 0 => P0.reduce(prod), + 1 => P1.reduce(prod), + 2 => P2.reduce(prod), + _ => unreachable!(), + } +} + +#[inline] +pub fn add_mod(a: Lane, b_val: Lane) -> Lane { + match PI { + 0 => P0.add(&a, &b_val), + 1 => P1.add(&a, &b_val), + 2 => P2.add(&a, &b_val), + _ => unreachable!(), + } +} + +#[inline] +pub fn sub_mod(a: Lane, b_val: Lane) -> Lane { + match PI { + 0 => P0.sub(&a, &b_val), + 1 => P1.sub(&a, &b_val), + 2 => P2.sub(&a, &b_val), + _ => unreachable!(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_primes_proth_form() { + assert_eq!(MODULI[0], 7u32 * (1u32 << 26) + 1); + assert_eq!(MODULI[1], 15u32 * (1u32 << 27) + 1); + assert_eq!(MODULI[2], 17u32 * (1u32 << 27) + 1); + } + + #[test] + fn test_primes_v2() { + for &p in &MODULI { + let v2 = (p - 1).trailing_zeros(); + assert!(v2 >= MAX_LOG_N, "v2(p-1) = {v2} < MAX_LOG_N"); + } + } + + #[test] + fn test_omega_order() { + for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { + let p = MODULI[pi]; + let sqr = |w: Lane| -> Lane { + match pi { + 0 => P0.reduce((w as u64) * (w as u64)), + 1 => P1.reduce((w as u64) * (w as u64)), + 2 => P2.reduce((w as u64) * (w as u64)), + _ => unreachable!(), + } + }; + + let mut w = match pi { + 0 => to_monty::<0>(omega_max), + 1 => to_monty::<1>(omega_max), + 2 => to_monty::<2>(omega_max), + _ => unreachable!(), + }; + for _ in 0..MAX_LOG_N - 1 { + w = sqr(w); + } + let w_std = match pi { + 0 => from_monty::<0>(w), + 1 => from_monty::<1>(w), + 2 => from_monty::<2>(w), + _ => unreachable!(), + }; + assert_eq!(w_std, p - 1, "omega^(2^(MAX_LOG_N-1)) != -1 mod p for prime {pi}"); + w = sqr(w); + let one = match pi { + 0 => from_monty::<0>(w), + 1 => from_monty::<1>(w), + 2 => from_monty::<2>(w), + _ => unreachable!(), + }; + assert_eq!(one, 1, "omega^(2^MAX_LOG_N) != 1 mod p for prime {pi}"); + } + } +} diff --git a/integer/src/arch/generic_64_bit/mod.rs b/integer/src/arch/generic_64_bit/mod.rs index eaab456f..fd98d73d 100644 --- a/integer/src/arch/generic_64_bit/mod.rs +++ b/integer/src/arch/generic_64_bit/mod.rs @@ -4,4 +4,5 @@ pub(crate) mod add; #[path = "../generic/digits.rs"] pub(crate) mod digits; +pub(crate) mod ntt; pub(crate) mod word; diff --git a/integer/src/arch/generic_64_bit/ntt.rs b/integer/src/arch/generic_64_bit/ntt.rs new file mode 100644 index 00000000..e0bd8de2 --- /dev/null +++ b/integer/src/arch/generic_64_bit/ntt.rs @@ -0,0 +1,151 @@ +//! NTT primes and constants for 64-bit Word targets. +//! +//! Uses Proth primes of the form `K * 2^N + 1`. +//! All constants computed by `integer/src/mul/ntt/compute_constants.py`. + +use num_modular::{FixedProth64, Reducer}; + +// Proth reducer instances — each with a different (N, K) pair. +pub const P0: FixedProth64<57, 29> = FixedProth64::<57, 29>; +pub const P1: FixedProth64<57, 71> = FixedProth64::<57, 71>; +pub const P2: FixedProth64<57, 75> = FixedProth64::<57, 75>; + +pub const K: usize = 3; +pub const MAX_LOG_N: u32 = 57; +pub const B_PACK_MIN: u32 = 16; +pub const B_PACK_CANDIDATES: &[u32] = &[64, 32, 16]; + +pub type Lane = u64; +pub type DoubleLane = u128; + +/// Primitive `MAX_LOG_N`-th roots of unity for each prime: +/// `omega_max[i]` = `g^{(p_i-1) / 2^MAX_LOG_N} mod p_i`. +pub const OMEGA_MAX: [Lane; K] = [ + 0x00003e6b41437d93, // P0 + 0x2f754195e85edc63, // P1 + 0x75544cac36cebb29, // P2 +]; + +pub const CRT_INV_IJ: [[Lane; K]; K] = [ + [0, 0x3979e79e79e79e7c, 0x8c37a6f4de9bd37d], + [0, 0, 0x2580000000000013], + [0, 0, 0], +]; + +#[inline] +pub fn to_monty(val: Lane) -> Lane { + match PI { + 0 => P0.transform(val), + 1 => P1.transform(val), + 2 => P2.transform(val), + _ => unreachable!(), + } +} + +#[inline] +pub fn from_monty(val: Lane) -> Lane { + match PI { + 0 => P0.residue(val), + 1 => P1.residue(val), + 2 => P2.residue(val), + _ => unreachable!(), + } +} + +#[inline] +pub fn mul_mod(a: Lane, b_val: Lane) -> Lane { + let prod = (a as DoubleLane) * (b_val as DoubleLane); + match PI { + 0 => P0.reduce(prod), + 1 => P1.reduce(prod), + 2 => P2.reduce(prod), + _ => unreachable!(), + } +} + +#[inline] +pub fn add_mod(a: Lane, b_val: Lane) -> Lane { + match PI { + 0 => P0.add(&a, &b_val), + 1 => P1.add(&a, &b_val), + 2 => P2.add(&a, &b_val), + _ => unreachable!(), + } +} + +#[inline] +pub fn sub_mod(a: Lane, b_val: Lane) -> Lane { + match PI { + 0 => P0.sub(&a, &b_val), + 1 => P1.sub(&a, &b_val), + 2 => P2.sub(&a, &b_val), + _ => unreachable!(), + } +} + +/// Prime moduli indexed by PI. +pub const MODULI: [Lane; K] = [ + FixedProth64::<57, 29>::MODULUS, + FixedProth64::<57, 71>::MODULUS, + FixedProth64::<57, 75>::MODULUS, +]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_primes_proth_form() { + assert_eq!(MODULI[0], 29u64 * (1u64 << 57) + 1); + assert_eq!(MODULI[1], 71u64 * (1u64 << 57) + 1); + assert_eq!(MODULI[2], 75u64 * (1u64 << 57) + 1); + } + + #[test] + fn test_primes_v2() { + for &p in &MODULI { + let v2 = (p - 1).trailing_zeros(); + assert!(v2 >= MAX_LOG_N, "v2(p-1) = {v2} < MAX_LOG_N"); + } + } + + #[test] + fn test_omega_order() { + for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { + let p = MODULI[pi]; + let sqr = |w: Lane| -> Lane { + match pi { + 0 => P0.reduce((w as u128) * (w as u128)), + 1 => P1.reduce((w as u128) * (w as u128)), + 2 => P2.reduce((w as u128) * (w as u128)), + _ => unreachable!(), + } + }; + + let mut w = match pi { + 0 => to_monty::<0>(omega_max), + 1 => to_monty::<1>(omega_max), + 2 => to_monty::<2>(omega_max), + _ => unreachable!(), + }; + for _ in 0..MAX_LOG_N - 1 { + w = sqr(w); + } + let w_std = match pi { + 0 => from_monty::<0>(w), + 1 => from_monty::<1>(w), + 2 => from_monty::<2>(w), + _ => unreachable!(), + }; + assert_eq!(w_std, p - 1, "omega^(2^(MAX_LOG_N-1)) != -1 mod p for prime {pi}"); + w = sqr(w); + let one = match pi { + 0 => from_monty::<0>(w), + 1 => from_monty::<1>(w), + 2 => from_monty::<2>(w), + _ => unreachable!(), + }; + assert_eq!(one, 1, "omega^(2^MAX_LOG_N) != 1 mod p for prime {pi}"); + } + } +} diff --git a/integer/src/arch/mod.rs b/integer/src/arch/mod.rs index 3bcd88fa..dd11b686 100644 --- a/integer/src/arch/mod.rs +++ b/integer/src/arch/mod.rs @@ -6,6 +6,9 @@ pub(crate) use arch_impl::add; pub(crate) use arch_impl::digits; pub(crate) use arch_impl::word; +#[cfg(not(any(force_bits = "16", target_pointer_width = "16")))] +pub(crate) use arch_impl::ntt; + // Architecture choice. The logic works like this: // 1. If the configuration option force_bits is set to 16, 32 or 64, use generic__bit. // 2. Otherwise if target_arch is known, select that architecture. diff --git a/integer/src/arch/x86/mod.rs b/integer/src/arch/x86/mod.rs index bf7fccf5..c370f391 100644 --- a/integer/src/arch/x86/mod.rs +++ b/integer/src/arch/x86/mod.rs @@ -3,5 +3,8 @@ pub(crate) mod add; #[path = "../generic/digits.rs"] pub(crate) mod digits; +#[path = "../generic_32_bit/ntt.rs"] +pub(crate) mod ntt; + #[path = "../generic_32_bit/word.rs"] pub(crate) mod word; diff --git a/integer/src/arch/x86_64/mod.rs b/integer/src/arch/x86_64/mod.rs index 57d34fbe..fef5d136 100644 --- a/integer/src/arch/x86_64/mod.rs +++ b/integer/src/arch/x86_64/mod.rs @@ -3,5 +3,8 @@ pub(crate) mod add; #[path = "../generic/digits.rs"] pub(crate) mod digits; +#[path = "../generic_64_bit/ntt.rs"] +pub(crate) mod ntt; + #[path = "../generic_64_bit/word.rs"] pub(crate) mod word; diff --git a/integer/src/mul/mod.rs b/integer/src/mul/mod.rs index 59ea6962..261f53aa 100644 --- a/integer/src/mul/mod.rs +++ b/integer/src/mul/mod.rs @@ -27,25 +27,19 @@ const_assert!(THRESHOLD_KARATSUBA_DEFAULT + 1 >= toom_3::MIN_LEN); /// If smaller operand length > this, NTT multiplication will be used. #[cfg(not(any( force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" + target_pointer_width = "16" )))] const THRESHOLD_NTT_DEFAULT: usize = ntt::THRESHOLD_NTT; /// NTT unavailable on 16/32-bit word targets — use `usize::MAX` so dispatch never /// routes to the NTT path. #[cfg(any( force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" + target_pointer_width = "16" ))] const THRESHOLD_NTT_DEFAULT: usize = usize::MAX; #[cfg(not(any( force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" + target_pointer_width = "16" )))] const_assert!(THRESHOLD_NTT_DEFAULT + 1 >= toom_3::MIN_LEN); @@ -97,9 +91,7 @@ mod helpers; mod karatsuba; #[cfg(not(any( force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" + target_pointer_width = "16" )))] pub(crate) mod ntt; mod simple; @@ -260,12 +252,13 @@ pub fn memory_requirement_up_to(total_len: usize, smaller_len: usize) -> Layout ))] { let _ = (total_len, smaller_len); - unreachable!("NTT requires 64-bit Word"); + unreachable!("NTT unavailable on 16-bit targets"); } } } /// Temporary scratch space required for multiplication. +#[inline] pub fn memory_requirement_exact(total_len: usize, smaller_len: usize) -> Layout { memory_requirement_up_to(total_len, smaller_len) } @@ -318,7 +311,7 @@ pub fn add_signed_mul<'a>( ))] { let _ = (c, sign, a, b, memory); - unreachable!("NTT requires 64-bit Word"); + unreachable!("NTT unavailable on 16-bit targets"); } } } @@ -361,7 +354,7 @@ pub fn add_signed_mul_same_len( ))] { let _ = (c, sign, a, b, memory); - unreachable!("NTT requires 64-bit Word"); + unreachable!("NTT unavailable on 16-bit targets"); } } } diff --git a/integer/src/mul/ntt/compute_constants.py b/integer/src/mul/ntt/compute_constants.py new file mode 100644 index 00000000..dfc3e864 --- /dev/null +++ b/integer/src/mul/ntt/compute_constants.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +"""Compute omega_max and CRT constants for Proth NTT primes. + +Prints the computed constants in a readable format — does NOT generate Rust code. +Copy the values into the arch ntt.rs files by hand. +""" + +# --- 64-bit Proth primes --- +PRIMES_64 = [ + (0x3a00000000000001, 57, 29), # Proth(57, 29) + (0x8e00000000000001, 57, 71), # Proth(57, 71) + (0x9600000000000001, 57, 75), # Proth(57, 75) +] +MAX_LOG_N_64 = 57 + +# --- 32-bit Proth primes --- +PRIMES_32 = [ + (0x1c000001, 26, 7), # Proth(26, 7) + (0x78000001, 27, 15), # Proth(27, 15) + (0x88000001, 27, 17), # Proth(27, 17) +] +MAX_LOG_N_32 = 26 + + +def mod_pow(base, exp, mod): + """base**exp mod mod.""" + result = 1 + while exp > 0: + if exp & 1: + result = (result * base) % mod + base = (base * base) % mod + exp >>= 1 + return result + + +def mod_inv(a, mod): + """Inverse of a mod mod (mod is prime).""" + return mod_pow(a, mod - 2, mod) + + +def factorize(n): + """Return list of distinct prime factors of n.""" + factors = [] + d = 2 + m = n + while d * d <= m: + if m % d == 0: + factors.append(d) + while m % d == 0: + m //= d + d += 1 if d == 2 else 2 # skip even after 2 + if m > 1: + factors.append(m) + return factors + + +def is_primitive_root(g, p, factors_of_pm1): + """Check if g is a primitive root mod p.""" + for q in factors_of_pm1: + if mod_pow(g, (p - 1) // q, p) == 1: + return False + return True + + +def find_primitive_root(p): + """Find a primitive root mod p by brute force.""" + factors = factorize(p - 1) + for g in range(2, min(p, 2000)): + if is_primitive_root(g, p, factors): + return g + raise ValueError(f"No primitive root found for p = {p} (tried up to 2000)") + + +def compute_omega(p, g, max_log_n): + """omega = g^((p-1) / 2^max_log_n) mod p.""" + assert (p - 1) % (1 << max_log_n) == 0, \ + f"max_log_n={max_log_n} does not divide p-1 for p={p:#x}" + exp = (p - 1) >> max_log_n + return mod_pow(g, exp, p) + + +def verify_omega(omega, p, max_log_n): + """Verify omega^(2^(max_log_n-1)) == -1 and omega^(2^max_log_n) == 1.""" + w = omega + for _ in range(max_log_n - 1): + w = (w * w) % p + assert w == p - 1, f"omega^(2^{max_log_n-1}) != -1 mod p, got {w:#x}" + w = (w * w) % p + assert w == 1, f"omega^(2^{max_log_n}) != 1 mod p, got {w:#x}" + + +def compute_crt_constants(primes): + """Compute Garner CRT: inv(p_i mod p_j) mod p_j for i < j.""" + k = len(primes) + crt = [[0] * k for _ in range(k)] + for i in range(k): + for j in range(i + 1, k): + pi = primes[i] + pj = primes[j] + crt[i][j] = mod_inv(pi % pj, pj) + return crt + + +def print_results(name, primes_data, max_log_n): + """Pretty-print computed constants for one architecture.""" + primes = [p for p, _, _ in primes_data] + + print(f"===== {name} =====") + print(f" MAX_LOG_N = {max_log_n}") + print() + for i, (p, n, k) in enumerate(primes_data): + print(f" PI={i}: p = {p:#018x} ({k} * 2^{n} + 1)") + v2 = ((p - 1) & -(p - 1)).bit_length() - 1 # trailing zeros + print(f" v2(p-1) = {v2}") + + print() + + # Primitive roots & omega + for i, (p, n, k) in enumerate(primes_data): + print(f" PI={i}: finding primitive root...") + g = find_primitive_root(p) + omega = compute_omega(p, g, max_log_n) + verify_omega(omega, p, max_log_n) + print(f" g = {g}") + bit_width = 64 if max_log_n == 57 else 32 + print(f" omega_max = {omega:#0{bit_width//4 + 2}x}") + + print() + + # CRT constants + crt = compute_crt_constants(primes) + bit_width = 64 if max_log_n == 57 else 32 + print(f" CRT_INV_IJ:") + for i in range(len(primes)): + for j in range(len(primes)): + if crt[i][j] != 0: + print(f" inv(p{i} mod p{j}) = {crt[i][j]:#0{bit_width//4 + 2}x}") + + # Two-prime product for headroom checks + prod_01 = primes[0] * primes[1] + print(f"\n p0 * p1 = {prod_01:#x}") + + print() + + +# --- Main --- +if __name__ == "__main__": + print_results("64-bit", PRIMES_64, MAX_LOG_N_64) + print_results("32-bit", PRIMES_32, MAX_LOG_N_32) diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs index 976bb91a..b25a7f61 100644 --- a/integer/src/mul/ntt/crt.rs +++ b/integer/src/mul/ntt/crt.rs @@ -1,34 +1,31 @@ //! Garner CRT: combine `K` residues modulo `K` primes into a small integer. //! -//! Takes `num_modular::Reducer` implementations so the caller can supply -//! specialized Solinas reducers for the hot per-coefficient path. +//! Generic over the lane type — supports u64 (→[`U192`]) and u32 (→[`U128`]). #![allow(clippy::unnecessary_cast)] -use num_modular::Reducer; +use crate::arch::ntt::K; +use num_modular::ModularCoreOps; -/// Subset of `Reducer` that is object-safe (no `new` or other -/// non-`&self` methods). Implemented automatically for every -/// `Reducer` via a blanket impl. -pub trait ModOps { - fn sub(&self, lhs: &u64, rhs: &u64) -> u64; - fn mul(&self, lhs: &u64, rhs: &u64) -> u64; +/// Accumulator for Garner CRT — either [`U192`] (64-bit lanes) or [`U96`] (32-bit). +pub trait CrtAccum: Default + Copy { + type Lane: Copy + Into + for<'a> ModularCoreOps; + fn from_lane(v: Self::Lane) -> Self; + /// `self += t * factor` where `t` is a lane-sized coefficient. + fn add_product(&mut self, t: Self::Lane, factor: u128); + /// `self mod m` + fn rem_lane(&self, m: Self::Lane) -> Self::Lane; + /// Number of non-zero u64 words. + #[allow(dead_code)] + fn len_words(&self) -> u32; + /// View as `&[u64]`. + fn as_u64_slice(&self) -> &[u64]; } -impl> ModOps for T { - #[inline] - fn sub(&self, lhs: &u64, rhs: &u64) -> u64 { - Reducer::sub(self, lhs, rhs) - } - #[inline] - fn mul(&self, lhs: &u64, rhs: &u64) -> u64 { - Reducer::mul(self, lhs, rhs) - } -} +// ── U192 (64-bit lanes) ──────────────────────────────────────────────── /// A 192-bit unsigned integer (3 × u64, little-endian). /// -/// Used to hold Garner CRT results (which are bounded by the product of -/// three ≈2^64 primes, therefore < 2^192). +/// Used to hold Garner CRT results for three ≈2^64 primes (product < 2^192). #[derive(Clone, Copy, Debug, Default)] pub struct U192(pub [u64; 3]); @@ -40,6 +37,7 @@ impl U192 { /// `self += v` where `v` fits in 128 bits. #[inline] + #[allow(dead_code)] pub fn add_u128(&mut self, v: u128) { let lo = v as u64; let hi = (v >> 64) as u64; @@ -98,42 +96,161 @@ impl U192 { } } -use super::primes::{CRT_INV_IJ, PRIMES}; +impl CrtAccum for U192 { + type Lane = u64; + + #[inline] + fn from_lane(v: u64) -> Self { + U192::new(v) + } + + #[inline] + fn add_product(&mut self, t: u64, factor: u128) { + self.add_mul_u64_u128(t, factor); + } + + #[inline] + fn rem_lane(&self, m: u64) -> u64 { + self.rem_u64(m) + } + + #[inline] + fn len_words(&self) -> u32 { + self.len_words() + } -/// Combine `K` residues into a `U192` via Garner's algorithm. + #[inline] + fn as_u64_slice(&self) -> &[u64] { + &self.0[..self.len_words() as usize] + } +} + +// ── U96 (32-bit lanes) ───────────────────────────────────────────────── + +/// A value bounded by 2^96 (product of three ≈2^32 primes). +/// +/// Stored as `[u64; 2]` (128 bits) so [`as_u64_slice`] can return a +/// `&[u64]` for [`add_shifted_to_prod`]. The upper 32 bits of the +/// second limb are always zero. /// -/// All primes and precomputed inverses are hardcoded in [`super::primes`]. -/// `reducers[i]` must be a reducer for the i-th prime. -pub fn garner_combine(residues: &[u64], reducers: &[&dyn ModOps]) -> U192 { +/// [`add_shifted_to_prod`]: super::add_shifted_to_prod +#[derive(Clone, Copy, Debug, Default)] +pub struct U96(pub [u64; 2]); + +impl U96 { + /// `self += t * factor` where `t` < 2^32. + #[inline] + pub fn add_mul_u32_u96(&mut self, t: u32, factor: u128) { + let fac_lo = factor as u64; + let fac_hi = (factor >> 64) as u64; + + // t × fac_lo (max 2^32 × 2^64 = 2^96 → fits in u128) + let m_lo_full = (t as u128) * (fac_lo as u128); + let lo = m_lo_full as u64; + let m_lo_carry = (m_lo_full >> 64) as u64; + + // t × fac_hi + carry + let m_hi_full = (t as u64 as u128) * (fac_hi as u128) + m_lo_carry as u128; + let m_hi = m_hi_full as u64; + + let (r0, c0) = self.0[0].overflowing_add(lo); + self.0[0] = r0; + let (r1, c1) = self.0[1].overflowing_add(m_hi.wrapping_add(c0 as u64)); + self.0[1] = r1; + let _ = c1; + } + + /// `self mod m`, where `m` < 2^32. + #[inline] + pub fn rem_u32(&self, m: u32) -> u32 { + let m128 = m as u128; + let mut r: u128 = 0; + for &word in self.0.iter().rev() { + r = (r << 64) | (word as u128); + r %= m128; + } + r as u32 + } + + #[inline] + pub fn len_words(&self) -> u32 { + if self.0[1] != 0 { + 2 + } else { + 1 + } + } +} + +impl CrtAccum for U96 { + type Lane = u32; + + #[inline] + fn from_lane(v: u32) -> Self { + U96([v as u64, 0]) + } + + #[inline] + fn add_product(&mut self, t: u32, factor: u128) { + self.add_mul_u32_u96(t, factor); + } + + #[inline] + fn rem_lane(&self, m: u32) -> u32 { + self.rem_u32(m) + } + + #[inline] + fn len_words(&self) -> u32 { + self.len_words() + } + + #[inline] + fn as_u64_slice(&self) -> &[u64] { + &self.0[..self.len_words() as usize] + } +} + +// ── Garner combine ───────────────────────────────────────────────────── + +/// Combine `K` residues into a [`CrtAccum`] via Garner's algorithm. +/// +/// All arithmetic is standard-form (not Montgomery). `crt_inv_ij[i][j]` +/// holds `inv(p_i mod p_j) mod p_j` for `i < j`. +/// `primes` contains the prime values (only `primes[0..k]` are used). +pub fn garner_combine( + residues: &[A::Lane], + crt_inv_ij: &[[A::Lane; K]; K], + primes: &[A::Lane; K], +) -> A { let k = residues.len(); - assert!(k <= 3, "CRT supports up to 3 primes"); - assert!(reducers.len() >= k); + assert!(k <= K, "CRT supports up to {K} primes"); - let p0 = PRIMES[0].p; - let p1 = PRIMES[1].p; - let p2 = PRIMES[2].p; - let mut x = U192::new(residues[0]); + let p0 = primes[0]; + let p1 = primes[1]; + let p2 = primes[2]; + let mut x = A::from_lane(residues[0]); if k == 1 { return x; } // t_1 = (r_1 - x mod p1) * inv(p0 mod p1) mod p1 - let x_mod_p1 = x.0[0] % p1; - let diff1 = reducers[1].sub(&residues[1], &x_mod_p1); - let t1 = reducers[1].mul(&diff1, &CRT_INV_IJ[0][1]); - x.add_u128((t1 as u128) * (p0 as u128)); + let x_mod_p1 = x.rem_lane(p1); + let diff1 = residues[1].subm(x_mod_p1, &p1); + let t1 = diff1.mulm(crt_inv_ij[0][1], &p1); + x.add_product(t1, p0.into()); if k == 2 { return x; } // t_2 = (r_2 - x mod p2) * inv(p0*p1 mod p2) mod p2 - let x_mod_p2 = x.rem_u64(p2); - let diff2 = reducers[2].sub(&residues[2], &x_mod_p2); - let inv_prod = reducers[2].mul(&CRT_INV_IJ[0][2], &CRT_INV_IJ[1][2]); - let t2 = reducers[2].mul(&diff2, &inv_prod); - x.add_mul_u64_u128(t2, (p0 as u128) * (p1 as u128)); + let x_mod_p2 = x.rem_lane(p2); + let diff2 = residues[2].subm(x_mod_p2, &p2); + let inv_prod = crt_inv_ij[0][2].mulm(crt_inv_ij[1][2], &p2); + let t2 = diff2.mulm(inv_prod, &p2); + x.add_product(t2, p0.into() * p1.into()); x } @@ -145,30 +262,25 @@ mod tests { use alloc::vec; #[test] - fn test_garner_with_ntt_primes() { - use super::super::primes::PRIMES; - use num_modular::FixedTrinomialSolinas64; - - let p0 = PRIMES[0].p; - let p1 = PRIMES[1].p; - let p2 = PRIMES[2].p; + fn test_garner_with_u64_primes() { + use crate::arch::ntt::{CRT_INV_IJ, MODULI}; - let r0 = FixedTrinomialSolinas64::<64, 32, 1>::new(&p0); - let r1 = FixedTrinomialSolinas64::<64, 34, 1>::new(&p1); - let r2 = FixedTrinomialSolinas64::<64, 40, 1>::new(&p2); - let reducers: [&dyn ModOps; 3] = [&r0, &r1, &r2]; + let p0 = MODULI[0]; + let p1 = MODULI[1]; + let p2 = MODULI[2]; + let primes = [p0, p1, p2]; let residues = vec![12345u64, 67890u64, 11111u64]; - let x = garner_combine(&residues, &reducers); + let x = garner_combine::(&residues, &CRT_INV_IJ, &primes); assert_eq!(x.rem_u64(p0), residues[0]); assert_eq!(x.rem_u64(p1), residues[1]); assert_eq!(x.rem_u64(p2), residues[2]); - let x = garner_combine(&residues[..2], &reducers[..2]); + let x = garner_combine::(&residues[..2], &CRT_INV_IJ, &primes); assert_eq!(x.rem_u64(p0), residues[0]); assert_eq!(x.rem_u64(p1), residues[1]); - let x = garner_combine(&residues[..1], &reducers[..1]); + let x = garner_combine::(&residues[..1], &CRT_INV_IJ, &primes); assert_eq!(x.0[0], residues[0]); } } diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index f77624f7..c713d406 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -1,9 +1,8 @@ //! NTT-based multiplication for very large integers. //! -//! Uses Number Theoretic Transforms over several 64-bit primes of the form -//! `2^64 - 2^b + 1` combined with the Chinese Remainder Theorem (CRT). +//! Uses Number Theoretic Transforms over Proth primes of the form +//! `K * 2^N + 1` combined with the Chinese Remainder Theorem (CRT). -use crate::mul::ntt::crt::ModOps; use crate::{ add, arch::word::{SignedWord, Word}, @@ -11,37 +10,18 @@ use crate::{ Sign::{self, *}, }; use alloc::alloc::Layout; -use num_modular::{FixedTrinomialSolinas64, Reducer}; +use core::mem; mod crt; mod pack; -mod primes; mod transform; -use crate::mul::ntt::crt::U192; -pub use primes::{K, PRIMES}; +use crate::arch::ntt::{mul_mod, B_PACK_CANDIDATES, B_PACK_MIN, K, MAX_LOG_N, MODULI, OMEGA_MAX}; +use crate::mul::ntt::crt::{garner_combine, CrtAccum}; /// Minimum smaller-operand length (in words) for the NTT path. -/// -/// With `b_pack = 64` the crossover is at ~25 000 words (~1.6 M bits) on -/// Apple M4 Pro. N-doubling at 32 769 / 65 537 words creates narrow -/// regression windows; radix-4 will shrink the step size further. -/// Chosen at 40 000 words where NTT is ≥18% faster. pub const THRESHOLD_NTT: usize = 40_000; -/// Smallest admissible coefficient bit width (used for worst-case memory bound). -const B_PACK_MIN: u32 = 16; - -/// Preferred coefficient bit width. 32 bits gives 2 coeffs/word and halves -/// the transform length vs. 16 bits, while staying comfortably within the -/// ~2^128 headroom of the two smallest primes. -/// Coefficient bit widths to try, in descending preference. -/// 64 uses K_eff = 3 primes; 32 and 16 use K_eff = 2. -const B_PACK_CANDIDATES: &[u32] = &[64, 32, 16]; - -/// Maximum `log2(transform length)`, set by `min(v2) = 32` across all primes. -const MAX_LOG_N: u32 = 32; - /// Select NTT parameters for operands with the given word lengths. /// /// Returns `(b_pack, N, K_eff)`. @@ -49,7 +29,7 @@ pub fn select_params(la_words: usize, lb_words: usize) -> (u32, usize, usize) { let word_bits = Word::BITS; let la_bits = la_words as u64 * word_bits as u64; let lb_bits = lb_words as u64 * word_bits as u64; - let prod_2 = (PRIMES[0].p as u128) * (PRIMES[1].p as u128); + let prod_2 = (MODULI[0] as u128) * (MODULI[1] as u128); for &b_pack in B_PACK_CANDIDATES { let coeffs_a = (la_bits + b_pack as u64 - 1) / b_pack as u64; @@ -62,7 +42,7 @@ pub fn select_params(la_words: usize, lb_words: usize) -> (u32, usize, usize) { } // Compute max coefficient value, guarding against u128 overflow for - // b_pack = 64 where (2^64−1)^2 ≈ 2^128 and n/2 can push it past 2^128. + // b_pack = 64 where (2^64−1)² ≈ 2^128 and n/2 can push it past 2^128. let coeff_max = (1u128 << b_pack) - 1; let max_coeff = coeff_max .checked_mul(coeff_max) @@ -70,14 +50,15 @@ pub fn select_params(la_words: usize, lb_words: usize) -> (u32, usize, usize) { let k_eff = match max_coeff { Some(mc) if mc < prod_2 => 2, - // Overflow or exceeds two-prime product → need all three primes. - // Three-prime product ≈ 2^192 ≫ any max_coeff we can encounter. _ => K, }; return (b_pack, n, k_eff); } - unreachable!("b_pack = 16 always passes the headroom check") + unreachable!( + "b_pack = {} always passes the headroom check", + B_PACK_CANDIDATES.last().unwrap() + ) } /// Estimate bit length from a word slice (excludes leading zeros). @@ -99,20 +80,27 @@ fn coeff_count(bit_len: u64, b_pack: u32) -> usize { /// Worst-case scratch memory bound. pub fn memory_requirement_up_to(total_len: usize, _smaller_len: usize) -> Layout { + use crate::arch::ntt::Lane; + let word_bits = Word::BITS; let max_coeffs = (total_len as u64 * word_bits as u64 + B_PACK_MIN as u64 - 1) / B_PACK_MIN as u64; let n_max = ((max_coeffs + 1) as usize).next_power_of_two().max(2); - let lanes_u64 = 2 * n_max; // a_lane + b_lane - let residues_u64 = K * n_max; // per-prime inverse results - let twiddles_u64 = n_max; // fwd + inv twiddle tables (n_max/2 each, reused) - let product_u64 = total_len; + let lanes = 2 * n_max; + let residues = K * n_max; + let twiddles = n_max; + let product = total_len; + let lane_bytes = mem::size_of::(); let u64_bytes = 8usize; - let word_bytes = core::mem::size_of::(); - let factor = (u64_bytes + word_bytes - 1) / word_bytes; - let total_words = product_u64 + (lanes_u64 + residues_u64 + twiddles_u64) * factor; + let word_bytes = mem::size_of::(); + + let lanes_words = lanes * lane_bytes / word_bytes; + let residues_words = residues * lane_bytes / word_bytes; + let twiddles_words = twiddles * lane_bytes / word_bytes; + let product_words = product * u64_bytes / word_bytes; + let total_words = product_words + lanes_words + residues_words + twiddles_words; memory::array_layout::(total_words) } @@ -158,6 +146,8 @@ fn add_signed_mul_impl( b: &[Word], memory: &mut Memory, ) -> SignedWord { + use crate::arch::ntt::Lane; + let la = a.len(); let lb = b.len(); @@ -176,145 +166,237 @@ fn add_signed_mul_impl( let coeffs_b = coeff_count(lb_bits, b_pack); let output_coeffs = coeffs_a + coeffs_b - 1; - // Per-prime CRT reducers (no allocation) - let r0 = FixedTrinomialSolinas64::<64, 32, 1>::new(&PRIMES[0].p); - let r1 = FixedTrinomialSolinas64::<64, 34, 1>::new(&PRIMES[1].p); - let r2 = FixedTrinomialSolinas64::<64, 40, 1>::new(&PRIMES[2].p); - let crt_reducers: [&dyn ModOps; 3] = [&r0, &r1, &r2]; - // ---- Memory carve (longest-lived first) ---- - // 1. Product buffer + // 1. Product buffer (always u64) let prod_len = la + lb; let (prod, mut mem) = memory.allocate_slice_fill::(prod_len, 0); // 2. Residue storage (per-prime inverse results) let residues_len = k_eff * nn; - let (residues, mut mem) = mem.allocate_slice_fill::(residues_len, 0); + let (residues, mut mem) = mem.allocate_slice_fill::(residues_len, 0); // 3. Lane buffers (reused across primes) - let (a_lane, mut mem) = mem.allocate_slice_fill::(nn, 0); - let (b_lane, mut mem) = mem.allocate_slice_fill::(nn, 0); + let (a_lane, mut mem) = mem.allocate_slice_fill::(nn, 0); + let (b_lane, mut mem) = mem.allocate_slice_fill::(nn, 0); // 4. Twiddle tables (fwd + inv, reused per prime) - let (fwd_twiddles, mut mem) = mem.allocate_slice_fill::(nn / 2, 0); - let (inv_twiddles, _) = mem.allocate_slice_fill::(nn / 2, 0); + let (fwd_twiddles, mut mem) = mem.allocate_slice_fill::(nn / 2, 0); + let (inv_twiddles, _) = mem.allocate_slice_fill::(nn / 2, 0); // ---- Per-prime transforms (const-generic dispatch) ---- - for (pi, prime) in PRIMES[..k_eff].iter().enumerate() { + for pi in 0..k_eff { let mut ctx = TransformCtx { a_lane, b_lane, fwd_twiddles, inv_twiddles, - p: prime.p, - omega_2_32: prime.omega_2_32, + p: MODULI[pi], + omega_max: OMEGA_MAX[pi], nn, b_pack, residues, pi, }; - match prime.b { - 32 => process_prime::<32>(a, b, &mut ctx), - 34 => process_prime::<34>(a, b, &mut ctx), - 40 => process_prime::<40>(a, b, &mut ctx), + match pi { + 0 => process_prime::<0>(a, b, &mut ctx), + 1 => process_prime::<1>(a, b, &mut ctx), + 2 => process_prime::<2>(a, b, &mut ctx), _ => unreachable!(), } } // ---- CRT per coefficient + accumulate ---- + // Extract prime constants as both u64 and u32 so the Lane-size + // dispatch below type-checks correctly in both branches. + // The dead branch (wrong width) is eliminated by the compiler. + let primes_u64: [u64; K] = [MODULI[0], MODULI[1], MODULI[2]]; + let crt_inv_u64: [[u64; K]; K] = { + use crate::arch::ntt::CRT_INV_IJ; + let mut m = [[0u64; K]; K]; + for i in 0..K { + for j in 0..K { + m[i][j] = CRT_INV_IJ[i][j]; + } + } + m + }; + let primes_u32: [u32; K] = [ + MODULI[0] as u32, + MODULI[1] as u32, + MODULI[2] as u32, + ]; + let crt_inv_u32: [[u32; K]; K] = { + let mut m = [[0u32; K]; K]; + for i in 0..K { + for j in 0..K { + m[i][j] = crt_inv_u64[i][j] as u32; + } + } + m + }; + + #[allow(clippy::unnecessary_cast)] + if mem::size_of::() == 8 { + // SAFETY: mem::size_of::() == 8, so Lane = u64 and + // residues is backed by u64 elements. + let residues_u64: &[u64] = + unsafe { core::slice::from_raw_parts(residues.as_ptr() as *const u64, residues.len()) }; + do_crt_u64( + prod, + residues_u64, + k_eff, + nn, + output_coeffs, + b_pack, + &primes_u64, + &crt_inv_u64, + ); + } else { + #[allow(clippy::unnecessary_cast)] + // SAFETY: mem::size_of::() != 8, so Lane = u32 and + // residues is backed by u32 elements. + let residues_u32: &[u32] = + unsafe { core::slice::from_raw_parts(residues.as_ptr() as *const u32, residues.len()) }; + do_crt_u32( + prod, + residues_u32, + k_eff, + nn, + output_coeffs, + b_pack, + &primes_u32, + &crt_inv_u32, + ); + } + + // ---- Fold product into c with sign ---- let output_words = la + lb; + fold_prod_into_c(c, sign, prod, output_words) +} + +/// CRT + accumulate for 64-bit lanes (U192 accumulator). +#[allow(clippy::too_many_arguments)] +#[inline(never)] +fn do_crt_u64( + prod: &mut [u64], + residues: &[u64], + k_eff: usize, + nn: usize, + output_coeffs: usize, + b_pack: u32, + primes: &[u64; K], + crt_inv: &[[u64; K]; K], +) { + use crate::mul::ntt::crt::U192; + for k in 0..output_coeffs { let mut coeff_residues = [0u64; 3]; #[allow(clippy::needless_range_loop)] for pi in 0..k_eff { coeff_residues[pi] = residues[pi * nn + k]; } - let crt_val = crt::garner_combine(&coeff_residues[..k_eff], &crt_reducers[..k_eff]); - add_shifted_to_prod(prod, &crt_val, k, b_pack); + let crt_val = + garner_combine::(&coeff_residues[..k_eff], crt_inv, primes); + add_shifted_to_prod(prod, crt_val.as_u64_slice(), crt_val.len_words(), k, b_pack); } +} - // ---- Fold product into c with sign ---- - assert_eq!( - core::mem::size_of::(), - core::mem::size_of::(), - "NTT requires 64-bit Word" - ); - // SAFETY: Word and u64 have the same size (asserted above) and - // prod is allocated with u64 alignment, compatible with Word. - let prod_words: &[Word] = - unsafe { core::slice::from_raw_parts(prod.as_ptr() as *const Word, output_words) }; - match sign { - Positive => add::add_signed_in_place(c, Positive, prod_words), - Negative => add::add_signed_in_place(c, Negative, prod_words), +/// CRT + accumulate for 32-bit lanes (U96 accumulator). +#[allow(clippy::too_many_arguments)] +#[inline(never)] +fn do_crt_u32( + prod: &mut [u64], + residues: &[u32], + k_eff: usize, + nn: usize, + output_coeffs: usize, + b_pack: u32, + primes: &[u32; K], + crt_inv: &[[u32; K]; K], +) { + use crate::mul::ntt::crt::U96; + + for k in 0..output_coeffs { + let mut coeff_residues = [0u32; 3]; + #[allow(clippy::needless_range_loop)] + for pi in 0..k_eff { + coeff_residues[pi] = residues[pi * nn + k]; + } + let crt_val = + garner_combine::(&coeff_residues[..k_eff], crt_inv, primes); + add_shifted_to_prod(prod, crt_val.as_u64_slice(), crt_val.len_words(), k, b_pack); } } /// Scratch buffers and parameters for one prime's NTT pipeline. struct TransformCtx<'a> { - a_lane: &'a mut [u64], - b_lane: &'a mut [u64], - fwd_twiddles: &'a mut [u64], - inv_twiddles: &'a mut [u64], - p: u64, - omega_2_32: u64, + a_lane: &'a mut [crate::arch::ntt::Lane], + b_lane: &'a mut [crate::arch::ntt::Lane], + fwd_twiddles: &'a mut [crate::arch::ntt::Lane], + inv_twiddles: &'a mut [crate::arch::ntt::Lane], + p: crate::arch::ntt::Lane, + omega_max: crate::arch::ntt::Lane, nn: usize, b_pack: u32, - residues: &'a mut [u64], + residues: &'a mut [crate::arch::ntt::Lane], pi: usize, } -/// Per-prime NTT pipeline, monomorphized for a specific `B`. +/// Per-prime NTT pipeline, monomorphized for a specific prime index `PI`. #[inline(never)] -fn process_prime(a: &[Word], b: &[Word], ctx: &mut TransformCtx<'_>) { +fn process_prime( + a: &[Word], + b: &[Word], + ctx: &mut TransformCtx<'_>, +) { + use crate::arch::ntt::{to_monty, from_monty}; + pack::pack(ctx.a_lane, a, ctx.b_pack, ctx.nn); pack::pack(ctx.b_lane, b, ctx.b_pack, ctx.nn); - // For b_pack = 64 coefficients may reach 2^64−1, which can exceed p. - // Reduce each coefficient mod p (one conditional subtract suffices: - // c < 2^64 < 2p for all three primes). - if ctx.b_pack >= 64 { - for c in ctx.a_lane[..ctx.nn].iter_mut() { - if *c >= ctx.p { - *c -= ctx.p; - } - } - for c in ctx.b_lane[..ctx.nn].iter_mut() { - if *c >= ctx.p { - *c -= ctx.p; - } - } + // Convert standard-form coefficients to Montgomery form. + // transform() handles any value in [0, 2^BITS), no pre-reduction needed. + for c in ctx.a_lane[..ctx.nn].iter_mut() { + *c = to_monty::(*c); + } + for c in ctx.b_lane[..ctx.nn].iter_mut() { + *c = to_monty::(*c); } - transform::precompute_twiddles::(ctx.fwd_twiddles, ctx.nn, ctx.p, ctx.omega_2_32, false); - transform::precompute_twiddles::(ctx.inv_twiddles, ctx.nn, ctx.p, ctx.omega_2_32, true); + transform::precompute_twiddles::(ctx.fwd_twiddles, ctx.nn, ctx.p, ctx.omega_max, false); + transform::precompute_twiddles::(ctx.inv_twiddles, ctx.nn, ctx.p, ctx.omega_max, true); transform::bit_reverse(ctx.a_lane); transform::bit_reverse(ctx.b_lane); - transform::forward::(ctx.a_lane, ctx.fwd_twiddles, ctx.p); - transform::forward::(ctx.b_lane, ctx.fwd_twiddles, ctx.p); - transform::pointwise_mul::(ctx.a_lane, ctx.b_lane); - transform::inverse::(ctx.a_lane, ctx.inv_twiddles, ctx.p); + transform::forward::(ctx.a_lane, ctx.fwd_twiddles); + transform::forward::(ctx.b_lane, ctx.fwd_twiddles); + transform::pointwise_mul::(ctx.a_lane, ctx.b_lane); + transform::inverse::(ctx.a_lane, ctx.inv_twiddles, ctx.p); + + // Convert residues back from Montgomery to standard form. + for c in ctx.a_lane[..ctx.nn].iter_mut() { + *c = from_monty::(*c); + } let offset = ctx.pi * ctx.nn; ctx.residues[offset..offset + ctx.nn].copy_from_slice(ctx.a_lane); } -/// Add a CRT value to `prod`, shifted left by `k * b_pack` bits. -fn add_shifted_to_prod(prod: &mut [u64], val: &U192, k: usize, b_pack: u32) { - let count = val.len_words() as usize; +/// Add a CRT value (as u64 words) to `prod`, shifted left by `k * b_pack` bits. +fn add_shifted_to_prod(prod: &mut [u64], words: &[u64], count: u32, k: usize, b_pack: u32) { let shift_bits = (k as u32).wrapping_mul(b_pack); let word_idx = (shift_bits / 64) as usize; let bit_shift = shift_bits % 64; let mut carry: u64 = 0; #[allow(clippy::needless_range_loop)] - for vi in 0..count { + for vi in 0..(count as usize) { let idx = word_idx + vi; if idx >= prod.len() { return; } - let v = val.0[vi]; + let v = words[vi]; let v128 = v as u128; if bit_shift == 0 { @@ -340,7 +422,7 @@ fn add_shifted_to_prod(prod: &mut [u64], val: &U192, k: usize, b_pack: u32) { } } - let mut idx = word_idx + count; + let mut idx = word_idx + count as usize; while carry != 0 && idx < prod.len() { let (r, c) = prod[idx].overflowing_add(carry); prod[idx] = r; @@ -349,6 +431,71 @@ fn add_shifted_to_prod(prod: &mut [u64], val: &U192, k: usize, b_pack: u32) { } } +/// Fold the u64 product buffer into the Word output array `c`. +/// +/// For 64-bit Word targets, `u64` == `Word`, so a direct transmute suffices. +/// For 32-bit Word targets, each u64 is split into two u32 words with carry. +fn fold_prod_into_c(c: &mut [Word], sign: Sign, prod: &[u64], output_words: usize) -> SignedWord { + if mem::size_of::() == 8 { + // 64-bit Word: direct transmute + assert_eq!( + mem::size_of::(), + mem::size_of::(), + "NTT requires 64-bit Word" + ); + // SAFETY: Word and u64 have the same size (asserted above) and + // prod is allocated with u64 alignment, compatible with Word. + let prod_words: &[Word] = + unsafe { core::slice::from_raw_parts(prod.as_ptr() as *const Word, output_words) }; + match sign { + Positive => add::add_signed_in_place(c, Positive, prod_words), + Negative => add::add_signed_in_place(c, Negative, prod_words), + } + } else { + // 32-bit Word: split each u64 into two u32 words. + // For 64-bit Word targets this branch is dead (eliminated by the compiler) + // but must still type-check. + let mut carry: u32 = 0; + let double_output_words = output_words.min(prod.len() * 2); + for i in 0..double_output_words { + let prod_word = prod[i / 2]; + let lo_word = if i % 2 == 0 { + (prod_word as u32).wrapping_add(carry) + } else { + ((prod_word >> 32) as u32).wrapping_add(carry) + }; + // Carry out of the low-word addition + if i % 2 == 0 { + let lo_before = prod_word as u32; + carry = (prod_word >> 32) as u32; + if lo_word < lo_before { + carry = carry.wrapping_add(1); + } + } + if i < c.len() { + let c_val = c[i] as u32; + match sign { + Positive => { + let (sum, c_out) = c_val.overflowing_add(lo_word); + c[i] = sum as Word; + carry = carry.wrapping_add(c_out as u32); + } + Negative => { + let (diff, b) = c_val.overflowing_sub(lo_word); + c[i] = diff as Word; + carry = carry.wrapping_add(b as u32); + } + } + } + } + if sign == Positive { + carry as SignedWord + } else { + -(carry as SignedWord) + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -360,56 +507,37 @@ mod tests { #[test] fn test_select_params_small() { let (b_pack, n, k_eff) = select_params(10, 10); - assert_eq!(b_pack, 64); + // On 64-bit: B_PACK_CANDIDATES[0] = 64, needs K_eff = 3. + // On 32-bit: B_PACK_CANDIDATES[0] = 32, likely K_eff = 2. + assert!(b_pack >= 32); assert!(n >= 2 && n.is_power_of_two()); - // b_pack = 64 needs K_eff = 3 primes. - assert_eq!(k_eff, K); + assert!((2..=K).contains(&k_eff)); } #[test] fn test_select_params_large() { - let (b_pack, n, k_eff) = select_params(THRESHOLD_NTT, THRESHOLD_NTT); - assert_eq!(b_pack, 64); + let (b_pack, n, _k_eff) = select_params(THRESHOLD_NTT, THRESHOLD_NTT); + assert!(b_pack >= 32); assert!(n.is_power_of_two()); - assert_eq!(k_eff, K); - let coeffs_a = (THRESHOLD_NTT * Word::BITS as usize + 63) / 64; + let coeffs_a = + (THRESHOLD_NTT * Word::BITS as usize + b_pack as usize - 1) / b_pack as usize; let coeffs_b = coeffs_a; let min_n = (coeffs_a + coeffs_b).next_power_of_two().max(2); assert!(n >= min_n, "n={n} < min_n={min_n}"); } - #[test] - fn test_headroom_holds() { - let la = THRESHOLD_NTT; - let lb = THRESHOLD_NTT; - let (b_pack, n, _k_eff) = select_params(la, lb); - // For b_pack = 64 the product overflows u128; checked_mul in - // select_params handles this and falls back to K_eff = 3. - // Three-prime product ≈ 2^192 ≫ max_coeff for n ≤ 2^32. - let coeff_max = (1u128 << b_pack) - 1; - let overflow = coeff_max - .checked_mul(coeff_max) - .and_then(|sq| (n as u128 / 2).checked_mul(sq)) - .is_none(); - assert!(overflow || _k_eff == 2, "K_eff=2 only when max_coeff fits in u128"); - assert_eq!(b_pack, 64); - } - #[test] fn test_bit_len() { assert_eq!(bit_len(&[]), 0); assert_eq!(bit_len(&[0]), 0); assert_eq!(bit_len(&[1]), 1); - assert_eq!(bit_len(&[0, 1]), 65); - assert_eq!(bit_len(&[0xFF, 0]), 8); } #[test] fn test_ntt_multiply_one_word() { - // Simplest case: single-word operands let a: Vec = vec![3]; let b: Vec = vec![5]; - let mut c = vec![0u64; 2]; + let mut c = vec![0u64 as Word; 2]; let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); @@ -420,28 +548,44 @@ mod tests { } #[test] - fn test_ntt_multiply_two_words() { - // Two-word operands - let a: Vec = vec![Word::MAX, 1]; // 2^64 + (2^64-1) - let b: Vec = vec![2, 0]; // 2 - let expected = schoolbook_mul(&a, &b); - let mut c = vec![0u64; a.len() + b.len()]; + fn test_ntt_zero_operand() { + let a = vec![0xDEADu64 as Word; 30]; + let b = vec![0u64 as Word; 30]; + let mut c = vec![0u64 as Word; a.len() + b.len()]; let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); - add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); - assert_eq!(&c[..], &expected[..], "two-word mismatch"); + let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + assert_eq!(carry, 0); + assert!(c.iter().all(|&w| w == 0)); + } + + #[test] + fn test_ntt_sign_negative() { + let a: Vec = (0..30).map(|i| (i as Word + 1) * 100).collect(); + let b: Vec = (0..30).map(|i| (i as Word + 1) * 200).collect(); + + let mut c = vec![0u64 as Word; a.len() + b.len()]; + let layout = memory_requirement_up_to(c.len(), b.len()); + let mut alloc = crate::memory::MemoryAllocation::new(layout); + let mut memory = alloc.memory(); + let _ = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + + let layout2 = memory_requirement_up_to(c.len(), b.len()); + let mut alloc2 = crate::memory::MemoryAllocation::new(layout2); + let mut memory2 = alloc2.memory(); + let _ = add_signed_mul_impl(&mut c, Negative, &a, &b, &mut memory2); + + assert!(c.iter().all(|&w| w == 0)); } - const NTT_TEST_LEN: usize = 1024; + const NTT_TEST_LEN: usize = 512; #[test] fn test_ntt_multiply_small() { - // Smoke test: NTT multiply with operands large enough to exercise - // the full pipeline (pack, forward, pointwise, inverse, CRT, accumulate). - let a: Vec = vec![0xDEADBEEFu64; NTT_TEST_LEN]; - let b: Vec = vec![0xCAFEBABEu64; NTT_TEST_LEN]; - let mut c = vec![0u64; a.len() + b.len()]; + let a: Vec = vec![0xDEADBEEFu64 as Word; NTT_TEST_LEN]; + let b: Vec = vec![0xCAFEBABEu64 as Word; NTT_TEST_LEN]; + let mut c = vec![0u64 as Word; a.len() + b.len()]; let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); @@ -453,45 +597,41 @@ mod tests { /// Naive schoolbook multiplication for comparison. fn schoolbook_mul(a: &[Word], b: &[Word]) -> Vec { - let mut c = vec![0u64; a.len() + b.len()]; + let mut c = vec![0u64 as Word; a.len() + b.len()]; for (i, &ai) in a.iter().enumerate() { let mut carry: u128 = 0; for (j, &bj) in b.iter().enumerate() { let idx = i + j; let prod = (ai as u128) * (bj as u128) + (c[idx] as u128) + carry; - c[idx] = prod as u64; - carry = prod >> 64; + c[idx] = prod as Word; + carry = prod >> Word::BITS; } - // Propagate carry into higher words let mut k = i + b.len(); while carry != 0 { let sum = (c[k] as u128) + carry; - c[k] = sum as u64; - carry = sum >> 64; + c[k] = sum as Word; + carry = sum >> Word::BITS; k += 1; } } c } - /// Test NTT against schoolbook with moderate operand sizes. fn run_ntt_vs_schoolbook(la: usize, lb: usize) { - // Generate deterministic test data let a: Vec = (0..la) - .map(|i| (i as u64 + 1).wrapping_mul(0x9E3779B97F4A7C15)) + .map(|i| (i as Word + 1).wrapping_mul(0x9E3779B97F4A7C15u64 as Word)) .collect(); let b: Vec = (0..lb) - .map(|i| (i as u64 + 1).wrapping_mul(0xC6A4A7935BD1E995)) + .map(|i| (i as Word + 1).wrapping_mul(0xC6A4A7935BD1E995u64 as Word)) .collect(); let expected = schoolbook_mul(&a, &b); - let mut c = vec![0u64; a.len() + b.len()]; + let mut c = vec![0u64 as Word; a.len() + b.len()]; let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); assert_eq!(carry, 0, "carry should be 0"); - assert_eq!(&c[..], &expected[..], "NTT mismatch: la={la}, lb={lb}"); } @@ -511,7 +651,6 @@ mod tests { #[test] fn test_ntt_vs_schoolbook_asymmetric() { - // Very asymmetric sizes for &(la, lb) in &[(200, 30), (150, 20)] { run_ntt_vs_schoolbook(la, lb); } @@ -519,13 +658,12 @@ mod tests { #[test] fn test_ntt_all_ones() { - // All-ones operands stress the carry chain. for &len in &[20, 50] { let a = vec![Word::MAX; len]; let b = vec![Word::MAX; len]; let expected = schoolbook_mul(&a, &b); - let mut c = vec![0u64; a.len() + b.len()]; + let mut c = vec![0u64 as Word; a.len() + b.len()]; let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); @@ -534,55 +672,17 @@ mod tests { } } - #[test] - fn test_ntt_zero_operand() { - let a = vec![0xDEADu64; 30]; - let b = vec![0u64; 30]; - let mut c = vec![0u64; a.len() + b.len()]; - let layout = memory_requirement_up_to(c.len(), b.len()); - let mut alloc = crate::memory::MemoryAllocation::new(layout); - let mut memory = alloc.memory(); - let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); - assert_eq!(carry, 0); - assert!(c.iter().all(|&w| w == 0), "zero operand should give zero product"); - } - - #[test] - fn test_ntt_sign_negative() { - // Test that Negative sign works (c -= a * b) - let a: Vec = (0..30).map(|i| (i as u64 + 1) * 100).collect(); - let b: Vec = (0..30).map(|i| (i as u64 + 1) * 200).collect(); - let _expected = schoolbook_mul(&a, &b); - - // First add: c += a * b - let mut c = vec![0u64; a.len() + b.len()]; - let layout = memory_requirement_up_to(c.len(), b.len()); - let mut alloc = crate::memory::MemoryAllocation::new(layout); - let mut memory = alloc.memory(); - let _ = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); - - // Then subtract: c -= a * b - let layout2 = memory_requirement_up_to(c.len(), b.len()); - let mut alloc2 = crate::memory::MemoryAllocation::new(layout2); - let mut memory2 = alloc2.memory(); - let _ = add_signed_mul_impl(&mut c, Negative, &a, &b, &mut memory2); - - // Result should be zero - assert!(c.iter().all(|&w| w == 0), "add then subtract should give zero"); - } - #[test] fn test_ntt_high_low_zero_limbs() { - // Operands with leading/trailing zero limbs - let mut a = vec![0u64; 80]; - let mut b = vec![0u64; 80]; + let mut a = vec![0u64 as Word; 80]; + let mut b = vec![0u64 as Word; 80]; for i in 20..60 { - a[i] = (i as u64 + 1).wrapping_mul(0xDEADBEEF); - b[i] = (i as u64 + 1).wrapping_mul(0xCAFEBABE); + a[i] = (i as Word + 1).wrapping_mul(0xDEADBEEF); + b[i] = (i as Word + 1).wrapping_mul(0xCAFEBABE); } let expected = schoolbook_mul(&a, &b); - let mut c = vec![0u64; a.len() + b.len()]; + let mut c = vec![0u64 as Word; a.len() + b.len()]; let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); diff --git a/integer/src/mul/ntt/pack.rs b/integer/src/mul/ntt/pack.rs index 5efdd08e..39d570ae 100644 --- a/integer/src/mul/ntt/pack.rs +++ b/integer/src/mul/ntt/pack.rs @@ -1,24 +1,23 @@ //! Bit-level packing / unpacking of `b`-bit coefficients. -#![allow(clippy::unnecessary_cast)] +use crate::arch::ntt::Lane; use crate::arch::word::Word; -// TODO: shall we specialize the packing function? Since we only have three b_pack options. - /// Pack a big integer (given as `&[Word]`, little-endian) into `out`, /// producing `n` coefficients of `b_pack` bits each, zero-padded. /// /// Each coefficient `c_i` satisfies `0 ≤ c_i < 2^{b_pack}`. /// Panics if `out.len() < n`. -pub fn pack(out: &mut [u64], words: &[Word], b_pack: u32, n: usize) { +pub fn pack(out: &mut [Lane], words: &[Word], b_pack: u32, n: usize) { assert!(out.len() >= n); // Fast path: one coefficient per word, no bit shifting needed. if b_pack == Word::BITS { let len = words.len().min(n); - // SAFETY: NTT path requires Word = u64 (asserted by caller). - let words_u64 = unsafe { &*(words as *const [Word] as *const [u64]) }; - out[..len].copy_from_slice(&words_u64[..len]); + #[allow(clippy::unnecessary_cast)] + // SAFETY: NTT path requires Word and Lane have the same size. + let words_lane = unsafe { &*(words as *const [Word] as *const [Lane]) }; + out[..len].copy_from_slice(&words_lane[..len]); out[len..n].fill(0); return; } @@ -39,7 +38,7 @@ pub fn pack(out: &mut [u64], words: &[Word], b_pack: u32, n: usize) { } if bit_offset + b_pack <= word_bits { - *coeff = (words[word_idx] >> bit_offset) & mask; + *coeff = ((words[word_idx] >> bit_offset) & mask) as Lane; bit_offset += b_pack; if bit_offset == word_bits { bit_offset = 0; @@ -48,12 +47,14 @@ pub fn pack(out: &mut [u64], words: &[Word], b_pack: u32, n: usize) { } else { let bits_first = word_bits - bit_offset; let bits_second = b_pack - bits_first; - let mut val = (words[word_idx] >> bit_offset) & ((1u64 << bits_first) - 1); + let mut val = + (words[word_idx] >> bit_offset) & ((1u64 << bits_first) - 1); word_idx += 1; if word_idx < words.len() { - val |= (words[word_idx] & ((1u64 << bits_second) - 1)) << bits_first; + val |= (words[word_idx] & ((1u64 << bits_second) - 1)) + << bits_first; } - *coeff = val; + *coeff = val as Lane; bit_offset = bits_second; } } @@ -73,8 +74,8 @@ mod tests { /// Each coefficient `c_k` contributes `c_k << (k * b_pack)` bits to the /// output. `output` must have capacity for `c.len()` coefficients plus any /// carry overflow. - fn unpack_accumulate(output: &mut [Word], coeffs: &[u64], b_pack: u32, output_len: usize) { - let word_bits = Word::BITS as u32; + fn unpack_accumulate(output: &mut [Word], coeffs: &[Lane], b_pack: u32, output_len: usize) { + let word_bits = Word::BITS; for (k, &coeff) in coeffs.iter().enumerate().take(output_len) { if coeff == 0 { @@ -84,12 +85,16 @@ mod tests { let word_idx = (shift_bits / word_bits) as usize; let bit_shift = shift_bits % word_bits; - let lo = coeff as u64; + let lo = coeff; let mut carry: Word; let mut idx = word_idx; if bit_shift == 0 { - let (sum, c) = output.get(idx).copied().unwrap_or(0).overflowing_add(lo); + let (sum, c) = output + .get(idx) + .copied() + .unwrap_or(0) + .overflowing_add(lo); carry = Word::from(c); if idx < output.len() { output[idx] = sum; @@ -97,11 +102,7 @@ mod tests { idx += 1; } else { let lo_part = lo << bit_shift; - let hi_part = if bit_shift > 0 { - lo >> (64 - bit_shift) - } else { - 0 - }; + let hi_part = if bit_shift > 0 { lo >> (64 - bit_shift) } else { 0 }; let (sum, c1) = output .get(idx) @@ -143,20 +144,20 @@ mod tests { let coeffs_per_word = (Word::BITS / b_pack) as usize; let n = test_words.len() * coeffs_per_word; - let mut packed = vec![0u64; n]; + let mut packed = vec![0u64 as Lane; n]; pack(&mut packed, &test_words, b_pack, n); let output_len = test_words.len() + 1; - let mut output = vec![0u64; output_len]; + let mut output = vec![0u64 as Word; output_len]; unpack_accumulate(&mut output, &packed, b_pack, n); assert_eq!(&output[..test_words.len()], &test_words[..]); } #[test] fn test_pack_zero_pads() { - let words = vec![0xFFFFu64]; + let words = vec![0xFFFFu64 as Word]; let n = 32; - let mut packed = vec![0u64; n]; + let mut packed = vec![0u64 as Lane; n]; pack(&mut packed, &words, 16, n); assert_eq!(packed[0], 0xFFFF); for &c in packed.iter().skip(1) { @@ -166,14 +167,14 @@ mod tests { #[test] fn test_pack_empty_input() { - let mut packed = vec![0u64; 8]; + let mut packed = vec![0u64 as Lane; 8]; pack(&mut packed, &[], 16, 8); - assert_eq!(packed, vec![0u64; 8]); + assert_eq!(packed, vec![0u64 as Lane; 8]); } #[test] fn test_unpack_single_coeff() { - let mut output = vec![0u64; 2]; + let mut output = vec![0u64 as Word; 2]; unpack_accumulate(&mut output, &[0xABCD], 16, 1); assert_eq!(output[0], 0xABCD); assert_eq!(output[1], 0); @@ -182,7 +183,7 @@ mod tests { #[test] fn test_unpack_carry_propagation() { // Coefficient at k=4 (shift by 64 bits = 1 word) + carry - let mut output = vec![0u64; 3]; + let mut output = vec![0u64 as Word; 3]; unpack_accumulate(&mut output, &[0, 0, 0, 0, 1], 16, 5); assert_eq!(output[0], 0); assert_eq!(output[1], 1); diff --git a/integer/src/mul/ntt/primes.rs b/integer/src/mul/ntt/primes.rs deleted file mode 100644 index fd2c2fa0..00000000 --- a/integer/src/mul/ntt/primes.rs +++ /dev/null @@ -1,257 +0,0 @@ -//! NTT-friendly primes of the form `2^64 - 2^b + 1`. -//! -//! All three support shift-based reduction via the identity `2^64 ≡ 2^b - 1 (mod p)`. - -/// The number of primes in the fixed array. -pub const K: usize = 3; - -/// Precomputed data for one NTT-friendly prime. -#[derive(Clone, Copy, Debug)] -pub struct NttPrime { - /// The prime value `p = 2^64 - 2^b + 1`. - pub p: u64, - /// The exponent `b` in the Solinas form. - pub b: u32, - /// The exponent of 2 in `p - 1`: `v2(p - 1)`. - #[cfg(test)] - pub v2: u32, - /// A primitive root modulo `p` that generates the full multiplicative group. - #[cfg(test)] - pub g: u64, - /// A primitive `2^32`-th root of unity: `ω = g^{(p-1) / 2^32} mod p`. - pub omega_2_32: u64, -} - -/// The three chosen primes, all of the form `2^64 - 2^b + 1` with `b ∈ {32, 34, 40}`. -/// -/// | name | `b` | `p` | `v2(p-1)` | gen `g` | `2^32`-th root ω | -/// |------|-----|---------------------|-----------|---------|----------------------------| -/// | GL | 32 | `0xFFFFFFFF00000001` | 32 | 7 | `1753635133440165772` | -/// | P1 | 34 | `0xFFFFFFFC00000001` | 34 | 5 | `11315553352654630047` | -/// | P2 | 40 | `0xFFFFFF0000000001` | 40 | 19 | `551857376737322389` | -#[cfg(not(test))] -pub const PRIMES: [NttPrime; K] = [ - NttPrime { - p: 0xFFFFFFFF00000001, - b: 32, - omega_2_32: 1753635133440165772, - }, - NttPrime { - p: 0xFFFFFFFC00000001, - b: 34, - omega_2_32: 11315553352654630047, - }, - NttPrime { - p: 0xFFFFFF0000000001, - b: 40, - omega_2_32: 551857376737322389, - }, -]; - -#[cfg(test)] -pub const PRIMES: [NttPrime; K] = [ - NttPrime { - p: 0xFFFFFFFF00000001, - b: 32, - v2: 32, - g: 7, - omega_2_32: 1753635133440165772, - }, - NttPrime { - p: 0xFFFFFFFC00000001, - b: 34, - v2: 34, - g: 5, - omega_2_32: 11315553352654630047, - }, - NttPrime { - p: 0xFFFFFF0000000001, - b: 40, - v2: 40, - g: 19, - omega_2_32: 551857376737322389, - }, -]; - -/// Garner CRT constants: `inv(p_i mod p_j)` for i < j. -/// Computed offline via `pow(p_i % p_j, -1, p_j)`. -pub const CRT_INV_IJ: [[u64; 3]; 3] = [ - [0, 0xfffffffbaaaaaaad, 0xfffffefffefeff01], - [0, 0, 0xfffffefffefbefc1], - [0, 0, 0], -]; - -#[cfg(test)] -mod tests { - use super::*; - use num_modular::FixedTrinomialSolinas64; - - /// Deterministic Miller–Rabin for 64-bit integers with known bases. - /// Tests `n` against bases `[2, 325, 9375, 28178, 450775, 9780504, 1795265022]` - /// which together suffice for all `n < 2^64` (deterministic). - fn is_prime_u64(n: u64) -> bool { - if n < 2 { - return false; - } - if n % 2 == 0 { - return n == 2; - } - - // Write n-1 = d * 2^s - let d = (n - 1) >> (n - 1).trailing_zeros(); - let s = (n - 1).trailing_zeros(); - - let bases = [2u64, 325, 9375, 28178, 450775, 9780504, 1795265022]; - - 'next_base: for &a in &bases { - if a >= n { - continue; - } - let mut x = mod_pow_u64(a % n, d, n); - if x == 1 || x == n - 1 { - continue 'next_base; - } - for _ in 1..s { - x = ((x as u128 * x as u128) % (n as u128)) as u64; - if x == n - 1 { - continue 'next_base; - } - } - return false; - } - true - } - - fn mod_pow_u64(mut base: u64, mut exp: u64, modulus: u64) -> u64 { - let mut result = 1u64; - while exp > 0 { - if exp & 1 != 0 { - result = ((result as u128 * base as u128) % (modulus as u128)) as u64; - } - base = ((base as u128 * base as u128) % (modulus as u128)) as u64; - exp >>= 1; - } - result - } - - #[test] - fn verify_primes() { - for &NttPrime { - p, - b, - v2, - g, - omega_2_32, - } in &PRIMES - { - // 1. Correct form: p == 2^64 - 2^b + 1 - let expected_p = (1u128 << 64) - (1u128 << b) + 1; - assert!(expected_p < (1u128 << 64), "p must fit in 64 bits"); - assert_eq!(p as u128, expected_p, "p = 0x{p:X} does not match 2^64 - 2^{b} + 1"); - assert!(p > 0, "p must be positive"); - - // 2. Primality - assert!(is_prime_u64(p), "p = 0x{p:X} is not prime"); - - // 3. v2(p-1) is at least 32 - let actual_v2 = (p - 1).trailing_zeros(); - assert!(actual_v2 >= 32, "v2(p-1) = {actual_v2} < 32 for p = 0x{p:X}"); - assert_eq!(actual_v2, v2, "stored v2 mismatch for p = 0x{p:X}"); - - // 4. g generates the full multiplicative group mod p. - // g^((p-1)/2) mod p ≠ 1 (g is a quadratic non-residue) - let g_order_half = mod_pow_u64(g, (p - 1) / 2, p); - assert_ne!(g_order_half, 1, "g = {g} is a quadratic residue mod p = 0x{p:X}"); - - // g^(p-1) ≡ 1 - let g_full = mod_pow_u64(g, p - 1, p); - assert_eq!(g_full, 1, "g^(p-1) != 1 mod p = 0x{p:X}"); - - // 5. ω has exact order 2^32 - let mut omega_pow = omega_2_32; - for _ in 0..31 { - omega_pow = ((omega_pow as u128 * omega_pow as u128) % (p as u128)) as u64; - } - // After 31 squarings: ω^{2^31} mod p - // Should be -1 mod p (order is exactly 2^32) - assert_eq!(omega_pow, p - 1, "omega^(2^31) != -1 mod p = 0x{p:X}, order not 2^32"); - - // ω^{2^32} ≡ 1 - let omega_full = mod_pow_u64(omega_2_32, 1u64 << 32, p); - assert_eq!(omega_full, 1, "omega^(2^32) != 1 mod p = 0x{p:X}"); - - // 6. Reduction identity: 2^64 ≡ 2^b - 1 (mod p) - let two_64_mod_p = ((1u128 << 64) % (p as u128)) as u64; - let expected = (if b == 0 { - 0 - } else { - (1u64 << (b - 1)).wrapping_mul(2) - }) - 1; - assert_eq!(two_64_mod_p, expected, "2^64 mod p != 2^(b) - 1 for p = 0x{p:X}"); - } - } - - #[test] - fn test_reduction_identity_per_prime() { - // Verify reduction works for all three primes using the actual reducer types. - // GL: b=32 - { - type Reducer = FixedTrinomialSolinas64<64, 32, 1>; - let p = Reducer::MODULUS; - assert_eq!(p, PRIMES[0].p); - - // Test reduce_double - let v = (p as u128) * 3; // 3p → should reduce to 0 - let r = Reducer::reduce_double(v); - assert!(r < p); - assert_eq!((r as u128) % (p as u128), v % (p as u128)); - } - // P1: b=34 - { - type Reducer = FixedTrinomialSolinas64<64, 34, 1>; - let p = Reducer::MODULUS; - assert_eq!(p, PRIMES[1].p); - - let v = (p as u128) * 3; - let r = Reducer::reduce_double(v); - assert!(r < p); - assert_eq!((r as u128) % (p as u128), v % (p as u128)); - } - // P2: b=40 - { - type Reducer = FixedTrinomialSolinas64<64, 40, 1>; - let p = Reducer::MODULUS; - assert_eq!(p, PRIMES[2].p); - - let v = (p as u128) * 3; - let r = Reducer::reduce_double(v); - assert!(r < p); - assert_eq!((r as u128) % (p as u128), v % (p as u128)); - } - } - - #[test] - fn test_pow_inv_roundtrip() { - // pow and inv round-trip checks using the trait API - use num_modular::{ModularPow, ModularUnaryOps}; - - for &NttPrime { p, g, .. } in &PRIMES { - // inv(x)·x ≡ 1 - let x = 123456789u64; - let inv = x.invm(&p); - if let Some(inv) = inv { - let prod = ((x as u128 * inv as u128) % (p as u128)) as u64; - assert_eq!(prod, 1, "inv round-trip failed for p = 0x{p:X}"); - } - - // pow(g, p-1) ≡ 1 - let g_pow = g.powm(&(p - 1), &p); - assert_eq!(g_pow, 1, "g^(p-1) != 1 mod p = 0x{p:X}"); - - // inv(g) · g ≡ 1 - let g_inv = g.invm(&p).unwrap(); - let prod = ((g as u128 * g_inv as u128) % (p as u128)) as u64; - assert_eq!(prod, 1, "inv(g) round-trip failed for p = 0x{p:X}"); - } - } -} diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs index 4bb63e65..11d0fc61 100644 --- a/integer/src/mul/ntt/transform.rs +++ b/integer/src/mul/ntt/transform.rs @@ -1,59 +1,26 @@ -//! Iterative in-place radix-2 NTT over primes of the form `2^64 - 2^b + 1`. +//! Iterative in-place radix-2 NTT over Proth primes `K * 2^N + 1`. //! -//! All functions are const-generic over `B` (the Solinas exponent, one of -//! `{32, 34, 40}`) so the compiler monomorphizes each prime's hot path. -//! Modular arithmetic delegates to `num_modular::FixedTrinomialSolinas64`. - -use num_modular::{FixedTrinomialSolinas64, ModularPow, ModularUnaryOps, Reducer}; - -// ---- dispatch helpers (the match is optimized away since B is const) ---- - -#[inline] -fn mul_mod(a: u64, b_val: u64) -> u64 { - let prod = (a as u128) * (b_val as u128); - match B { - 32 => FixedTrinomialSolinas64::<64, 32, 1>::reduce_double(prod), - 34 => FixedTrinomialSolinas64::<64, 34, 1>::reduce_double(prod), - 40 => FixedTrinomialSolinas64::<64, 40, 1>::reduce_double(prod), - _ => unreachable!(), - } -} +//! All functions are const-generic over `PI` (the prime index `0..K`). +//! Modular arithmetic delegates to `crate::arch::ntt`. -#[inline] -fn add_mod(a: u64, b_val: u64, p: u64) -> u64 { - match B { - 32 => FixedTrinomialSolinas64::<64, 32, 1>::new(&p).add(&a, &b_val), - 34 => FixedTrinomialSolinas64::<64, 34, 1>::new(&p).add(&a, &b_val), - 40 => FixedTrinomialSolinas64::<64, 40, 1>::new(&p).add(&a, &b_val), - _ => unreachable!(), - } -} - -#[inline] -fn sub_mod(a: u64, b_val: u64, p: u64) -> u64 { - match B { - 32 => FixedTrinomialSolinas64::<64, 32, 1>::new(&p).sub(&a, &b_val), - 34 => FixedTrinomialSolinas64::<64, 34, 1>::new(&p).sub(&a, &b_val), - 40 => FixedTrinomialSolinas64::<64, 40, 1>::new(&p).sub(&a, &b_val), - _ => unreachable!(), - } -} +use crate::arch::ntt::{add_mod, mul_mod, sub_mod, to_monty, Lane, MAX_LOG_N}; +use num_modular::{ModularPow, ModularUnaryOps}; // ---- public API ---- -/// Fill `out[0..n/2]` with twiddle factors `omega_n^k`. +/// Fill `out[0..n/2]` with twiddle factors `omega_n^k` in Montgomery form. /// /// Panics if `out.len() < n / 2`. -pub fn precompute_twiddles( - out: &mut [u64], +pub fn precompute_twiddles( + out: &mut [Lane], n: usize, - p: u64, - omega_2_32: u64, + p: Lane, + omega_max: Lane, inverse: bool, ) { assert!(out.len() >= n / 2); - let shift = 32 - n.trailing_zeros(); - let omega_n = omega_2_32.powm(&(1u64 << shift), &p); + let shift = MAX_LOG_N - n.trailing_zeros(); + let omega_n = omega_max.powm(&((1u64 as Lane) << shift), &p); let base = if inverse { omega_n.invm(&p).expect("omega_n not invertible") @@ -61,14 +28,16 @@ pub fn precompute_twiddles( omega_n }; - out[0] = 1; + // Convert base and 1 to Montgomery form + let base_mont = to_monty::(base); + out[0] = to_monty::(1); for k in 1..(n / 2) { - out[k] = mul_mod::(out[k - 1], base); + out[k] = mul_mod::(out[k - 1], base_mont); } } /// Bit-reverse `a` in place. Length must be a power of two. -pub fn bit_reverse(a: &mut [u64]) { +pub fn bit_reverse(a: &mut [Lane]) { let n = a.len(); assert!(n.is_power_of_two()); let log_n = n.trailing_zeros(); @@ -81,28 +50,31 @@ pub fn bit_reverse(a: &mut [u64]) { } /// Forward NTT in place (decimation-in-time, radix-2). -pub fn forward(a: &mut [u64], twiddles: &[u64], p: u64) { - ntt_core::(a, twiddles, p); +pub fn forward(a: &mut [Lane], twiddles: &[Lane]) { + ntt_core::(a, twiddles); } /// Inverse NTT in place. /// -/// Computed as `bit_reverse → forward(ω^{-1}) → scale`, producing output +/// Computed as `bit_reverse → forward(ω⁻¹) → scale`, producing output /// in **natural order**. /// /// `twiddles` must have been precomputed with `inverse = true`. -pub fn inverse(a: &mut [u64], twiddles: &[u64], p: u64) { +pub fn inverse(a: &mut [Lane], twiddles: &[Lane], p: Lane) { let n = a.len(); bit_reverse(a); - ntt_core::(a, twiddles, p); - let n_inv = (n as u64).invm(&p).expect("n not invertible mod p"); + ntt_core::(a, twiddles); + let n_val = n as Lane; + let n_inv = n_val.invm(&p).expect("n not invertible mod p"); + // Convert n⁻¹ to Montgomery form so the result stays in Montgomery form. + let n_inv_mont = to_monty::(n_inv); for x in a.iter_mut() { - *x = mul_mod::(*x, n_inv); + *x = mul_mod::(*x, n_inv_mont); } } /// In-place radix-2 DIT NTT (Cooley–Tukey). -fn ntt_core(a: &mut [u64], twiddles: &[u64], p: u64) { +fn ntt_core(a: &mut [Lane], twiddles: &[Lane]) { let n = a.len(); debug_assert!(n.is_power_of_two() && twiddles.len() == n / 2); @@ -114,9 +86,9 @@ fn ntt_core(a: &mut [u64], twiddles: &[u64], p: u64) { for i in (0..n).step_by(sub_len) { for j in 0..half { let u = a[i + j]; - let v = mul_mod::(a[i + j + half], twiddles[j * step]); - a[i + j] = add_mod::(u, v, p); - a[i + j + half] = sub_mod::(u, v, p); + let v = mul_mod::(a[i + j + half], twiddles[j * step]); + a[i + j] = add_mod::(u, v); + a[i + j + half] = sub_mod::(u, v); } } @@ -125,51 +97,52 @@ fn ntt_core(a: &mut [u64], twiddles: &[u64], p: u64) { } /// Pointwise multiply of two transformed vectors in place. -pub fn pointwise_mul(a_hat: &mut [u64], b_hat: &[u64]) { +pub fn pointwise_mul(a_hat: &mut [Lane], b_hat: &[Lane]) { assert_eq!(a_hat.len(), b_hat.len()); for (a, &b_val) in a_hat.iter_mut().zip(b_hat.iter()) { - *a = mul_mod::(*a, b_val); + *a = mul_mod::(*a, b_val); } } #[cfg(test)] mod tests { use super::*; - use crate::mul::ntt::primes::PRIMES; + use crate::arch::ntt::{from_monty, to_monty, K, MODULI, OMEGA_MAX}; #[cfg(not(feature = "std"))] use alloc::vec; #[cfg(not(feature = "std"))] use alloc::vec::Vec; - fn assert_all_eq(a: &[u64], b_val: &[u64]) { - assert_eq!(a.len(), b_val.len()); + fn assert_all_eq(a: &[Lane], b_val: &[Lane], context: &str) { + assert_eq!(a.len(), b_val.len(), "{context}: length mismatch"); for (i, (x, y)) in a.iter().zip(b_val.iter()).enumerate() { - assert_eq!(x, y, "mismatch at index {i}: {x} != {y}"); + assert_eq!(x, y, "{context}: mismatch at index {i}: {x} != {y}"); } } macro_rules! for_each_prime { - ($b:ident, $p:ident, $omega:ident, $body:block) => { - for prime in &PRIMES { - let $b = prime.b; - match $b { - 32 => { - let $p = prime.p; - let $omega = prime.omega_2_32; - fn go($p: u64, $omega: u64) $body - go::<32>($p, $omega); + ($pi:ident, $p:ident, $omega:ident, $body:block) => { + for idx in 0..K { + let $p = MODULI[idx]; + let $omega = OMEGA_MAX[idx]; + match idx { + 0 => { + let $pi: usize = 0; + let _ = $pi; + fn go($p: Lane, $omega: Lane) $body + go::<0>($p, $omega); } - 34 => { - let $p = prime.p; - let $omega = prime.omega_2_32; - fn go($p: u64, $omega: u64) $body - go::<34>($p, $omega); + 1 => { + let $pi: usize = 1; + let _ = $pi; + fn go($p: Lane, $omega: Lane) $body + go::<1>($p, $omega); } - 40 => { - let $p = prime.p; - let $omega = prime.omega_2_32; - fn go($p: u64, $omega: u64) $body - go::<40>($p, $omega); + 2 => { + let $pi: usize = 2; + let _ = $pi; + fn go($p: Lane, $omega: Lane) $body + go::<2>($p, $omega); } _ => unreachable!(), } @@ -179,65 +152,80 @@ mod tests { #[test] fn test_forward_inverse_roundtrip() { - for_each_prime!(b, p, omega, { + for_each_prime!(pi, p, omega, { for &n in &[2, 4, 8, 16, 32, 64, 128, 256, 512] { - let mut fwd_twiddles = alloc::vec![0u64; n / 2]; - let mut inv_twiddles = alloc::vec![0u64; n / 2]; - precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); - precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); + let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; + let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; + precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); + precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); - let mut a: Vec = (0..n) - .map(|i| ((i as u64 + 1).wrapping_mul(123456789)) % p) + let mut a: Vec = (0..n) + .map(|i| ((i as Lane + 1).wrapping_mul(123456789)) % p) .collect(); + // Convert to Montgomery form for the NTT pipeline + for val in a.iter_mut() { + *val = to_monty::(*val); + } let orig = a.clone(); bit_reverse(&mut a); - forward::(&mut a, &fwd_twiddles, p); - inverse::(&mut a, &inv_twiddles, p); + forward::(&mut a, &fwd_twiddles); + inverse::(&mut a, &inv_twiddles, p); - assert_all_eq(&a, &orig); + assert_all_eq(&a, &orig, "roundtrip failed for n={n}"); } }); } #[test] fn test_convolution_via_ntt() { - for_each_prime!(b, p, omega, { + for_each_prime!(pi, p, omega, { for len_a in [1, 2, 3, 5] { for len_b in [1, 2, 3, 5] { let conv_len: usize = len_a + len_b - 1; let n = conv_len.next_power_of_two().max(2); - let a: Vec = (0..len_a).map(|i| ((i + 1) as u64 * 12345) % p).collect(); - let b_vec: Vec = - (0..len_b).map(|i| ((i + 1) as u64 * 67890) % p).collect(); + let a: Vec = + (0..len_a).map(|i| ((i + 1) as Lane * 12345) % p).collect(); + let b_vec: Vec = + (0..len_b).map(|i| ((i + 1) as Lane * 67890) % p).collect(); - let mut expected = vec![0u64; conv_len]; + // Compute expected convolution in standard form + let mut expected = vec![0u64 as Lane; conv_len]; for (i, &ai) in a.iter().enumerate() { for (j, &bj) in b_vec.iter().enumerate() { - expected[i + j] = - add_mod::(expected[i + j], mul_mod::(ai, bj), p); + let prod = (ai as u128 * bj as u128 % p as u128) as Lane; + expected[i + j] = add_mod::(expected[i + j], prod); } } - let mut fwd_twiddles = alloc::vec![0u64; n / 2]; - let mut inv_twiddles = alloc::vec![0u64; n / 2]; - precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); - precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); + let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; + let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; + precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); + precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); - let mut a_pad = vec![0u64; n]; - let mut b_pad = vec![0u64; n]; - a_pad[..len_a].copy_from_slice(&a); - b_pad[..len_b].copy_from_slice(&b_vec); + // Convert inputs to Montgomery form + let mut a_pad = vec![0u64 as Lane; n]; + let mut b_pad = vec![0u64 as Lane; n]; + for i in 0..len_a { + a_pad[i] = to_monty::(a[i]); + } + for i in 0..len_b { + b_pad[i] = to_monty::(b_vec[i]); + } bit_reverse(&mut a_pad); bit_reverse(&mut b_pad); - forward::(&mut a_pad, &fwd_twiddles, p); - forward::(&mut b_pad, &fwd_twiddles, p); - pointwise_mul::(&mut a_pad, &b_pad); - inverse::(&mut a_pad, &inv_twiddles, p); + forward::(&mut a_pad, &fwd_twiddles); + forward::(&mut b_pad, &fwd_twiddles); + pointwise_mul::(&mut a_pad, &b_pad); + inverse::(&mut a_pad, &inv_twiddles, p); + // Convert results back to standard form + for val in a_pad[..conv_len].iter_mut() { + *val = from_monty::(*val); + } - assert_all_eq(&a_pad[..conv_len], &expected); + assert_all_eq(&a_pad[..conv_len], &expected, "convolution mismatch"); } } }); @@ -245,47 +233,56 @@ mod tests { #[test] fn test_bit_reverse() { - let mut a: Vec = (0..8).collect(); + let mut a: Vec = (0..8).map(|i| i as Lane).collect(); bit_reverse(&mut a); assert_eq!(a, vec![0, 4, 2, 6, 1, 5, 3, 7]); } + /// Naive O(n²) NTT using standard-form modular arithmetic. #[allow(clippy::needless_range_loop)] - fn ntt_naive(x: &[u64], omega_n: u64, p: u64) -> Vec { + fn ntt_naive_std(x: &[Lane], omega_n: Lane, p: Lane) -> Vec { let n = x.len(); - let mut result = vec![0u64; n]; + let mut result = vec![0u64 as Lane; n]; for k in 0..n { - let mut acc = 0u64; + let mut acc: u128 = 0; for j in 0..n { let twiddle = if k == 0 || j == 0 { 1 } else { - omega_n.powm(&((k * j) as u64), &p) + omega_n.powm(&((k * j) as Lane), &p) as u128 }; - acc = add_mod::(acc, mul_mod::(x[j], twiddle), p); + acc = (acc + x[j] as u128 * twiddle) % p as u128; } - result[k] = acc; + result[k] = acc as Lane; } result } #[test] fn test_forward_correctness() { - for_each_prime!(b, p, omega, { + for_each_prime!(pi, p, omega, { for &n in &[2usize, 4, 8] { - let x: Vec = (0..n).map(|i| ((i + 1) as u64 * 11111) % p).collect(); + let x: Vec = (0..n).map(|i| ((i + 1) as Lane * 11111) % p).collect(); + - let shift = 32 - n.trailing_zeros(); - let omega_n = omega.powm(&(1u64 << shift), &p); + let shift = MAX_LOG_N - n.trailing_zeros(); + let omega_n = omega.powm(&((1u64 as Lane) << shift), &p); - let mut a = x.clone(); + // Convert to Montgomery form + let mut a: Vec = x.iter().map(|&v| to_monty::(v)).collect(); bit_reverse(&mut a); - let mut fwd_twiddles = alloc::vec![0u64; n / 2]; - precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); - forward::(&mut a, &fwd_twiddles, p); + let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; + precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); + forward::(&mut a, &fwd_twiddles); + + // Convert forward output back to standard form for comparison + for val in a.iter_mut() { + *val = from_monty::(*val); + } - let expected = ntt_naive::(&x, omega_n, p); + // expected: compute naive NTT in standard form + let expected = ntt_naive_std::(&x, omega_n, p); assert_eq!(a, expected, "forward NTT mismatch"); } }); @@ -293,56 +290,71 @@ mod tests { #[test] fn test_convolution_debug() { - let prime = &PRIMES[0]; // GL: b=32 - let p = prime.p; + let p = MODULI[0]; + let omega = OMEGA_MAX[0]; - let a = vec![12345u64 % p]; - let b_vec = vec![67890u64 % p, 135780u64 % p, 203670u64 % p]; + let a = [12345u64 as Lane % p]; + let b_vec = [67890u64 as Lane % p, + 135780u64 as Lane % p, + 203670u64 as Lane % p]; let conv_len = a.len() + b_vec.len() - 1; let n = 4; - let mut expected = vec![0u64; conv_len]; + // Expected values in standard form + let mut expected = vec![0u64 as Lane; conv_len]; for (i, &ai) in a.iter().enumerate() { for (j, &bj) in b_vec.iter().enumerate() { - expected[i + j] = add_mod::<32>(expected[i + j], mul_mod::<32>(ai, bj), p); + let prod = (ai as u128 * bj as u128 % p as u128) as Lane; + expected[i + j] = add_mod::<0>(expected[i + j], prod); } } - let mut fwd_twiddles = alloc::vec![0u64; n / 2]; - let mut inv_twiddles = alloc::vec![0u64; n / 2]; - precompute_twiddles::<32>(&mut fwd_twiddles, n, p, prime.omega_2_32, false); - precompute_twiddles::<32>(&mut inv_twiddles, n, p, prime.omega_2_32, true); + let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; + let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; + precompute_twiddles::<0>(&mut fwd_twiddles, n, p, omega, false); + precompute_twiddles::<0>(&mut inv_twiddles, n, p, omega, true); - let mut a_pad = vec![0u64; n]; - let mut b_pad = vec![0u64; n]; - a_pad[..a.len()].copy_from_slice(&a); - b_pad[..b_vec.len()].copy_from_slice(&b_vec); + // Convert to Montgomery form + let mut a_pad = vec![0u64 as Lane; n]; + let mut b_pad = vec![0u64 as Lane; n]; + for i in 0..a.len() { + a_pad[i] = to_monty::<0>(a[i]); + } + for i in 0..b_vec.len() { + b_pad[i] = to_monty::<0>(b_vec[i]); + } bit_reverse(&mut a_pad); bit_reverse(&mut b_pad); - forward::<32>(&mut a_pad, &fwd_twiddles, p); - forward::<32>(&mut b_pad, &fwd_twiddles, p); - pointwise_mul::<32>(&mut a_pad, &b_pad); - inverse::<32>(&mut a_pad, &inv_twiddles, p); + forward::<0>(&mut a_pad, &fwd_twiddles); + forward::<0>(&mut b_pad, &fwd_twiddles); + pointwise_mul::<0>(&mut a_pad, &b_pad); + inverse::<0>(&mut a_pad, &inv_twiddles, p); + // Convert back to standard form + for val in a_pad[..conv_len].iter_mut() { + *val = from_monty::<0>(*val); + } assert_eq!(&a_pad[..conv_len], &expected[..]); } #[test] fn test_length_two_edge_case() { - for_each_prime!(b, p, omega, { + for_each_prime!(pi, p, omega, { let n = 2; - let mut fwd_twiddles = alloc::vec![0u64; n / 2]; - let mut inv_twiddles = alloc::vec![0u64; n / 2]; - precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); - precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); - - let a_orig = vec![1u64 % p, 2u64 % p]; + let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; + let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; + precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); + precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); + + let a_std = [1u64 as Lane % p, 2u64 as Lane % p]; + // Convert to Montgomery form + let a_orig: Vec = a_std.iter().map(|&v| to_monty::(v)).collect(); let mut a = a_orig.clone(); bit_reverse(&mut a); - forward::(&mut a, &fwd_twiddles, p); - inverse::(&mut a, &inv_twiddles, p); - assert_all_eq(&a, &a_orig); + forward::(&mut a, &fwd_twiddles); + inverse::(&mut a, &inv_twiddles, p); + assert_all_eq(&a, &a_orig, "length two roundtrip"); }); } } From 7c653b12defcd67742dbcb6fd109124d6922ee68 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 11:03:02 +0800 Subject: [PATCH 11/19] Tidy up num-modular usage --- integer/src/arch/generic_32_bit/ntt.rs | 110 +++----- integer/src/arch/generic_32_bit/word.rs | 4 + integer/src/arch/generic_64_bit/ntt.rs | 110 +++----- integer/src/arch/generic_64_bit/word.rs | 4 + integer/src/mul/mod.rs | 62 +---- integer/src/mul/ntt/crt.rs | 349 +++++++++--------------- integer/src/mul/ntt/mod.rs | 259 +++++------------- integer/src/mul/ntt/pack.rs | 19 +- integer/src/mul/ntt/transform.rs | 203 ++++++-------- 9 files changed, 380 insertions(+), 740 deletions(-) diff --git a/integer/src/arch/generic_32_bit/ntt.rs b/integer/src/arch/generic_32_bit/ntt.rs index ce55b608..9f339b6a 100644 --- a/integer/src/arch/generic_32_bit/ntt.rs +++ b/integer/src/arch/generic_32_bit/ntt.rs @@ -3,13 +3,18 @@ //! Uses Proth primes of the form `K * 2^N + 1`. //! All constants computed by `integer/src/mul/ntt/compute_constants.py`. -use num_modular::{FixedProth32, Reducer}; +use num_modular::FixedProth32; // Proth reducer instances — each with a different (N, K) pair. pub const P0: FixedProth32<26, 7> = FixedProth32::<26, 7>; pub const P1: FixedProth32<27, 15> = FixedProth32::<27, 15>; pub const P2: FixedProth32<27, 17> = FixedProth32::<27, 17>; +// Type aliases needed by for_each_prime! macro in transform tests. +pub type Rp0 = FixedProth32<26, 7>; +pub type Rp1 = FixedProth32<27, 15>; +pub type Rp2 = FixedProth32<27, 17>; + pub const K: usize = 3; pub const MAX_LOG_N: u32 = 26; pub const B_PACK_MIN: u32 = 8; @@ -38,60 +43,10 @@ pub const MODULI: [Lane; K] = [ FixedProth32::<27, 17>::MODULUS, ]; -#[inline] -pub fn to_monty(val: Lane) -> Lane { - match PI { - 0 => P0.transform(val), - 1 => P1.transform(val), - 2 => P2.transform(val), - _ => unreachable!(), - } -} - -#[inline] -pub fn from_monty(val: Lane) -> Lane { - match PI { - 0 => P0.residue(val), - 1 => P1.residue(val), - 2 => P2.residue(val), - _ => unreachable!(), - } -} - -#[inline] -pub fn mul_mod(a: Lane, b_val: Lane) -> Lane { - let prod = (a as DoubleLane) * (b_val as DoubleLane); - match PI { - 0 => P0.reduce(prod), - 1 => P1.reduce(prod), - 2 => P2.reduce(prod), - _ => unreachable!(), - } -} - -#[inline] -pub fn add_mod(a: Lane, b_val: Lane) -> Lane { - match PI { - 0 => P0.add(&a, &b_val), - 1 => P1.add(&a, &b_val), - 2 => P2.add(&a, &b_val), - _ => unreachable!(), - } -} - -#[inline] -pub fn sub_mod(a: Lane, b_val: Lane) -> Lane { - match PI { - 0 => P0.sub(&a, &b_val), - 1 => P1.sub(&a, &b_val), - 2 => P2.sub(&a, &b_val), - _ => unreachable!(), - } -} - #[cfg(test)] mod tests { use super::*; + use num_modular::Reducer; #[test] fn test_primes_proth_form() { @@ -112,39 +67,36 @@ mod tests { fn test_omega_order() { for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { let p = MODULI[pi]; - let sqr = |w: Lane| -> Lane { - match pi { - 0 => P0.reduce((w as u64) * (w as u64)), - 1 => P1.reduce((w as u64) * (w as u64)), - 2 => P2.reduce((w as u64) * (w as u64)), - _ => unreachable!(), - } - }; - - let mut w = match pi { - 0 => to_monty::<0>(omega_max), - 1 => to_monty::<1>(omega_max), - 2 => to_monty::<2>(omega_max), + let (sqr, to_m, from_m): ( + fn(Lane) -> Lane, + fn(Lane) -> Lane, + fn(Lane) -> Lane, + ) = match pi { + 0 => ( + |w| P0.reduce((w as u64) * (w as u64)), + |v| P0.transform(v), + |v| P0.residue(v), + ), + 1 => ( + |w| P1.reduce((w as u64) * (w as u64)), + |v| P1.transform(v), + |v| P1.residue(v), + ), + 2 => ( + |w| P2.reduce((w as u64) * (w as u64)), + |v| P2.transform(v), + |v| P2.residue(v), + ), _ => unreachable!(), }; + + let mut w = to_m(omega_max); for _ in 0..MAX_LOG_N - 1 { w = sqr(w); } - let w_std = match pi { - 0 => from_monty::<0>(w), - 1 => from_monty::<1>(w), - 2 => from_monty::<2>(w), - _ => unreachable!(), - }; - assert_eq!(w_std, p - 1, "omega^(2^(MAX_LOG_N-1)) != -1 mod p for prime {pi}"); + assert_eq!(from_m(w), p - 1, "omega^(2^(MAX_LOG_N-1)) != -1 mod p for prime {pi}"); w = sqr(w); - let one = match pi { - 0 => from_monty::<0>(w), - 1 => from_monty::<1>(w), - 2 => from_monty::<2>(w), - _ => unreachable!(), - }; - assert_eq!(one, 1, "omega^(2^MAX_LOG_N) != 1 mod p for prime {pi}"); + assert_eq!(from_m(w), 1, "omega^(2^MAX_LOG_N) != 1 mod p for prime {pi}"); } } } diff --git a/integer/src/arch/generic_32_bit/word.rs b/integer/src/arch/generic_32_bit/word.rs index 4b3769d3..2cb44b61 100644 --- a/integer/src/arch/generic_32_bit/word.rs +++ b/integer/src/arch/generic_32_bit/word.rs @@ -9,3 +9,7 @@ pub type DoubleWord = u64; /// Signed double machine word. pub type SignedDoubleWord = i64; + +/// Accumulator for the product of three primes (3 × 2^32 ≈ 2^96). +#[derive(Clone, Copy, Debug, Default)] +pub struct TripleWord(pub [u32; 3]); diff --git a/integer/src/arch/generic_64_bit/ntt.rs b/integer/src/arch/generic_64_bit/ntt.rs index e0bd8de2..2e20d9a0 100644 --- a/integer/src/arch/generic_64_bit/ntt.rs +++ b/integer/src/arch/generic_64_bit/ntt.rs @@ -3,13 +3,18 @@ //! Uses Proth primes of the form `K * 2^N + 1`. //! All constants computed by `integer/src/mul/ntt/compute_constants.py`. -use num_modular::{FixedProth64, Reducer}; +use num_modular::FixedProth64; // Proth reducer instances — each with a different (N, K) pair. pub const P0: FixedProth64<57, 29> = FixedProth64::<57, 29>; pub const P1: FixedProth64<57, 71> = FixedProth64::<57, 71>; pub const P2: FixedProth64<57, 75> = FixedProth64::<57, 75>; +// Type aliases needed by for_each_prime! macro in transform tests. +pub type Rp0 = FixedProth64<57, 29>; +pub type Rp1 = FixedProth64<57, 71>; +pub type Rp2 = FixedProth64<57, 75>; + pub const K: usize = 3; pub const MAX_LOG_N: u32 = 57; pub const B_PACK_MIN: u32 = 16; @@ -32,57 +37,6 @@ pub const CRT_INV_IJ: [[Lane; K]; K] = [ [0, 0, 0], ]; -#[inline] -pub fn to_monty(val: Lane) -> Lane { - match PI { - 0 => P0.transform(val), - 1 => P1.transform(val), - 2 => P2.transform(val), - _ => unreachable!(), - } -} - -#[inline] -pub fn from_monty(val: Lane) -> Lane { - match PI { - 0 => P0.residue(val), - 1 => P1.residue(val), - 2 => P2.residue(val), - _ => unreachable!(), - } -} - -#[inline] -pub fn mul_mod(a: Lane, b_val: Lane) -> Lane { - let prod = (a as DoubleLane) * (b_val as DoubleLane); - match PI { - 0 => P0.reduce(prod), - 1 => P1.reduce(prod), - 2 => P2.reduce(prod), - _ => unreachable!(), - } -} - -#[inline] -pub fn add_mod(a: Lane, b_val: Lane) -> Lane { - match PI { - 0 => P0.add(&a, &b_val), - 1 => P1.add(&a, &b_val), - 2 => P2.add(&a, &b_val), - _ => unreachable!(), - } -} - -#[inline] -pub fn sub_mod(a: Lane, b_val: Lane) -> Lane { - match PI { - 0 => P0.sub(&a, &b_val), - 1 => P1.sub(&a, &b_val), - 2 => P2.sub(&a, &b_val), - _ => unreachable!(), - } -} - /// Prime moduli indexed by PI. pub const MODULI: [Lane; K] = [ FixedProth64::<57, 29>::MODULUS, @@ -93,6 +47,7 @@ pub const MODULI: [Lane; K] = [ #[cfg(test)] mod tests { use super::*; + use num_modular::Reducer; #[test] fn test_primes_proth_form() { @@ -113,39 +68,36 @@ mod tests { fn test_omega_order() { for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { let p = MODULI[pi]; - let sqr = |w: Lane| -> Lane { - match pi { - 0 => P0.reduce((w as u128) * (w as u128)), - 1 => P1.reduce((w as u128) * (w as u128)), - 2 => P2.reduce((w as u128) * (w as u128)), - _ => unreachable!(), - } - }; - - let mut w = match pi { - 0 => to_monty::<0>(omega_max), - 1 => to_monty::<1>(omega_max), - 2 => to_monty::<2>(omega_max), + let (sqr, to_m, from_m): ( + fn(Lane) -> Lane, + fn(Lane) -> Lane, + fn(Lane) -> Lane, + ) = match pi { + 0 => ( + |w| P0.reduce((w as u128) * (w as u128)), + |v| P0.transform(v), + |v| P0.residue(v), + ), + 1 => ( + |w| P1.reduce((w as u128) * (w as u128)), + |v| P1.transform(v), + |v| P1.residue(v), + ), + 2 => ( + |w| P2.reduce((w as u128) * (w as u128)), + |v| P2.transform(v), + |v| P2.residue(v), + ), _ => unreachable!(), }; + + let mut w = to_m(omega_max); for _ in 0..MAX_LOG_N - 1 { w = sqr(w); } - let w_std = match pi { - 0 => from_monty::<0>(w), - 1 => from_monty::<1>(w), - 2 => from_monty::<2>(w), - _ => unreachable!(), - }; - assert_eq!(w_std, p - 1, "omega^(2^(MAX_LOG_N-1)) != -1 mod p for prime {pi}"); + assert_eq!(from_m(w), p - 1, "omega^(2^(MAX_LOG_N-1)) != -1 mod p for prime {pi}"); w = sqr(w); - let one = match pi { - 0 => from_monty::<0>(w), - 1 => from_monty::<1>(w), - 2 => from_monty::<2>(w), - _ => unreachable!(), - }; - assert_eq!(one, 1, "omega^(2^MAX_LOG_N) != 1 mod p for prime {pi}"); + assert_eq!(from_m(w), 1, "omega^(2^MAX_LOG_N) != 1 mod p for prime {pi}"); } } } diff --git a/integer/src/arch/generic_64_bit/word.rs b/integer/src/arch/generic_64_bit/word.rs index fefd7e7f..a5f15ed9 100644 --- a/integer/src/arch/generic_64_bit/word.rs +++ b/integer/src/arch/generic_64_bit/word.rs @@ -9,3 +9,7 @@ pub type DoubleWord = u128; /// Signed double machine word. pub type SignedDoubleWord = i128; + +/// Accumulator for the product of three primes (3 × 2^64 = 2^192). +#[derive(Clone, Copy, Debug, Default)] +pub struct TripleWord(pub [u64; 3]); diff --git a/integer/src/mul/mod.rs b/integer/src/mul/mod.rs index 261f53aa..8e97f208 100644 --- a/integer/src/mul/mod.rs +++ b/integer/src/mul/mod.rs @@ -25,22 +25,13 @@ const THRESHOLD_KARATSUBA_DEFAULT: usize = 96; const_assert!(THRESHOLD_KARATSUBA_DEFAULT + 1 >= toom_3::MIN_LEN); /// If smaller operand length > this, NTT multiplication will be used. -#[cfg(not(any( - force_bits = "16", - target_pointer_width = "16" -)))] +#[cfg(not(any(force_bits = "16", target_pointer_width = "16")))] const THRESHOLD_NTT_DEFAULT: usize = ntt::THRESHOLD_NTT; /// NTT unavailable on 16/32-bit word targets — use `usize::MAX` so dispatch never /// routes to the NTT path. -#[cfg(any( - force_bits = "16", - target_pointer_width = "16" -))] +#[cfg(any(force_bits = "16", target_pointer_width = "16"))] const THRESHOLD_NTT_DEFAULT: usize = usize::MAX; -#[cfg(not(any( - force_bits = "16", - target_pointer_width = "16" -)))] +#[cfg(not(any(force_bits = "16", target_pointer_width = "16")))] const_assert!(THRESHOLD_NTT_DEFAULT + 1 >= toom_3::MIN_LEN); /// Environment-variable overrides for multiplication thresholds. @@ -89,10 +80,7 @@ mod threshold { mod helpers; mod karatsuba; -#[cfg(not(any( - force_bits = "16", - target_pointer_width = "16" -)))] +#[cfg(not(any(force_bits = "16", target_pointer_width = "16")))] pub(crate) mod ntt; mod simple; pub(crate) mod toom_3; @@ -235,21 +223,11 @@ pub fn memory_requirement_up_to(total_len: usize, smaller_len: usize) -> Layout toom_3::memory_requirement_up_to(smaller_len) } else { // NTT path — only available on 64-bit word targets. - #[cfg(not(any( - force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" - )))] + #[cfg(not(any(force_bits = "16", target_pointer_width = "16")))] { ntt::memory_requirement_up_to(total_len, smaller_len) } - #[cfg(any( - force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" - ))] + #[cfg(any(force_bits = "16", target_pointer_width = "16"))] { let _ = (total_len, smaller_len); unreachable!("NTT unavailable on 16-bit targets"); @@ -294,21 +272,11 @@ pub fn add_signed_mul<'a>( } else if b.len() <= threshold::ntt() { toom_3::add_signed_mul(c, sign, a, b, memory) } else { - #[cfg(not(any( - force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" - )))] + #[cfg(not(any(force_bits = "16", target_pointer_width = "16")))] { ntt::add_signed_mul(c, sign, a, b, memory) } - #[cfg(any( - force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" - ))] + #[cfg(any(force_bits = "16", target_pointer_width = "16"))] { let _ = (c, sign, a, b, memory); unreachable!("NTT unavailable on 16-bit targets"); @@ -337,21 +305,11 @@ pub fn add_signed_mul_same_len( } else if n <= threshold::ntt() { toom_3::add_signed_mul_same_len(c, sign, a, b, memory) } else { - #[cfg(not(any( - force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" - )))] + #[cfg(not(any(force_bits = "16", target_pointer_width = "16")))] { ntt::add_signed_mul_same_len(c, sign, a, b, memory) } - #[cfg(any( - force_bits = "16", - force_bits = "32", - target_pointer_width = "16", - target_pointer_width = "32" - ))] + #[cfg(any(force_bits = "16", target_pointer_width = "16"))] { let _ = (c, sign, a, b, memory); unreachable!("NTT unavailable on 16-bit targets"); diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs index b25a7f61..926029f2 100644 --- a/integer/src/mul/ntt/crt.rs +++ b/integer/src/mul/ntt/crt.rs @@ -1,221 +1,32 @@ //! Garner CRT: combine `K` residues modulo `K` primes into a small integer. -//! -//! Generic over the lane type — supports u64 (→[`U192`]) and u32 (→[`U128`]). -#![allow(clippy::unnecessary_cast)] use crate::arch::ntt::K; use num_modular::ModularCoreOps; -/// Accumulator for Garner CRT — either [`U192`] (64-bit lanes) or [`U96`] (32-bit). +/// Accumulator for Garner CRT. +/// +/// Implemented by [`crate::arch::word::TripleWord`] (192 bits on 64-bit +/// targets, 96 bits on 32-bit targets). pub trait CrtAccum: Default + Copy { - type Lane: Copy + Into + for<'a> ModularCoreOps; + type Lane: Copy + + Default + + Into + + for<'a> ModularCoreOps; fn from_lane(v: Self::Lane) -> Self; - /// `self += t * factor` where `t` is a lane-sized coefficient. + /// `self += t * factor` fn add_product(&mut self, t: Self::Lane, factor: u128); /// `self mod m` fn rem_lane(&self, m: Self::Lane) -> Self::Lane; - /// Number of non-zero u64 words. - #[allow(dead_code)] - fn len_words(&self) -> u32; - /// View as `&[u64]`. - fn as_u64_slice(&self) -> &[u64]; -} - -// ── U192 (64-bit lanes) ──────────────────────────────────────────────── - -/// A 192-bit unsigned integer (3 × u64, little-endian). -/// -/// Used to hold Garner CRT results for three ≈2^64 primes (product < 2^192). -#[derive(Clone, Copy, Debug, Default)] -pub struct U192(pub [u64; 3]); - -impl U192 { - #[inline] - pub fn new(lo: u64) -> Self { - U192([lo, 0, 0]) - } - - /// `self += v` where `v` fits in 128 bits. - #[inline] - #[allow(dead_code)] - pub fn add_u128(&mut self, v: u128) { - let lo = v as u64; - let hi = (v >> 64) as u64; - let (r0, c0) = self.0[0].overflowing_add(lo); - self.0[0] = r0; - let (r1, c1) = self.0[1].overflowing_add(hi.wrapping_add(c0 as u64)); - self.0[1] = r1; - self.0[2] = self.0[2].wrapping_add(c1 as u64); - } - - /// `self += t * factor` where `t` < 2^64, `factor` < 2^128. - #[inline] - pub fn add_mul_u64_u128(&mut self, t: u64, factor: u128) { - let fac_lo = factor as u64; - let fac_hi = (factor >> 64) as u64; - - let m_lo_full = (t as u128) * (fac_lo as u128); - let lo = m_lo_full as u64; - let m_lo = (m_lo_full >> 64) as u64; - - let m_hi_full = (t as u128) * (fac_hi as u128); - let m_hi = m_hi_full as u64; - let hi = (m_hi_full >> 64) as u64; - - let (mid, c) = m_lo.overflowing_add(m_hi); - let hi_word = hi.wrapping_add(c as u64); - - let (r0, c0) = self.0[0].overflowing_add(lo); - self.0[0] = r0; - let (r1, c1) = self.0[1].overflowing_add(mid.wrapping_add(c0 as u64)); - self.0[1] = r1; - self.0[2] = self.0[2].wrapping_add(hi_word.wrapping_add(c1 as u64)); - } - - /// `self mod m`, where `m` < 2^64. - #[inline] - pub fn rem_u64(&self, m: u64) -> u64 { - let m128 = m as u128; - let mut r: u128 = 0; - for &word in self.0.iter().rev() { - r = (r << 64) | (word as u128); - r %= m128; - } - r as u64 - } - - #[inline] - pub fn len_words(&self) -> u32 { - if self.0[2] != 0 { - 3 - } else if self.0[1] != 0 { - 2 - } else { - 1 - } - } -} - -impl CrtAccum for U192 { - type Lane = u64; - - #[inline] - fn from_lane(v: u64) -> Self { - U192::new(v) - } - - #[inline] - fn add_product(&mut self, t: u64, factor: u128) { - self.add_mul_u64_u128(t, factor); - } - - #[inline] - fn rem_lane(&self, m: u64) -> u64 { - self.rem_u64(m) - } - - #[inline] - fn len_words(&self) -> u32 { - self.len_words() - } - - #[inline] - fn as_u64_slice(&self) -> &[u64] { - &self.0[..self.len_words() as usize] - } -} - -// ── U96 (32-bit lanes) ───────────────────────────────────────────────── - -/// A value bounded by 2^96 (product of three ≈2^32 primes). -/// -/// Stored as `[u64; 2]` (128 bits) so [`as_u64_slice`] can return a -/// `&[u64]` for [`add_shifted_to_prod`]. The upper 32 bits of the -/// second limb are always zero. -/// -/// [`add_shifted_to_prod`]: super::add_shifted_to_prod -#[derive(Clone, Copy, Debug, Default)] -pub struct U96(pub [u64; 2]); - -impl U96 { - /// `self += t * factor` where `t` < 2^32. - #[inline] - pub fn add_mul_u32_u96(&mut self, t: u32, factor: u128) { - let fac_lo = factor as u64; - let fac_hi = (factor >> 64) as u64; - - // t × fac_lo (max 2^32 × 2^64 = 2^96 → fits in u128) - let m_lo_full = (t as u128) * (fac_lo as u128); - let lo = m_lo_full as u64; - let m_lo_carry = (m_lo_full >> 64) as u64; - - // t × fac_hi + carry - let m_hi_full = (t as u64 as u128) * (fac_hi as u128) + m_lo_carry as u128; - let m_hi = m_hi_full as u64; - - let (r0, c0) = self.0[0].overflowing_add(lo); - self.0[0] = r0; - let (r1, c1) = self.0[1].overflowing_add(m_hi.wrapping_add(c0 as u64)); - self.0[1] = r1; - let _ = c1; - } - - /// `self mod m`, where `m` < 2^32. - #[inline] - pub fn rem_u32(&self, m: u32) -> u32 { - let m128 = m as u128; - let mut r: u128 = 0; - for &word in self.0.iter().rev() { - r = (r << 64) | (word as u128); - r %= m128; - } - r as u32 - } - - #[inline] - pub fn len_words(&self) -> u32 { - if self.0[1] != 0 { - 2 - } else { - 1 - } - } -} - -impl CrtAccum for U96 { - type Lane = u32; - - #[inline] - fn from_lane(v: u32) -> Self { - U96([v as u64, 0]) - } - - #[inline] - fn add_product(&mut self, t: u32, factor: u128) { - self.add_mul_u32_u96(t, factor); - } - - #[inline] - fn rem_lane(&self, m: u32) -> u32 { - self.rem_u32(m) - } - - #[inline] - fn len_words(&self) -> u32 { - self.len_words() - } - - #[inline] - fn as_u64_slice(&self) -> &[u64] { - &self.0[..self.len_words() as usize] - } + /// Write the value into `out` as little-endian `Word` values, + /// returning the number of non-zero words written. + fn write_words(&self, out: &mut [crate::arch::word::Word; 6]) -> u32; } // ── Garner combine ───────────────────────────────────────────────────── /// Combine `K` residues into a [`CrtAccum`] via Garner's algorithm. /// -/// All arithmetic is standard-form (not Montgomery). `crt_inv_ij[i][j]` +/// All arithmetic is standard-form. `crt_inv_ij[i][j]` /// holds `inv(p_i mod p_j) mod p_j` for `i < j`. /// `primes` contains the prime values (only `primes[0..k]` are used). pub fn garner_combine( @@ -255,32 +66,144 @@ pub fn garner_combine( x } +// ── TripleWord impls (cfg-gated per arch) ───────────────────────────── + +/// 64-bit: 3 × u64 = 192 bits. +#[cfg(not(any(force_bits = "32", target_pointer_width = "32")))] +mod triple_impl { + use super::CrtAccum; + use crate::arch::word::TripleWord; + + impl CrtAccum for TripleWord { + type Lane = u64; + + #[inline] + fn from_lane(v: u64) -> Self { + TripleWord([v, 0, 0]) + } + + #[inline] + fn add_product(&mut self, t: u64, factor: u128) { + let fac_lo = factor as u64; + let fac_hi = (factor >> 64) as u64; + let m_lo_full = (t as u128) * (fac_lo as u128); + let lo = m_lo_full as u64; + let m_lo = (m_lo_full >> 64) as u64; + let m_hi_full = (t as u128) * (fac_hi as u128); + let m_hi = m_hi_full as u64; + let hi = (m_hi_full >> 64) as u64; + let (mid, c) = m_lo.overflowing_add(m_hi); + let hi_word = hi.wrapping_add(c as u64); + let (r0, c0) = self.0[0].overflowing_add(lo); + self.0[0] = r0; + let (r1, c1) = self.0[1].overflowing_add(mid.wrapping_add(c0 as u64)); + self.0[1] = r1; + self.0[2] = self.0[2].wrapping_add(hi_word.wrapping_add(c1 as u64)); + } + + #[inline] + fn rem_lane(&self, m: u64) -> u64 { + let m128 = m as u128; + let mut r: u128 = 0; + for &word in self.0.iter().rev() { + r = (r << 64) | (word as u128); + r %= m128; + } + r as u64 + } + + #[inline] + fn write_words(&self, out: &mut [crate::arch::word::Word; 6]) -> u32 { + out[0] = self.0[0]; + out[1] = self.0[1]; + out[2] = self.0[2]; + if self.0[2] != 0 { 3 } else if self.0[1] != 0 { 2 } else { 1 } + } + } +} + +/// 32-bit: 3 × u32 = 96 bits. +#[cfg(any(force_bits = "32", target_pointer_width = "32"))] +mod triple_impl { + use super::CrtAccum; + use crate::arch::word::TripleWord; + + impl CrtAccum for TripleWord { + type Lane = u32; + + #[inline] + fn from_lane(v: u32) -> Self { + TripleWord([v, 0, 0]) + } + + #[inline] + fn add_product(&mut self, t: u32, factor: u128) { + let factor_lo = factor as u32; + let factor_hi = (factor >> 32) as u32; + let m_lo = (t as u64) * (factor_lo as u64); + let lo = m_lo as u32; + let m_mid = (m_lo >> 32) as u32; + let m_hi = (t as u64) * (factor_hi as u64) + m_mid as u64; + let mid = m_hi as u32; + let hi = (m_hi >> 32) as u32; + let (r0, c0) = self.0[0].overflowing_add(lo); + self.0[0] = r0; + let (r1, c1) = self.0[1].overflowing_add(mid.wrapping_add(c0 as u32)); + self.0[1] = r1; + self.0[2] = self.0[2].wrapping_add(hi.wrapping_add(c1 as u32)); + } + + #[inline] + fn rem_lane(&self, m: u32) -> u32 { + let m64 = m as u64; + let mut r: u64 = 0; + for &word in self.0.iter().rev() { + r = (r << 32) | (word as u64); + r %= m64; + } + r as u32 + } + + #[inline] + fn write_words(&self, out: &mut [crate::arch::word::Word; 6]) -> u32 { + out[0] = self.0[0]; + out[1] = self.0[1]; + out[2] = self.0[2]; + if self.0[2] != 0 { 3 } else if self.0[1] != 0 { 2 } else { 1 } + } + } +} + #[cfg(test)] mod tests { use super::*; + use crate::arch::word::TripleWord; #[cfg(not(feature = "std"))] use alloc::vec; #[test] - fn test_garner_with_u64_primes() { + fn test_garner_roundtrip() { use crate::arch::ntt::{CRT_INV_IJ, MODULI}; + type Lane = ::Lane; let p0 = MODULI[0]; let p1 = MODULI[1]; let p2 = MODULI[2]; let primes = [p0, p1, p2]; - let residues = vec![12345u64, 67890u64, 11111u64]; - let x = garner_combine::(&residues, &CRT_INV_IJ, &primes); - assert_eq!(x.rem_u64(p0), residues[0]); - assert_eq!(x.rem_u64(p1), residues[1]); - assert_eq!(x.rem_u64(p2), residues[2]); + let residues = vec![12345u64 as Lane, 67890u64 as Lane, 11111u64 as Lane]; + let x = garner_combine::(&residues, &CRT_INV_IJ, &primes); + assert_eq!(x.rem_lane(p0), residues[0]); + assert_eq!(x.rem_lane(p1), residues[1]); + assert_eq!(x.rem_lane(p2), residues[2]); - let x = garner_combine::(&residues[..2], &CRT_INV_IJ, &primes); - assert_eq!(x.rem_u64(p0), residues[0]); - assert_eq!(x.rem_u64(p1), residues[1]); + let x = garner_combine::(&residues[..2], &CRT_INV_IJ, &primes); + assert_eq!(x.rem_lane(p0), residues[0]); + assert_eq!(x.rem_lane(p1), residues[1]); - let x = garner_combine::(&residues[..1], &CRT_INV_IJ, &primes); - assert_eq!(x.0[0], residues[0]); + let x = garner_combine::(&residues[..1], &CRT_INV_IJ, &primes); + let mut buf = [crate::arch::word::Word::default(); 6]; + x.write_words(&mut buf); + assert_eq!(buf[0] as u64, residues[0] as u64); } } diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index c713d406..e6e2f7bc 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -12,12 +12,13 @@ use crate::{ use alloc::alloc::Layout; use core::mem; -mod crt; +pub(crate) mod crt; mod pack; mod transform; -use crate::arch::ntt::{mul_mod, B_PACK_CANDIDATES, B_PACK_MIN, K, MAX_LOG_N, MODULI, OMEGA_MAX}; +use crate::arch::ntt::{B_PACK_CANDIDATES, B_PACK_MIN, K, MAX_LOG_N, MODULI, OMEGA_MAX, P0, P1, P2}; use crate::mul::ntt::crt::{garner_combine, CrtAccum}; +use num_modular::Reducer; /// Minimum smaller-operand length (in words) for the NTT path. pub const THRESHOLD_NTT: usize = 40_000; @@ -93,14 +94,12 @@ pub fn memory_requirement_up_to(total_len: usize, _smaller_len: usize) -> Layout let product = total_len; let lane_bytes = mem::size_of::(); - let u64_bytes = 8usize; let word_bytes = mem::size_of::(); let lanes_words = lanes * lane_bytes / word_bytes; let residues_words = residues * lane_bytes / word_bytes; let twiddles_words = twiddles * lane_bytes / word_bytes; - let product_words = product * u64_bytes / word_bytes; - let total_words = product_words + lanes_words + residues_words + twiddles_words; + let total_words = product + lanes_words + residues_words + twiddles_words; memory::array_layout::(total_words) } @@ -168,9 +167,9 @@ fn add_signed_mul_impl( // ---- Memory carve (longest-lived first) ---- - // 1. Product buffer (always u64) + // 1. Product buffer (Word-sized, CRT splits u64 words into Word limbs) let prod_len = la + lb; - let (prod, mut mem) = memory.allocate_slice_fill::(prod_len, 0); + let (prod, mut mem) = memory.allocate_slice_fill::(prod_len, 0); // 2. Residue storage (per-prime inverse results) let residues_len = k_eff * nn; @@ -184,14 +183,13 @@ fn add_signed_mul_impl( let (fwd_twiddles, mut mem) = mem.allocate_slice_fill::(nn / 2, 0); let (inv_twiddles, _) = mem.allocate_slice_fill::(nn / 2, 0); - // ---- Per-prime transforms (const-generic dispatch) ---- + // ---- Per-prime transforms (monomorphized per reducer) ---- for pi in 0..k_eff { let mut ctx = TransformCtx { a_lane, b_lane, fwd_twiddles, inv_twiddles, - p: MODULI[pi], omega_max: OMEGA_MAX[pi], nn, b_pack, @@ -199,9 +197,9 @@ fn add_signed_mul_impl( pi, }; match pi { - 0 => process_prime::<0>(a, b, &mut ctx), - 1 => process_prime::<1>(a, b, &mut ctx), - 2 => process_prime::<2>(a, b, &mut ctx), + 0 => process_prime(a, b, &mut ctx, &P0), + 1 => process_prime(a, b, &mut ctx, &P1), + 2 => process_prime(a, b, &mut ctx, &P2), _ => unreachable!(), } } @@ -210,22 +208,18 @@ fn add_signed_mul_impl( // Extract prime constants as both u64 and u32 so the Lane-size // dispatch below type-checks correctly in both branches. // The dead branch (wrong width) is eliminated by the compiler. - let primes_u64: [u64; K] = [MODULI[0], MODULI[1], MODULI[2]]; + let primes_u64: [u64; K] = [MODULI[0] as u64, MODULI[1] as u64, MODULI[2] as u64]; let crt_inv_u64: [[u64; K]; K] = { use crate::arch::ntt::CRT_INV_IJ; let mut m = [[0u64; K]; K]; for i in 0..K { for j in 0..K { - m[i][j] = CRT_INV_IJ[i][j]; + m[i][j] = CRT_INV_IJ[i][j] as u64; } } m }; - let primes_u32: [u32; K] = [ - MODULI[0] as u32, - MODULI[1] as u32, - MODULI[2] as u32, - ]; + let primes_u32: [u32; K] = [MODULI[0] as u32, MODULI[1] as u32, MODULI[2] as u32]; let crt_inv_u32: [[u32; K]; K] = { let mut m = [[0u32; K]; K]; for i in 0..K { @@ -236,96 +230,56 @@ fn add_signed_mul_impl( m }; - #[allow(clippy::unnecessary_cast)] - if mem::size_of::() == 8 { - // SAFETY: mem::size_of::() == 8, so Lane = u64 and - // residues is backed by u64 elements. + // CRT dispatch: one branch per Word size, gated by cfg so only + // one compiles — no dummy types needed in dead branches. + #[cfg(not(any(force_bits = "32", target_pointer_width = "32")))] + { let residues_u64: &[u64] = unsafe { core::slice::from_raw_parts(residues.as_ptr() as *const u64, residues.len()) }; - do_crt_u64( - prod, - residues_u64, - k_eff, - nn, - output_coeffs, - b_pack, - &primes_u64, - &crt_inv_u64, + do_crt::( + prod, residues_u64, k_eff, nn, output_coeffs, b_pack, &primes_u64, &crt_inv_u64, ); - } else { - #[allow(clippy::unnecessary_cast)] - // SAFETY: mem::size_of::() != 8, so Lane = u32 and - // residues is backed by u32 elements. + } + #[cfg(any(force_bits = "32", target_pointer_width = "32"))] + { let residues_u32: &[u32] = unsafe { core::slice::from_raw_parts(residues.as_ptr() as *const u32, residues.len()) }; - do_crt_u32( - prod, - residues_u32, - k_eff, - nn, - output_coeffs, - b_pack, - &primes_u32, - &crt_inv_u32, + do_crt::( + prod, residues_u32, k_eff, nn, output_coeffs, b_pack, &primes_u32, &crt_inv_u32, ); } // ---- Fold product into c with sign ---- let output_words = la + lb; - fold_prod_into_c(c, sign, prod, output_words) -} - -/// CRT + accumulate for 64-bit lanes (U192 accumulator). -#[allow(clippy::too_many_arguments)] -#[inline(never)] -fn do_crt_u64( - prod: &mut [u64], - residues: &[u64], - k_eff: usize, - nn: usize, - output_coeffs: usize, - b_pack: u32, - primes: &[u64; K], - crt_inv: &[[u64; K]; K], -) { - use crate::mul::ntt::crt::U192; - - for k in 0..output_coeffs { - let mut coeff_residues = [0u64; 3]; - #[allow(clippy::needless_range_loop)] - for pi in 0..k_eff { - coeff_residues[pi] = residues[pi * nn + k]; - } - let crt_val = - garner_combine::(&coeff_residues[..k_eff], crt_inv, primes); - add_shifted_to_prod(prod, crt_val.as_u64_slice(), crt_val.len_words(), k, b_pack); + match sign { + Positive => add::add_signed_in_place(c, Positive, &prod[..output_words]), + Negative => add::add_signed_in_place(c, Negative, &prod[..output_words]), } } -/// CRT + accumulate for 32-bit lanes (U96 accumulator). +/// CRT + accumulate, generic over the accumulator type. #[allow(clippy::too_many_arguments)] #[inline(never)] -fn do_crt_u32( - prod: &mut [u64], - residues: &[u32], +fn do_crt( + prod: &mut [Word], + residues: &[A::Lane], k_eff: usize, nn: usize, output_coeffs: usize, b_pack: u32, - primes: &[u32; K], - crt_inv: &[[u32; K]; K], + primes: &[A::Lane; K], + crt_inv: &[[A::Lane; K]; K], ) { - use crate::mul::ntt::crt::U96; - for k in 0..output_coeffs { - let mut coeff_residues = [0u32; 3]; + let mut coeff_residues = [A::Lane::default(); 3]; #[allow(clippy::needless_range_loop)] for pi in 0..k_eff { coeff_residues[pi] = residues[pi * nn + k]; } - let crt_val = - garner_combine::(&coeff_residues[..k_eff], crt_inv, primes); - add_shifted_to_prod(prod, crt_val.as_u64_slice(), crt_val.len_words(), k, b_pack); + let crt_val = garner_combine::(&coeff_residues[..k_eff], crt_inv, primes); + let mut crt_buf = [Word::default(); 6]; + let crt_n = crt_val.write_words(&mut crt_buf); + add_shifted_to_prod(prod, &crt_buf[..crt_n as usize], crt_n, k, b_pack); } } @@ -335,7 +289,6 @@ struct TransformCtx<'a> { b_lane: &'a mut [crate::arch::ntt::Lane], fwd_twiddles: &'a mut [crate::arch::ntt::Lane], inv_twiddles: &'a mut [crate::arch::ntt::Lane], - p: crate::arch::ntt::Lane, omega_max: crate::arch::ntt::Lane, nn: usize, b_pack: u32, @@ -343,159 +296,95 @@ struct TransformCtx<'a> { pi: usize, } -/// Per-prime NTT pipeline, monomorphized for a specific prime index `PI`. +/// Per-prime NTT pipeline, monomorphized for a specific reducer `R`. #[inline(never)] -fn process_prime( +fn process_prime>( a: &[Word], b: &[Word], ctx: &mut TransformCtx<'_>, + r: &R, ) { - use crate::arch::ntt::{to_monty, from_monty}; - pack::pack(ctx.a_lane, a, ctx.b_pack, ctx.nn); pack::pack(ctx.b_lane, b, ctx.b_pack, ctx.nn); // Convert standard-form coefficients to Montgomery form. - // transform() handles any value in [0, 2^BITS), no pre-reduction needed. for c in ctx.a_lane[..ctx.nn].iter_mut() { - *c = to_monty::(*c); + *c = r.transform(*c); } for c in ctx.b_lane[..ctx.nn].iter_mut() { - *c = to_monty::(*c); + *c = r.transform(*c); } - transform::precompute_twiddles::(ctx.fwd_twiddles, ctx.nn, ctx.p, ctx.omega_max, false); - transform::precompute_twiddles::(ctx.inv_twiddles, ctx.nn, ctx.p, ctx.omega_max, true); + transform::precompute_twiddles( + ctx.fwd_twiddles, ctx.nn, ctx.omega_max, false, r, + ); + transform::precompute_twiddles( + ctx.inv_twiddles, ctx.nn, ctx.omega_max, true, r, + ); transform::bit_reverse(ctx.a_lane); transform::bit_reverse(ctx.b_lane); - transform::forward::(ctx.a_lane, ctx.fwd_twiddles); - transform::forward::(ctx.b_lane, ctx.fwd_twiddles); - transform::pointwise_mul::(ctx.a_lane, ctx.b_lane); - transform::inverse::(ctx.a_lane, ctx.inv_twiddles, ctx.p); + transform::forward(ctx.a_lane, ctx.fwd_twiddles, r); + transform::forward(ctx.b_lane, ctx.fwd_twiddles, r); + transform::pointwise_mul(ctx.a_lane, ctx.b_lane, r); + transform::inverse(ctx.a_lane, ctx.inv_twiddles, r); // Convert residues back from Montgomery to standard form. for c in ctx.a_lane[..ctx.nn].iter_mut() { - *c = from_monty::(*c); + *c = r.residue(*c); } let offset = ctx.pi * ctx.nn; ctx.residues[offset..offset + ctx.nn].copy_from_slice(ctx.a_lane); } -/// Add a CRT value (as u64 words) to `prod`, shifted left by `k * b_pack` bits. -fn add_shifted_to_prod(prod: &mut [u64], words: &[u64], count: u32, k: usize, b_pack: u32) { +/// Add a CRT value (as `Word`-sized limbs) to `prod`, shifted left by +/// `k * b_pack` bits. +fn add_shifted_to_prod(prod: &mut [Word], words: &[Word], count: u32, k: usize, b_pack: u32) { let shift_bits = (k as u32).wrapping_mul(b_pack); - let word_idx = (shift_bits / 64) as usize; - let bit_shift = shift_bits % 64; + let word_bits = Word::BITS; + let start_idx = (shift_bits / word_bits) as usize; + let bit_shift = shift_bits % word_bits; + + let mut carry: Word = 0; - let mut carry: u64 = 0; - #[allow(clippy::needless_range_loop)] for vi in 0..(count as usize) { - let idx = word_idx + vi; + let limb = words[vi].wrapping_add(carry); + let idx = start_idx + vi; if idx >= prod.len() { return; } - let v = words[vi]; - let v128 = v as u128; if bit_shift == 0 { - let sum = v128.wrapping_add(carry as u128); - let (r, c) = prod[idx].overflowing_add(sum as u64); + let (r, c) = prod[idx].overflowing_add(limb); prod[idx] = r; - carry = (sum >> 64) as u64 + c as u64; + carry = Word::from(c); } else { - let lo = v128 << bit_shift; - let sum = lo.wrapping_add(carry as u128); - let lo_carry = (sum >> 64) as u64; - let lo_word = sum as u64; + let val = (limb as u128) << bit_shift; + let lo = val as Word; + let hi = (val >> word_bits) as Word; - let (r, c1) = prod[idx].overflowing_add(lo_word); + let (r, c1) = prod[idx].overflowing_add(lo); prod[idx] = r; - carry = lo_carry + c1 as u64; + carry = Word::from(c1).wrapping_add(hi); if idx + 1 < prod.len() && carry != 0 { let (r2, c2) = prod[idx + 1].overflowing_add(carry); prod[idx + 1] = r2; - carry = c2 as u64; + carry = Word::from(c2); } } } - let mut idx = word_idx + count as usize; + let mut idx = start_idx + count as usize; while carry != 0 && idx < prod.len() { let (r, c) = prod[idx].overflowing_add(carry); prod[idx] = r; - carry = c as u64; + carry = Word::from(c); idx += 1; } } -/// Fold the u64 product buffer into the Word output array `c`. -/// -/// For 64-bit Word targets, `u64` == `Word`, so a direct transmute suffices. -/// For 32-bit Word targets, each u64 is split into two u32 words with carry. -fn fold_prod_into_c(c: &mut [Word], sign: Sign, prod: &[u64], output_words: usize) -> SignedWord { - if mem::size_of::() == 8 { - // 64-bit Word: direct transmute - assert_eq!( - mem::size_of::(), - mem::size_of::(), - "NTT requires 64-bit Word" - ); - // SAFETY: Word and u64 have the same size (asserted above) and - // prod is allocated with u64 alignment, compatible with Word. - let prod_words: &[Word] = - unsafe { core::slice::from_raw_parts(prod.as_ptr() as *const Word, output_words) }; - match sign { - Positive => add::add_signed_in_place(c, Positive, prod_words), - Negative => add::add_signed_in_place(c, Negative, prod_words), - } - } else { - // 32-bit Word: split each u64 into two u32 words. - // For 64-bit Word targets this branch is dead (eliminated by the compiler) - // but must still type-check. - let mut carry: u32 = 0; - let double_output_words = output_words.min(prod.len() * 2); - for i in 0..double_output_words { - let prod_word = prod[i / 2]; - let lo_word = if i % 2 == 0 { - (prod_word as u32).wrapping_add(carry) - } else { - ((prod_word >> 32) as u32).wrapping_add(carry) - }; - // Carry out of the low-word addition - if i % 2 == 0 { - let lo_before = prod_word as u32; - carry = (prod_word >> 32) as u32; - if lo_word < lo_before { - carry = carry.wrapping_add(1); - } - } - if i < c.len() { - let c_val = c[i] as u32; - match sign { - Positive => { - let (sum, c_out) = c_val.overflowing_add(lo_word); - c[i] = sum as Word; - carry = carry.wrapping_add(c_out as u32); - } - Negative => { - let (diff, b) = c_val.overflowing_sub(lo_word); - c[i] = diff as Word; - carry = carry.wrapping_add(b as u32); - } - } - } - } - if sign == Positive { - carry as SignedWord - } else { - -(carry as SignedWord) - } - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/integer/src/mul/ntt/pack.rs b/integer/src/mul/ntt/pack.rs index 39d570ae..0989beeb 100644 --- a/integer/src/mul/ntt/pack.rs +++ b/integer/src/mul/ntt/pack.rs @@ -38,7 +38,7 @@ pub fn pack(out: &mut [Lane], words: &[Word], b_pack: u32, n: usize) { } if bit_offset + b_pack <= word_bits { - *coeff = ((words[word_idx] >> bit_offset) & mask) as Lane; + *coeff = ((words[word_idx] as u64 >> bit_offset) & mask) as Lane; bit_offset += b_pack; if bit_offset == word_bits { bit_offset = 0; @@ -47,12 +47,10 @@ pub fn pack(out: &mut [Lane], words: &[Word], b_pack: u32, n: usize) { } else { let bits_first = word_bits - bit_offset; let bits_second = b_pack - bits_first; - let mut val = - (words[word_idx] >> bit_offset) & ((1u64 << bits_first) - 1); + let mut val = (words[word_idx] as u64 >> bit_offset) & ((1u64 << bits_first) - 1); word_idx += 1; if word_idx < words.len() { - val |= (words[word_idx] & ((1u64 << bits_second) - 1)) - << bits_first; + val |= (words[word_idx] as u64 & ((1u64 << bits_second) - 1)) << bits_first; } *coeff = val as Lane; bit_offset = bits_second; @@ -90,19 +88,16 @@ mod tests { let mut idx = word_idx; if bit_shift == 0 { - let (sum, c) = output - .get(idx) - .copied() - .unwrap_or(0) - .overflowing_add(lo); + let (sum, c) = output.get(idx).copied().unwrap_or(0).overflowing_add(lo); carry = Word::from(c); if idx < output.len() { output[idx] = sum; } idx += 1; } else { - let lo_part = lo << bit_shift; - let hi_part = if bit_shift > 0 { lo >> (64 - bit_shift) } else { 0 }; + let val = (lo as u128) << bit_shift; + let lo_part = val as Word; + let hi_part = (val >> word_bits) as Word; let (sum, c1) = output .get(idx) diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs index 11d0fc61..15a98de2 100644 --- a/integer/src/mul/ntt/transform.rs +++ b/integer/src/mul/ntt/transform.rs @@ -1,38 +1,37 @@ //! Iterative in-place radix-2 NTT over Proth primes `K * 2^N + 1`. //! -//! All functions are const-generic over `PI` (the prime index `0..K`). -//! Modular arithmetic delegates to `crate::arch::ntt`. +//! All functions are generic over `R: Reducer` so each prime's +//! reducer is monomorphized at the call site. -use crate::arch::ntt::{add_mod, mul_mod, sub_mod, to_monty, Lane, MAX_LOG_N}; -use num_modular::{ModularPow, ModularUnaryOps}; +use crate::arch::ntt::{Lane, MAX_LOG_N}; +use num_modular::Reducer; // ---- public API ---- /// Fill `out[0..n/2]` with twiddle factors `omega_n^k` in Montgomery form. /// /// Panics if `out.len() < n / 2`. -pub fn precompute_twiddles( +pub fn precompute_twiddles>( out: &mut [Lane], n: usize, - p: Lane, omega_max: Lane, inverse: bool, + r: &R, ) { assert!(out.len() >= n / 2); let shift = MAX_LOG_N - n.trailing_zeros(); - let omega_n = omega_max.powm(&((1u64 as Lane) << shift), &p); + let omega_max_mont = r.transform(omega_max); + let omega_n_mont = r.pow(omega_max_mont, &((1u64 as Lane) << shift)); - let base = if inverse { - omega_n.invm(&p).expect("omega_n not invertible") + let base_mont = if inverse { + r.inv(omega_n_mont).expect("omega_n not invertible") } else { - omega_n + omega_n_mont }; - // Convert base and 1 to Montgomery form - let base_mont = to_monty::(base); - out[0] = to_monty::(1); + out[0] = r.transform(1); for k in 1..(n / 2) { - out[k] = mul_mod::(out[k - 1], base_mont); + out[k] = r.mul(&out[k - 1], &base_mont); } } @@ -50,8 +49,8 @@ pub fn bit_reverse(a: &mut [Lane]) { } /// Forward NTT in place (decimation-in-time, radix-2). -pub fn forward(a: &mut [Lane], twiddles: &[Lane]) { - ntt_core::(a, twiddles); +pub fn forward>(a: &mut [Lane], twiddles: &[Lane], r: &R) { + ntt_core(a, twiddles, r); } /// Inverse NTT in place. @@ -60,21 +59,19 @@ pub fn forward(a: &mut [Lane], twiddles: &[Lane]) { /// in **natural order**. /// /// `twiddles` must have been precomputed with `inverse = true`. -pub fn inverse(a: &mut [Lane], twiddles: &[Lane], p: Lane) { +pub fn inverse>(a: &mut [Lane], twiddles: &[Lane], r: &R) { let n = a.len(); bit_reverse(a); - ntt_core::(a, twiddles); - let n_val = n as Lane; - let n_inv = n_val.invm(&p).expect("n not invertible mod p"); - // Convert n⁻¹ to Montgomery form so the result stays in Montgomery form. - let n_inv_mont = to_monty::(n_inv); + ntt_core(a, twiddles, r); + let n_mont = r.transform(n as Lane); + let n_inv_mont = r.inv(n_mont).expect("n not invertible mod p"); for x in a.iter_mut() { - *x = mul_mod::(*x, n_inv_mont); + *x = r.mul(x, &n_inv_mont); } } /// In-place radix-2 DIT NTT (Cooley–Tukey). -fn ntt_core(a: &mut [Lane], twiddles: &[Lane]) { +fn ntt_core>(a: &mut [Lane], twiddles: &[Lane], r: &R) { let n = a.len(); debug_assert!(n.is_power_of_two() && twiddles.len() == n / 2); @@ -86,9 +83,9 @@ fn ntt_core(a: &mut [Lane], twiddles: &[Lane]) { for i in (0..n).step_by(sub_len) { for j in 0..half { let u = a[i + j]; - let v = mul_mod::(a[i + j + half], twiddles[j * step]); - a[i + j] = add_mod::(u, v); - a[i + j + half] = sub_mod::(u, v); + let v = r.mul(&a[i + j + half], &twiddles[j * step]); + a[i + j] = r.add(&u, &v); + a[i + j + half] = r.sub(&u, &v); } } @@ -97,17 +94,18 @@ fn ntt_core(a: &mut [Lane], twiddles: &[Lane]) { } /// Pointwise multiply of two transformed vectors in place. -pub fn pointwise_mul(a_hat: &mut [Lane], b_hat: &[Lane]) { +pub fn pointwise_mul>(a_hat: &mut [Lane], b_hat: &[Lane], r: &R) { assert_eq!(a_hat.len(), b_hat.len()); for (a, &b_val) in a_hat.iter_mut().zip(b_hat.iter()) { - *a = mul_mod::(*a, b_val); + *a = r.mul(a, &b_val); } } #[cfg(test)] mod tests { use super::*; - use crate::arch::ntt::{from_monty, to_monty, K, MODULI, OMEGA_MAX}; + use crate::arch::ntt::{K, MODULI, OMEGA_MAX, P0, P1, P2}; + use num_modular::{ModularCoreOps, ModularPow, ModularUnaryOps}; #[cfg(not(feature = "std"))] use alloc::vec; #[cfg(not(feature = "std"))] @@ -121,28 +119,22 @@ mod tests { } macro_rules! for_each_prime { - ($pi:ident, $p:ident, $omega:ident, $body:block) => { + ($r:ident, $p:ident, $omega:ident, $body:block) => { for idx in 0..K { let $p = MODULI[idx]; let $omega = OMEGA_MAX[idx]; match idx { 0 => { - let $pi: usize = 0; - let _ = $pi; - fn go($p: Lane, $omega: Lane) $body - go::<0>($p, $omega); + fn go>($r: &R, $p: Lane, $omega: Lane) $body + go::(&P0, $p, $omega); } 1 => { - let $pi: usize = 1; - let _ = $pi; - fn go($p: Lane, $omega: Lane) $body - go::<1>($p, $omega); + fn go>($r: &R, $p: Lane, $omega: Lane) $body + go::(&P1, $p, $omega); } 2 => { - let $pi: usize = 2; - let _ = $pi; - fn go($p: Lane, $omega: Lane) $body - go::<2>($p, $omega); + fn go>($r: &R, $p: Lane, $omega: Lane) $body + go::(&P2, $p, $omega); } _ => unreachable!(), } @@ -152,25 +144,22 @@ mod tests { #[test] fn test_forward_inverse_roundtrip() { - for_each_prime!(pi, p, omega, { + for_each_prime!(r, p, omega, { for &n in &[2, 4, 8, 16, 32, 64, 128, 256, 512] { let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; - precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); - precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); + precompute_twiddles(&mut fwd_twiddles, n, omega, false, r); + precompute_twiddles(&mut inv_twiddles, n, omega, true, r); let mut a: Vec = (0..n) .map(|i| ((i as Lane + 1).wrapping_mul(123456789)) % p) .collect(); - // Convert to Montgomery form for the NTT pipeline - for val in a.iter_mut() { - *val = to_monty::(*val); - } + for val in a.iter_mut() { *val = r.transform(*val); } let orig = a.clone(); bit_reverse(&mut a); - forward::(&mut a, &fwd_twiddles); - inverse::(&mut a, &inv_twiddles, p); + forward(&mut a, &fwd_twiddles, r); + inverse(&mut a, &inv_twiddles, r); assert_all_eq(&a, &orig, "roundtrip failed for n={n}"); } @@ -179,7 +168,7 @@ mod tests { #[test] fn test_convolution_via_ntt() { - for_each_prime!(pi, p, omega, { + for_each_prime!(r, p, omega, { for len_a in [1, 2, 3, 5] { for len_b in [1, 2, 3, 5] { let conv_len: usize = len_a + len_b - 1; @@ -190,40 +179,31 @@ mod tests { let b_vec: Vec = (0..len_b).map(|i| ((i + 1) as Lane * 67890) % p).collect(); - // Compute expected convolution in standard form let mut expected = vec![0u64 as Lane; conv_len]; for (i, &ai) in a.iter().enumerate() { for (j, &bj) in b_vec.iter().enumerate() { let prod = (ai as u128 * bj as u128 % p as u128) as Lane; - expected[i + j] = add_mod::(expected[i + j], prod); + expected[i + j] = r.add(&expected[i + j], &prod); } } - let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; + let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; - precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); - precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); + precompute_twiddles(&mut fwd_twiddles, n, omega, false, r); + precompute_twiddles(&mut inv_twiddles, n, omega, true, r); - // Convert inputs to Montgomery form let mut a_pad = vec![0u64 as Lane; n]; let mut b_pad = vec![0u64 as Lane; n]; - for i in 0..len_a { - a_pad[i] = to_monty::(a[i]); - } - for i in 0..len_b { - b_pad[i] = to_monty::(b_vec[i]); - } + for i in 0..len_a { a_pad[i] = r.transform(a[i]); } + for i in 0..len_b { b_pad[i] = r.transform(b_vec[i]); } bit_reverse(&mut a_pad); bit_reverse(&mut b_pad); - forward::(&mut a_pad, &fwd_twiddles); - forward::(&mut b_pad, &fwd_twiddles); - pointwise_mul::(&mut a_pad, &b_pad); - inverse::(&mut a_pad, &inv_twiddles, p); - // Convert results back to standard form - for val in a_pad[..conv_len].iter_mut() { - *val = from_monty::(*val); - } + forward(&mut a_pad, &fwd_twiddles, r); + forward(&mut b_pad, &fwd_twiddles, r); + pointwise_mul(&mut a_pad, &b_pad, r); + inverse(&mut a_pad, &inv_twiddles, r); + for val in a_pad[..conv_len].iter_mut() { *val = r.residue(*val); } assert_all_eq(&a_pad[..conv_len], &expected, "convolution mismatch"); } @@ -240,49 +220,43 @@ mod tests { /// Naive O(n²) NTT using standard-form modular arithmetic. #[allow(clippy::needless_range_loop)] - fn ntt_naive_std(x: &[Lane], omega_n: Lane, p: Lane) -> Vec { + fn ntt_naive_std(x: &[Lane], omega_n: Lane, p: Lane) -> Vec { + use num_modular::ModularCoreOps; let n = x.len(); let mut result = vec![0u64 as Lane; n]; for k in 0..n { - let mut acc: u128 = 0; + let mut acc = 0u64 as Lane; for j in 0..n { let twiddle = if k == 0 || j == 0 { 1 } else { - omega_n.powm(&((k * j) as Lane), &p) as u128 + omega_n.powm(&((k * j) as Lane), &p) }; - acc = (acc + x[j] as u128 * twiddle) % p as u128; + acc = acc.addm(x[j].mulm(twiddle, &p), &p); } - result[k] = acc as Lane; + result[k] = acc; } result } #[test] fn test_forward_correctness() { - for_each_prime!(pi, p, omega, { + for_each_prime!(r, p, omega, { for &n in &[2usize, 4, 8] { let x: Vec = (0..n).map(|i| ((i + 1) as Lane * 11111) % p).collect(); - let shift = MAX_LOG_N - n.trailing_zeros(); let omega_n = omega.powm(&((1u64 as Lane) << shift), &p); - // Convert to Montgomery form - let mut a: Vec = x.iter().map(|&v| to_monty::(v)).collect(); + let mut a: Vec = x.iter().map(|&v| r.transform(v)).collect(); bit_reverse(&mut a); let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; - precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); - forward::(&mut a, &fwd_twiddles); + precompute_twiddles(&mut fwd_twiddles, n, omega, false, r); + forward(&mut a, &fwd_twiddles, r); - // Convert forward output back to standard form for comparison - for val in a.iter_mut() { - *val = from_monty::(*val); - } - - // expected: compute naive NTT in standard form - let expected = ntt_naive_std::(&x, omega_n, p); + for val in a.iter_mut() { *val = r.residue(*val); } + let expected = ntt_naive_std(&x, omega_n, p); assert_eq!(a, expected, "forward NTT mismatch"); } }); @@ -292,68 +266,57 @@ mod tests { fn test_convolution_debug() { let p = MODULI[0]; let omega = OMEGA_MAX[0]; + let r = &P0; let a = [12345u64 as Lane % p]; - let b_vec = [67890u64 as Lane % p, - 135780u64 as Lane % p, - 203670u64 as Lane % p]; + let b_vec = [67890u64 as Lane % p, 135780u64 as Lane % p, 203670u64 as Lane % p]; let conv_len = a.len() + b_vec.len() - 1; let n = 4; - // Expected values in standard form let mut expected = vec![0u64 as Lane; conv_len]; for (i, &ai) in a.iter().enumerate() { for (j, &bj) in b_vec.iter().enumerate() { let prod = (ai as u128 * bj as u128 % p as u128) as Lane; - expected[i + j] = add_mod::<0>(expected[i + j], prod); + expected[i + j] = r.add(&expected[i + j], &prod); } } let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; - precompute_twiddles::<0>(&mut fwd_twiddles, n, p, omega, false); - precompute_twiddles::<0>(&mut inv_twiddles, n, p, omega, true); + precompute_twiddles(&mut fwd_twiddles, n, omega, false, r); + precompute_twiddles(&mut inv_twiddles, n, omega, true, r); - // Convert to Montgomery form let mut a_pad = vec![0u64 as Lane; n]; let mut b_pad = vec![0u64 as Lane; n]; - for i in 0..a.len() { - a_pad[i] = to_monty::<0>(a[i]); - } - for i in 0..b_vec.len() { - b_pad[i] = to_monty::<0>(b_vec[i]); - } + for i in 0..a.len() { a_pad[i] = r.transform(a[i]); } + for i in 0..b_vec.len() { b_pad[i] = r.transform(b_vec[i]); } bit_reverse(&mut a_pad); bit_reverse(&mut b_pad); - forward::<0>(&mut a_pad, &fwd_twiddles); - forward::<0>(&mut b_pad, &fwd_twiddles); - pointwise_mul::<0>(&mut a_pad, &b_pad); - inverse::<0>(&mut a_pad, &inv_twiddles, p); - // Convert back to standard form - for val in a_pad[..conv_len].iter_mut() { - *val = from_monty::<0>(*val); - } + forward(&mut a_pad, &fwd_twiddles, r); + forward(&mut b_pad, &fwd_twiddles, r); + pointwise_mul(&mut a_pad, &b_pad, r); + inverse(&mut a_pad, &inv_twiddles, r); + for val in a_pad[..conv_len].iter_mut() { *val = r.residue(*val); } assert_eq!(&a_pad[..conv_len], &expected[..]); } #[test] fn test_length_two_edge_case() { - for_each_prime!(pi, p, omega, { + for_each_prime!(r, p, omega, { let n = 2; let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; - precompute_twiddles::(&mut fwd_twiddles, n, p, omega, false); - precompute_twiddles::(&mut inv_twiddles, n, p, omega, true); + precompute_twiddles(&mut fwd_twiddles, n, omega, false, r); + precompute_twiddles(&mut inv_twiddles, n, omega, true, r); let a_std = [1u64 as Lane % p, 2u64 as Lane % p]; - // Convert to Montgomery form - let a_orig: Vec = a_std.iter().map(|&v| to_monty::(v)).collect(); + let a_orig: Vec = a_std.iter().map(|&v| r.transform(v)).collect(); let mut a = a_orig.clone(); bit_reverse(&mut a); - forward::(&mut a, &fwd_twiddles); - inverse::(&mut a, &inv_twiddles, p); + forward(&mut a, &fwd_twiddles, r); + inverse(&mut a, &inv_twiddles, r); assert_all_eq(&a, &a_orig, "length two roundtrip"); }); } From c78700cc4824db2a3835c156f76df92b35c507f0 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 11:30:49 +0800 Subject: [PATCH 12/19] Tune the NTT threshold --- integer/src/arch/generic_32_bit/ntt.rs | 47 +++++++++++--------------- integer/src/arch/generic_64_bit/ntt.rs | 41 +++++++++++----------- integer/src/mul/mod.rs | 4 +-- integer/src/mul/ntt/crt.rs | 16 +++++++-- integer/src/mul/ntt/mod.rs | 35 +++++++++++++------ integer/src/mul/ntt/transform.rs | 43 ++++++++++++++++------- 6 files changed, 111 insertions(+), 75 deletions(-) diff --git a/integer/src/arch/generic_32_bit/ntt.rs b/integer/src/arch/generic_32_bit/ntt.rs index 9f339b6a..6e6cbfb6 100644 --- a/integer/src/arch/generic_32_bit/ntt.rs +++ b/integer/src/arch/generic_32_bit/ntt.rs @@ -30,11 +30,7 @@ pub const OMEGA_MAX: [Lane; K] = [ 0x1aa0ab5e, // P2 ]; -pub const CRT_INV_IJ: [[Lane; K]; K] = [ - [0, 0x4e42c85b, 0x5fb425ef], - [0, 0, 0x44000009], - [0, 0, 0], -]; +pub const CRT_INV_IJ: [[Lane; K]; K] = [[0, 0x4e42c85b, 0x5fb425ef], [0, 0, 0x44000009], [0, 0, 0]]; /// Prime moduli indexed by PI. pub const MODULI: [Lane; K] = [ @@ -67,28 +63,25 @@ mod tests { fn test_omega_order() { for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { let p = MODULI[pi]; - let (sqr, to_m, from_m): ( - fn(Lane) -> Lane, - fn(Lane) -> Lane, - fn(Lane) -> Lane, - ) = match pi { - 0 => ( - |w| P0.reduce((w as u64) * (w as u64)), - |v| P0.transform(v), - |v| P0.residue(v), - ), - 1 => ( - |w| P1.reduce((w as u64) * (w as u64)), - |v| P1.transform(v), - |v| P1.residue(v), - ), - 2 => ( - |w| P2.reduce((w as u64) * (w as u64)), - |v| P2.transform(v), - |v| P2.residue(v), - ), - _ => unreachable!(), - }; + let (sqr, to_m, from_m): (fn(Lane) -> Lane, fn(Lane) -> Lane, fn(Lane) -> Lane) = + match pi { + 0 => ( + |w| P0.reduce((w as u64) * (w as u64)), + |v| P0.transform(v), + |v| P0.residue(v), + ), + 1 => ( + |w| P1.reduce((w as u64) * (w as u64)), + |v| P1.transform(v), + |v| P1.residue(v), + ), + 2 => ( + |w| P2.reduce((w as u64) * (w as u64)), + |v| P2.transform(v), + |v| P2.residue(v), + ), + _ => unreachable!(), + }; let mut w = to_m(omega_max); for _ in 0..MAX_LOG_N - 1 { diff --git a/integer/src/arch/generic_64_bit/ntt.rs b/integer/src/arch/generic_64_bit/ntt.rs index 2e20d9a0..86f6b159 100644 --- a/integer/src/arch/generic_64_bit/ntt.rs +++ b/integer/src/arch/generic_64_bit/ntt.rs @@ -68,28 +68,25 @@ mod tests { fn test_omega_order() { for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { let p = MODULI[pi]; - let (sqr, to_m, from_m): ( - fn(Lane) -> Lane, - fn(Lane) -> Lane, - fn(Lane) -> Lane, - ) = match pi { - 0 => ( - |w| P0.reduce((w as u128) * (w as u128)), - |v| P0.transform(v), - |v| P0.residue(v), - ), - 1 => ( - |w| P1.reduce((w as u128) * (w as u128)), - |v| P1.transform(v), - |v| P1.residue(v), - ), - 2 => ( - |w| P2.reduce((w as u128) * (w as u128)), - |v| P2.transform(v), - |v| P2.residue(v), - ), - _ => unreachable!(), - }; + let (sqr, to_m, from_m): (fn(Lane) -> Lane, fn(Lane) -> Lane, fn(Lane) -> Lane) = + match pi { + 0 => ( + |w| P0.reduce((w as u128) * (w as u128)), + |v| P0.transform(v), + |v| P0.residue(v), + ), + 1 => ( + |w| P1.reduce((w as u128) * (w as u128)), + |v| P1.transform(v), + |v| P1.residue(v), + ), + 2 => ( + |w| P2.reduce((w as u128) * (w as u128)), + |v| P2.transform(v), + |v| P2.residue(v), + ), + _ => unreachable!(), + }; let mut w = to_m(omega_max); for _ in 0..MAX_LOG_N - 1 { diff --git a/integer/src/mul/mod.rs b/integer/src/mul/mod.rs index 8e97f208..6f438957 100644 --- a/integer/src/mul/mod.rs +++ b/integer/src/mul/mod.rs @@ -437,8 +437,8 @@ mod threshold_tests { use std::time::Instant; let sizes: &[usize] = &[ - 5_000, 10_000, 20_000, 30_000, 40_000, 50_000, 60_000, 80_000, 100_000, 120_000, - 131_000, + 1_000, 2_000, 3_000, 4_000, 5_000, 6_000, 7_000, 8_000, 9_000, 10_000, 20_000, 40_000, + 80_000, ]; println!( diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs index 926029f2..045f5074 100644 --- a/integer/src/mul/ntt/crt.rs +++ b/integer/src/mul/ntt/crt.rs @@ -117,7 +117,13 @@ mod triple_impl { out[0] = self.0[0]; out[1] = self.0[1]; out[2] = self.0[2]; - if self.0[2] != 0 { 3 } else if self.0[1] != 0 { 2 } else { 1 } + if self.0[2] != 0 { + 3 + } else if self.0[1] != 0 { + 2 + } else { + 1 + } } } } @@ -169,7 +175,13 @@ mod triple_impl { out[0] = self.0[0]; out[1] = self.0[1]; out[2] = self.0[2]; - if self.0[2] != 0 { 3 } else if self.0[1] != 0 { 2 } else { 1 } + if self.0[2] != 0 { + 3 + } else if self.0[1] != 0 { + 2 + } else { + 1 + } } } } diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index e6e2f7bc..aef7b2c0 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -16,12 +16,17 @@ pub(crate) mod crt; mod pack; mod transform; -use crate::arch::ntt::{B_PACK_CANDIDATES, B_PACK_MIN, K, MAX_LOG_N, MODULI, OMEGA_MAX, P0, P1, P2}; +use crate::arch::ntt::{ + B_PACK_CANDIDATES, B_PACK_MIN, K, MAX_LOG_N, MODULI, OMEGA_MAX, P0, P1, P2, +}; use crate::mul::ntt::crt::{garner_combine, CrtAccum}; use num_modular::Reducer; /// Minimum smaller-operand length (in words) for the NTT path. -pub const THRESHOLD_NTT: usize = 40_000; +/// +/// Crossover with Toom-3 lies at ~3 200 words; chosen at 4 000 where +/// NTT is a clear 30%+ faster. +pub const THRESHOLD_NTT: usize = 4_000; /// Select NTT parameters for operands with the given word lengths. /// @@ -237,7 +242,14 @@ fn add_signed_mul_impl( let residues_u64: &[u64] = unsafe { core::slice::from_raw_parts(residues.as_ptr() as *const u64, residues.len()) }; do_crt::( - prod, residues_u64, k_eff, nn, output_coeffs, b_pack, &primes_u64, &crt_inv_u64, + prod, + residues_u64, + k_eff, + nn, + output_coeffs, + b_pack, + &primes_u64, + &crt_inv_u64, ); } #[cfg(any(force_bits = "32", target_pointer_width = "32"))] @@ -245,7 +257,14 @@ fn add_signed_mul_impl( let residues_u32: &[u32] = unsafe { core::slice::from_raw_parts(residues.as_ptr() as *const u32, residues.len()) }; do_crt::( - prod, residues_u32, k_eff, nn, output_coeffs, b_pack, &primes_u32, &crt_inv_u32, + prod, + residues_u32, + k_eff, + nn, + output_coeffs, + b_pack, + &primes_u32, + &crt_inv_u32, ); } @@ -315,12 +334,8 @@ fn process_prime>( *c = r.transform(*c); } - transform::precompute_twiddles( - ctx.fwd_twiddles, ctx.nn, ctx.omega_max, false, r, - ); - transform::precompute_twiddles( - ctx.inv_twiddles, ctx.nn, ctx.omega_max, true, r, - ); + transform::precompute_twiddles(ctx.fwd_twiddles, ctx.nn, ctx.omega_max, false, r); + transform::precompute_twiddles(ctx.inv_twiddles, ctx.nn, ctx.omega_max, true, r); transform::bit_reverse(ctx.a_lane); transform::bit_reverse(ctx.b_lane); diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs index 15a98de2..631932c6 100644 --- a/integer/src/mul/ntt/transform.rs +++ b/integer/src/mul/ntt/transform.rs @@ -105,11 +105,11 @@ pub fn pointwise_mul>(a_hat: &mut [Lane], b_hat: &[Lane], r: &R mod tests { use super::*; use crate::arch::ntt::{K, MODULI, OMEGA_MAX, P0, P1, P2}; - use num_modular::{ModularCoreOps, ModularPow, ModularUnaryOps}; #[cfg(not(feature = "std"))] use alloc::vec; #[cfg(not(feature = "std"))] use alloc::vec::Vec; + use num_modular::{ModularCoreOps, ModularPow, ModularUnaryOps}; fn assert_all_eq(a: &[Lane], b_val: &[Lane], context: &str) { assert_eq!(a.len(), b_val.len(), "{context}: length mismatch"); @@ -154,7 +154,9 @@ mod tests { let mut a: Vec = (0..n) .map(|i| ((i as Lane + 1).wrapping_mul(123456789)) % p) .collect(); - for val in a.iter_mut() { *val = r.transform(*val); } + for val in a.iter_mut() { + *val = r.transform(*val); + } let orig = a.clone(); bit_reverse(&mut a); @@ -174,8 +176,7 @@ mod tests { let conv_len: usize = len_a + len_b - 1; let n = conv_len.next_power_of_two().max(2); - let a: Vec = - (0..len_a).map(|i| ((i + 1) as Lane * 12345) % p).collect(); + let a: Vec = (0..len_a).map(|i| ((i + 1) as Lane * 12345) % p).collect(); let b_vec: Vec = (0..len_b).map(|i| ((i + 1) as Lane * 67890) % p).collect(); @@ -194,8 +195,12 @@ mod tests { let mut a_pad = vec![0u64 as Lane; n]; let mut b_pad = vec![0u64 as Lane; n]; - for i in 0..len_a { a_pad[i] = r.transform(a[i]); } - for i in 0..len_b { b_pad[i] = r.transform(b_vec[i]); } + for i in 0..len_a { + a_pad[i] = r.transform(a[i]); + } + for i in 0..len_b { + b_pad[i] = r.transform(b_vec[i]); + } bit_reverse(&mut a_pad); bit_reverse(&mut b_pad); @@ -203,7 +208,9 @@ mod tests { forward(&mut b_pad, &fwd_twiddles, r); pointwise_mul(&mut a_pad, &b_pad, r); inverse(&mut a_pad, &inv_twiddles, r); - for val in a_pad[..conv_len].iter_mut() { *val = r.residue(*val); } + for val in a_pad[..conv_len].iter_mut() { + *val = r.residue(*val); + } assert_all_eq(&a_pad[..conv_len], &expected, "convolution mismatch"); } @@ -255,7 +262,9 @@ mod tests { precompute_twiddles(&mut fwd_twiddles, n, omega, false, r); forward(&mut a, &fwd_twiddles, r); - for val in a.iter_mut() { *val = r.residue(*val); } + for val in a.iter_mut() { + *val = r.residue(*val); + } let expected = ntt_naive_std(&x, omega_n, p); assert_eq!(a, expected, "forward NTT mismatch"); } @@ -269,7 +278,11 @@ mod tests { let r = &P0; let a = [12345u64 as Lane % p]; - let b_vec = [67890u64 as Lane % p, 135780u64 as Lane % p, 203670u64 as Lane % p]; + let b_vec = [ + 67890u64 as Lane % p, + 135780u64 as Lane % p, + 203670u64 as Lane % p, + ]; let conv_len = a.len() + b_vec.len() - 1; let n = 4; @@ -288,8 +301,12 @@ mod tests { let mut a_pad = vec![0u64 as Lane; n]; let mut b_pad = vec![0u64 as Lane; n]; - for i in 0..a.len() { a_pad[i] = r.transform(a[i]); } - for i in 0..b_vec.len() { b_pad[i] = r.transform(b_vec[i]); } + for i in 0..a.len() { + a_pad[i] = r.transform(a[i]); + } + for i in 0..b_vec.len() { + b_pad[i] = r.transform(b_vec[i]); + } bit_reverse(&mut a_pad); bit_reverse(&mut b_pad); @@ -297,7 +314,9 @@ mod tests { forward(&mut b_pad, &fwd_twiddles, r); pointwise_mul(&mut a_pad, &b_pad, r); inverse(&mut a_pad, &inv_twiddles, r); - for val in a_pad[..conv_len].iter_mut() { *val = r.residue(*val); } + for val in a_pad[..conv_len].iter_mut() { + *val = r.residue(*val); + } assert_eq!(&a_pad[..conv_len], &expected[..]); } From 54adac525abae3b6d49887b66c757462869911a5 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 15:49:48 +0800 Subject: [PATCH 13/19] Tidy up --- NTT_RADIX4.md | 191 ++++++++++ TODO_NTT.md | 149 ++------ integer/CHANGELOG.md | 9 +- integer/benches/primitive.rs | 26 ++ integer/src/arch/generic_32_bit/ntt.rs | 16 +- integer/src/arch/generic_64_bit/ntt.rs | 16 +- integer/src/mul/ntt/mod.rs | 463 ++++++++++++++++--------- integer/src/mul/ntt/transform.rs | 2 +- 8 files changed, 568 insertions(+), 304 deletions(-) create mode 100644 NTT_RADIX4.md diff --git a/NTT_RADIX4.md b/NTT_RADIX4.md new file mode 100644 index 00000000..77088445 --- /dev/null +++ b/NTT_RADIX4.md @@ -0,0 +1,191 @@ +# Radix-4 NTT for `dashu-int` + +## Context + +The NTT multiplication path in `integer/src/mul/ntt/` currently uses an iterative in-place **radix-2** DIT transform. `TODO_NTT.md` lists "Radix-4 or split-radix NTT" as optimisation opportunity #1 — roughly halving the number of memory passes (`log₄(N)` stages instead of `log₂(N)`), with a small reduction in arithmetic work. Memory traffic, not arithmetic, is the dominant cost on the Goldilocks-style primes used here, so the pass count is the win that matters. + +This plan implements radix-4 in place of the current radix-2 core, while keeping the rest of the pipeline (`pack`, `bit_reverse`, `pointwise_mul`, CRT, Garner, and the recently-introduced `NttGeometry` / `prepare_b_hat_and_twiddles` / `run_ntt_pipeline` / `process_prime` layering) untouched. A legacy copy of the radix-2 core is kept during development as a differential test oracle, then removed before the final commit. + +Work happens in a new worktree `ntt-radix4` branching off the current `ssa` branch. + +--- + +## Current state of the code (post-refactor) + +The NTT module has been re-organised so that the conv and chunked paths share a single symmetric pipeline, with geometry constants factored into a small value type: + +- `add_signed_mul` (mod.rs:143) dispatches: `a.len() > 2 * b.len()` → `add_signed_mul_chunked`; otherwise `add_signed_mul_conv`. +- **Both paths pre-transform `b` and precompute twiddles up front**, then call `run_ntt_pipeline`. There is no "raw b" vs "cached b" distinction — `b_hat`, `fwd_tw_cache`, `inv_tw_cache` are always populated by `prepare_b_hat_and_twiddles` (mod.rs:412) and passed in. +- `NttGeometry` (mod.rs:372) is a small value struct holding `nn`, `b_pack`, `k_eff`, `output_coeffs`. It's passed by reference into `prepare_b_hat_and_twiddles` and `run_ntt_pipeline`, and embedded inside `TransformCtx`. +- `TransformCtx` (mod.rs:380) is now just four scratch buffer slices (`a_lane`, `b_lane`, `fwd_twiddles`, `inv_twiddles`) plus a `geom: NttGeometry`. `prod` and `residues` are not in the ctx — they are allocated in `run_ntt_pipeline` (mod.rs:267–268) and passed as separate `&mut` arguments to `process_prime` and `do_crt`. +- `run_ntt_pipeline` (mod.rs:251) owns the per-call scratch allocation, runs the per-prime loop calling `process_prime`, then calls `do_crt` and signed-accumulates into `c_out`. +- `process_prime` (mod.rs:464) takes `(a, b_hat_slice, ctx, residues, pi, r)`. It transforms `a` from raw words, copies the pre-transformed `b_hat_slice` into `b_lane`, pointwise-multiplies, inverse-transforms, and writes residues. No twiddle precompute happens here. +- `transform_b_forward` (mod.rs:391) is the helper used by `prepare_b_hat_and_twiddles` to pack/Montgomery/bit-reverse/forward-transform `b`. + +**Implication for radix-4:** the transform-level changes (radix-4 butterfly, expanded twiddle table) live entirely in `transform::forward` / `transform::inverse` / `transform::precompute_twiddles` / `transform::ntt_core`. Because every code path reaches the transform through these, the speedup propagates everywhere for free. The only multi-site edits in `mod.rs` are the twiddle *allocation sizes* and the cache-offset arithmetic, which now live in a small number of well-defined places. + +--- + +## Math summary (verified) + +Radix-4 DIT takes bit-reversed input and produces natural-order output — same I/O contract as the existing radix-2 DIT, so `bit_reverse`, `inverse()`, and all callers are unchanged. + +For each butterfly on quad `(a0, a1, a2, a3)` at positions `(k, k+q, k+2q, k+3q)` within a length-`sub_len` group, with `q = sub_len/4` and `step = n/sub_len`: + +``` +b1 = a1 · ω_n^(k·step) +b2 = a2 · ω_n^(2k·step) +b3 = a3 · ω_n^(3k·step) +e0 = a0 + b2 +e1 = a0 − b2 +e2 = b1 + b3 +e3 = b1 − b3 // order matters: b1 − b3, not b3 − b1 +y0 = e0 + e2 +y1 = e1 + j·e3 // j = ω_n^(n/4), read from twiddles[n/4] +y2 = e0 − e2 +y3 = e1 − j·e3 +``` + +Iterative structure: +- If `log₂(n)` is even, stages run with `sub_len = 4, 16, 64, …, n` (pure radix-4). +- If `log₂(n)` is odd (n = 2·4^L), run **one** radix-2 stage with `sub_len = 2` (uses only `twiddles[0] = 1`), then radix-4 stages with `sub_len = 8, 32, …, n`. +- `n = 2` is a degenerate case — emit a single radix-2 butterfly with twiddle 1 and return early. + +The constant `j = ω_n^(n/4)` is read once from `twiddles[n/4]` at the top of `ntt_core`. For inverse transforms, `twiddles[n/4]` holds `ω_n^(−n/4) = −j`, which is the *other* primitive 4th root; the same butterfly formula applies with it (a sign flip on the `j·e3` terms). No special handling needed — the symmetry falls out naturally. + +The maximum twiddle index touched is `3k·step ≤ 3(n/4 − 1) ≈ 3n/4` at the final stage, which exceeds the current `n/2`-long table. **Fix: expand the twiddle table from `n/2` to `n` lanes.** Memory overhead ≈ +`n` lanes per table. + +--- + +## Files to modify + +### `integer/src/mul/ntt/transform.rs` (primary rewrite) + +1. **`precompute_twiddles`** — change `assert!(out.len() >= n / 2)` to `assert!(out.len() >= n)`, and extend the fill loop from `1..(n/2)` to `1..n`. Output is now `ω_n^k` for `k ∈ [0, n)`. + +2. **`ntt_core` (rewrite)** — replace with the radix-4 algorithm: + ```rust + fn ntt_core>(a: &mut [Lane], twiddles: &[Lane], r: &R) { + let n = a.len(); + debug_assert!(n.is_power_of_two() && twiddles.len() >= n); + if n == 1 { return; } + if n == 2 { + // Radix-2 fallback: twiddle = twiddles[0] = 1. + let u = a[0]; let v = a[1]; + a[0] = r.add(&u, &v); + a[1] = r.sub(&u, &v); + return; + } + let j_mont = twiddles[n / 4]; + + let log_n = n.trailing_zeros(); + let mut sub_len = if log_n & 1 == 1 { + // One radix-2 stage with step = n/2 (only k=0, twiddle = 1). + for i in (0..n).step_by(2) { + let u = a[i]; let v = a[i + 1]; + a[i] = r.add(&u, &v); + a[i + 1] = r.sub(&u, &v); + } + 8 + } else { + 4 + }; + + // Radix-4 stages. + while sub_len <= n { + let q = sub_len / 4; + let step = n / sub_len; + for i in (0..n).step_by(sub_len) { + // k = 0: twiddles are all 1, skip the multiplies. + butterfly_radix4(a, i, q, twiddles[0], twiddles[0], twiddles[0], j_mont, r); + for k in 1..q { + let t1 = twiddles[k * step]; + let t2 = twiddles[2 * k * step]; + let t3 = twiddles[3 * k * step]; + butterfly_radix4(a, i + k, q, t1, t2, t3, j_mont, r); + } + } + sub_len *= 4; + } + } + ``` + `butterfly_radix4` is a small `#[inline(always)]` helper that performs the four muls (`b1, b2, b3` + `j·e3`) and writes back to the four positions. Make sure reads of `a[idx]` happen before any writes. + +3. **Keep a private `ntt_core_radix2_legacy`** during development — the current body of `ntt_core`, renamed. It reads only `twiddles[0..n/2]` so it works fine on the expanded table. Used only by the cross-check test (below) and deleted before merge. + +4. **Tests** — update local allocations in tests from `n/2` to `n` lanes. Extend `test_forward_correctness` to cover `n ∈ {2, 4, 8, 16, 32}` (currently only `{2, 4, 8}`). Add: + - `test_radix4_matches_legacy` — for each prime and `n ∈ {2, 4, 8, 16, 32, 64, 128, 256}`, run forward via both `ntt_core` and `ntt_core_radix2_legacy` on identical random bit-reversed input, assert byte-equal output. Delete together with the legacy fn before merge. + +### `integer/src/mul/ntt/mod.rs` (memory layout — four logical sites) + +Because of the recent refactor, twiddle allocation is centralised. The current code uses `nn / 2` for twiddle lengths in a small number of well-defined places, all of which need to become `nn`: + +1. **`memory_requirement_up_to` (line 98)** — worst-case scratch bound. The `twiddles = n_max` constant assumes two tables of size `n/2` (forward + inverse). Bump to `2 * n_max`: + ```rust + let twiddles = 2 * n_max; // was n_max + ``` + +2. **`run_ntt_pipeline` (lines 271–272, 283, 284–287)** — per-call scratch allocation + cache slicing inside the per-prime loop: + ```rust + // line 271-272 + let (fwd_twiddles, mut m) = m.allocate_slice_fill::(nn, 0); // was nn / 2 + let (inv_twiddles, _) = m.allocate_slice_fill::(nn, 0); // was nn / 2 + // line 283 + let tw_off = pi * nn; // was pi * (nn / 2) + // line 284-287 + ctx.fwd_twiddles.copy_from_slice(&fwd_tw_cache[tw_off..tw_off + nn]); // was nn / 2 + ctx.inv_twiddles.copy_from_slice(&inv_tw_cache[tw_off..tw_off + nn]); // was nn / 2 + ``` + +3. **`prepare_b_hat_and_twiddles` (lines 411 docstring, 428–429, 452, 454–455)** — per-prime scratch during precompute + cache write-back: + ```rust + // docstring at line ~411: "fwd_tw_cache and inv_tw_cache each geom.k_eff * geom.nn" + // (was geom.k_eff * (geom.nn / 2)) + // line 428-429 + let (fwd_tw, mut rest) = rest.allocate_slice_fill::(nn, 0); // was nn / 2 + let (inv_tw, _) = rest.allocate_slice_fill::(nn, 0); // was nn / 2 + // line 452 + let tw_off = pi * nn; // was pi * (nn / 2) + // line 454-455 + fwd_tw_cache[tw_off..tw_off + nn].copy_from_slice(fwd_tw); // was nn / 2 + inv_tw_cache[tw_off..tw_off + nn].copy_from_slice(inv_tw); // was nn / 2 + ``` + +4. **Cache length computations** — two sites that derive the total cache size from `nn`: + - `add_signed_mul_chunked` line 184: `let twiddle_len = k_eff * nn_chunk;` (was `k_eff * (nn_chunk / 2)`) + - `add_signed_mul_conv` line 334: `let twiddle_len = k_eff * nn;` (was `k_eff * (nn / 2)`) + +`NttGeometry` itself does **not** store twiddle size — only `nn`, `b_pack`, `k_eff`, `output_coeffs`. The `nn / 2` → `nn` change is local to the four sites above; no field needs adding to the geometry struct. + +### `integer/CHANGELOG.md` + +Add under `## Unreleased` → `### Improve`: +> NTT inner transform rewritten as radix-4 DIT (with one radix-2 stage when N is not a power of 4), halving the number of passes over the coefficient array. ~20–30% faster large-integer multiplication above the NTT threshold. + +### `TODO_NTT.md` + +Mark section "1. Radix-4 or split-radix NTT" as completed (move from "Remaining" to "Implemented", or strike through with a dated note). Leave split-radix as a possible future improvement. + +### No changes + +- `integer/src/arch/generic_64_bit/ntt.rs` and `integer/src/arch/generic_32_bit/ntt.rs` — primes and `OMEGA_MAX` are already sufficient. `MAX_LOG_N ≥ 2` is all radix-4 needs. +- `integer/src/mul/ntt/pack.rs`, `crt.rs` — unaffected. +- All callers in `mod.rs` (`add_signed_mul_conv`, `add_signed_mul_chunked` closure, `prepare_b_hat_and_twiddles`, `transform_b_forward`, `process_prime`, `do_crt`, `NttGeometry`, `TransformCtx`) — unaffected because `forward`/`inverse` signatures are unchanged. + +--- + +## Implementation sequence + +1. **Enter the worktree** `ntt-radix4` off the current `ssa` HEAD (via the worktree tool). +2. **Commit 1 — Twiddle table expansion.** Bump allocation sizes and cache-offset arithmetic in `transform.rs::precompute_twiddles`, and at all the mod.rs sites listed above. Also bump test-local allocations. The old radix-2 still works correctly on the now-oversized table. Run `cargo test -p dashu-int mul::ntt` — everything should pass. +3. **Commit 2 — Radix-4 core + legacy differential test.** Rename the existing `ntt_core` body to `ntt_core_radix2_legacy`. Write the new `ntt_core` (radix-4). Add `test_radix4_matches_legacy`. Extend `test_forward_correctness` to n ∈ {2, 4, 8, 16, 32}. Run the full NTT test suite; the schoolbook comparison tests in `mod.rs` are the strongest correctness gate. +4. **Commit 3 — Cleanup + docs.** Delete `ntt_core_radix2_legacy` and `test_radix4_matches_legacy`. Update `CHANGELOG.md` and `TODO_NTT.md`. +5. **Verification.** Run the existing `crossover_ntt` ignored test (`cargo test -p dashu-int --release -- mul::threshold_tests::crossover_ntt --ignored --nocapture`) and compare timings before/after by checking out `ssa` temporarily. A 20–30% drop in the NTT column at sizes ≥ 4096 words confirms the optimisation landed. + +## Verification end-to-end + +- `cargo check --all-features --tests` +- `cargo test --workspace --exclude dashu-python` +- `cargo clippy --all-features --all-targets --workspace --exclude dashu-python -- -D warnings` +- `cargo fmt --all -- --check` +- `cargo test -p dashu-int --release -- mul::threshold_tests::crossover_ntt --ignored --nocapture` — for end-to-end timing sanity. diff --git a/TODO_NTT.md b/TODO_NTT.md index 295a6b4a..537634a4 100644 --- a/TODO_NTT.md +++ b/TODO_NTT.md @@ -2,131 +2,56 @@ ## Implemented -### Core algorithm (Phases 1–6) - -- **Primes.** Three Goldilocks-style Solinas primes `p = 2^64 − 2^b + 1` with - `b ∈ {32, 34, 40}`. All support shift-based reduction via `2^64 ≡ 2^b − 1`. - Stored in `integer/src/mul/ntt/primes.rs` with a full `verify_primes()` test - (Miller–Rabin, exact root order, reduction-identity self-check). - -- **Modular arithmetic.** Lane arithmetic delegates to - `num_modular::FixedTrinomialSolinas64` for `add`/`sub`/`mul`/`reduce_double`. - `add`/`sub`/`mul` are monomorphized per `B` via const generics so the - compiler fully inlines each prime's hot path. - -- **NTT transforms.** Iterative in-place radix-2 decimation-in-time - (`integer/src/mul/ntt/transform.rs`). Forward transform: `bit_reverse → - forward(ω)`. Inverse transform: `bit_reverse → forward(ω⁻¹) → scale by - N⁻¹`. Twiddle tables precomputed once per prime per call. Pointwise - multiply in the transform domain. - -- **Packing / unpacking.** Bit-level `pack` slices `&[Word]` into `b_pack`-bit - coefficients (`integer/src/mul/ntt/pack.rs`). CRT-recovered coefficients are - accumulated into the output limb array via shifted addition - (`add_shifted_to_prod`). - -- **CRT.** Garner's algorithm combining `K` residues modulo `K` primes into a - `U192` (3 × u64) integer (`integer/src/mul/ntt/crt.rs`). All Garner - precomputed constants are hardcoded in `primes.rs` (`CRT_INV_IJ`). Uses an - object-safe `ModOps` trait (subset of `Reducer`) for dynamic dispatch - over the per-prime reducers. - -- **Dispatch.** Multiplication above `THRESHOLD_NTT` words routes to the NTT - path (`integer/src/mul/mod.rs`). `add_signed_mul` (unequal lengths) and - `add_signed_mul_same_len` (equal lengths) share a single `add_signed_mul_impl` - that does one NTT convolution — no chunking for equal/similar lengths. - -- **Memory.** Scratch space carved from the linear `Memory` arena. Worst-case - bound computed in `memory_requirement_up_to` using `B_PACK_MIN = 16` - (largest possible N for a given operand size). - -### Phase 7 optimisations (completed) - -- **K_eff = 2 auto-selection.** `select_params` checks headroom against - `P0·P1` (≈2^128). For `b_pack = 16`, `max_coeff < 2^63 ≪ 2^128`, so two - primes always suffice. Third-prime fallback (`K_eff = 3`) is retained as a - safety net for larger `b_pack`. - -- **Threshold calibrated.** `THRESHOLD_NTT = 40 000` words (~2.6 M bits), - the first measured crossover where NTT beats pure toom-3 on Apple M4 Pro. - -- **Env-var overrides.** `DASHU_THRESHOLD_SIMPLE`, `DASHU_THRESHOLD_KARATSUBA`, - `DASHU_THRESHOLD_NTT` override the compile-time defaults at runtime. Gated - behind the `tuning` feature (implies `std`). - -- **Crossover benchmark.** `#[ignore]` test `crossover()` in - `integer/src/mul/ntt/mod.rs` compares NTT against toom-3 at key sizes. - Run with `DASHU_THRESHOLD_NTT=99999999` to force pure toom-3. +### Core algorithm ---- - -## Remaining optimisation opportunities +- **Primes.** Three Proth primes of the form `K·2^N + 1` per word size: + 64-bit: `29·2^57+1`, `71·2^57+1`, `75·2^57+1` (MAX_LOG_N=57). + 32-bit: `7·2^26+1`, `15·2^27+1`, `17·2^27+1` (MAX_LOG_N=26). + Defined in `integer/src/arch/generic_{64,32}_bit/ntt.rs` with per-prime + reducer instances `P0`/`P1`/`P2`. -### 1. Radix-4 or split-radix NTT (~25–33% fewer twiddle multiplies) +- **Modular arithmetic.** Delegates to `num_modular::Reducer` (Proth + Montgomery reduction). Transform functions are generic over + `R: Reducer`, monomorphized per prime at the `process_prime` call + site. No per-prime wrapper functions. -Radix-4 processes 4 elements per butterfly with 3 twiddle multiplies and -`log₄(N)` stages (half as many passes through memory). Split-radix pushes -the savings closer to 33%. +- **NTT transforms.** Iterative in-place radix-2 DIT. Forward: + `bit_reverse → forward(ω)`. Inverse: `bit_reverse → forward(ω⁻¹) → scale`. + All arithmetic in Montgomery form; conversion at pipeline boundaries via + `r.transform`/`r.residue`. -**Work items:** -- Rewrite `ntt_core` in `transform.rs` with a radix-4 butterfly. -- Handle N that is a power of 2 but not a power of 4: do one radix-2 stage - followed by radix-4 stages. -- Update twiddle indexing; the twiddle table layout changes. -- A primitive 4-th root `j = ω_N^{N/4}` is needed for the butterfly core; - derive it from the existing `ω_2_32` root. - -### 2. Harvey lazy-reduction butterflies (~10–15%) +- **CRT.** Garner's algorithm combining `K` residues into a `TripleWord` + (`[u64;3]` on 64-bit, `[u32;3]` on 32-bit). `CrtAccum` trait in + `crt.rs`, impl gated by `#[cfg]` per word size. Standard-form arithmetic + via `num_modular::ModularCoreOps::subm`/`mulm`. -Currently every `add_mod` / `sub_mod` fully normalizes to `[0, p)`. Harvey's -approach keeps values in `[0, 2p)` across multiple butterfly stages, deferring -the conditional subtract to the end (or to the next `mul_mod`). This replaces -a branch + subtract with a no-op in the inner loop. +- **Dispatch.** `THRESHOLD_NTT = 4 000` words (256 kbits). Asymmetric + chunking (`a > 2·b`): pre-transforms `b` once, reuses `b̂` across chunks of + `a`. Shared entry point `process_prime(a, b: BInput<'_>, ctx, r)` handles + both raw `b` and cached `b̂`. -**Work items:** -- Change `add_mod` / `sub_mod` to allow `[0, 2p)` outputs. -- Add a normalization pass at the end of `pointwise_mul` and `inverse`. -- Verify no overflow in the radix-2 structure (each stage at most doubles the - dynamic range, so worst-case after log₂(N) stages is `[0, N·p)` — we need a - cleanup before it overflows `u64`). +### Architecture -### 3. Merge `bit_reverse` with `pack` (~5–10%) +- NTT constants and reducer instances live in arch-specific `ntt.rs` files, + re-exported through `arch/mod.rs` → `crate::arch::ntt`. +- 16-bit Word targets excluded at compile time (`#[cfg]` on `pub(crate) mod ntt`). -Currently `pack` writes coefficients in natural order, then `bit_reverse` -permutes them in a second pass. Write packed coefficients directly to their -bit-reversed positions, saving one full array pass. +### Benchmarking -### 4. Shift-expressible twiddle factors (stage-dependent) +- `ubig_mul_asymmetric` in `integer/benches/primitive.rs` — fixed `b` (500 kbits), + varying `a` (1 kbit – 5 Mbits). Exercises all chunked-mul code paths. -In Goldilocks primes, `2^k mod p = 2^k` when `2^k < p`. The first few NTT -stages have twiddle factors that are pure powers of 2, so `mul_mod(t, 2^k)` -reduces to a shift + conditional subtract — no `u128` multiply needed. - -### 5. Specialize `b = 32` lane (~5%) +--- -The `b = 32` prime (`0xFFFFFFFF00000001`) has the cleanest reduction identity -(splits a `u128` product exactly into 32-bit limbs). A dedicated code path -for this prime alone could squeeze out a few more cycles vs. the generic -`match B` dispatch in `mul_mod`. +## Remaining optimisation opportunities -### 6. Asymmetric operand chunking (conditional) +### 1. Radix-4 NTT (~25–33% fewer twiddle multiplies) -When `a ≫ b`, chunk the long operand, forward-transform the short operand -once, and reuse `b̂` (the transformed short operand) across all chunks. Only -matters for extremely lopsided inputs. +Attempted but reverted — the 4-point DFT output ordering within DIT/DIF +interacts non-trivially with bit-reversal. Needs careful re-analysis. -### 7. u32-word support via u32 Solinas primes +### 2. Harvey lazy-reduction butterflies (~10–15%) -The NTT path currently requires `Word = u64` and uses three 64-bit Solinas -primes. For 32-bit (and potentially 16-bit) targets, we need a separate set -of u32-friendly Solinas primes of the form `2^32 − 2^b + 1`. +Bypass `r.add`/`r.sub` normalization in `ntt_core`, deferring to a cleanup +pass. Trickiest due to interaction with `num_modular`'s `Reducer` API. -**Work items:** -- Find 2–3 primes `p = 2^32 − 2^b + 1` with `v2(p-1) ≥ 16` (enough for N up - to 2^16) and distinct `b` values. -- Implement `FixedTrinomialSolinas32` (or equivalent) in `num-modular`, or - hand-roll the 32-bit reduction inline. -- Generalize the NTT pipeline over `Word` size: the packing, transform, and - CRT layers need to work with `u32` coefficients instead of `u64`. -- Assert `Word = u32` or `Word = u64` at entry and dispatch to the appropriate - prime set. diff --git a/integer/CHANGELOG.md b/integer/CHANGELOG.md index d2e4a013..cb29c33b 100644 --- a/integer/CHANGELOG.md +++ b/integer/CHANGELOG.md @@ -3,7 +3,8 @@ ## Unreleased ### Add -- NTT-based multiplication for very large integers (above 40000 words / ~2.6M bits), using two or three Proth primes of the form `K·2^N + 1` combined with the Chinese Remainder Theorem. Supports both 64-bit and 32-bit Word targets. +- NTT-based multiplication using Proth primes (`K·2^N + 1`), combined via Garner CRT. Supports 64-bit and 32-bit Word targets. Threshold at 4 000 words (~256 kbits). +- Asymmetric NTT chunking: when one operand is much larger than the other, the shorter operand is forward-transformed once and reused across chunks. - `UBig::from_u64` and `IBig::from_i64`, const on 32-bit and 64-bit targets. ### Improve @@ -14,6 +15,12 @@ - NTT multiplication auto-selects `K_eff = 2` primes when headroom allows, skipping the third prime. - Multiplication thresholds can be overridden at runtime via `DASHU_THRESHOLD_SIMPLE`, `DASHU_THRESHOLD_KARATSUBA`, and `DASHU_THRESHOLD_NTT` environment variables (requires `tuning` feature). +### Change +- NTT multiplication now uses Proth primes (`K·2^N + 1`) instead of Solinas primes, improving modular reduction speed. +- NTT threshold lowered from 40 000 to 4 000 words. +- NTT enabled for 32-bit Word targets. +- Arch-specific NTT prime definitions under `arch/generic_{32,64}_bit/ntt.rs`. + ## 0.4.2 - Add `UBig::ones`. diff --git a/integer/benches/primitive.rs b/integer/benches/primitive.rs index 2958fc54..6bcee837 100644 --- a/integer/benches/primitive.rs +++ b/integer/benches/primitive.rs @@ -150,6 +150,31 @@ fn ubig_ilog_large(criterion: &mut Criterion) { group.finish(); } +fn ubig_mul_asymmetric(criterion: &mut Criterion) { + let mut rng = StdRng::seed_from_u64(SEED); + let mut group = criterion.benchmark_group("ubig_mul_asymmetric"); + group.plot_config(PlotConfiguration::default().summary_scale(AxisScale::Logarithmic)); + + // b just above the NTT threshold (4 000 words = 256 kbits → use 500 kbits). + let b_bits = 500_000; + let b = random_ubig(b_bits, &mut rng); + + // a ranges from 1 kbit (below Karatsuba threshold) to heavily + // asymmetric (10×), exercising all chunked-mul code paths. + for &a_bits in &[ + 1_000, 10_000, 100_000, 500_000, 1_000_000, 2_000_000, 5_000_000, + ] { + let a = random_ubig(a_bits, &mut rng); + group.bench_with_input( + BenchmarkId::from_parameter(format!("{a_bits}/{b_bits}")), + &(a, &b), + |bencher, (ta, tb)| bencher.iter(|| ta * *tb), + ); + } + + group.finish(); +} + criterion_group!( benches, ubig_add, @@ -163,6 +188,7 @@ criterion_group!( ubig_modulo_pow, ubig_pow_large_base, ubig_ilog_large, + ubig_mul_asymmetric, ); criterion_main!(benches); diff --git a/integer/src/arch/generic_32_bit/ntt.rs b/integer/src/arch/generic_32_bit/ntt.rs index 6e6cbfb6..a2218dd1 100644 --- a/integer/src/arch/generic_32_bit/ntt.rs +++ b/integer/src/arch/generic_32_bit/ntt.rs @@ -6,22 +6,20 @@ use num_modular::FixedProth32; // Proth reducer instances — each with a different (N, K) pair. -pub const P0: FixedProth32<26, 7> = FixedProth32::<26, 7>; -pub const P1: FixedProth32<27, 15> = FixedProth32::<27, 15>; -pub const P2: FixedProth32<27, 17> = FixedProth32::<27, 17>; - -// Type aliases needed by for_each_prime! macro in transform tests. pub type Rp0 = FixedProth32<26, 7>; pub type Rp1 = FixedProth32<27, 15>; pub type Rp2 = FixedProth32<27, 17>; +pub const P0: Rp0 = FixedProth32::<26, 7>; +pub const P1: Rp1 = FixedProth32::<27, 15>; +pub const P2: Rp2 = FixedProth32::<27, 17>; + pub const K: usize = 3; pub const MAX_LOG_N: u32 = 26; pub const B_PACK_MIN: u32 = 8; pub const B_PACK_CANDIDATES: &[u32] = &[32, 16, 8]; pub type Lane = u32; -pub type DoubleLane = u64; /// Primitive `MAX_LOG_N`-th roots of unity for each prime. pub const OMEGA_MAX: [Lane; K] = [ @@ -33,11 +31,7 @@ pub const OMEGA_MAX: [Lane; K] = [ pub const CRT_INV_IJ: [[Lane; K]; K] = [[0, 0x4e42c85b, 0x5fb425ef], [0, 0, 0x44000009], [0, 0, 0]]; /// Prime moduli indexed by PI. -pub const MODULI: [Lane; K] = [ - FixedProth32::<26, 7>::MODULUS, - FixedProth32::<27, 15>::MODULUS, - FixedProth32::<27, 17>::MODULUS, -]; +pub const MODULI: [Lane; K] = [Rp0::MODULUS, Rp1::MODULUS, Rp2::MODULUS]; #[cfg(test)] mod tests { diff --git a/integer/src/arch/generic_64_bit/ntt.rs b/integer/src/arch/generic_64_bit/ntt.rs index 86f6b159..0bf25969 100644 --- a/integer/src/arch/generic_64_bit/ntt.rs +++ b/integer/src/arch/generic_64_bit/ntt.rs @@ -6,22 +6,20 @@ use num_modular::FixedProth64; // Proth reducer instances — each with a different (N, K) pair. -pub const P0: FixedProth64<57, 29> = FixedProth64::<57, 29>; -pub const P1: FixedProth64<57, 71> = FixedProth64::<57, 71>; -pub const P2: FixedProth64<57, 75> = FixedProth64::<57, 75>; - -// Type aliases needed by for_each_prime! macro in transform tests. pub type Rp0 = FixedProth64<57, 29>; pub type Rp1 = FixedProth64<57, 71>; pub type Rp2 = FixedProth64<57, 75>; +pub const P0: Rp0 = FixedProth64::<57, 29>; +pub const P1: Rp1 = FixedProth64::<57, 71>; +pub const P2: Rp2 = FixedProth64::<57, 75>; + pub const K: usize = 3; pub const MAX_LOG_N: u32 = 57; pub const B_PACK_MIN: u32 = 16; pub const B_PACK_CANDIDATES: &[u32] = &[64, 32, 16]; pub type Lane = u64; -pub type DoubleLane = u128; /// Primitive `MAX_LOG_N`-th roots of unity for each prime: /// `omega_max[i]` = `g^{(p_i-1) / 2^MAX_LOG_N} mod p_i`. @@ -38,11 +36,7 @@ pub const CRT_INV_IJ: [[Lane; K]; K] = [ ]; /// Prime moduli indexed by PI. -pub const MODULI: [Lane; K] = [ - FixedProth64::<57, 29>::MODULUS, - FixedProth64::<57, 71>::MODULUS, - FixedProth64::<57, 75>::MODULUS, -]; +pub const MODULI: [Lane; K] = [Rp0::MODULUS, Rp1::MODULUS, Rp2::MODULUS]; #[cfg(test)] mod tests { diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index aef7b2c0..2a615afa 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -17,7 +17,7 @@ mod pack; mod transform; use crate::arch::ntt::{ - B_PACK_CANDIDATES, B_PACK_MIN, K, MAX_LOG_N, MODULI, OMEGA_MAX, P0, P1, P2, + B_PACK_CANDIDATES, B_PACK_MIN, CRT_INV_IJ, K, MAX_LOG_N, MODULI, OMEGA_MAX, P0, P1, P2, }; use crate::mul::ntt::crt::{garner_combine, CrtAccum}; use num_modular::Reducer; @@ -113,6 +113,7 @@ pub fn memory_requirement_up_to(total_len: usize, _smaller_len: usize) -> Layout /// /// Returns carry. #[must_use] +#[inline] pub fn add_signed_mul_same_len( c: &mut [Word], sign: Sign, @@ -122,13 +123,23 @@ pub fn add_signed_mul_same_len( ) -> SignedWord { let n = a.len(); debug_assert!(b.len() == n && c.len() == 2 * n); - add_signed_mul_impl(c, sign, a, b, memory) + add_signed_mul_conv(c, sign, a, b, memory) } /// `c += sign * a * b` (general, a may be longer than b). /// +/// When `a ≫ b` the implementation forks: +/// - If `b` is below [`THRESHOLD_NTT`], dispatch already routes to +/// `toom_3::add_signed_mul` (which uses +/// `add_signed_mul_split_into_chunks` from +/// [`helpers`](crate::mul::helpers)). +/// - If `b` is above [`THRESHOLD_NTT`], this function pre-transforms +/// `b` once per prime and reuses `b̂` across chunks of `a` via +/// [`add_signed_mul_chunked`]. +/// /// Returns carry. #[must_use] +#[inline] pub fn add_signed_mul( c: &mut [Word], sign: Sign, @@ -137,13 +148,18 @@ pub fn add_signed_mul( memory: &mut Memory, ) -> SignedWord { debug_assert!(a.len() >= b.len() && c.len() == a.len() + b.len()); - add_signed_mul_impl(c, sign, a, b, memory) + if a.len() > 2 * b.len() { + return add_signed_mul_chunked(c, sign, a, b, memory); + } + add_signed_mul_conv(c, sign, a, b, memory) } -/// Core implementation: c += sign * a * b. +/// NTT multiplication with asymmetric chunking. /// -/// Does a single NTT convolution of the full operands (no chunking). -fn add_signed_mul_impl( +/// When `la > 2 * lb`, transform `b` once and reuse `b̂` across chunks +/// of `a`, reducing total transform work from O((la+lb)·log(la+lb)) +/// to O(la + lb·log(lb)). +fn add_signed_mul_chunked( c: &mut [Word], sign: Sign, a: &[Word], @@ -151,206 +167,330 @@ fn add_signed_mul_impl( memory: &mut Memory, ) -> SignedWord { use crate::arch::ntt::Lane; + use crate::mul::helpers::add_signed_mul_split_into_chunks; - let la = a.len(); let lb = b.len(); + let chunk_len = lb * 2; + + // Parameters for chunk-sized transforms. + let (b_pack, nn_chunk, k_eff) = select_params(chunk_len, lb); // a_chunk ≈ 2*lb + + // ---- Allocate long-lived buffers ---- + + // Per-prime forward-transformed b̂ and cached twiddles (fwd + inv). + // Twiddles depend only on (pi, nn_chunk, omega_max) — precompute + // once so the per-chunk closure can copy instead of recomputing. + let b_hat_len = k_eff * nn_chunk; + let twiddle_len = k_eff * (nn_chunk / 2); + let (b_hat, mut mem) = memory.allocate_slice_fill::(b_hat_len, 0); + let (fwd_tw_cache, mut mem) = mem.allocate_slice_fill::(twiddle_len, 0); + let (inv_tw_cache, mut mem) = mem.allocate_slice_fill::(twiddle_len, 0); + + // ---- Transform b once per prime; also precompute twiddles ---- + let geom = NttGeometry { + nn: nn_chunk, + b_pack, + k_eff, + output_coeffs: 0, // unused by prepare_b_hat_and_twiddles + }; + prepare_b_hat_and_twiddles(b_hat, fwd_tw_cache, inv_tw_cache, b, &geom, &mut mem); - if la == 0 || lb == 0 { - return 0; - } - - let (b_pack, nn, k_eff) = select_params(la, lb); - let la_bits = bit_len(a); + // ---- Setup for the closure ---- let lb_bits = bit_len(b); - if la_bits == 0 || lb_bits == 0 { - return 0; - } - - let coeffs_a = coeff_count(la_bits, b_pack); let coeffs_b = coeff_count(lb_bits, b_pack); - let output_coeffs = coeffs_a + coeffs_b - 1; - // ---- Memory carve (longest-lived first) ---- + // ---- Chunked multiply ---- + let carry = add_signed_mul_split_into_chunks( + c, + sign, + a, + b, + chunk_len, + &mut mem, + |c_slice, sign, a_chunk, b, mem| { + let a_bits = bit_len(a_chunk); + if a_bits == 0 { + return 0; + } + let coeffs_a = coeff_count(a_bits, b_pack); + let output_coeffs = coeffs_a + coeffs_b - 1; + let out_words = a_chunk.len() + b.len(); + + let geom = NttGeometry { + nn: nn_chunk, + b_pack, + k_eff, + output_coeffs, + }; + run_ntt_pipeline( + a_chunk, + b_hat, + fwd_tw_cache, + inv_tw_cache, + &geom, + out_words, + c_slice, + sign, + mem, + ) + }, + ); + + carry +} - // 1. Product buffer (Word-sized, CRT splits u64 words into Word limbs) - let prod_len = la + lb; - let (prod, mut mem) = memory.allocate_slice_fill::(prod_len, 0); +/// Run the full NTT pipeline: allocate → per-prime transform → CRT → fold into `c_out`. +/// +/// `b_hat`, `fwd_tw_cache`, and `inv_tw_cache` must have been precomputed by the +/// caller (pack + Montgomery convert + bit-reverse + forward-transform for `b_hat`; +/// forward/inverse twiddle tables for the caches). See `transform_b_forward` and +/// `transform::precompute_twiddles`. +/// +/// Shared body of `add_signed_mul_conv` and the per-chunk callback in +/// `add_signed_mul_chunked`. +fn run_ntt_pipeline( + a: &[Word], + b_hat: &[crate::arch::ntt::Lane], + fwd_tw_cache: &[crate::arch::ntt::Lane], + inv_tw_cache: &[crate::arch::ntt::Lane], + geom: &NttGeometry, + out_words: usize, + c_out: &mut [Word], + sign: Sign, + mem: &mut Memory, +) -> SignedWord { + use crate::arch::ntt::Lane; - // 2. Residue storage (per-prime inverse results) - let residues_len = k_eff * nn; - let (residues, mut mem) = mem.allocate_slice_fill::(residues_len, 0); + let nn = geom.nn; + let k_eff = geom.k_eff; + + let (prod, mut m) = mem.allocate_slice_fill::(out_words, 0); + let (residues, mut m) = m.allocate_slice_fill::(k_eff * nn, 0); + let (a_lane, mut m) = m.allocate_slice_fill::(nn, 0); + let (b_lane, mut m) = m.allocate_slice_fill::(nn, 0); + let (fwd_twiddles, mut m) = m.allocate_slice_fill::(nn / 2, 0); + let (inv_twiddles, _) = m.allocate_slice_fill::(nn / 2, 0); + + let mut ctx = TransformCtx { + a_lane, + b_lane, + fwd_twiddles, + inv_twiddles, + geom: NttGeometry { ..*geom }, + }; - // 3. Lane buffers (reused across primes) - let (a_lane, mut mem) = mem.allocate_slice_fill::(nn, 0); - let (b_lane, mut mem) = mem.allocate_slice_fill::(nn, 0); + for pi in 0..k_eff { + let tw_off = pi * (nn / 2); + ctx.fwd_twiddles + .copy_from_slice(&fwd_tw_cache[tw_off..tw_off + nn / 2]); + ctx.inv_twiddles + .copy_from_slice(&inv_tw_cache[tw_off..tw_off + nn / 2]); - // 4. Twiddle tables (fwd + inv, reused per prime) - let (fwd_twiddles, mut mem) = mem.allocate_slice_fill::(nn / 2, 0); - let (inv_twiddles, _) = mem.allocate_slice_fill::(nn / 2, 0); + let b_hat_slice = &b_hat[pi * nn..(pi + 1) * nn]; - // ---- Per-prime transforms (monomorphized per reducer) ---- - for pi in 0..k_eff { - let mut ctx = TransformCtx { - a_lane, - b_lane, - fwd_twiddles, - inv_twiddles, - omega_max: OMEGA_MAX[pi], - nn, - b_pack, - residues, - pi, - }; match pi { - 0 => process_prime(a, b, &mut ctx, &P0), - 1 => process_prime(a, b, &mut ctx, &P1), - 2 => process_prime(a, b, &mut ctx, &P2), + 0 => process_prime(a, b_hat_slice, &mut ctx, residues, pi, &P0), + 1 => process_prime(a, b_hat_slice, &mut ctx, residues, pi, &P1), + 2 => process_prime(a, b_hat_slice, &mut ctx, residues, pi, &P2), _ => unreachable!(), } } - // ---- CRT per coefficient + accumulate ---- - // Extract prime constants as both u64 and u32 so the Lane-size - // dispatch below type-checks correctly in both branches. - // The dead branch (wrong width) is eliminated by the compiler. - let primes_u64: [u64; K] = [MODULI[0] as u64, MODULI[1] as u64, MODULI[2] as u64]; - let crt_inv_u64: [[u64; K]; K] = { - use crate::arch::ntt::CRT_INV_IJ; - let mut m = [[0u64; K]; K]; - for i in 0..K { - for j in 0..K { - m[i][j] = CRT_INV_IJ[i][j] as u64; - } - } - m - }; - let primes_u32: [u32; K] = [MODULI[0] as u32, MODULI[1] as u32, MODULI[2] as u32]; - let crt_inv_u32: [[u32; K]; K] = { - let mut m = [[0u32; K]; K]; - for i in 0..K { - for j in 0..K { - m[i][j] = crt_inv_u64[i][j] as u32; - } - } - m - }; - - // CRT dispatch: one branch per Word size, gated by cfg so only - // one compiles — no dummy types needed in dead branches. - #[cfg(not(any(force_bits = "32", target_pointer_width = "32")))] - { - let residues_u64: &[u64] = - unsafe { core::slice::from_raw_parts(residues.as_ptr() as *const u64, residues.len()) }; - do_crt::( - prod, - residues_u64, - k_eff, - nn, - output_coeffs, - b_pack, - &primes_u64, - &crt_inv_u64, - ); - } - #[cfg(any(force_bits = "32", target_pointer_width = "32"))] - { - let residues_u32: &[u32] = - unsafe { core::slice::from_raw_parts(residues.as_ptr() as *const u32, residues.len()) }; - do_crt::( - prod, - residues_u32, - k_eff, - nn, - output_coeffs, - b_pack, - &primes_u32, - &crt_inv_u32, - ); - } + do_crt::(prod, residues, &ctx, &MODULI, &CRT_INV_IJ); - // ---- Fold product into c with sign ---- - let output_words = la + lb; match sign { - Positive => add::add_signed_in_place(c, Positive, &prod[..output_words]), - Negative => add::add_signed_in_place(c, Negative, &prod[..output_words]), + Positive => add::add_signed_in_place(&mut c_out[..out_words], Positive, &prod[..out_words]), + Negative => add::add_signed_in_place(&mut c_out[..out_words], Negative, &prod[..out_words]), } } +/// Core implementation: c += sign * a * b. +/// +/// Does a single NTT convolution of the full operands (no chunking). +fn add_signed_mul_conv( + c: &mut [Word], + sign: Sign, + a: &[Word], + b: &[Word], + memory: &mut Memory, +) -> SignedWord { + use crate::arch::ntt::Lane; + + let la = a.len(); + let lb = b.len(); + + debug_assert!(la > 0 && lb > 0); + let (b_pack, nn, k_eff) = select_params(la, lb); + let la_bits = bit_len(a); + let lb_bits = bit_len(b); + debug_assert!(la_bits > 0 && lb_bits > 0); + + let coeffs_a = coeff_count(la_bits, b_pack); + let coeffs_b = coeff_count(lb_bits, b_pack); + let output_coeffs = coeffs_a + coeffs_b - 1; + + // Pre-transform b and precompute twiddles. + let b_hat_len = k_eff * nn; + let twiddle_len = k_eff * (nn / 2); + let (b_hat, mut mem) = memory.allocate_slice_fill::(b_hat_len, 0); + let (fwd_tw_cache, mut mem) = mem.allocate_slice_fill::(twiddle_len, 0); + let (inv_tw_cache, mut mem) = mem.allocate_slice_fill::(twiddle_len, 0); + + let geom = NttGeometry { + nn, + b_pack, + k_eff, + output_coeffs, + }; + prepare_b_hat_and_twiddles(b_hat, fwd_tw_cache, inv_tw_cache, b, &geom, &mut mem); + run_ntt_pipeline(a, b_hat, fwd_tw_cache, inv_tw_cache, &geom, la + lb, c, sign, &mut mem) +} + /// CRT + accumulate, generic over the accumulator type. -#[allow(clippy::too_many_arguments)] -#[inline(never)] fn do_crt( prod: &mut [Word], residues: &[A::Lane], - k_eff: usize, - nn: usize, - output_coeffs: usize, - b_pack: u32, + ctx: &TransformCtx<'_>, primes: &[A::Lane; K], crt_inv: &[[A::Lane; K]; K], ) { - for k in 0..output_coeffs { + let g = &ctx.geom; + for k in 0..g.output_coeffs { let mut coeff_residues = [A::Lane::default(); 3]; #[allow(clippy::needless_range_loop)] - for pi in 0..k_eff { - coeff_residues[pi] = residues[pi * nn + k]; + for pi in 0..g.k_eff { + coeff_residues[pi] = residues[pi * g.nn + k]; } - let crt_val = garner_combine::(&coeff_residues[..k_eff], crt_inv, primes); + let crt_val = garner_combine::(&coeff_residues[..g.k_eff], crt_inv, primes); let mut crt_buf = [Word::default(); 6]; let crt_n = crt_val.write_words(&mut crt_buf); - add_shifted_to_prod(prod, &crt_buf[..crt_n as usize], crt_n, k, b_pack); + add_shifted_to_prod(prod, &crt_buf[..crt_n as usize], crt_n, k, g.b_pack); } } -/// Scratch buffers and parameters for one prime's NTT pipeline. +/// Geometry constants for an NTT pipeline invocation. +struct NttGeometry { + nn: usize, + b_pack: u32, + k_eff: usize, + output_coeffs: usize, +} + +/// Scratch buffers and geometry for the per-prime NTT pipeline. struct TransformCtx<'a> { a_lane: &'a mut [crate::arch::ntt::Lane], b_lane: &'a mut [crate::arch::ntt::Lane], fwd_twiddles: &'a mut [crate::arch::ntt::Lane], inv_twiddles: &'a mut [crate::arch::ntt::Lane], - omega_max: crate::arch::ntt::Lane, + geom: NttGeometry, +} + +/// Transform `b` and leave the result in `b_lane` (forward-transformed, +/// Montgomery form). `fwd_twiddles` must already be precomputed. + +fn transform_b_forward>( + b_lane: &mut [crate::arch::ntt::Lane], + b: &[Word], nn: usize, b_pack: u32, - residues: &'a mut [crate::arch::ntt::Lane], - pi: usize, + fwd_twiddles: &[crate::arch::ntt::Lane], + r: &R, +) { + pack::pack(b_lane, b, b_pack, nn); + for c in b_lane[..nn].iter_mut() { + *c = r.transform(*c); + } + transform::bit_reverse(&mut b_lane[..nn]); + transform::forward(&mut b_lane[..nn], fwd_twiddles, r); } -/// Per-prime NTT pipeline, monomorphized for a specific reducer `R`. -#[inline(never)] +/// Pre-transform `b` and precompute twiddles, storing results into the +/// pre-allocated cache slices. +/// +/// `b_hat` must have length `geom.k_eff * geom.nn`, `fwd_tw_cache` and +/// `inv_tw_cache` each `geom.k_eff * (geom.nn / 2)`. +fn prepare_b_hat_and_twiddles( + b_hat: &mut [crate::arch::ntt::Lane], + fwd_tw_cache: &mut [crate::arch::ntt::Lane], + inv_tw_cache: &mut [crate::arch::ntt::Lane], + b: &[Word], + geom: &NttGeometry, + mem: &mut Memory, +) { + use crate::arch::ntt::Lane; + + let nn = geom.nn; + let b_pack = geom.b_pack; + let k_eff = geom.k_eff; + + for pi in 0..k_eff { + let (b_lane, mut rest) = mem.allocate_slice_fill::(nn, 0); + let (fwd_tw, mut rest) = rest.allocate_slice_fill::(nn / 2, 0); + let (inv_tw, _) = rest.allocate_slice_fill::(nn / 2, 0); + + let omega = OMEGA_MAX[pi]; + match pi { + 0 => { + transform::precompute_twiddles(fwd_tw, nn, omega, false, &P0); + transform::precompute_twiddles(inv_tw, nn, omega, true, &P0); + transform_b_forward(b_lane, b, nn, b_pack, fwd_tw, &P0); + } + 1 => { + transform::precompute_twiddles(fwd_tw, nn, omega, false, &P1); + transform::precompute_twiddles(inv_tw, nn, omega, true, &P1); + transform_b_forward(b_lane, b, nn, b_pack, fwd_tw, &P1); + } + 2 => { + transform::precompute_twiddles(fwd_tw, nn, omega, false, &P2); + transform::precompute_twiddles(inv_tw, nn, omega, true, &P2); + transform_b_forward(b_lane, b, nn, b_pack, fwd_tw, &P2); + } + _ => unreachable!(), + } + + let b_off = pi * nn; + let tw_off = pi * (nn / 2); + b_hat[b_off..b_off + nn].copy_from_slice(b_lane); + fwd_tw_cache[tw_off..tw_off + nn / 2].copy_from_slice(fwd_tw); + inv_tw_cache[tw_off..tw_off + nn / 2].copy_from_slice(inv_tw); + } +} + +/// Per-prime NTT pipeline. +/// +/// `b_hat_slice` must already be forward-transformed (packed, Montgomery +/// form, bit-reversed). `ctx.fwd_twiddles` and `ctx.inv_twiddles` must +/// already be precomputed. fn process_prime>( a: &[Word], - b: &[Word], + b_hat_slice: &[crate::arch::ntt::Lane], ctx: &mut TransformCtx<'_>, + residues: &mut [crate::arch::ntt::Lane], + pi: usize, r: &R, ) { - pack::pack(ctx.a_lane, a, ctx.b_pack, ctx.nn); - pack::pack(ctx.b_lane, b, ctx.b_pack, ctx.nn); + let nn = ctx.geom.nn; + let b_pack = ctx.geom.b_pack; - // Convert standard-form coefficients to Montgomery form. - for c in ctx.a_lane[..ctx.nn].iter_mut() { - *c = r.transform(*c); - } - for c in ctx.b_lane[..ctx.nn].iter_mut() { + // Transform a + pack::pack(ctx.a_lane, a, b_pack, nn); + for c in ctx.a_lane[..nn].iter_mut() { *c = r.transform(*c); } + transform::bit_reverse(&mut ctx.a_lane[..nn]); + transform::forward(&mut ctx.a_lane[..nn], ctx.fwd_twiddles, r); - transform::precompute_twiddles(ctx.fwd_twiddles, ctx.nn, ctx.omega_max, false, r); - transform::precompute_twiddles(ctx.inv_twiddles, ctx.nn, ctx.omega_max, true, r); + // Copy pre-transformed b + ctx.b_lane[..nn].copy_from_slice(b_hat_slice); - transform::bit_reverse(ctx.a_lane); - transform::bit_reverse(ctx.b_lane); - transform::forward(ctx.a_lane, ctx.fwd_twiddles, r); - transform::forward(ctx.b_lane, ctx.fwd_twiddles, r); - transform::pointwise_mul(ctx.a_lane, ctx.b_lane, r); - transform::inverse(ctx.a_lane, ctx.inv_twiddles, r); - - // Convert residues back from Montgomery to standard form. - for c in ctx.a_lane[..ctx.nn].iter_mut() { + transform::pointwise_mul(&mut ctx.a_lane[..nn], &ctx.b_lane[..nn], r); + transform::inverse(&mut ctx.a_lane[..nn], ctx.inv_twiddles, r); + for c in ctx.a_lane[..nn].iter_mut() { *c = r.residue(*c); } - let offset = ctx.pi * ctx.nn; - ctx.residues[offset..offset + ctx.nn].copy_from_slice(ctx.a_lane); + let offset = pi * nn; + residues[offset..offset + nn].copy_from_slice(&ctx.a_lane[..nn]); } /// Add a CRT value (as `Word`-sized limbs) to `prod`, shifted left by @@ -445,25 +585,12 @@ mod tests { let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); - let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + let carry = add_signed_mul_conv(&mut c, Positive, &a, &b, &mut memory); assert_eq!(carry, 0); assert_eq!(c[0], 15); assert_eq!(c[1], 0); } - #[test] - fn test_ntt_zero_operand() { - let a = vec![0xDEADu64 as Word; 30]; - let b = vec![0u64 as Word; 30]; - let mut c = vec![0u64 as Word; a.len() + b.len()]; - let layout = memory_requirement_up_to(c.len(), b.len()); - let mut alloc = crate::memory::MemoryAllocation::new(layout); - let mut memory = alloc.memory(); - let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); - assert_eq!(carry, 0); - assert!(c.iter().all(|&w| w == 0)); - } - #[test] fn test_ntt_sign_negative() { let a: Vec = (0..30).map(|i| (i as Word + 1) * 100).collect(); @@ -473,12 +600,12 @@ mod tests { let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); - let _ = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + let _ = add_signed_mul_conv(&mut c, Positive, &a, &b, &mut memory); let layout2 = memory_requirement_up_to(c.len(), b.len()); let mut alloc2 = crate::memory::MemoryAllocation::new(layout2); let mut memory2 = alloc2.memory(); - let _ = add_signed_mul_impl(&mut c, Negative, &a, &b, &mut memory2); + let _ = add_signed_mul_conv(&mut c, Negative, &a, &b, &mut memory2); assert!(c.iter().all(|&w| w == 0)); } @@ -494,7 +621,7 @@ mod tests { let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); - let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + let carry = add_signed_mul_conv(&mut c, Positive, &a, &b, &mut memory); assert_eq!(carry, 0); assert!(c.iter().any(|&w| w != 0)); } @@ -534,7 +661,7 @@ mod tests { let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); - let carry = add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + let carry = add_signed_mul_conv(&mut c, Positive, &a, &b, &mut memory); assert_eq!(carry, 0, "carry should be 0"); assert_eq!(&c[..], &expected[..], "NTT mismatch: la={la}, lb={lb}"); } @@ -571,7 +698,7 @@ mod tests { let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); - add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + add_signed_mul_conv(&mut c, Positive, &a, &b, &mut memory); assert_eq!(&c[..], &expected[..], "all-ones mismatch len={len}"); } } @@ -590,7 +717,7 @@ mod tests { let layout = memory_requirement_up_to(c.len(), b.len()); let mut alloc = crate::memory::MemoryAllocation::new(layout); let mut memory = alloc.memory(); - add_signed_mul_impl(&mut c, Positive, &a, &b, &mut memory); + add_signed_mul_conv(&mut c, Positive, &a, &b, &mut memory); assert_eq!(&c[..], &expected[..], "sparse operand mismatch"); } } diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs index 631932c6..43c769da 100644 --- a/integer/src/mul/ntt/transform.rs +++ b/integer/src/mul/ntt/transform.rs @@ -109,7 +109,7 @@ mod tests { use alloc::vec; #[cfg(not(feature = "std"))] use alloc::vec::Vec; - use num_modular::{ModularCoreOps, ModularPow, ModularUnaryOps}; + use num_modular::ModularPow; fn assert_all_eq(a: &[Lane], b_val: &[Lane], context: &str) { assert_eq!(a.len(), b_val.len(), "{context}: length mismatch"); From 382245270db398788e6feb6152540e5ab6601f95 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 18:01:44 +0800 Subject: [PATCH 14/19] Remove todos --- NTT_RADIX4.md | 191 -------------------------------------------------- TODO.md | 69 ------------------ TODO_NTT.md | 57 --------------- 3 files changed, 317 deletions(-) delete mode 100644 NTT_RADIX4.md delete mode 100644 TODO.md delete mode 100644 TODO_NTT.md diff --git a/NTT_RADIX4.md b/NTT_RADIX4.md deleted file mode 100644 index 77088445..00000000 --- a/NTT_RADIX4.md +++ /dev/null @@ -1,191 +0,0 @@ -# Radix-4 NTT for `dashu-int` - -## Context - -The NTT multiplication path in `integer/src/mul/ntt/` currently uses an iterative in-place **radix-2** DIT transform. `TODO_NTT.md` lists "Radix-4 or split-radix NTT" as optimisation opportunity #1 — roughly halving the number of memory passes (`log₄(N)` stages instead of `log₂(N)`), with a small reduction in arithmetic work. Memory traffic, not arithmetic, is the dominant cost on the Goldilocks-style primes used here, so the pass count is the win that matters. - -This plan implements radix-4 in place of the current radix-2 core, while keeping the rest of the pipeline (`pack`, `bit_reverse`, `pointwise_mul`, CRT, Garner, and the recently-introduced `NttGeometry` / `prepare_b_hat_and_twiddles` / `run_ntt_pipeline` / `process_prime` layering) untouched. A legacy copy of the radix-2 core is kept during development as a differential test oracle, then removed before the final commit. - -Work happens in a new worktree `ntt-radix4` branching off the current `ssa` branch. - ---- - -## Current state of the code (post-refactor) - -The NTT module has been re-organised so that the conv and chunked paths share a single symmetric pipeline, with geometry constants factored into a small value type: - -- `add_signed_mul` (mod.rs:143) dispatches: `a.len() > 2 * b.len()` → `add_signed_mul_chunked`; otherwise `add_signed_mul_conv`. -- **Both paths pre-transform `b` and precompute twiddles up front**, then call `run_ntt_pipeline`. There is no "raw b" vs "cached b" distinction — `b_hat`, `fwd_tw_cache`, `inv_tw_cache` are always populated by `prepare_b_hat_and_twiddles` (mod.rs:412) and passed in. -- `NttGeometry` (mod.rs:372) is a small value struct holding `nn`, `b_pack`, `k_eff`, `output_coeffs`. It's passed by reference into `prepare_b_hat_and_twiddles` and `run_ntt_pipeline`, and embedded inside `TransformCtx`. -- `TransformCtx` (mod.rs:380) is now just four scratch buffer slices (`a_lane`, `b_lane`, `fwd_twiddles`, `inv_twiddles`) plus a `geom: NttGeometry`. `prod` and `residues` are not in the ctx — they are allocated in `run_ntt_pipeline` (mod.rs:267–268) and passed as separate `&mut` arguments to `process_prime` and `do_crt`. -- `run_ntt_pipeline` (mod.rs:251) owns the per-call scratch allocation, runs the per-prime loop calling `process_prime`, then calls `do_crt` and signed-accumulates into `c_out`. -- `process_prime` (mod.rs:464) takes `(a, b_hat_slice, ctx, residues, pi, r)`. It transforms `a` from raw words, copies the pre-transformed `b_hat_slice` into `b_lane`, pointwise-multiplies, inverse-transforms, and writes residues. No twiddle precompute happens here. -- `transform_b_forward` (mod.rs:391) is the helper used by `prepare_b_hat_and_twiddles` to pack/Montgomery/bit-reverse/forward-transform `b`. - -**Implication for radix-4:** the transform-level changes (radix-4 butterfly, expanded twiddle table) live entirely in `transform::forward` / `transform::inverse` / `transform::precompute_twiddles` / `transform::ntt_core`. Because every code path reaches the transform through these, the speedup propagates everywhere for free. The only multi-site edits in `mod.rs` are the twiddle *allocation sizes* and the cache-offset arithmetic, which now live in a small number of well-defined places. - ---- - -## Math summary (verified) - -Radix-4 DIT takes bit-reversed input and produces natural-order output — same I/O contract as the existing radix-2 DIT, so `bit_reverse`, `inverse()`, and all callers are unchanged. - -For each butterfly on quad `(a0, a1, a2, a3)` at positions `(k, k+q, k+2q, k+3q)` within a length-`sub_len` group, with `q = sub_len/4` and `step = n/sub_len`: - -``` -b1 = a1 · ω_n^(k·step) -b2 = a2 · ω_n^(2k·step) -b3 = a3 · ω_n^(3k·step) -e0 = a0 + b2 -e1 = a0 − b2 -e2 = b1 + b3 -e3 = b1 − b3 // order matters: b1 − b3, not b3 − b1 -y0 = e0 + e2 -y1 = e1 + j·e3 // j = ω_n^(n/4), read from twiddles[n/4] -y2 = e0 − e2 -y3 = e1 − j·e3 -``` - -Iterative structure: -- If `log₂(n)` is even, stages run with `sub_len = 4, 16, 64, …, n` (pure radix-4). -- If `log₂(n)` is odd (n = 2·4^L), run **one** radix-2 stage with `sub_len = 2` (uses only `twiddles[0] = 1`), then radix-4 stages with `sub_len = 8, 32, …, n`. -- `n = 2` is a degenerate case — emit a single radix-2 butterfly with twiddle 1 and return early. - -The constant `j = ω_n^(n/4)` is read once from `twiddles[n/4]` at the top of `ntt_core`. For inverse transforms, `twiddles[n/4]` holds `ω_n^(−n/4) = −j`, which is the *other* primitive 4th root; the same butterfly formula applies with it (a sign flip on the `j·e3` terms). No special handling needed — the symmetry falls out naturally. - -The maximum twiddle index touched is `3k·step ≤ 3(n/4 − 1) ≈ 3n/4` at the final stage, which exceeds the current `n/2`-long table. **Fix: expand the twiddle table from `n/2` to `n` lanes.** Memory overhead ≈ +`n` lanes per table. - ---- - -## Files to modify - -### `integer/src/mul/ntt/transform.rs` (primary rewrite) - -1. **`precompute_twiddles`** — change `assert!(out.len() >= n / 2)` to `assert!(out.len() >= n)`, and extend the fill loop from `1..(n/2)` to `1..n`. Output is now `ω_n^k` for `k ∈ [0, n)`. - -2. **`ntt_core` (rewrite)** — replace with the radix-4 algorithm: - ```rust - fn ntt_core>(a: &mut [Lane], twiddles: &[Lane], r: &R) { - let n = a.len(); - debug_assert!(n.is_power_of_two() && twiddles.len() >= n); - if n == 1 { return; } - if n == 2 { - // Radix-2 fallback: twiddle = twiddles[0] = 1. - let u = a[0]; let v = a[1]; - a[0] = r.add(&u, &v); - a[1] = r.sub(&u, &v); - return; - } - let j_mont = twiddles[n / 4]; - - let log_n = n.trailing_zeros(); - let mut sub_len = if log_n & 1 == 1 { - // One radix-2 stage with step = n/2 (only k=0, twiddle = 1). - for i in (0..n).step_by(2) { - let u = a[i]; let v = a[i + 1]; - a[i] = r.add(&u, &v); - a[i + 1] = r.sub(&u, &v); - } - 8 - } else { - 4 - }; - - // Radix-4 stages. - while sub_len <= n { - let q = sub_len / 4; - let step = n / sub_len; - for i in (0..n).step_by(sub_len) { - // k = 0: twiddles are all 1, skip the multiplies. - butterfly_radix4(a, i, q, twiddles[0], twiddles[0], twiddles[0], j_mont, r); - for k in 1..q { - let t1 = twiddles[k * step]; - let t2 = twiddles[2 * k * step]; - let t3 = twiddles[3 * k * step]; - butterfly_radix4(a, i + k, q, t1, t2, t3, j_mont, r); - } - } - sub_len *= 4; - } - } - ``` - `butterfly_radix4` is a small `#[inline(always)]` helper that performs the four muls (`b1, b2, b3` + `j·e3`) and writes back to the four positions. Make sure reads of `a[idx]` happen before any writes. - -3. **Keep a private `ntt_core_radix2_legacy`** during development — the current body of `ntt_core`, renamed. It reads only `twiddles[0..n/2]` so it works fine on the expanded table. Used only by the cross-check test (below) and deleted before merge. - -4. **Tests** — update local allocations in tests from `n/2` to `n` lanes. Extend `test_forward_correctness` to cover `n ∈ {2, 4, 8, 16, 32}` (currently only `{2, 4, 8}`). Add: - - `test_radix4_matches_legacy` — for each prime and `n ∈ {2, 4, 8, 16, 32, 64, 128, 256}`, run forward via both `ntt_core` and `ntt_core_radix2_legacy` on identical random bit-reversed input, assert byte-equal output. Delete together with the legacy fn before merge. - -### `integer/src/mul/ntt/mod.rs` (memory layout — four logical sites) - -Because of the recent refactor, twiddle allocation is centralised. The current code uses `nn / 2` for twiddle lengths in a small number of well-defined places, all of which need to become `nn`: - -1. **`memory_requirement_up_to` (line 98)** — worst-case scratch bound. The `twiddles = n_max` constant assumes two tables of size `n/2` (forward + inverse). Bump to `2 * n_max`: - ```rust - let twiddles = 2 * n_max; // was n_max - ``` - -2. **`run_ntt_pipeline` (lines 271–272, 283, 284–287)** — per-call scratch allocation + cache slicing inside the per-prime loop: - ```rust - // line 271-272 - let (fwd_twiddles, mut m) = m.allocate_slice_fill::(nn, 0); // was nn / 2 - let (inv_twiddles, _) = m.allocate_slice_fill::(nn, 0); // was nn / 2 - // line 283 - let tw_off = pi * nn; // was pi * (nn / 2) - // line 284-287 - ctx.fwd_twiddles.copy_from_slice(&fwd_tw_cache[tw_off..tw_off + nn]); // was nn / 2 - ctx.inv_twiddles.copy_from_slice(&inv_tw_cache[tw_off..tw_off + nn]); // was nn / 2 - ``` - -3. **`prepare_b_hat_and_twiddles` (lines 411 docstring, 428–429, 452, 454–455)** — per-prime scratch during precompute + cache write-back: - ```rust - // docstring at line ~411: "fwd_tw_cache and inv_tw_cache each geom.k_eff * geom.nn" - // (was geom.k_eff * (geom.nn / 2)) - // line 428-429 - let (fwd_tw, mut rest) = rest.allocate_slice_fill::(nn, 0); // was nn / 2 - let (inv_tw, _) = rest.allocate_slice_fill::(nn, 0); // was nn / 2 - // line 452 - let tw_off = pi * nn; // was pi * (nn / 2) - // line 454-455 - fwd_tw_cache[tw_off..tw_off + nn].copy_from_slice(fwd_tw); // was nn / 2 - inv_tw_cache[tw_off..tw_off + nn].copy_from_slice(inv_tw); // was nn / 2 - ``` - -4. **Cache length computations** — two sites that derive the total cache size from `nn`: - - `add_signed_mul_chunked` line 184: `let twiddle_len = k_eff * nn_chunk;` (was `k_eff * (nn_chunk / 2)`) - - `add_signed_mul_conv` line 334: `let twiddle_len = k_eff * nn;` (was `k_eff * (nn / 2)`) - -`NttGeometry` itself does **not** store twiddle size — only `nn`, `b_pack`, `k_eff`, `output_coeffs`. The `nn / 2` → `nn` change is local to the four sites above; no field needs adding to the geometry struct. - -### `integer/CHANGELOG.md` - -Add under `## Unreleased` → `### Improve`: -> NTT inner transform rewritten as radix-4 DIT (with one radix-2 stage when N is not a power of 4), halving the number of passes over the coefficient array. ~20–30% faster large-integer multiplication above the NTT threshold. - -### `TODO_NTT.md` - -Mark section "1. Radix-4 or split-radix NTT" as completed (move from "Remaining" to "Implemented", or strike through with a dated note). Leave split-radix as a possible future improvement. - -### No changes - -- `integer/src/arch/generic_64_bit/ntt.rs` and `integer/src/arch/generic_32_bit/ntt.rs` — primes and `OMEGA_MAX` are already sufficient. `MAX_LOG_N ≥ 2` is all radix-4 needs. -- `integer/src/mul/ntt/pack.rs`, `crt.rs` — unaffected. -- All callers in `mod.rs` (`add_signed_mul_conv`, `add_signed_mul_chunked` closure, `prepare_b_hat_and_twiddles`, `transform_b_forward`, `process_prime`, `do_crt`, `NttGeometry`, `TransformCtx`) — unaffected because `forward`/`inverse` signatures are unchanged. - ---- - -## Implementation sequence - -1. **Enter the worktree** `ntt-radix4` off the current `ssa` HEAD (via the worktree tool). -2. **Commit 1 — Twiddle table expansion.** Bump allocation sizes and cache-offset arithmetic in `transform.rs::precompute_twiddles`, and at all the mod.rs sites listed above. Also bump test-local allocations. The old radix-2 still works correctly on the now-oversized table. Run `cargo test -p dashu-int mul::ntt` — everything should pass. -3. **Commit 2 — Radix-4 core + legacy differential test.** Rename the existing `ntt_core` body to `ntt_core_radix2_legacy`. Write the new `ntt_core` (radix-4). Add `test_radix4_matches_legacy`. Extend `test_forward_correctness` to n ∈ {2, 4, 8, 16, 32}. Run the full NTT test suite; the schoolbook comparison tests in `mod.rs` are the strongest correctness gate. -4. **Commit 3 — Cleanup + docs.** Delete `ntt_core_radix2_legacy` and `test_radix4_matches_legacy`. Update `CHANGELOG.md` and `TODO_NTT.md`. -5. **Verification.** Run the existing `crossover_ntt` ignored test (`cargo test -p dashu-int --release -- mul::threshold_tests::crossover_ntt --ignored --nocapture`) and compare timings before/after by checking out `ssa` temporarily. A 20–30% drop in the NTT column at sizes ≥ 4096 words confirms the optimisation landed. - -## Verification end-to-end - -- `cargo check --all-features --tests` -- `cargo test --workspace --exclude dashu-python` -- `cargo clippy --all-features --all-targets --workspace --exclude dashu-python -- -D warnings` -- `cargo fmt --all -- --check` -- `cargo test -p dashu-int --release -- mul::threshold_tests::crossover_ntt --ignored --nocapture` — for end-to-end timing sanity. diff --git a/TODO.md b/TODO.md deleted file mode 100644 index 9d459cf1..00000000 --- a/TODO.md +++ /dev/null @@ -1,69 +0,0 @@ -## dashu-int Improvements - -### High impact - -- **`submul_1` fused primitive** — Multiply-and-subtract in one pass for the division inner loop correction step. Currently done as separate mul + sub, doubling memory passes. Reference: `ramp/src/ll/mul.rs:134-182`. - -- **`div_preinv` / 3-by-2 division with pre-inverted divisor** — Ramp computes `invert_pi(d1, d0)` for fast approximate quotient without the x86 `div` instruction, with separate `divrem_1` (single-limb) and `divrem_2` (two-limb) fast paths. Could replace the `num-modular` dependency. Affects many hot paths: modular arithmetic, formatting, base conversion. Reference: `ramp/src/ll/limb.rs:768-783`, `ramp/src/ll/div.rs:208-253`. - -- **Dedicated Toom-2 squaring** — `sqr_toom2` exploits symmetry: uses `z1 = x0*x1` instead of `(x0-x0)*(y1-y0)`, eliminating subtraction operations — only 3 sub-products instead of 4, and the cross term is `2*z1` without signed arithmetic. dashu has Karatsuba and Toom-3 general multiplication but no squaring-specific variant that takes advantage of `x == y`. dashu only uses specialized squaring up to 30 words. Reference: `ramp/src/ll/mul.rs:473-512`. - -- **Toom-22 as intermediate multiplication** — Ramp uses Toom-22 above 20 limbs before falling back to unbalanced mul. dashu goes Karatsuba → Toom-3 at 192 words. Toom-22 could fill the 24–192 word gap. Reference: `ramp/src/ll/mul.rs:243-390`. - -### Medium impact - -- **Trailing zero stripping in GCD loop** — Strip trailing zeros after each subtraction, not just at initialization. Helps for random inputs where intermediate results often gain trailing zeros. Reference: `ramp/src/ll/gcd.rs:20-86`. - -- **Trailing zero stripping in pow** — Factor out `(m * 2^k)^exp = m^exp * 2^(k*exp)` to reduce operand size. Reference: `ramp/src/ll/pow.rs:41-118`. - -### Low impact / ergonomics - -- **Build-time BASES table** — Pre-compute `digits_per_limb` and `big_base` per base via `build.rs` so base-10 conversion avoids repeated division. dashu's `integer/src/fmt/non_power_two.rs` uses simpler chunking (`CHUNK_LEN = 16`) without precomputed powers. Reference: `ramp/src/ll/base.rs:31-40`, `ramp/build.rs`. - -- **Scratch allocator improvements** — Ramp's `TmpAllocator` uses a linked list of dynamic allocations freed on drop, vs. dashu's pre-computed layout approach. Might be simpler for algorithms with hard-to-predict memory needs. Reference: `ramp/src/mem.rs`. - -## dashu-ratio Improvements - -- GCD: An idea of fast gcd check for rational number: don't do gcd reduction after every operation. - For small numerators or denominators, we can directly do a gcd, otherwise, we first do gcd with a primorial that - fits in a word (min is u16), and only remove these small divisors. - Further improvement: store a const divisor for the prime factors in the primorial, thus supports a fast factorial of - the gcd result, and then divide with these const divisor. - -## dashu-float Improvements - -- **Trig (`sin`/`cos`) — current baseline** — `float/src/math/trig.rs` uses dynamic guard-digit - work precision, simple `x mod (π/2)` range reduction, and Taylor series on the reduced - argument `r`. Adequate for moderate precision and moderate `|x|`; items below target large - arguments and very high precision. Reference: MPFR `mpfr_sin` / `mpfr_sin_cos`. - -### High impact - -- **Payne–Hanek range reduction** — For large `|x|`, replace `k = round(x/(π/2)); r = x - k·(π/2)` with - multiplication by precomputed blocks of `2/π`, extracting the integer part without a full high-precision - division. Avoids catastrophic cancellation that currently forces `work_precision ≈ precision + log|x| + guards` - (`compute_work_context`). This is the main gap vs. MPFR for huge arguments. Reference: `float/src/math/trig.rs`. - -- **Binary splitting for Taylor core** — At high precision, evaluate the `sin`/`cos` series via binary splitting - (same technique as Chudnovsky π in `float/src/math/consts.rs`) instead of naive term-by-term accumulation. - Reduces cost from O(p²) to roughly O(M(p) log p) for p-bit results. - -### Medium impact - -- **Remez minimax polynomial + Clenshaw (low/medium p)** — For `p ≲ 512`, use a fixed-degree minimax polynomial - on `[-π/4, π/4]` evaluated with Clenshaw recurrence instead of Taylor. MPFR uses this for its fast path; switch - to series/binary splitting only when p is large. - -- **Cody–Waite π/2 split** — Represent `π/2 = hi + lo` and compute `r = ((x - k·hi) - k·lo)` to reduce guard - digit pressure for moderate `|x|` before Payne–Hanek is needed. Complements the existing `reduce_to_quadrant`. - -- **Argument shrinking for `|r| > π/4`** — Use `sin(r) = cos(π/2 - r)` (and the cosine analogue) so the Taylor - series runs on a smaller interval, needing fewer terms when `r` is near ±π/2. - -### Low impact / ergonomics - -- **Cache π at common precisions** — Avoid recomputing Chudnovsky π on every trig call when work precision repeats. - TODO already noted in `float/src/math/consts.rs`. - -- **Precomputed `2/π` block table** — Storage for Payne–Hanek: blocks of `2/π` bits (e.g. 32/64 bits per entry), - generated once or lazily on first use at a given precision. diff --git a/TODO_NTT.md b/TODO_NTT.md deleted file mode 100644 index 537634a4..00000000 --- a/TODO_NTT.md +++ /dev/null @@ -1,57 +0,0 @@ -# NTT multiplication for `UBig` — status & remaining work - -## Implemented - -### Core algorithm - -- **Primes.** Three Proth primes of the form `K·2^N + 1` per word size: - 64-bit: `29·2^57+1`, `71·2^57+1`, `75·2^57+1` (MAX_LOG_N=57). - 32-bit: `7·2^26+1`, `15·2^27+1`, `17·2^27+1` (MAX_LOG_N=26). - Defined in `integer/src/arch/generic_{64,32}_bit/ntt.rs` with per-prime - reducer instances `P0`/`P1`/`P2`. - -- **Modular arithmetic.** Delegates to `num_modular::Reducer` (Proth - Montgomery reduction). Transform functions are generic over - `R: Reducer`, monomorphized per prime at the `process_prime` call - site. No per-prime wrapper functions. - -- **NTT transforms.** Iterative in-place radix-2 DIT. Forward: - `bit_reverse → forward(ω)`. Inverse: `bit_reverse → forward(ω⁻¹) → scale`. - All arithmetic in Montgomery form; conversion at pipeline boundaries via - `r.transform`/`r.residue`. - -- **CRT.** Garner's algorithm combining `K` residues into a `TripleWord` - (`[u64;3]` on 64-bit, `[u32;3]` on 32-bit). `CrtAccum` trait in - `crt.rs`, impl gated by `#[cfg]` per word size. Standard-form arithmetic - via `num_modular::ModularCoreOps::subm`/`mulm`. - -- **Dispatch.** `THRESHOLD_NTT = 4 000` words (256 kbits). Asymmetric - chunking (`a > 2·b`): pre-transforms `b` once, reuses `b̂` across chunks of - `a`. Shared entry point `process_prime(a, b: BInput<'_>, ctx, r)` handles - both raw `b` and cached `b̂`. - -### Architecture - -- NTT constants and reducer instances live in arch-specific `ntt.rs` files, - re-exported through `arch/mod.rs` → `crate::arch::ntt`. -- 16-bit Word targets excluded at compile time (`#[cfg]` on `pub(crate) mod ntt`). - -### Benchmarking - -- `ubig_mul_asymmetric` in `integer/benches/primitive.rs` — fixed `b` (500 kbits), - varying `a` (1 kbit – 5 Mbits). Exercises all chunked-mul code paths. - ---- - -## Remaining optimisation opportunities - -### 1. Radix-4 NTT (~25–33% fewer twiddle multiplies) - -Attempted but reverted — the 4-point DFT output ordering within DIT/DIF -interacts non-trivially with bit-reversal. Needs careful re-analysis. - -### 2. Harvey lazy-reduction butterflies (~10–15%) - -Bypass `r.add`/`r.sub` normalization in `ntt_core`, deferring to a cleanup -pass. Trickiest due to interaction with `num_modular`'s `Reducer` API. - From 1f9d93d9565bd81752cd787fd892a3ecb726c820 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 18:23:09 +0800 Subject: [PATCH 15/19] Fix 32-bit and clippy CI failures in NTT module Use native Word/Lane types throughout pack.rs (instead of u64/u32) so the same source compiles cleanly on both 32-bit and 64-bit targets. Resolve the remaining clippy warnings (unnecessary_cast, let_and_return, too_many_arguments, needless_range_loop, type_complexity) that were failing the -D warnings CI run. Co-Authored-By: Claude Opus 4.7 --- integer/CHANGELOG.md | 6 ++++ integer/src/arch/generic_32_bit/ntt.rs | 4 ++- integer/src/arch/generic_64_bit/ntt.rs | 4 ++- integer/src/mul/ntt/crt.rs | 2 +- integer/src/mul/ntt/mod.rs | 15 ++++----- integer/src/mul/ntt/pack.rs | 43 ++++++++++++++------------ 6 files changed, 42 insertions(+), 32 deletions(-) diff --git a/integer/CHANGELOG.md b/integer/CHANGELOG.md index cb29c33b..17e4b422 100644 --- a/integer/CHANGELOG.md +++ b/integer/CHANGELOG.md @@ -21,6 +21,12 @@ - NTT enabled for 32-bit Word targets. - Arch-specific NTT prime definitions under `arch/generic_{32,64}_bit/ntt.rs`. +### Fix +- `pack.rs` test used 64-bit literals that overflowed `Word` (`u32`) on 32-bit targets, breaking the test build. +- `pack.rs` now uses native `Word`/`Lane` types throughout instead of `u64`/`u32`, fixing clippy `unnecessary_cast` warnings on 64-bit. +- `test_unpack_carry_propagation` had a hardcoded 64-bit shift assumption; now derived from `Word::BITS` so it works on 32-bit. +- Various clippy warnings (`let_and_return`, `too_many_arguments`, `needless_range_loop`, `type_complexity`) resolved across the NTT module. + ## 0.4.2 - Add `UBig::ones`. diff --git a/integer/src/arch/generic_32_bit/ntt.rs b/integer/src/arch/generic_32_bit/ntt.rs index a2218dd1..04682b95 100644 --- a/integer/src/arch/generic_32_bit/ntt.rs +++ b/integer/src/arch/generic_32_bit/ntt.rs @@ -38,6 +38,8 @@ mod tests { use super::*; use num_modular::Reducer; + type ReducerFns = (fn(Lane) -> Lane, fn(Lane) -> Lane, fn(Lane) -> Lane); + #[test] fn test_primes_proth_form() { assert_eq!(MODULI[0], 7u32 * (1u32 << 26) + 1); @@ -57,7 +59,7 @@ mod tests { fn test_omega_order() { for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { let p = MODULI[pi]; - let (sqr, to_m, from_m): (fn(Lane) -> Lane, fn(Lane) -> Lane, fn(Lane) -> Lane) = + let (sqr, to_m, from_m): ReducerFns = match pi { 0 => ( |w| P0.reduce((w as u64) * (w as u64)), diff --git a/integer/src/arch/generic_64_bit/ntt.rs b/integer/src/arch/generic_64_bit/ntt.rs index 0bf25969..24942c76 100644 --- a/integer/src/arch/generic_64_bit/ntt.rs +++ b/integer/src/arch/generic_64_bit/ntt.rs @@ -43,6 +43,8 @@ mod tests { use super::*; use num_modular::Reducer; + type ReducerFns = (fn(Lane) -> Lane, fn(Lane) -> Lane, fn(Lane) -> Lane); + #[test] fn test_primes_proth_form() { assert_eq!(MODULI[0], 29u64 * (1u64 << 57) + 1); @@ -62,7 +64,7 @@ mod tests { fn test_omega_order() { for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { let p = MODULI[pi]; - let (sqr, to_m, from_m): (fn(Lane) -> Lane, fn(Lane) -> Lane, fn(Lane) -> Lane) = + let (sqr, to_m, from_m): ReducerFns = match pi { 0 => ( |w| P0.reduce((w as u128) * (w as u128)), diff --git a/integer/src/mul/ntt/crt.rs b/integer/src/mul/ntt/crt.rs index 045f5074..af6ba991 100644 --- a/integer/src/mul/ntt/crt.rs +++ b/integer/src/mul/ntt/crt.rs @@ -216,6 +216,6 @@ mod tests { let x = garner_combine::(&residues[..1], &CRT_INV_IJ, &primes); let mut buf = [crate::arch::word::Word::default(); 6]; x.write_words(&mut buf); - assert_eq!(buf[0] as u64, residues[0] as u64); + assert_eq!(buf[0], residues[0]); } } diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index 2a615afa..0b5df34a 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -200,7 +200,7 @@ fn add_signed_mul_chunked( let coeffs_b = coeff_count(lb_bits, b_pack); // ---- Chunked multiply ---- - let carry = add_signed_mul_split_into_chunks( + add_signed_mul_split_into_chunks( c, sign, a, @@ -234,9 +234,7 @@ fn add_signed_mul_chunked( mem, ) }, - ); - - carry + ) } /// Run the full NTT pipeline: allocate → per-prime transform → CRT → fold into `c_out`. @@ -248,6 +246,7 @@ fn add_signed_mul_chunked( /// /// Shared body of `add_signed_mul_conv` and the per-chunk callback in /// `add_signed_mul_chunked`. +#[allow(clippy::too_many_arguments)] fn run_ntt_pipeline( a: &[Word], b_hat: &[crate::arch::ntt::Lane], @@ -387,7 +386,6 @@ struct TransformCtx<'a> { /// Transform `b` and leave the result in `b_lane` (forward-transformed, /// Montgomery form). `fwd_twiddles` must already be precomputed. - fn transform_b_forward>( b_lane: &mut [crate::arch::ntt::Lane], b: &[Word], @@ -423,12 +421,11 @@ fn prepare_b_hat_and_twiddles( let b_pack = geom.b_pack; let k_eff = geom.k_eff; - for pi in 0..k_eff { + for (pi, &omega) in OMEGA_MAX.iter().enumerate().take(k_eff) { let (b_lane, mut rest) = mem.allocate_slice_fill::(nn, 0); let (fwd_tw, mut rest) = rest.allocate_slice_fill::(nn / 2, 0); let (inv_tw, _) = rest.allocate_slice_fill::(nn / 2, 0); - let omega = OMEGA_MAX[pi]; match pi { 0 => { transform::precompute_twiddles(fwd_tw, nn, omega, false, &P0); @@ -503,8 +500,8 @@ fn add_shifted_to_prod(prod: &mut [Word], words: &[Word], count: u32, k: usize, let mut carry: Word = 0; - for vi in 0..(count as usize) { - let limb = words[vi].wrapping_add(carry); + for (vi, &word) in words.iter().enumerate().take(count as usize) { + let limb = word.wrapping_add(carry); let idx = start_idx + vi; if idx >= prod.len() { return; diff --git a/integer/src/mul/ntt/pack.rs b/integer/src/mul/ntt/pack.rs index 0989beeb..05111224 100644 --- a/integer/src/mul/ntt/pack.rs +++ b/integer/src/mul/ntt/pack.rs @@ -14,20 +14,20 @@ pub fn pack(out: &mut [Lane], words: &[Word], b_pack: u32, n: usize) { // Fast path: one coefficient per word, no bit shifting needed. if b_pack == Word::BITS { let len = words.len().min(n); - #[allow(clippy::unnecessary_cast)] // SAFETY: NTT path requires Word and Lane have the same size. + #[allow(clippy::unnecessary_cast)] let words_lane = unsafe { &*(words as *const [Word] as *const [Lane]) }; out[..len].copy_from_slice(&words_lane[..len]); out[len..n].fill(0); return; } - let mask = if b_pack < Word::BITS { - (1u64 << b_pack) - 1 + let word_bits = Word::BITS; + let mask: Word = if b_pack < word_bits { + (1 << b_pack) - 1 } else { - u64::MAX + Word::MAX }; - let word_bits = Word::BITS; let mut word_idx = 0usize; let mut bit_offset = 0u32; @@ -38,7 +38,7 @@ pub fn pack(out: &mut [Lane], words: &[Word], b_pack: u32, n: usize) { } if bit_offset + b_pack <= word_bits { - *coeff = ((words[word_idx] as u64 >> bit_offset) & mask) as Lane; + *coeff = (words[word_idx] >> bit_offset) & mask; bit_offset += b_pack; if bit_offset == word_bits { bit_offset = 0; @@ -47,12 +47,12 @@ pub fn pack(out: &mut [Lane], words: &[Word], b_pack: u32, n: usize) { } else { let bits_first = word_bits - bit_offset; let bits_second = b_pack - bits_first; - let mut val = (words[word_idx] as u64 >> bit_offset) & ((1u64 << bits_first) - 1); + let mut val = (words[word_idx] >> bit_offset) & ((1 << bits_first) - 1); word_idx += 1; if word_idx < words.len() { - val |= (words[word_idx] as u64 & ((1u64 << bits_second) - 1)) << bits_first; + val |= (words[word_idx] & ((1 << bits_second) - 1)) << bits_first; } - *coeff = val as Lane; + *coeff = val; bit_offset = bits_second; } } @@ -135,24 +135,24 @@ mod tests { #[test] fn test_pack_unpack_roundtrip() { let b_pack = 16u32; - let test_words: Vec = vec![0xDEADBEEF_CAFEBABE, 0x12345678_9ABCDEF0]; + let test_words: Vec = vec![0xDEADBEEF, 0x12345678]; let coeffs_per_word = (Word::BITS / b_pack) as usize; let n = test_words.len() * coeffs_per_word; - let mut packed = vec![0u64 as Lane; n]; + let mut packed: Vec = vec![0; n]; pack(&mut packed, &test_words, b_pack, n); let output_len = test_words.len() + 1; - let mut output = vec![0u64 as Word; output_len]; + let mut output: Vec = vec![0; output_len]; unpack_accumulate(&mut output, &packed, b_pack, n); assert_eq!(&output[..test_words.len()], &test_words[..]); } #[test] fn test_pack_zero_pads() { - let words = vec![0xFFFFu64 as Word]; + let words: Vec = vec![0xFFFF]; let n = 32; - let mut packed = vec![0u64 as Lane; n]; + let mut packed: Vec = vec![0; n]; pack(&mut packed, &words, 16, n); assert_eq!(packed[0], 0xFFFF); for &c in packed.iter().skip(1) { @@ -162,14 +162,14 @@ mod tests { #[test] fn test_pack_empty_input() { - let mut packed = vec![0u64 as Lane; 8]; + let mut packed: Vec = vec![0; 8]; pack(&mut packed, &[], 16, 8); - assert_eq!(packed, vec![0u64 as Lane; 8]); + assert_eq!(packed, vec![0; 8]); } #[test] fn test_unpack_single_coeff() { - let mut output = vec![0u64 as Word; 2]; + let mut output: Vec = vec![0; 2]; unpack_accumulate(&mut output, &[0xABCD], 16, 1); assert_eq!(output[0], 0xABCD); assert_eq!(output[1], 0); @@ -177,9 +177,12 @@ mod tests { #[test] fn test_unpack_carry_propagation() { - // Coefficient at k=4 (shift by 64 bits = 1 word) + carry - let mut output = vec![0u64 as Word; 3]; - unpack_accumulate(&mut output, &[0, 0, 0, 0, 1], 16, 5); + // Coefficient at k = Word::BITS/16 shifts by exactly Word::BITS bits = 1 word. + let k = (Word::BITS / 16) as usize; + let mut coeffs: Vec = vec![0; k + 1]; + coeffs[k] = 1; + let mut output: Vec = vec![0; 3]; + unpack_accumulate(&mut output, &coeffs, 16, k + 1); assert_eq!(output[0], 0); assert_eq!(output[1], 1); assert_eq!(output[2], 0); From ed2379cd28944225cec14f43c10a6a656336959d Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 18:27:57 +0800 Subject: [PATCH 16/19] CI: run clippy on 32-bit Word target Adds a second clippy step to the existing Clippy job that runs with --cfg force_bits="32", catching Word-width-dependent lints like the unnecessary_cast warnings recently fixed in the NTT module. Co-Authored-By: Claude Opus 4.7 --- .github/workflows/tests.yml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index a40ff963..1abdb09f 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -162,4 +162,9 @@ jobs: with: toolchain: stable components: clippy - - run: cargo clippy --all-features --all-targets --workspace --exclude dashu-python -- -D warnings \ No newline at end of file + - name: Clippy (default / 64-bit Word) + run: cargo clippy --all-features --all-targets --workspace --exclude dashu-python -- -D warnings + - name: Clippy (32-bit Word) + env: + RUSTFLAGS: --cfg force_bits="32" + run: cargo clippy --all-features --all-targets --workspace --exclude dashu-python -- -D warnings \ No newline at end of file From a355d577816be5a8bb4f045de5cd28416ef95639 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 18:43:58 +0800 Subject: [PATCH 17/19] Fix clippy warnings on 32-bit Word in dashu-float Extract radix once in from_str_native to avoid repeated B as u32 casts, and silence the identity try_into() in num_traits::Num::from_str_radix on 32-bit Word targets. Co-Authored-By: Claude Opus 4.7 --- float/src/parse.rs | 12 ++++++--- float/src/third_party/num_traits.rs | 2 +- integer/src/arch/generic_32_bit/ntt.rs | 31 +++++++++------------ integer/src/arch/generic_64_bit/ntt.rs | 37 +++++++++++++------------- 4 files changed, 40 insertions(+), 42 deletions(-) diff --git a/float/src/parse.rs b/float/src/parse.rs index 1b164cee..3bf106c5 100644 --- a/float/src/parse.rs +++ b/float/src/parse.rs @@ -27,6 +27,12 @@ impl Repr { pub fn from_str_native(mut src: &str) -> Result<(Self, usize), ParseError> { assert!(MIN_RADIX as Word <= B && B <= MAX_RADIX as Word); + // B is guaranteed to be in 2..=36 by the assert above; the cast to u32 + // is needed because `from_str_radix` takes a u32 radix. On 32-bit Word + // targets the cast is a no-op. + #[allow(clippy::unnecessary_cast)] + let radix: u32 = B as u32; + // parse and remove the sign let sign = match src.strip_prefix('-') { Some(s) => { @@ -100,14 +106,14 @@ impl Repr { return Err(ParseError::UnsupportedRadix); } else { let digits = int_str.len() - int_str.matches('_').count(); - (UBig::from_str_radix(&src[..dot], B as u32)?, digits, B as u32) + (UBig::from_str_radix(&src[..dot], radix)?, digits, radix) } } else { if pmarker { // prefix is required for using `p` as scale marker return Err(ParseError::UnsupportedRadix); } - (UBig::ZERO, 0, B as u32) + (UBig::ZERO, 0, radix) }; // parse fractional part @@ -139,7 +145,7 @@ impl Repr { return Err(ParseError::UnsupportedRadix); } else { ndigits = src.len() - src.matches('_').count(); - UBig::from_str_radix(src, B as u32)? + UBig::from_str_radix(src, radix)? } }; diff --git a/float/src/third_party/num_traits.rs b/float/src/third_party/num_traits.rs index 1d5e3dbe..0c93e7ef 100644 --- a/float/src/third_party/num_traits.rs +++ b/float/src/third_party/num_traits.rs @@ -133,7 +133,7 @@ impl num_traits::Num for FBig { #[inline] fn from_str_radix(s: &str, radix: u32) -> Result { // the conversion might a fail with 16-bit words. - #[allow(clippy::unnecessary_fallible_conversions)] + #[allow(clippy::unnecessary_fallible_conversions, clippy::useless_conversion)] let r: Word = radix.try_into().map_err(|_| ParseError::UnsupportedRadix)?; if r == B { #[allow(deprecated)] // TODO(v0.5): remove after from_str_native is made private. diff --git a/integer/src/arch/generic_32_bit/ntt.rs b/integer/src/arch/generic_32_bit/ntt.rs index 04682b95..746bc188 100644 --- a/integer/src/arch/generic_32_bit/ntt.rs +++ b/integer/src/arch/generic_32_bit/ntt.rs @@ -59,25 +59,18 @@ mod tests { fn test_omega_order() { for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { let p = MODULI[pi]; - let (sqr, to_m, from_m): ReducerFns = - match pi { - 0 => ( - |w| P0.reduce((w as u64) * (w as u64)), - |v| P0.transform(v), - |v| P0.residue(v), - ), - 1 => ( - |w| P1.reduce((w as u64) * (w as u64)), - |v| P1.transform(v), - |v| P1.residue(v), - ), - 2 => ( - |w| P2.reduce((w as u64) * (w as u64)), - |v| P2.transform(v), - |v| P2.residue(v), - ), - _ => unreachable!(), - }; + let (sqr, to_m, from_m): ReducerFns = match pi { + 0 => { + (|w| P0.reduce((w as u64) * (w as u64)), |v| P0.transform(v), |v| P0.residue(v)) + } + 1 => { + (|w| P1.reduce((w as u64) * (w as u64)), |v| P1.transform(v), |v| P1.residue(v)) + } + 2 => { + (|w| P2.reduce((w as u64) * (w as u64)), |v| P2.transform(v), |v| P2.residue(v)) + } + _ => unreachable!(), + }; let mut w = to_m(omega_max); for _ in 0..MAX_LOG_N - 1 { diff --git a/integer/src/arch/generic_64_bit/ntt.rs b/integer/src/arch/generic_64_bit/ntt.rs index 24942c76..03896b5b 100644 --- a/integer/src/arch/generic_64_bit/ntt.rs +++ b/integer/src/arch/generic_64_bit/ntt.rs @@ -64,25 +64,24 @@ mod tests { fn test_omega_order() { for (pi, &omega_max) in OMEGA_MAX.iter().enumerate() { let p = MODULI[pi]; - let (sqr, to_m, from_m): ReducerFns = - match pi { - 0 => ( - |w| P0.reduce((w as u128) * (w as u128)), - |v| P0.transform(v), - |v| P0.residue(v), - ), - 1 => ( - |w| P1.reduce((w as u128) * (w as u128)), - |v| P1.transform(v), - |v| P1.residue(v), - ), - 2 => ( - |w| P2.reduce((w as u128) * (w as u128)), - |v| P2.transform(v), - |v| P2.residue(v), - ), - _ => unreachable!(), - }; + let (sqr, to_m, from_m): ReducerFns = match pi { + 0 => ( + |w| P0.reduce((w as u128) * (w as u128)), + |v| P0.transform(v), + |v| P0.residue(v), + ), + 1 => ( + |w| P1.reduce((w as u128) * (w as u128)), + |v| P1.transform(v), + |v| P1.residue(v), + ), + 2 => ( + |w| P2.reduce((w as u128) * (w as u128)), + |v| P2.transform(v), + |v| P2.residue(v), + ), + _ => unreachable!(), + }; let mut w = to_m(omega_max); for _ in 0..MAX_LOG_N - 1 { From b16f1a7fdf98953e8eae8c46db4cda6cca30f6d2 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 19:01:05 +0800 Subject: [PATCH 18/19] Remove development-only NTT tests Tests removed are fully subsumed by the schoolbook-comparison and roundtrip tests added later in development. Co-Authored-By: Claude Opus 4.7 --- integer/src/mul/ntt/mod.rs | 30 -------- integer/src/mul/ntt/transform.rs | 115 ------------------------------- 2 files changed, 145 deletions(-) diff --git a/integer/src/mul/ntt/mod.rs b/integer/src/mul/ntt/mod.rs index 0b5df34a..a0f21200 100644 --- a/integer/src/mul/ntt/mod.rs +++ b/integer/src/mul/ntt/mod.rs @@ -574,20 +574,6 @@ mod tests { assert_eq!(bit_len(&[1]), 1); } - #[test] - fn test_ntt_multiply_one_word() { - let a: Vec = vec![3]; - let b: Vec = vec![5]; - let mut c = vec![0u64 as Word; 2]; - let layout = memory_requirement_up_to(c.len(), b.len()); - let mut alloc = crate::memory::MemoryAllocation::new(layout); - let mut memory = alloc.memory(); - let carry = add_signed_mul_conv(&mut c, Positive, &a, &b, &mut memory); - assert_eq!(carry, 0); - assert_eq!(c[0], 15); - assert_eq!(c[1], 0); - } - #[test] fn test_ntt_sign_negative() { let a: Vec = (0..30).map(|i| (i as Word + 1) * 100).collect(); @@ -607,22 +593,6 @@ mod tests { assert!(c.iter().all(|&w| w == 0)); } - const NTT_TEST_LEN: usize = 512; - - #[test] - fn test_ntt_multiply_small() { - let a: Vec = vec![0xDEADBEEFu64 as Word; NTT_TEST_LEN]; - let b: Vec = vec![0xCAFEBABEu64 as Word; NTT_TEST_LEN]; - let mut c = vec![0u64 as Word; a.len() + b.len()]; - - let layout = memory_requirement_up_to(c.len(), b.len()); - let mut alloc = crate::memory::MemoryAllocation::new(layout); - let mut memory = alloc.memory(); - let carry = add_signed_mul_conv(&mut c, Positive, &a, &b, &mut memory); - assert_eq!(carry, 0); - assert!(c.iter().any(|&w| w != 0)); - } - /// Naive schoolbook multiplication for comparison. fn schoolbook_mul(a: &[Word], b: &[Word]) -> Vec { let mut c = vec![0u64 as Word; a.len() + b.len()]; diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs index 43c769da..6d7d7e52 100644 --- a/integer/src/mul/ntt/transform.rs +++ b/integer/src/mul/ntt/transform.rs @@ -109,7 +109,6 @@ mod tests { use alloc::vec; #[cfg(not(feature = "std"))] use alloc::vec::Vec; - use num_modular::ModularPow; fn assert_all_eq(a: &[Lane], b_val: &[Lane], context: &str) { assert_eq!(a.len(), b_val.len(), "{context}: length mismatch"); @@ -225,118 +224,4 @@ mod tests { assert_eq!(a, vec![0, 4, 2, 6, 1, 5, 3, 7]); } - /// Naive O(n²) NTT using standard-form modular arithmetic. - #[allow(clippy::needless_range_loop)] - fn ntt_naive_std(x: &[Lane], omega_n: Lane, p: Lane) -> Vec { - use num_modular::ModularCoreOps; - let n = x.len(); - let mut result = vec![0u64 as Lane; n]; - for k in 0..n { - let mut acc = 0u64 as Lane; - for j in 0..n { - let twiddle = if k == 0 || j == 0 { - 1 - } else { - omega_n.powm(&((k * j) as Lane), &p) - }; - acc = acc.addm(x[j].mulm(twiddle, &p), &p); - } - result[k] = acc; - } - result - } - - #[test] - fn test_forward_correctness() { - for_each_prime!(r, p, omega, { - for &n in &[2usize, 4, 8] { - let x: Vec = (0..n).map(|i| ((i + 1) as Lane * 11111) % p).collect(); - - let shift = MAX_LOG_N - n.trailing_zeros(); - let omega_n = omega.powm(&((1u64 as Lane) << shift), &p); - - let mut a: Vec = x.iter().map(|&v| r.transform(v)).collect(); - bit_reverse(&mut a); - - let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; - precompute_twiddles(&mut fwd_twiddles, n, omega, false, r); - forward(&mut a, &fwd_twiddles, r); - - for val in a.iter_mut() { - *val = r.residue(*val); - } - let expected = ntt_naive_std(&x, omega_n, p); - assert_eq!(a, expected, "forward NTT mismatch"); - } - }); - } - - #[test] - fn test_convolution_debug() { - let p = MODULI[0]; - let omega = OMEGA_MAX[0]; - let r = &P0; - - let a = [12345u64 as Lane % p]; - let b_vec = [ - 67890u64 as Lane % p, - 135780u64 as Lane % p, - 203670u64 as Lane % p, - ]; - let conv_len = a.len() + b_vec.len() - 1; - let n = 4; - - let mut expected = vec![0u64 as Lane; conv_len]; - for (i, &ai) in a.iter().enumerate() { - for (j, &bj) in b_vec.iter().enumerate() { - let prod = (ai as u128 * bj as u128 % p as u128) as Lane; - expected[i + j] = r.add(&expected[i + j], &prod); - } - } - - let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; - let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; - precompute_twiddles(&mut fwd_twiddles, n, omega, false, r); - precompute_twiddles(&mut inv_twiddles, n, omega, true, r); - - let mut a_pad = vec![0u64 as Lane; n]; - let mut b_pad = vec![0u64 as Lane; n]; - for i in 0..a.len() { - a_pad[i] = r.transform(a[i]); - } - for i in 0..b_vec.len() { - b_pad[i] = r.transform(b_vec[i]); - } - - bit_reverse(&mut a_pad); - bit_reverse(&mut b_pad); - forward(&mut a_pad, &fwd_twiddles, r); - forward(&mut b_pad, &fwd_twiddles, r); - pointwise_mul(&mut a_pad, &b_pad, r); - inverse(&mut a_pad, &inv_twiddles, r); - for val in a_pad[..conv_len].iter_mut() { - *val = r.residue(*val); - } - - assert_eq!(&a_pad[..conv_len], &expected[..]); - } - - #[test] - fn test_length_two_edge_case() { - for_each_prime!(r, p, omega, { - let n = 2; - let mut fwd_twiddles = alloc::vec![0u64 as Lane; n / 2]; - let mut inv_twiddles = alloc::vec![0u64 as Lane; n / 2]; - precompute_twiddles(&mut fwd_twiddles, n, omega, false, r); - precompute_twiddles(&mut inv_twiddles, n, omega, true, r); - - let a_std = [1u64 as Lane % p, 2u64 as Lane % p]; - let a_orig: Vec = a_std.iter().map(|&v| r.transform(v)).collect(); - let mut a = a_orig.clone(); - bit_reverse(&mut a); - forward(&mut a, &fwd_twiddles, r); - inverse(&mut a, &inv_twiddles, r); - assert_all_eq(&a, &a_orig, "length two roundtrip"); - }); - } } From 687f58c9425f75dcda7666bc534c0e02fad956d7 Mon Sep 17 00:00:00 2001 From: Jacob Zhong Date: Sun, 14 Jun 2026 22:33:47 +0800 Subject: [PATCH 19/19] Fix fmt --- integer/src/mul/ntt/transform.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/integer/src/mul/ntt/transform.rs b/integer/src/mul/ntt/transform.rs index 6d7d7e52..4b8866d6 100644 --- a/integer/src/mul/ntt/transform.rs +++ b/integer/src/mul/ntt/transform.rs @@ -223,5 +223,4 @@ mod tests { bit_reverse(&mut a); assert_eq!(a, vec![0, 4, 2, 6, 1, 5, 3, 7]); } - }