refactor(wordram): consolidate square-root kernels and assess power specializations

Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandExec.lean

@@ -78,7 +78,7 @@ theorem bandTime_le (L : ℕ) : bandTime L + 3 ≤ 512 * (L + 1) * 10 ^ L := by
 /-- The sum of the previously proved native component charges fits the generic bound. -/
 theorem execution_time_le (L m N : ℕ) :
     (16 * L + 49) + ((7 * inputEnd m N + 7) +
-      (((6 * L + 16) * 2 ^ L + 14) + (((6 * L + 17) * 2 ^ L + 4) +
+      (((6 * L + 16) * 2 ^ L + 16) + (((6 * L + 17) * 2 ^ L + 4) +
         ((numBands L m N * (bandTime L + 3) + 3) +
           ((numBands L m N * (bandTime L + 3) + 3) + (7 * inputEnd m N + 7)))))) ≤
       timeBound L m N := by

Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandFlat.lean

@@ -60,7 +60,8 @@ theorem program_length_le : (program.program catalogue).length ≤ 16384 := by
   simp only [Code.program, List.length_append, Code.length_assemble,
     List.length_cons, List.length_nil] at hl hr ⊢
   simp only [program, header, headerLoads, headerPowers, headerLayout, powerCode,
-    saveInput, countCode, countMasks, countBody, populationCode, rootCode, rootBody,
+    saveInput, countCode, countMasks, countBody, populationCode, rootCode,
+    NatSqrt.program, NatSqrt.body,
     bandCount, tablesCode, tableBody, tableWrite, encodeBands, bandBody, prepareBand,
     prepareHead, prepareBody, splitCode, splitHead, splitBody, readCell, collectBand,
     collectHead, restoreInput, Code.size, Code.size_rename, Power.program, Power.body,

Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandParameters.lean

@@ -86,7 +86,7 @@ theorem local_fits {w L m N : ℕ} (hfit : Fits w (wordBudget L m N)) :
 /-- Compute the root and band counts with the actual native loops. -/
 theorem exec_countCode {w L m N : ℕ} (x : Data 96)
     (hx : HeaderPost L m N x) (hfit : Fits w (wordBudget L m N)) :
-    ∃ y, ExecWithin w countCode x y ((6 * L + 16) * 2 ^ L + 14) ∧
+    ∃ y, ExecWithin w countCode x y ((6 * L + 16) * 2 ^ L + 16) ∧
       BandParameters L m N y ∧ y.mem = x.mem := by
   have hl := local_fits hfit
   have hK : K L m ≤ 2 ^ L := Nat.choose_le_two_pow _ _
@@ -101,17 +101,16 @@ theorem exec_countCode {w L m N : ℕ} (x : Data 96)
       simp only [countMasks, countBody, populationCode, Code.Preserves, Op.dest, Op.bin]
       decide
     exact hp i hi)
-  obtain ⟨r, hr, hroot, hsquare, hrmem, _⟩ := exec_root (K L m) c hcount hcp.fixed.one
+  obtain ⟨r, hr, hroot, hsquare, hrmem, hrframe⟩ := exec_root (K L m) c hcount hcp.fixed.one
     (hl.of_le (by
       nlinarith only [hK, Nat.zero_le L, Nat.zero_le N,
         Nat.zero_le (N * D m), Nat.zero_le (2 ^ L), Nat.zero_le (10 ^ L)]))
   have hrp : HeaderPost L m N r := headerPost_reframe hcp (by
     intro i hi
-    apply hr.regs_eq
-    have hp : ∀ i ∈ headerRegisters, rootCode.Preserves i := by
-      simp only [rootCode, rootBody, Code.Preserves, Op.dest, Op.bin]
-      decide
-    exact hp i hi)
+    have hdisj : ∀ i ∈ headerRegisters,
+        i ≠ root ∧ i ≠ square ∧ i ≠ temporary ∧ i ≠ limit := by decide
+    obtain ⟨h1, h2, h3, h4⟩ := hdisj i hi
+    exact hrframe i h1 h2 h3 h4)
   have hB : bandSize L m ≤ 7 ^ L := bandSize_le_seven_pow L m
   have h7 : 7 ^ L ≤ 10 ^ L := Nat.pow_le_pow_left (by decide) _
   have hbfit : Fits w (bandSize L m) := hl.of_le (by
@@ -132,16 +131,15 @@ theorem exec_countCode {w L m N : ℕ} (x : Data 96)
     ram_eval [bandCount, y, hroot0, hrp.outerSize, hrp.fixed.rows, hrp.fixed.one,
       bandSize, numBands, hbfit, hnfit, hnsubfit, hdivfit,
       temporary, root, outerSize, bands, rows, one, Fin.mk.injEq]
-  have he : ExecWithin w countCode x y ((6 * L + 16) * 2 ^ L + 14) := by
+  have he : ExecWithin w countCode x y ((6 * L + 16) * 2 ^ L + 16) := by
     apply (hc.seq (hr.seq hy)).mono
     have hsqrt : Nat.sqrt (K L m) ≤ 2 ^ L := (Nat.sqrt_le_self _).trans hK
     nlinarith
-  have hyp : HeaderPost L m N y := headerPost_reframe hx (by
+  have hyp : HeaderPost L m N y := headerPost_reframe hrp (by
     intro i hi
-    apply he.regs_eq
-    have hp : ∀ i ∈ headerRegisters, countCode.Preserves i := by
-      simp only [countCode, countMasks, countBody, populationCode, rootCode, rootBody,
-        bandCount, Code.Preserves, Op.dest, Op.bin]
+    apply hy.regs_eq
+    have hp : ∀ i ∈ headerRegisters, bandCount.Preserves i := by
+      simp only [bandCount, Code.Preserves, Op.dest, Op.bin]
       decide
     exact hp i hi)
   refine ⟨y, he, ⟨hyp, ?_, ?_, ?_⟩, hrmem.trans hcmem⟩

Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandProgram.lean

@@ -9,6 +9,7 @@ module
 
 public import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Flat
 public import Complexity.Algorithms.MatrixMultiplication.Schonhage.ThinProduct
+public import Complexity.Algorithms.WordRAM.NatSqrt
 public import Complexity.Computability.WordRAM.Power
 
 /-!
@@ -198,16 +199,13 @@ def countBody : Code .standard 96 :=
 def countMasks : Code .standard 96 :=
   .seq (ram[ mask ← #0; count ← #0 ]) (.loop (.lt mask masks) countBody)
 
-/-- One explicitly squared trial in the integer square-root search. -/
-def rootBody : Code .standard 96 := ram[
-  root ← root + one; temporary ← root + one; square ← temporary * temporary ]
+/-- Embed the six-register integer square-root kernel into the band-preparation register bank. -/
+def rootRegisters : Fin 6 → Fin 96 := ![count, root, temporary, square, one, limit]
 
-/-- Compute the integer square root by squared comparisons. No square-root instruction
-is added to the native basis. -/
+/-- Compute the integer square root and its square by renaming `NatSqrt.program` and charging
+one final multiplication. No square-root instruction is added to the native basis. -/
 def rootCode : Code .standard 96 :=
-  .seq (ram[
-  root ← #0; square ← #1; limit ← count + one
-  ]) (.seq (.loop (.lt square limit) rootBody) (ram[ square ← root * root ]))
+  .seq (NatSqrt.program.rename rootRegisters) (ram[ square ← root * root ])
 
 /-- Compute the ceiling quotient giving the number of bands. -/
 def bandCount : Code .standard 96 := ram[

Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandRoot.lean

@@ -7,15 +7,16 @@ Authors: Erik Demaine
 module
 -- upstream: candidate
 
-public import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.BandProgram
+public import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.BandCount
 
 /-!
-# Native integer square root for band parameters
+# Native integer square root for the band block size
 
 ## Main results
 
-* `BandPreparation.exec_root`: successive squared comparisons compute the ordinary natural
-  square root within `5 * sqrt k + 6` instructions, with every trial square proved to fit.
+* `BandPreparation.exec_root`: reuse `NatSqrt.exec_rename` and one final multiplication to
+  compute the ordinary natural square root and its square within `5 * Nat.sqrt k + 8`
+  native instructions.
 -/
 
 @[expose] public section
@@ -24,78 +25,49 @@ namespace Cslib.Algorithms.WordRAM.MatrixMultiplication.Schonhage.BandPreparatio
 
 open Cslib.WordRAM
 
-/-- The root loop uses only increment, multiplication and unsigned comparison. The bound on
-trial squares is derived from the original count, including the first failing trial. -/
+/-- The six-register square-root assignment is injective. -/
+theorem rootRegisters_injective : Function.Injective rootRegisters := by decide
+
+private theorem rootRegisters_preserves (i : Fin 96)
+    (hr : i ≠ root) (hs : i ≠ square) (ht : i ≠ temporary) (hl : i ≠ limit) (hone : i ≠ one) :
+    (NatSqrt.program.rename rootRegisters).Preserves i := by
+  simp [NatSqrt.program, NatSqrt.body, Code.rename, Code.Preserves, Op.rename, Op.dest, Op.bin,
+    rootRegisters, NatSqrt.arg, NatSqrt.root, NatSqrt.next, NatSqrt.square, NatSqrt.one,
+    NatSqrt.bound, Ne.symm hr, Ne.symm hs, Ne.symm ht, Ne.symm hl, Ne.symm hone]
+
+/-- Compute the floor square root and its square by renaming `NatSqrt.program` and charging one
+final multiplication. The only intermediate bound is `4*k+4`, which covers `(sqrt k + 1)²`. -/
 theorem exec_root {w : ℕ} (k : ℕ) (x : Data 96)
     (hk : x.regs count = k) (ho : x.regs one = 1) (hfit : Fits w (4 * k + 4)) :
-    ∃ y, ExecWithin w rootCode x y (5 * Nat.sqrt k + 6) ∧
-      y.regs root = Nat.sqrt k ∧ y.regs square = Nat.sqrt k * Nat.sqrt k ∧
-      y.mem = x.mem ∧
+    ∃ y, ExecWithin w rootCode x y (5 * Nat.sqrt k + 8) ∧
+      y.regs root = Nat.sqrt k ∧ y.regs square = Nat.sqrt k * Nat.sqrt k ∧ y.mem = x.mem ∧
       (∀ i, i ≠ root → i ≠ square → i ≠ temporary → i ≠ limit → y.regs i = x.regs i) := by
-  have h0 : Fits w 0 := hfit.of_le (by omega)
-  have h1 : Fits w 1 := hfit.of_le (by omega)
-  have hnext : Fits w (k + 1) := hfit.of_le (by omega)
-  let start := ((x.set root 0).set square 1).set limit (k + 1)
-  have hs : ExecWithin w (ram[ root ← #0; square ← #1; limit ← count + one ])
-      x start 3 := by
-    ram_eval [start, hk, ho, h0, h1, hnext, root, square, limit, count, one, Fin.mk.injEq]
-  let I (y : Data 96) : Prop :=
-    y.mem = x.mem ∧ y.regs one = 1 ∧ y.regs limit = k + 1 ∧
-      y.regs root ≤ Nat.sqrt k ∧
-      y.regs square = (y.regs root + 1) * (y.regs root + 1)
-  have hi : I start := by simp [I, start, ho, root, square, limit, one]
-  have hb (y : Data 96) (hy : I y) (hc : (Cond.lt square limit).Holds y) :
-      ∃ z, ExecWithin w rootBody y z 3 ∧ I z ∧
-        Nat.sqrt k - z.regs root < Nat.sqrt k - y.regs root := by
-    obtain ⟨hm, h1, hlimit, hr, hsq⟩ := hy
-    have htrial : (y.regs root + 1) * (y.regs root + 1) ≤ k := by
-      simpa only [Cond.holds_lt, hlimit, hsq, Nat.lt_add_one_iff] using hc
-    have hrnext : y.regs root + 1 ≤ Nat.sqrt k := Nat.le_sqrt.mpr htrial
-    have hlarge : (y.regs root + 2) * (y.regs root + 2) ≤ 4 * k + 4 := by
-      nlinarith only [htrial]
-    have hrfit : Fits w (y.regs root + 1) :=
-      hfit.of_le (by have := Nat.sqrt_le_self k; omega)
-    have htfit : Fits w (y.regs root + 2) :=
-      hfit.of_le (by have := Nat.sqrt_le_self k; omega)
-    have hsfit : Fits w ((y.regs root + 2) * (y.regs root + 2)) := hfit.of_le hlarge
-    let z := ((y.set root (y.regs root + 1)).set temporary (y.regs root + 2)).set
-      square ((y.regs root + 2) * (y.regs root + 2))
-    have he : ExecWithin w rootBody y z 3 := by
-      ram_eval [rootBody, z, h1, hrfit, htfit, hsfit,
-        root, temporary, square, one, Fin.mk.injEq, Nat.add_assoc]
-    refine ⟨z, he, ?_, ?_⟩
-    · simpa +decide [I, z, root, temporary, square, one, limit, Function.update_apply,
-        Nat.add_assoc] using
-        (show y.mem = x.mem ∧ y.regs one = 1 ∧ y.regs limit = k + 1 ∧
-          y.regs root + 1 ≤ Nat.sqrt k ∧
-          (y.regs root + 2) * (y.regs root + 2) =
-            (y.regs root + 1 + 1) * (y.regs root + 1 + 1) from
-          ⟨hm, h1, hlimit, hrnext, by ring⟩)
-    · have hz : z.regs root = y.regs root + 1 := by simp [z, root, square, temporary]
-      rw [hz]
-      omega
-  obtain ⟨y, he, hy, hc⟩ := ExecWithin.loop_invariant I
-    (fun y => Nat.sqrt k - y.regs root) hb hi
-  have hr : y.regs root = Nat.sqrt k := by
-    have hstop : k < (y.regs root + 1) * (y.regs root + 1) := by
-      simp only [Cond.holds_lt, hy.2.2.1, hy.2.2.2.2] at hc
-      omega
-    exact Nat.eq_sqrt.mpr ⟨(Nat.le_sqrt.mp hy.2.2.2.1), hstop⟩
-  let z := y.set square (Nat.sqrt k * Nat.sqrt k)
-  have ht : ExecWithin w (ram[ square ← root * root ]) y z 1 := by
-    ram_eval [z, hr, hfit.of_le (Nat.sqrt_le k |>.trans (by omega)), square, root]
-  have hwhole : ExecWithin w rootCode x z (5 * Nat.sqrt k + 6) := by
-    apply (hs.seq (he.seq ht)).mono
-    have hsroot : start.regs root = 0 := by simp [start, root, square, limit]
-    rw [hsroot, Nat.sub_zero]
-    omega
-  refine ⟨z, hwhole, ?_, ?_, hy.1, ?_⟩
-  · simpa [z, square, root] using hr
-  · simp [z]
+  obtain ⟨z, hz, hzroot, _, hzmem, hzone⟩ :=
+    NatSqrt.exec_rename rootRegisters rootRegisters_injective x
+      (by simpa only [rootRegisters, NatSqrt.arg, Matrix.cons_val_zero, hk] using hfit)
+  simp only [rootRegisters, NatSqrt.arg, NatSqrt.root, NatSqrt.one,
+    Matrix.cons_val_zero, Matrix.cons_val_one, Matrix.cons_val] at hz hzroot hzone
+  rw [hk] at hz hzroot
+  have hfinal_fit : Fits w (z.regs root * z.regs root) := by
+    rw [hzroot]
+    have hle := Nat.sqrt_le k
+    exact hfit.of_le (by omega)
+  have hfinal : ExecWithin w (ram[ square ← root * root ]) z
+      (z.set square (Nat.sqrt k * Nat.sqrt k)) 1 := by
+    have hmul : Fits w (Nat.sqrt k * Nat.sqrt k) := by simpa only [hzroot] using hfinal_fit
+    ram_eval [hzroot, hmul]
+  have hwhole : ExecWithin w rootCode x (z.set square (Nat.sqrt k * Nat.sqrt k))
+      (5 * Nat.sqrt k + 8) :=
+    (hz.seq hfinal).mono (by omega)
+  refine ⟨z.set square (Nat.sqrt k * Nat.sqrt k), hwhole, ?_, ?_, hzmem, ?_⟩
+  · simpa +decide [Data.set, square, root] using hzroot
+  · simp [Data.set]
   · intro i hr hs ht hl
-    apply hwhole.regs_eq
-    have hf : ∀ i : Fin 96, i ≠ root → i ≠ square → i ≠ temporary → i ≠ limit →
-        rootCode.Preserves i := by decide
-    exact hf i hr hs ht hl
+    have his : (z.set square (Nat.sqrt k * Nat.sqrt k)).regs i = z.regs i := by
+      simp [Data.set, hs]
+    rw [his]
+    by_cases hone : i = one
+    · subst hone; exact hzone.trans ho.symm
+    · exact hz.regs_eq (rootRegisters_preserves i hr hs ht hl hone)
 
 end Cslib.Algorithms.WordRAM.MatrixMultiplication.Schonhage.BandPreparation

Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Encoding.lean

@@ -213,7 +213,10 @@ def setupHead : Code .standard 32 := ram[
   address ← #1; depth ← *address;
   inputSize ← #1; capacity ← #1; stage ← #0 ]
 
-/-- One multiplication by each radix and one counter increment. -/
+/-- One multiplication by each radix and one counter increment. Computing `7^L` and `10^L`
+simultaneously in a single `stage < depth` loop charges `5 * L + 18` setup instructions and
+shares one counter register, whereas two sequential `WordRAM.Power.program` loops would
+duplicate the loop control and counter reset. -/
 def powerBody : Code .standard 32 := ram[
   inputSize ← inputSize * seven; capacity ← capacity * ten; stage ← stage + one ]
 

Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Product/Parameters.lean

@@ -45,15 +45,15 @@ public theorem exec_parameters {w L m N Q : ℕ} (hm : m ≤ L) (x : Data 160)
     (hL : x.mem 1 = L) (hmemory : x.mem 2 = m) (hN : x.mem 3 = N)
     (hfit : Fits w (wordBudget L m N Q)) :
     ∃ y, ExecWithin w parameters x y
-        ((16 * L + 49) + ((6 * L + 16) * 2 ^ L + 14)) ∧
+        ((16 * L + 49) + ((6 * L + 16) * 2 ^ L + 16)) ∧
       BandPreparation.BandParameters L m N (y.comap bank) ∧ y.mem = x.mem ∧
       (∀ i, 96 ≤ i.val → y.regs i = x.regs i) := by
   obtain ⟨h, hh, hp, hm⟩ := BandPreparation.exec_header hm (x.comap bank) hL hmemory hN
     (producer_fits hfit).1
   obtain ⟨c, hc, cp, cm⟩ := BandPreparation.exec_countCode h hp (producer_fits hfit).1
   have he := (hh.seq hc).rename bank_injective
   change ExecWithin w parameters x (c.extend bank x)
-    ((16 * L + 49) + ((6 * L + 16) * 2 ^ L + 14)) at he
+    ((16 * L + 49) + ((6 * L + 16) * 2 ^ L + 16)) at he
   refine ⟨c.extend bank x, he, ?_, cm.trans hm, ?_⟩
   · simpa only [Data.comap_extend bank_injective] using cp
   · intro i hi

Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Product/Routing.lean

@@ -30,7 +30,7 @@ open Cslib.Algorithms.MatrixMultiplication.Schonhage
 
 /-- The sum of the actual preparation phases, before matrix adaptation and encoding. -/
 @[expose] public def routingTimeBound (L m N Q : ℕ) : ℕ :=
-  ((16 * L + 49) + ((6 * L + 16) * 2 ^ L + 14)) + 50 +
+  ((16 * L + 49) + ((6 * L + 16) * 2 ^ L + 16)) + 50 +
     (7 * inputEnd m N Q + 10) + (14 * Q + 10) + Wanted.timeBound L m N Q +
       (7 * numBands L m N ^ 2 + 14 * Q + 37)