Add specialized implementations for squaring

AGENTS.md

@@ -44,6 +44,8 @@ Note: always `--exclude dashu-python` when running workspace-wide commands, sinc
 - Doc comments use `# Examples` sections with runnable code — every public function on primitive number types must include a usage example
 - Modules are organized by operation (add, div, mul, cmp, convert, etc.)
 - Third-party trait implementations go in a `third_party/` module per crate, feature-gated
+- When borrowing an algorithm idea from GMP (or any other library), do **not** reference its function names in our docstrings or comments. Describe the algorithm in our own terms and use our own function names (e.g. write `add_mul_dword_same_len_in_place`, never `addmul_2` / `mpn_addmul_2`). External function names must not appear anywhere in the repo.
+- Tests for a specific algorithm/kernel belong in the same source file as the implementation, as a `#[cfg(test)] mod tests` block at the bottom — not in a separate integration test file under `tests/`. Reserve `tests/` for cross-cutting or public-API tests.
 
 ## Feature flags
 

integer/CHANGELOG.md

@@ -6,22 +6,30 @@
 - 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.
+- Specialized Karatsuba squaring: uses 3 recursive squarings instead of 3 multiplications, with simplified diff handling.
+- Specialized Toom-Cook-3 squaring: evaluates a single polynomial instead of two, 5 recursive squarings instead of multiplications.
+- Specialized NTT squaring: single forward transform instead of two, pointwise square instead of multiply.
+- Squaring thresholds can be overridden at runtime via `DASHU_THRESHOLD_SIMPLE_SQR`, `DASHU_THRESHOLD_KARATSUBA_SQR`, and `DASHU_THRESHOLD_NTT_SQR` environment variables (requires `tuning` feature).
 
 ### 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.
+- Basecase (schoolbook) multiplication now uses an dword mult inner kernel (two multiplier words per sweep over the accumulator, via the `add_mul_dword_same_len_in_place` and `sub_mul_dword_same_len_in_place` kernels), roughly halving accumulator memory traffic.
+- Basecase (schoolbook) squaring's off-diagonal phase now uses the same two-word kernel as multiplication, pairing consecutive limbs `(a[i], a[i+1])` against their shared suffix `a[i+2..]`. This halves the accumulator traffic of the basecase, speeding up squaring ~25% in the schoolbook range (≤30 words) and ~12-17% through the Karatsuba/Toom-3 bands that recurse into it; `ubig_pow` improves likewise since exponentiation is squaring-dominated.
 - 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.
 - 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).
+- Multiplication thresholds can be overridden at runtime via `DASHU_THRESHOLD_SIMPLE_MUL`, `DASHU_THRESHOLD_KARATSUBA_MUL`, and `DASHU_THRESHOLD_NTT_MUL` environment variables (requires `tuning` feature).
 
 ### Change
+- Multiplication threshold env vars renamed with `_MUL` suffix: `DASHU_THRESHOLD_SIMPLE_MUL`, `DASHU_THRESHOLD_KARATSUBA_MUL`, `DASHU_THRESHOLD_NTT_MUL` (was without suffix).
 - 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`.
 
 ### Fix
+- Modular exponentiation and `Reduced::sqr` under-allocated scratch memory for squaring (they sized it using the multiplication budget), which could exhaust the scratch allocator mid-recursion for moduli in the Karatsuba band (e.g. the Mersenne-prime `test_pow` case). The pow path now reserves `max(mul, sqr)` scratch and `Reduced::sqr` uses the dedicated squaring budget.
+- Unused imports (`Sign`, `debug_assert_zero`) in `sqr/mod.rs` on 16-bit Word targets, where the NTT arm is compiled out.
 - `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.

integer/benches/primitive.rs

@@ -175,6 +175,27 @@ fn ubig_mul_asymmetric(criterion: &mut Criterion) {
     group.finish();
 }
 
+fn ubig_sqr(criterion: &mut Criterion) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let mut group = criterion.benchmark_group("ubig_sqr");
+    group.plot_config(PlotConfiguration::default().summary_scale(AxisScale::Logarithmic));
+
+    for log_bits in 1..=6 {
+        if log_bits >= 5 {
+            group.sample_size(10);
+        }
+        let bits = 10usize.pow(log_bits);
+        let a = random_ubig(bits, &mut rng);
+        group.bench_with_input(
+            BenchmarkId::from_parameter(format!("1e{}", log_bits)),
+            &a,
+            |bencher, ta| bencher.iter(|| ta.sqr()),
+        );
+    }
+
+    group.finish();
+}
+
 criterion_group!(
     benches,
     ubig_add,
@@ -189,6 +210,7 @@ criterion_group!(
     ubig_pow_large_base,
     ubig_ilog_large,
     ubig_mul_asymmetric,
+    ubig_sqr,
 );
 
 criterion_main!(benches);

integer/src/modular/mul.rs

@@ -106,7 +106,7 @@ impl<'a> Reduced<'a> {
             }
             ReducedRepr::Large(raw, ring) => {
                 let mut result = raw.clone();
-                let memory_requirement = mul_memory_requirement(ring);
+                let memory_requirement = sqr::sqr_memory_requirement(ring.normalized_divisor.len());
                 let mut allocation = MemoryAllocation::new(memory_requirement);
                 sqr_in_place(ring, &mut result, &mut allocation.memory());
                 Reduced::from_large(result, ring)

integer/src/modular/pow.rs

@@ -142,7 +142,10 @@ mod large {
 
         let memory_requirement = memory::add_layout(
             memory::array_layout::<Word>(table_words),
-            mul_memory_requirement(ring),
+            // pow performs both multiplications and squarings, so size for the
+            // larger of the two (squaring needs more scratch than mul in the
+            // Karatsuba band).
+            memory::max_layout(mul_memory_requirement(ring), crate::sqr::sqr_memory_requirement(n)),
         );
         let mut allocation = MemoryAllocation::new(memory_requirement);
         let mut memory = allocation.memory();

integer/src/mul/mod.rs

@@ -36,15 +36,15 @@ 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
+/// When the `tuning` feature is active the user may set `DASHU_THRESHOLD_SIMPLE_MUL`,
+/// `DASHU_THRESHOLD_KARATSUBA_MUL` or `DASHU_THRESHOLD_NTT_MUL` 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(s) = std::env::var("DASHU_THRESHOLD_SIMPLE_MUL") {
                 if let Ok(v) = s.parse() {
                     return v;
                 }
@@ -56,7 +56,7 @@ mod threshold {
     pub fn karatsuba() -> usize {
         #[cfg(feature = "tuning")]
         {
-            if let Ok(s) = std::env::var("DASHU_THRESHOLD_KARATSUBA") {
+            if let Ok(s) = std::env::var("DASHU_THRESHOLD_KARATSUBA_MUL") {
                 if let Ok(v) = s.parse() {
                     return v;
                 }
@@ -68,7 +68,7 @@ mod threshold {
     pub fn ntt() -> usize {
         #[cfg(feature = "tuning")]
         {
-            if let Ok(s) = std::env::var("DASHU_THRESHOLD_NTT") {
+            if let Ok(s) = std::env::var("DASHU_THRESHOLD_NTT_MUL") {
                 if let Ok(v) = s.parse() {
                     return v;
                 }
@@ -416,7 +416,7 @@ mod threshold_tests {
     ///
     /// 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 \
+    /// DASHU_THRESHOLD_NTT_MUL=99999999 cargo test -p dashu-int --features tuning --release \
     ///   -- mul::threshold_tests::crossover_ntt --ignored --nocapture
     /// ```
     ///
@@ -467,7 +467,7 @@ mod threshold_tests {
             let warmup = 2;
             let iters = 5;
 
-            // toom-3 (may use NTT internally depending on DASHU_THRESHOLD_NTT)
+            // toom-3 (may use NTT internally depending on DASHU_THRESHOLD_NTT_MUL)
             let t_toom = {
                 let mut best = f64::MAX;
                 for _ in 0..warmup {

integer/src/mul/ntt/mod.rs

@@ -12,9 +12,9 @@ use crate::{
 use alloc::alloc::Layout;
 use core::mem;
 
-pub(crate) mod crt;
-mod pack;
-mod transform;
+pub mod crt;
+pub mod pack;
+pub mod transform;
 
 use crate::arch::ntt::{
     B_PACK_CANDIDATES, B_PACK_MIN, CRT_INV_IJ, K, MAX_LOG_N, MODULI, OMEGA_MAX, P0, P1, P2,
@@ -68,7 +68,7 @@ pub fn select_params(la_words: usize, lb_words: usize) -> (u32, usize, usize) {
 }
 
 /// Estimate bit length from a word slice (excludes leading zeros).
-fn bit_len(words: &[Word]) -> u64 {
+pub 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 {
@@ -80,7 +80,7 @@ fn bit_len(words: &[Word]) -> u64 {
 }
 
 /// Count number of coefficients needed for a given bit length.
-fn coeff_count(bit_len: u64, b_pack: u32) -> usize {
+pub fn coeff_count(bit_len: u64, b_pack: u32) -> usize {
     ((bit_len + b_pack as u64 - 1) / b_pack as u64) as usize
 }
 
@@ -295,7 +295,7 @@ fn run_ntt_pipeline(
         }
     }
 
-    do_crt::<crate::arch::word::TripleWord>(prod, residues, &ctx, &MODULI, &CRT_INV_IJ);
+    do_crt::<crate::arch::word::TripleWord>(prod, residues, &ctx.geom, &MODULI, &CRT_INV_IJ);
 
     match sign {
         Positive => add::add_signed_in_place(&mut c_out[..out_words], Positive, &prod[..out_words]),
@@ -346,14 +346,14 @@ fn add_signed_mul_conv(
 }
 
 /// CRT + accumulate, generic over the accumulator type.
-fn do_crt<A: CrtAccum>(
+pub fn do_crt<A: CrtAccum>(
     prod: &mut [Word],
     residues: &[A::Lane],
-    ctx: &TransformCtx<'_>,
+    geom: &NttGeometry,
     primes: &[A::Lane; K],
     crt_inv: &[[A::Lane; K]; K],
 ) {
-    let g = &ctx.geom;
+    let g = geom;
     for k in 0..g.output_coeffs {
         let mut coeff_residues = [A::Lane::default(); 3];
         #[allow(clippy::needless_range_loop)]
@@ -368,11 +368,11 @@ fn do_crt<A: CrtAccum>(
 }
 
 /// Geometry constants for an NTT pipeline invocation.
-struct NttGeometry {
-    nn: usize,
-    b_pack: u32,
-    k_eff: usize,
-    output_coeffs: usize,
+pub struct NttGeometry {
+    pub nn: usize,
+    pub b_pack: u32,
+    pub k_eff: usize,
+    pub output_coeffs: usize,
 }
 
 /// Scratch buffers and geometry for the per-prime NTT pipeline.
@@ -492,7 +492,7 @@ fn process_prime<R: Reducer<crate::arch::ntt::Lane>>(
 
 /// 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) {
+pub 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_bits = Word::BITS;
     let start_idx = (shift_bits / word_bits) as usize;

integer/src/mul/simple.rs

@@ -85,7 +85,7 @@ fn add_signed_mul_chunk(
 /// It consumes two multiplier words per sweep over `rhs`, so the accumulator
 /// word `words[k]` is loaded and stored once per two multiplier words instead
 /// of once per word. This halves the memory traffic on `words` and exposes two
-/// independent multiply chains, mirroring GMP's `mpn_addmul_2`.
+/// independent multiply chains.
 ///
 /// Only `words[..n]` is modified. The two extra high words of the product
 /// (the carries out of columns `n` and `n + 1`) are returned as
@@ -181,7 +181,7 @@ fn sub_mul_chunk(c: &mut [Word], a: &[Word], b: &[Word]) -> bool {
     let mut borrow_out = false;
     let mut i = 0;
 
-    // Consume the multiplier two words at a time via the submul_2 kernel.
+    // Consume the multiplier two words at a time via the sub_mul_dword_same_len_in_place kernel.
     let mut pairs = b.chunks_exact(2);
     for pair in &mut pairs {
         let (carry_lo, carry_hi) =

integer/src/sqr/karatsuba.rs

@@ -0,0 +1,90 @@
+//! Karatsuba squaring algorithm.
+
+use crate::{
+    add,
+    arch::word::{SignedWord, Word},
+    helper_macros::debug_assert_zero,
+    math,
+    memory::{self, Memory},
+    sqr,
+    Sign::*,
+};
+use alloc::alloc::Layout;
+
+// Same constraint as Karatsuba multiplication: 3 * floor((n+1)/2) <= 2n for n >= 3.
+/// Minimum supported length.
+pub const MIN_LEN: usize = 3;
+
+/// Temporary memory required for squaring.
+///
+/// n bounds the operand length in words.
+pub fn memory_requirement_up_to(n: usize) -> Layout {
+    // 3n + 2 ceil_log2 n  (vs 2n + 2 ceil_log2 n for mul).
+    //
+    // The extra n is the diff_sq temp buffer (2·mid words) that holds the
+    // cross term (a_lo − a_hi)² before it is subtracted from the output.
+    //
+    // Multiplication avoids this temp by accumulating its difference product
+    // straight into the output (mul::add_signed_mul_same_len), but squaring
+    // cannot do so cheaply. The symmetric basecase (simple::square) adds the
+    // off-diagonal products once and finishes with an O(n) in-place doubling,
+    // which requires a *zeroed* target. Accumulating into a non-zero output
+    // would force a full-schoolbook basecase (~2× the multiply work), which
+    // measures ~25% slower across the Toom-3 range where Karatsuba is used
+    // recursively. The temp is therefore kept deliberately: it is the price
+    // of the efficient symmetric (in-place) squaring basecase.
+    let num_words = 3 * n + 2 * (math::ceil_log2(n) as usize);
+    memory::array_layout::<Word>(num_words)
+}
+
+/// b = a², b must be filled with zeros. n >= MIN_LEN.
+///
+/// a² = a_lo² + (a_lo² + a_hi² − (a_lo−a_hi)²)·B^mid + a_hi²·B^(2·mid)
+pub fn square(b: &mut [Word], a: &[Word], memory: &mut Memory) {
+    let n = a.len();
+    debug_assert!(n >= MIN_LEN && b.len() == 2 * n);
+
+    let mid = (n + 1) / 2;
+    let (a_lo, a_hi) = a.split_at(mid);
+
+    let mut carry: SignedWord = 0;
+    let mut carry_c0: SignedWord = 0; // at 2*mid
+    let mut carry_c1: SignedWord = 0; // at 3*mid
+
+    {
+        // P0 = sqr(a_lo)
+        let (p0, mut mem) = memory.allocate_slice_fill::<Word>(2 * mid, 0);
+        sqr::sqr(p0, a_lo, &mut mem);
+        carry_c0 += add::add_signed_same_len_in_place(&mut b[..2 * mid], Positive, p0);
+        carry_c1 += add::add_signed_same_len_in_place(&mut b[mid..3 * mid], Positive, p0);
+    }
+    {
+        // P2 = sqr(a_hi)
+        let p2_len = 2 * (n - mid);
+        let (p2, mut mem) = memory.allocate_slice_fill::<Word>(p2_len, 0);
+        sqr::sqr(p2, a_hi, &mut mem);
+        carry += add::add_signed_same_len_in_place(&mut b[2 * mid..], Positive, p2);
+        carry_c1 += add::add_signed_in_place(&mut b[mid..3 * mid], Positive, p2);
+    }
+    {
+        // diff_sq = (|a_lo − a_hi|)²   (always non-negative)
+        let (diff, mut mem) = memory.allocate_slice_copy(a_lo);
+        let diff_sign = add::sub_in_place_with_sign(diff, a_hi);
+        if diff_sign == Negative {
+            // |a_lo − a_hi| = a_hi − a_lo
+            diff[..(n - mid)].copy_from_slice(a_hi);
+            diff[(n - mid)..].fill(0);
+            debug_assert_zero!(add::sub_in_place(diff, a_lo));
+        }
+        let (diff_sq, mut mem) = mem.allocate_slice_fill::<Word>(2 * mid, 0);
+        sqr::sqr(diff_sq, diff, &mut mem);
+        // c[mid..3*mid] -= diff_sq  (always subtract, diff_sq is non-negative)
+        carry_c1 += add::add_signed_same_len_in_place(&mut b[mid..3 * mid], Negative, diff_sq);
+    }
+
+    // Propagate carries.
+    carry_c1 += add::add_signed_word_in_place(&mut b[2 * mid..3 * mid], carry_c0);
+    carry += add::add_signed_word_in_place(&mut b[3 * mid..], carry_c1);
+
+    debug_assert!(carry.abs() <= 1);
+}