From 69207d42981d5e51fe127e51b7b190d10280c596 Mon Sep 17 00:00:00 2001 From: Samuel Schlesinger Date: Fri, 30 Oct 2026 22:29:45 +0000 Subject: [PATCH] refactor(wordram): consolidate square-root and power kernels (#1553) --- .../Schonhage/BandExec.lean | 2 +- .../Schonhage/BandFlat.lean | 3 +- .../Schonhage/BandParameters.lean | 24 ++-- .../Schonhage/BandProgram.lean | 14 +-- .../Schonhage/BandRoot.lean | 118 +++++++----------- .../Schonhage/Encoding.lean | 5 +- .../Schonhage/Product/Parameters.lean | 4 +- .../Schonhage/Product/Routing.lean | 2 +- .../Schonhage/Pruned/Powers.lean | 108 +++++++--------- .../Schonhage/Pruned/Program.lean | 15 ++- .../Schonhage/Wanted/Cost.lean | 2 +- .../Schonhage/Wanted/Setup.lean | 2 +- .../WordRAM/ThreeSum/ChanHe/Parameters.lean | 5 +- ComplexityTests/SchonhageBandPreparation.lean | 34 +++++ ComplexityTests/SchonhagePruned.lean | 23 ++++ UPSTREAM.md | 23 ++-- 16 files changed, 204 insertions(+), 180 deletions(-) diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandExec.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandExec.lean index c299842d7..d715a3b79 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandExec.lean +++ b/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 diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandFlat.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandFlat.lean index 1b356a8f0..c3e07ee30 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandFlat.lean +++ b/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, diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandParameters.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandParameters.lean index 519cb3b84..916f7013a 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandParameters.lean +++ b/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⟩ diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandProgram.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandProgram.lean index 127663f3c..01a742253 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandProgram.lean +++ b/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[ diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandRoot.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandRoot.lean index 48ca52819..077ab5810 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/BandRoot.lean +++ b/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 diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Encoding.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Encoding.lean index d9cf2eb32..6c6dd0e0f 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Encoding.lean +++ b/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 ] diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Product/Parameters.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Product/Parameters.lean index 4aabeed0a..29d0e7024 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Product/Parameters.lean +++ b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Product/Parameters.lean @@ -45,7 +45,7 @@ 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 @@ -53,7 +53,7 @@ public theorem exec_parameters {w L m N Q : ℕ} (hm : m ≤ L) (x : Data 160) 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 diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Product/Routing.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Product/Routing.lean index 6d5f61239..0954d8f38 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Product/Routing.lean +++ b/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) diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Pruned/Powers.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Pruned/Powers.lean index 93023bad4..4fbbf5ea4 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Pruned/Powers.lean +++ b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Pruned/Powers.lean @@ -14,8 +14,8 @@ public import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Prune ## Main results -* `PrunedQuery.exec_powers`: compute `10^L` using the actual native multiplication loop, - preserving all caller parameters and memory. +* `PrunedQuery.exec_powers`: compute `10^L` by renaming `Power.program` after the five + retained constant initializations, preserving all caller parameters and memory. -/ @[expose] public section @@ -24,7 +24,7 @@ namespace Cslib.Algorithms.WordRAM.MatrixMultiplication.Schonhage.PrunedQuery open Cslib.WordRAM -/-- The scalar registers initialized or advanced by the power loop. -/ +/-- The scalar registers initialized or advanced by the power routine. -/ def PowersLocal (i : Fin registers) : Prop := i = nil ∨ i = one ∨ i = two ∨ i = ten ∨ i = nine ∨ i = frameWidth ∨ i = span ∨ i = index @@ -32,11 +32,15 @@ instance (i : Fin registers) : Decidable (PowersLocal i) := inferInstanceAs (Decidable (i = nil ∨ i = one ∨ i = two ∨ i = ten ∨ i = nine ∨ i = frameWidth ∨ i = span ∨ i = index)) +/-- The five-register power assignment is injective. -/ +theorem powerRegisters_injective : Function.Injective powerRegisters := by decide + private theorem powers_preserves (i : Fin registers) : ¬PowersLocal i → powers.Preserves i := by decide +kernel +revert -/-- The power loop charges its eight initializations, both body instructions, loop tests and -jumps. The only width premise is a bound on ordinary input-derived natural numbers. -/ +/-- The power routine charges its five constant initializations and the renamed `4*L+5` +`Power.program` execution. The only width premise is a bound on ordinary input-derived +natural numbers. -/ theorem exec_powers {w L Q : ℕ} (x : Data registers) (hL : x.regs inputDepth = L) (hbound : Fits w (addressBound L Q)) : ∃ y, ExecWithin w powers x y (4 * L + 10) ∧ y.mem = x.mem ∧ @@ -53,65 +57,47 @@ theorem exec_powers {w L Q : ℕ} (x : Data registers) (hL : x.regs inputDepth = nlinarith have hbase := hbound.of_le hscalar have h0 : Fits w 0 := fits_zero w - have h1 : Fits w 1 := hbase.of_le (by omega) have h2 : Fits w 2 := hbase.of_le (by omega) have h4 : Fits w 4 := hbase.of_le (by omega) have h9 : Fits w 9 := hbase.of_le (by omega) have h10 : Fits w 10 := hbase.of_le (by omega) - have hpow (j : ℕ) (hj : j ≤ L) : Fits w (10 ^ j) := by - have hle := Nat.pow_le_pow_right (by decide : 1 ≤ (10 : ℕ)) hj - exact hbase.of_le (by omega) - let I (y : Data registers) : Prop := - y.mem = x.mem ∧ y.regs inputDepth = L ∧ y.regs span = 10 ^ y.regs index ∧ - y.regs index ≤ L ∧ y.regs one = 1 ∧ y.regs ten = 10 - have hstep (y : Data registers) (hy : I y) (hc : (Cond.lt index inputDepth).Holds y) : - ∃ z, ExecWithin w powerStep y z 2 ∧ I z ∧ - L - z.regs index < L - y.regs index := by - obtain ⟨hm, hd, hp, hj, ho, ht⟩ := hy - have hlt : y.regs index < L := by simpa only [Cond.holds_lt, hd] using hc - have hmul : Fits w (y.regs span * 10) := by - rw [hp, ← pow_succ] - exact hpow _ (by omega) - have hinc : Fits w (y.regs index + 1) := hbase.of_le (by omega) - let z := (y.set span (y.regs span * 10)).set index (y.regs index + 1) - have he : ExecWithin w powerStep y z 2 := by - ram_eval [powerStep, z, ho, ht, hmul, hinc, span, index, ten, one, Fin.mk.injEq] - refine ⟨z, he, ?_, ?_⟩ - · refine ⟨hm, ?_, ?_, ?_, ?_, ?_⟩ - · simpa +decide [z, inputDepth, span, index, Function.update_apply] using hd - · simp +decide [z, span, index, hp, pow_succ] - · simpa [z] using Nat.succ_le_of_lt hlt - · simpa +decide [z, one, span, index, Function.update_apply] using ho - · simpa +decide [z, ten, span, index, Function.update_apply] using ht - · simp only [z, Data.set_regs, Function.update_self] - omega - let x0 := ((((((((x.set nil 0).set one 1).set two 2).set ten 10).set nine 9).set - frameWidth 4).set span 1).set index 0) - have hsetup : ExecWithin w powerSetup x x0 8 := by - ram_eval [powerSetup, x0, h0, h1, h2, h4, h9, h10] - have hI : I x0 := by - simp +decide [I, x0, hL, inputDepth, nil, one, two, ten, nine, frameWidth, span, index] - obtain ⟨y, hy, hIy, hstop⟩ := ExecWithin.loop_invariant I (fun y => L - y.regs index) - hstep hI - have he : ExecWithin w powers x y (4 * L + 10) := by - exact (hsetup.seq hy).mono (by simp [x0]; omega) - obtain ⟨hm, hd, hp, hj, ho, ht⟩ := hIy - have hidx : y.regs index = L := by - simp only [Cond.holds_lt, hd, not_lt] at hstop - omega - refine ⟨y, he, hm, hd, hp.trans (congrArg (fun n => 10 ^ n) hidx), hidx, - ?_, ho, ?_, ht, ?_, ?_, fun i hi => he.regs_eq (powers_preserves i hi)⟩ - · simpa +decide [x0, nil, one, two, ten, nine, frameWidth, span, index] using - hy.regs_eq (by decide : - (Code.loop (.lt index inputDepth) powerStep).Preserves nil) - · simpa +decide [x0, nil, one, two, ten, nine, frameWidth, span, index] using - hy.regs_eq (by decide : - (Code.loop (.lt index inputDepth) powerStep).Preserves two) - · simpa +decide [x0, nil, one, two, ten, nine, frameWidth, span, index] using - hy.regs_eq (by decide : - (Code.loop (.lt index inputDepth) powerStep).Preserves nine) - · simpa +decide [x0, nil, one, two, ten, nine, frameWidth, span, index] using - hy.regs_eq (by decide : - (Code.loop (.lt index inputDepth) powerStep).Preserves frameWidth) + have hfit_pow : Fits w (10 ^ L + L + 2) := hbase.of_le (by omega) + let x0 := ((((x.set nil 0).set two 2).set ten 10).set nine 9).set frameWidth 4 + have hsetup : ExecWithin w powerSetup x x0 5 := by + ram_eval [powerSetup, x0, h0, h2, h4, h9, h10] + have hx0_ten : (x0.comap powerRegisters).regs Power.base = 10 := by + simp +decide [x0, powerRegisters, Power.base, nil, two, ten, nine, frameWidth] + have hx0_depth : (x0.comap powerRegisters).regs Power.exponent = L := by + simpa +decide [x0, powerRegisters, Power.exponent, inputDepth, nil, two, ten, nine, + frameWidth] using hL + obtain ⟨z, hz, hv, hc, ho, hm, hb', hd'⟩ := + Power.exec 10 L (by decide) (x0.comap powerRegisters) hx0_ten hx0_depth hfit_pow + have hpow : ExecWithin w (Power.program.rename powerRegisters) x0 + (z.extend powerRegisters x0) (4 * L + 5) := + hz.rename powerRegisters_injective + let y := z.extend powerRegisters x0 + have he : ExecWithin w powers x y (4 * L + 10) := + (hsetup.seq hpow).mono (by omega) + have hce := Data.comap_extend powerRegisters_injective z x0 + refine ⟨y, he, hm, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, ?_, + fun i hi => he.regs_eq (powers_preserves i hi)⟩ + · have h := congrArg (fun d : Data 5 => d.regs Power.exponent) hce + exact h.trans hd' + · have h := congrArg (fun d : Data 5 => d.regs Power.value) hce + exact h.trans hv + · have h := congrArg (fun d : Data 5 => d.regs Power.counter) hce + exact h.trans hc + · simpa +decide [x0, nil, two, ten, nine, frameWidth] using + hpow.regs_eq (by decide : (Power.program.rename powerRegisters).Preserves nil) + · have h := congrArg (fun d : Data 5 => d.regs Power.one) hce + exact h.trans ho + · simpa +decide [x0, nil, two, ten, nine, frameWidth] using + hpow.regs_eq (by decide : (Power.program.rename powerRegisters).Preserves two) + · have h := congrArg (fun d : Data 5 => d.regs Power.base) hce + exact h.trans hb' + · simpa +decide [x0, nil, two, ten, nine, frameWidth] using + hpow.regs_eq (by decide : (Power.program.rename powerRegisters).Preserves nine) + · simpa +decide [x0, nil, two, ten, nine, frameWidth] using + hpow.regs_eq (by decide : (Power.program.rename powerRegisters).Preserves frameWidth) end Cslib.Algorithms.WordRAM.MatrixMultiplication.Schonhage.PrunedQuery diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Pruned/Program.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Pruned/Program.lean index 37537e76f..a081d006e 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Pruned/Program.lean +++ b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Pruned/Program.lean @@ -8,6 +8,7 @@ module -- upstream: candidate public import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Pruned.Codes +public import Complexity.Computability.WordRAM.Power public import Complexity.Computability.WordRAM.Signed.Arithmetic /-! @@ -134,18 +135,16 @@ abbrev nil : Fin registers := 52 /-- Raw unsigned constant two, initialized with the fixed scalar constants. -/ abbrev two : Fin registers := 53 -/-- Constant initialization, before the depth-controlled power loop. -/ +/-- Constant initialization, before the generic depth-controlled power routine. -/ def powerSetup : Code .standard registers := ram[ - nil ← #0; one ← #1; two ← #2; ten ← #10; nine ← #9; frameWidth ← #4; - span ← #1; index ← #0 ] + nil ← #0; two ← #2; ten ← #10; nine ← #9; frameWidth ← #4 ] -/-- One charged power multiplication and counter increment. -/ -def powerStep : Code .standard registers := ram[ - span ← span * ten; index ← index + one ] +/-- Embed the five-register native power routine into the pruned-query register bank. -/ +def powerRegisters : Fin 5 → Fin registers := ![span, index, ten, inputDepth, one] /-- Initialize constants and compute the actual encoding-array length `10^L`. -/ -def powers : Code .standard registers := ram[ - ⟪powerSetup⟫; while index < inputDepth do [ ⟪powerStep⟫ ] ] +def powers : Code .standard registers := + .seq powerSetup (Power.program.rename powerRegisters) /-- Separate stack frames, query lists, value directories and membership timestamps. -/ def locate : Code .standard registers := ram[ diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Wanted/Cost.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Wanted/Cost.lean index 5bba08969..7caa27602 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Wanted/Cost.lean +++ b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Wanted/Cost.lean @@ -33,7 +33,7 @@ open Cslib.Algorithms.MatrixMultiplication.Schonhage /-- Prefix instruction budget, before the final flat halt. All initialization and loops count. -/ def timeBound (L m N Q : ℕ) : ℕ := - 16 * L + (12 * L + 33) * 2 ^ L + 15 * numBands L m N ^ 2 + (15 * L + 57) * Q + 98 + 16 * L + (12 * L + 33) * 2 ^ L + 15 * numBands L m N ^ 2 + (15 * L + 57) * Q + 100 /-- Including halt, routing is polynomial in depth times the original query, tile and subset-table sizes. In particular it does not scan `10^L` cells per query or tile. -/ diff --git a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Wanted/Setup.lean b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Wanted/Setup.lean index 7c3b293d2..e43d4e44c 100644 --- a/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Wanted/Setup.lean +++ b/Complexity/Algorithms/WordRAM/MatrixMultiplication/Schonhage/Wanted/Setup.lean @@ -132,7 +132,7 @@ theorem exec_tables {w L m N Q : ℕ} (hm : m ≤ L) (x : Data 96) (hfit : Fits w (wordBudget L m N Q)) : ∃ y, ExecWithin w (.seq BandPreparation.header (.seq BandPreparation.countCode (.seq setup BandPreparation.tablesCode))) x y - (16 * L + (12 * L + 33) * 2 ^ L + 83) ∧ Parameters L m N Q y ∧ + (16 * L + (12 * L + 33) * 2 ^ L + 85) ∧ Parameters L m N Q y ∧ (∀ gh : Fin (K0 L m) × Fin (K0 L m), y.mem (tableAddress L Q + (MaskTable.pairRank gh).val) = (MaskTable.selectedMask gh).val) ∧ diff --git a/Complexity/Algorithms/WordRAM/ThreeSum/ChanHe/Parameters.lean b/Complexity/Algorithms/WordRAM/ThreeSum/ChanHe/Parameters.lean index e72e896ab..66de0a5da 100644 --- a/Complexity/Algorithms/WordRAM/ThreeSum/ChanHe/Parameters.lean +++ b/Complexity/Algorithms/WordRAM/ThreeSum/ChanHe/Parameters.lean @@ -129,7 +129,10 @@ def setup : Code .standard 24 := ram[ quotient ← #0; len ← *quotient; one ← #1; two ← #2; three ← #3; four ← #4; five ← #5; eight ← #8; base ← len + two ] -/-- Unroll only the fixed exponent, never the input length or word width. -/ +/-- Unroll only the compile-time fixed exponent `κ`, never the input length or word width. +Unlike `WordRAM.Power.program`, which loops over a runtime exponent register using counter and +constant-one scratch registers, `power κ` emits `κ + 1` straight-line instructions over `magnitude` +and `base` alone and requires no `Fits w (κ + 1)` budget premise at empty input (`n = 0`). -/ def power : ℕ → Code .standard 24 | 0 => ram[ magnitude ← #1 ] | κ + 1 => .seq (power κ) (ram[ magnitude ← magnitude * base ]) diff --git a/ComplexityTests/SchonhageBandPreparation.lean b/ComplexityTests/SchonhageBandPreparation.lean index 6ee815d44..5eadcebc6 100644 --- a/ComplexityTests/SchonhageBandPreparation.lean +++ b/ComplexityTests/SchonhageBandPreparation.lean @@ -94,6 +94,40 @@ private def observeAsymmetricRead (side : Bool) (a : ℕ) : /-- The original numerical width premise has a checked strict boundary. -/ example : ¬Fits 13 (wordBudget 1 0 4) ∧ Fits 14 (wordBudget 1 0 4) := by decide +kernel +private def rootTestInput (k : ℕ) : Data 96 where + regs i := if i = count then k else if i = one then 1 else if i = rows then 37 else 0 + mem a := if a = 5 then 73 else 0 + +/-- Direct kernel-checked `exec_root` regressions covering zero (`k = 0`), nonsquare (`k = 6`), +and exact-square (`k = 9`) arguments, verifying the `5 * Nat.sqrt k + 8` step bound, `root`, +`square`, memory preservation, and register frame. -/ +example : + (∃ y, ExecWithin 6 rootCode (rootTestInput 0) y 8 ∧ + y.regs root = 0 ∧ y.regs square = 0 ∧ y.mem 5 = 73 ∧ y.regs rows = 37) ∧ + (∃ y, ExecWithin 6 rootCode (rootTestInput 6) y 18 ∧ + y.regs root = 2 ∧ y.regs square = 4 ∧ y.mem 5 = 73 ∧ y.regs rows = 37) ∧ + (∃ y, ExecWithin 6 rootCode (rootTestInput 9) y 23 ∧ + y.regs root = 3 ∧ y.regs square = 9 ∧ y.mem 5 = 73 ∧ y.regs rows = 37) := by + refine ⟨?_, ?_, ?_⟩ + · obtain ⟨y, he, hr, hs, hm, hf⟩ := + exec_root (w := 6) 0 (rootTestInput 0) rfl rfl (by decide) + have h0 : Nat.sqrt 0 = 0 := by norm_num + rw [h0] at he hr hs + exact ⟨y, he, hr, hs, congrFun hm 5, + (hf rows (by decide) (by decide) (by decide) (by decide)).trans rfl⟩ + · obtain ⟨y, he, hr, hs, hm, hf⟩ := + exec_root (w := 6) 6 (rootTestInput 6) rfl rfl (by decide) + have h6 : Nat.sqrt 6 = 2 := by decide +kernel + rw [h6] at he hr hs + exact ⟨y, he, hr, hs, congrFun hm 5, + (hf rows (by decide) (by decide) (by decide) (by decide)).trans rfl⟩ + · obtain ⟨y, he, hr, hs, hm, hf⟩ := + exec_root (w := 6) 9 (rootTestInput 9) rfl rfl (by decide) + have h9 : Nat.sqrt 9 = 3 := by decide +kernel + rw [h9] at he hr hs + exact ⟨y, he, hr, hs, congrFun hm 5, + (hf rows (by decide) (by decide) (by decide) (by decide)).trans rfl⟩ + /-- A negative zero-depth initialized client. The actual final memory belongs to the charged prefix run, which is followed by its charged halt. -/ example : ∃ y : Data 96, ∃ n, n ≤ timeBound 0 0 1 ∧ diff --git a/ComplexityTests/SchonhagePruned.lean b/ComplexityTests/SchonhagePruned.lean index ead284f2e..b9e3669f2 100644 --- a/ComplexityTests/SchonhagePruned.lean +++ b/ComplexityTests/SchonhagePruned.lean @@ -5,6 +5,7 @@ Authors: Erik Demaine -/ import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Pruned.Memory +import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Pruned.Powers import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Pruned.Root import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Pruned.Report import Complexity.Algorithms.WordRAM.MatrixMultiplication.Schonhage.Pruned.EmptyChild @@ -134,6 +135,28 @@ example : controller = ram[ while nil < active do [ ⟪step⟫ ] ] := rfl example : queryAt (L := 2) (90 : Fin 100) = ![.z 0 0, .z0] := by decide +kernel example : (requested (L := 2) ([90, 9, 90] : List (Fin 100))).card = 2 := by decide +kernel +private def powersTestInput (L : ℕ) : Data registers where + regs i := if i = inputDepth then L else if i = queryCount then 19 else 0 + mem a := if a = 11 then 43 else 0 + +-- Direct kernel-checked `exec_powers` regressions at depth zero and depth two, verifying the +-- `4 * L + 10` step bound from `powerSetup` plus `Power.program.rename powerRegisters`, the +-- computed `10^L` span, initialized constants, memory preservation, and `¬PowersLocal` frame. +example : + (∃ y, ExecWithin 10 powers (powersTestInput 0) y 10 ∧ + y.regs span = 1 ∧ y.regs index = 0 ∧ y.regs one = 1 ∧ y.regs frameWidth = 4 ∧ + y.mem 11 = 43 ∧ y.regs queryCount = 19) ∧ + (∃ y, ExecWithin 17 powers (powersTestInput 2) y 18 ∧ + y.regs span = 100 ∧ y.regs index = 2 ∧ y.regs one = 1 ∧ y.regs frameWidth = 4 ∧ + y.mem 11 = 43 ∧ y.regs queryCount = 19) := by + refine ⟨?_, ?_⟩ + · obtain ⟨y, he, hm, _, hspan, hidx, _, hone, _, _, _, hfw, hf⟩ := + exec_powers (w := 10) (L := 0) (Q := 0) (powersTestInput 0) rfl (by decide) + exact ⟨y, he, hspan, hidx, hone, hfw, congrFun hm 11, (hf queryCount (by decide)).trans rfl⟩ + · obtain ⟨y, he, hm, _, hspan, hidx, _, hone, _, _, _, hfw, hf⟩ := + exec_powers (w := 17) (L := 2) (Q := 0) (powersTestInput 2) rfl (by decide) + exact ⟨y, he, hspan, hidx, hone, hfw, congrFun hm 11, (hf queryCount (by decide)).trans rfl⟩ + private def observe (w L fuel : ℕ) (a b : Fin (10 ^ L) → ℤ) (qs : List (Fin (10 ^ L))) := (program.eval w fuel (State.init (input L a b qs) : State registers).data).map fun (y, time) => ((List.range qs.length).map fun i => decode (y.mem (outputBase L qs.length + i)), time) diff --git a/UPSTREAM.md b/UPSTREAM.md index d2c654807..75d6bd509 100644 --- a/UPSTREAM.md +++ b/UPSTREAM.md @@ -167,14 +167,18 @@ Do not add a ledger row for an ordinary candidate or repeat its source marker he `w+2`. The two extra bits and actual arithmetic-block costs come from the existing native signed kernels, without a copied source machine or a whole-program simulation hypothesis. The existing array-copy interface is retained, with a narrow generalization for disjoint - blocks in either address order or identical blocks. + blocks in either address order or identical blocks. `Encoding.setup` and `Encoding.powerBody` + retain their simultaneous `7^L`/`10^L` loop sharing the `stage` counter (`5*L+18` setup + instructions) rather than running two sequential `WordRAM.Power.program` loops. This is the encoder component only; band/tile preparation and full native composition remain in [#1394](https://github.com/cslib-community/complexitylib/issues/1394). * **Native shared-band preparation.** The native shared-band preparation component ([#1523](https://github.com/cslib-community/complexitylib/issues/1523)) computes a deterministic `Layout` and both subset tables by native mask/popcount scans, computes the square-root block - count by a charged loop, and decodes each band-array index into the original row/column and - thin-dimension coordinates. Real matrix bounds are checked before row-major multiplication. + count and its square by renaming `NatSqrt.program` (`BandPreparation.rootRegisters`) followed by + one multiplication (`5*sqrt(k)+8` steps), and decodes each band-array index into the original + row/column and thin-dimension coordinates. Real matrix bounds are checked before row-major + multiplication. Each shared row/column band calls the unchanged encoder once; tiles refer to its catalogue block in the existing leaf order. The fixed 96-register standard code starts from actual `State.init` scalar/matrix input, charges snapshot and restoration of the complete original @@ -451,11 +455,14 @@ Do not add a ledger row for an ordinary candidate or repeat its source marker he callers use the flat initialized runs, width theorem or pointer/epoch kernel. The narrow `WordRAM.ArrayAppend` and `WordRAM.InputArray` topic modules keep append/input dependencies and interface changes away from the widely imported `Allocation` core. - Post-integration [#1553](https://github.com/cslib-community/complexitylib/issues/1553) includes - replacing pruned `powers` with sibling #1536's generic `Power.program`, preserving the - `4L+10` budget via five retained constant initializations plus its `4L+5` run. Its exact - register map, affected execution/width/cost clients and sibling integration order are - recorded there. This documentation/test repair does not change either sibling's program. + Post-integration [#1553](https://github.com/cslib-community/complexitylib/issues/1553) replaces + pruned `powers` with five retained constant initializations (`powerSetup`) followed by + `Power.program.rename powerRegisters`, preserving the `4L+10` step bound and `PowersLocal` + frame, and replaces `BandPreparation.rootCode` with `NatSqrt.program.rename rootRegisters` + followed by `square ← root * root` (`5*sqrt(k)+8` steps). `ChanHe.Parameters.power` retains + its compile-time fixed-exponent unrolling (`κ+1` straight-line instructions over `magnitude` + and `base` without a runtime exponent/counter register or `Fits w (κ+1)` premise at `n=0`), + and `Encoding.setup`/`powerBody` retains its simultaneous `7^L`/`10^L` loop sharing `stage`. * **Typed RAM upstream core.** [CSLib #1158](https://github.com/leanprover/cslib/pull/1158) at `beea60a15eddc91d93e92fa5af9fb965c351d0d5` proposes `Cslib/Computability/Machines/RAM/Basic.lean`, in namespace `Cslib.RAM.Typed`.