Skip to content

Add specialized implementations for squaring - #76

Merged
cmpute merged 11 commits into
masterfrom
sqr
Jun 16, 2026
Merged

Add specialized implementations for squaring#76
cmpute merged 11 commits into
masterfrom
sqr

Conversation

@cmpute

@cmpute cmpute commented Jun 14, 2026

Copy link
Copy Markdown
Owner

No description provided.

@cmpute

cmpute commented Jun 14, 2026

Copy link
Copy Markdown
Owner Author

A blocking issue: squaring is slower than multiplying in certain sizes:

  ┌──────┬───────┬──────────┬──────────┬───────┬──────────────┐
  │ Size │ Words │ ubig_sqr │ ubig_mul │ Ratio │  Algorithm   │
  ├──────┼───────┼──────────┼──────────┼───────┼──────────────┤
  │ 1e1  │ 1     │ 4.09 ns  │ 6.04 ns  │ 0.68x │ inline dword │
  ├──────┼───────┼──────────┼──────────┼───────┼──────────────┤
  │ 1e2  │ 2     │ 64.3 ns  │ 36.2 ns  │ 1.78x │ simple       │
  ├──────┼───────┼──────────┼──────────┼───────┼──────────────┤
  │ 1e3  │ 16    │ 552.7 ns │ 376.7 ns │ 1.47x │ simple       │
  ├──────┼───────┼──────────┼──────────┼───────┼──────────────┤
  │ 1e4  │ 157   │ 15.62 µs │ 17.17 µs │ 0.91x │ karatsuba    │
  ├──────┼───────┼──────────┼──────────┼───────┼──────────────┤
  │ 1e5  │ 1563  │ 487.2 µs │ 628.0 µs │ 0.78x │ toom-3       │
  ├──────┼───────┼──────────┼──────────┼───────┼──────────────┤
  │ 1e6  │ 15625 │ 9.46 ms  │ 12.47 ms │ 0.76x │ NTT          │
  └──────┴───────┴──────────┴──────────┴───────┴──────────────┘

Jacob Zhong and others added 5 commits June 15, 2026 21:33
- Karatsuba: 3 recursive squarings instead of multiplications, simplified diff
- Toom-3: single polynomial evaluation, 5 squarings, V(-1) always positive
- NTT: single forward transform, pointwise square instead of multiply
- Separate _SQR env var thresholds from _MUL multiplication thresholds

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
- Karatsuba: 3n + 2·ceil_log2(n) (vs 2n for mul) — the diff_sq
  temp buffer adds ~n words of scratch pressure at peak.
- Toom-3: revert to 4n + 13·ceil_log2(n) (same as mul) as the
  earlier 3n + 10 derivation missed buffer overlap at recursion.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Move add_signed_sqr_same_len, add_signed_sqr_conv, and
process_prime_square to sqr/ntt.rs. Expose NTT internals
(pack, transform, NttGeometry, do_crt, bit_len, coeff_count,
add_shifted_to_prod) as pub(crate) for the new location.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Remove pointwise_square from mul/ntt/transform.rs and inline the
Montgomery squaring directly in process_prime_square. It was only
used by the NTT squaring path.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Tests 10^1 through 10^6 bit random operands, exercising simple,
Karatsuba, Toom-3, and NTT squaring paths.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Jacob Zhong and others added 6 commits June 15, 2026 23:50
Explain at the memory_requirement_up_to site that the diff_sq temp is
kept deliberately: accumulating the cross term straight into the output
(as multiplication does) would force a full-schoolbook basecase, since
the symmetric in-place squaring needs a zeroed target. The temp trades
scratch memory for the efficient symmetric basecase (~25% faster in the
Toom-3 range, which recurses into Karatsuba).

Co-Authored-By: Claude <noreply@anthropic.com>
The schoolbook squaring basecase (simple::square) computed its off-diagonal
triangle one source limb at a time via add_mul_word_same_len_in_place,
while the multiplication basecase had already moved to a two-word
(mpn_addmul_2-style) kernel that halves accumulator memory traffic.

Bring squaring to parity: pair consecutive limbs (a[i], a[i+1]) and multiply
their shared suffix a[i+2..] by the double word (a[i] + a[i+1]*B) through
two interleaved mul-accumulate chains (mul_add_2carry), folding the lone
corner product a[i]*a[i+1] into the kernel's seed carry. Carries propagate
in place — each pair's high words land at b[n+i], b[n+i+1] and the next
pair's sweep re-reads b[n+i+1], so only a single pending carry escapes a
sweep (and it resolves to 0, since the triangle < B^(2n-1)).

simple::square is the basecase every algorithm recurses into, so this lifts
the whole stack: ubig_sqr/1e3 (schoolbook range) ~1.25x, ~1.15x through the
Karatsuba/Toom-3 bands (1e4-1e5), flat at 1e6 (NTT-bound). ubig_pow improves
likewise (~1.17x at 1e3-1e4) since exponentiation is squaring-dominated.

Adds a sqr()==mul() regression test across sizes 2..=128 plus adversarial
(all-ones, single high bit, alternating) inputs.

Co-Authored-By: Claude <noreply@anthropic.com>
Review feedback:
- Drop all GMP function-name references from docstrings/comments (e.g.
  mpn_addmul_2 -> our add_mul_dword_same_len_in_place); none remain in the
  repo. Recorded as a guideline in AGENTS.md.
- Move the basecase-squaring kernel tests from tests/sqr_simple_kernel.rs
  into a #[cfg(test)] mod at the bottom of src/sqr/simple.rs (kernel tests
  belong with their implementation). Recorded as a guideline in AGENTS.md.

CI fixes:
- Modular exponentiation and Reduced::sqr sized their scratch with the
  multiplication budget, but squaring needs more scratch than multiplication
  in the Karatsuba band (the cross-term temporary), so the bump allocator
  was exhausted mid-recursion for moduli around 30-96 words (e.g. the
  Mersenne-prime test_pow). Add sqr_memory_requirement and size the pow path
  with max(mul, sqr) and Reduced::sqr with the squaring budget.
- Gate the Sign/debug_assert_zero imports in sqr/mod.rs behind the not-16-bit
  cfg so they aren't unused (build-breaking under -D warnings) on 16-bit Word
  targets where the NTT arm is compiled out.

Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
- Move sqr_memory_requirement from modular/mul.rs to sqr/mod.rs (it's a
  squaring concern). It now takes the operand length n (so sqr stays
  decoupled from the modular ring type); the modular callers pass
  ring.normalized_divisor.len().
- Make the mul::ntt helpers pub instead of pub(crate) (bit_len, coeff_count,
  do_crt, NttGeometry, select_params, the crt/pack/transform submodules,
  etc.). mul is a private module, so these cannot reach the public API —
  verified with cargo doc (the private mul module is not documented and none
  of the items leak into any public page).
- In sqr/ntt.rs, move the function-local `use crate::...` statements to the
  top of the file (and use the imported `Lane` short name in signatures).
- cargo fmt across the touched files (also picks up pre-existing rustfmt
  drift in sqr/ntt.rs, sqr/toom_3.rs, mul/ntt/mod.rs).

Co-Authored-By: Claude <noreply@anthropic.com>
@cmpute

cmpute commented Jun 16, 2026

Copy link
Copy Markdown
Owner Author

It turns out the previous benchmark compares the result against the master branch with additional improvements that are not present in the sqr branch.

@cmpute
cmpute merged commit 07802df into master Jun 16, 2026
13 checks passed
@cmpute
cmpute deleted the sqr branch June 16, 2026 06:51
CokieMiner pushed a commit to CokieMiner/dashu that referenced this pull request Jun 25, 2026
* Add specialized squaring for Karatsuba, Toom-3, and NTT

- Karatsuba: 3 recursive squarings instead of multiplications, simplified diff
- Toom-3: single polynomial evaluation, 5 squarings, V(-1) always positive
- NTT: single forward transform, pointwise square instead of multiply
- Separate _SQR env var thresholds from _MUL multiplication thresholds

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* Fix memory formulas for Karatsuba and Toom-3 squaring

- Karatsuba: 3n + 2·ceil_log2(n) (vs 2n for mul) — the diff_sq
  temp buffer adds ~n words of scratch pressure at peak.
- Toom-3: revert to 4n + 13·ceil_log2(n) (same as mul) as the
  earlier 3n + 10 derivation missed buffer overlap at recursion.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* Move NTT squaring code from mul/ntt to sqr/ntt

Move add_signed_sqr_same_len, add_signed_sqr_conv, and
process_prime_square to sqr/ntt.rs. Expose NTT internals
(pack, transform, NttGeometry, do_crt, bit_len, coeff_count,
add_shifted_to_prod) as pub(crate) for the new location.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* Inline pointwise_square into sqr/ntt.rs

Remove pointwise_square from mul/ntt/transform.rs and inline the
Montgomery squaring directly in process_prime_square. It was only
used by the NTT squaring path.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* Add ubig_sqr benchmark covering all algorithm thresholds

Tests 10^1 through 10^6 bit random operands, exercising simple,
Karatsuba, Toom-3, and NTT squaring paths.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* Document why Karatsuba squaring keeps its extra n-word temp

Explain at the memory_requirement_up_to site that the diff_sq temp is
kept deliberately: accumulating the cross term straight into the output
(as multiplication does) would force a full-schoolbook basecase, since
the symmetric in-place squaring needs a zeroed target. The temp trades
scratch memory for the efficient symmetric basecase (~25% faster in the
Toom-3 range, which recurses into Karatsuba).

Co-Authored-By: Claude <noreply@anthropic.com>

* Speed up basecase squaring with a two-word off-diagonal kernel

The schoolbook squaring basecase (simple::square) computed its off-diagonal
triangle one source limb at a time via add_mul_word_same_len_in_place,
while the multiplication basecase had already moved to a two-word
(mpn_addmul_2-style) kernel that halves accumulator memory traffic.

Bring squaring to parity: pair consecutive limbs (a[i], a[i+1]) and multiply
their shared suffix a[i+2..] by the double word (a[i] + a[i+1]*B) through
two interleaved mul-accumulate chains (mul_add_2carry), folding the lone
corner product a[i]*a[i+1] into the kernel's seed carry. Carries propagate
in place — each pair's high words land at b[n+i], b[n+i+1] and the next
pair's sweep re-reads b[n+i+1], so only a single pending carry escapes a
sweep (and it resolves to 0, since the triangle < B^(2n-1)).

simple::square is the basecase every algorithm recurses into, so this lifts
the whole stack: ubig_sqr/1e3 (schoolbook range) ~1.25x, ~1.15x through the
Karatsuba/Toom-3 bands (1e4-1e5), flat at 1e6 (NTT-bound). ubig_pow improves
likewise (~1.17x at 1e3-1e4) since exponentiation is squaring-dominated.

Adds a sqr()==mul() regression test across sizes 2..=128 plus adversarial
(all-ones, single high bit, alternating) inputs.

Co-Authored-By: Claude <noreply@anthropic.com>

* Address review feedback and fix modular pow/sqr scratch sizing

Review feedback:
- Drop all GMP function-name references from docstrings/comments (e.g.
  mpn_addmul_2 -> our add_mul_dword_same_len_in_place); none remain in the
  repo. Recorded as a guideline in AGENTS.md.
- Move the basecase-squaring kernel tests from tests/sqr_simple_kernel.rs
  into a #[cfg(test)] mod at the bottom of src/sqr/simple.rs (kernel tests
  belong with their implementation). Recorded as a guideline in AGENTS.md.

CI fixes:
- Modular exponentiation and Reduced::sqr sized their scratch with the
  multiplication budget, but squaring needs more scratch than multiplication
  in the Karatsuba band (the cross-term temporary), so the bump allocator
  was exhausted mid-recursion for moduli around 30-96 words (e.g. the
  Mersenne-prime test_pow). Add sqr_memory_requirement and size the pow path
  with max(mul, sqr) and Reduced::sqr with the squaring budget.
- Gate the Sign/debug_assert_zero imports in sqr/mod.rs behind the not-16-bit
  cfg so they aren't unused (build-breaking under -D warnings) on 16-bit Word
  targets where the NTT arm is compiled out.

Co-Authored-By: Claude <noreply@anthropic.com>

* cargo fmt: wrap modular::pow import list

Co-Authored-By: Claude <noreply@anthropic.com>

* Fix no-std build of sqr/simple.rs tests: import alloc::vec/Vec

Co-Authored-By: Claude <noreply@anthropic.com>

* Reorganize squaring/NTT modules per review

- Move sqr_memory_requirement from modular/mul.rs to sqr/mod.rs (it's a
  squaring concern). It now takes the operand length n (so sqr stays
  decoupled from the modular ring type); the modular callers pass
  ring.normalized_divisor.len().
- Make the mul::ntt helpers pub instead of pub(crate) (bit_len, coeff_count,
  do_crt, NttGeometry, select_params, the crt/pack/transform submodules,
  etc.). mul is a private module, so these cannot reach the public API —
  verified with cargo doc (the private mul module is not documented and none
  of the items leak into any public page).
- In sqr/ntt.rs, move the function-local `use crate::...` statements to the
  top of the file (and use the imported `Lane` short name in signatures).
- cargo fmt across the touched files (also picks up pre-existing rustfmt
  drift in sqr/ntt.rs, sqr/toom_3.rs, mul/ntt/mod.rs).

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: Jacob Zhong <jacob@rimbot.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant