κ-improvement development, part 19: item-language primitives

-- OpenAI-preprint / κ-improvement development (integer-multiplication-kappa), generated by skeleton subtraction.

import Definitions.Def_IntMul_Kappa_K_Defs18
import Definitions.Def_IntMul_HvdH_BApprox
import Definitions.Def_IntMul_HvdH_DiagApprox
import Definitions.Def_IntMul_HvdH_EApprox
import Definitions.Def_IntMul_HvdH_FPApprox
import Definitions.Def_IntMul_HvdH_JApprox
import Definitions.Def_IntMul_HvdH_SApprox

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KFNSel
/-!
# KFN: selected-bit operators by record strides

A selected-bit operator on a record array only needs, for every selected bit `i`, its record
stride `st i`: bit `i` of record `j` is `j / st i % 2` (`fbit`), and setting the selected bits to
`y` moves to record `fset st j y = j + Σ_i (y_i - bit_i) st_i`.  `selOpF st n T xs` is the
selected-bit operator in these terms.

* `SLaws st n`: the strides are those of a genuine selected-bit structure on `n` records
  (`fset` stays below `n` and reads back `y`).
* `selOpF_mul`, `selOpF_one`, `selOpF_diagonal`: the operator calculus.
* `StrideSel Sb st`: a `SelBits` structure whose bits/sets are given by the strides;
  `selOp_eq_selOpF` and `slaws_of_strideSel`.
* `strideSel_layer`: the layer layout `[P] × [2^K]^D × [S]` has strides `stT K D S ρ i`
  (independent of `P`).
* `subOp g T` (`T` on the bits `g κ`, identity elsewhere) and **`selOpF_subOp`**:
  `selOpF st n (subOp g T) = selOpF (st ∘ g) n T`.  This one lemma gives reindexing
  (`g` an equivalence, `subOp_tens`), the child slots (`slotOp_eq_subOp`) and the pieces
  (`embed_eq_subOp`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open AModel Matrix Finset

section Basic

variable {ι : Type*} [Fintype ι] [DecidableEq ι]

/-- Bit `i` of record `j`. -/
def fbit (st : ι → ℕ) (j : ℕ) : ι → Fin 2 := fun i => ⟨j / st i % 2, Nat.mod_lt _ (by norm_num)⟩

/-- The record with the selected bits set to `y`, as an integer. -/
def fsetZ (st : ι → ℕ) (j : ℕ) (y : ι → Fin 2) : ℤ :=
  (j : ℤ) + ∑ i, (((y i : ℕ) : ℤ) - ((fbit st j i : ℕ) : ℤ)) * (st i : ℤ)

/-- The record with the selected bits set to `y`. -/
def fset (st : ι → ℕ) (j : ℕ) (y : ι → Fin 2) : ℕ := (fsetZ st j y).toNat

/-- The strides describe a selected-bit structure on `n` records. -/
structure SLaws (st : ι → ℕ) (n : ℕ) : Prop where
  nonneg : ∀ j < n, ∀ y, 0 ≤ fsetZ st j y
  lt : ∀ j < n, ∀ y, fset st j y < n
  bit_set : ∀ j < n, ∀ y, fbit st (fset st j y) = y

variable {M : Type*} [AddCommMonoid M] [Module ℂ M]

/-- The selected-bit operator `T` on `n` records with strides `st`. -/
noncomputable def selOpF (st : ι → ℕ) (n : ℕ) (T : Matrix (ι → Fin 2) (ι → Fin 2) ℂ)
    (xs : List M) : List M :=
  (List.range n).map fun j => ∑ y, T (fbit st j) y • xs.getD (fset st j y) 0

end Basic

/-! ## The bridge to `SelBits` -/

section Bridge

variable {sh : Shape} {ι : Type*} [Fintype ι] [DecidableEq ι]

/-- A `SelBits` structure described by record strides. -/
def StrideSel (Sb : Butterfly.SelBits sh ι) (st : ι → ℕ) : Prop :=
  ∀ j < vol sh, Sb.bits (decode sh j) = fbit st j ∧
    ∀ y, ((encode sh (Sb.set (decode sh j) y) : ℕ) : ℤ) = fsetZ st j y

variable {M : Type*} [AddCommMonoid M] [Module ℂ M]

end Bridge

/-! ## The layer layout -/

section Layer

variable {P K D S ρ : ℕ}

end Layer

/-! ## Operators on a subfamily of the selected bits -/

section Sub

variable {ι κ : Type*} [Fintype ι] [DecidableEq ι] [Fintype κ] [DecidableEq κ]

/-- `T` on the selected bits `g κ`, identity on the others. -/
def subOp (g : κ → ι) (T : Matrix (κ → Fin 2) (κ → Fin 2) ℂ) : Matrix (ι → Fin 2) (ι → Fin 2) ℂ :=
  fun x y => if ∀ i, (∀ k, g k ≠ i) → x i = y i then T (x ∘ g) (y ∘ g) else 0

/-- Extension of bits on `g κ` by the bits of `x` elsewhere. -/
noncomputable def extB (g : κ → ι) (x : ι → Fin 2) (z : κ → Fin 2) : ι → Fin 2 := by
  classical
  exact fun i => if h : ∃ k, g k = i then z (Classical.choose h) else x i

variable {M : Type*} [AddCommMonoid M] [Module ℂ M]

end Sub

/-! ## Slots and pieces as subfamilies -/

end NLogN.KFN

end

section
-- Source module: Solutions.KFNRole
/-!
# KFN: the role streams of a network invocation

A network state on `W` roles is stored role-major: role `w` is a list `X w` of `n` records, and
record `j` of a role carries the group address `fbit st j : Addr m κ` (selected bits with record
strides `st`, `KFNSel`).  `roleOpF st n A X` is the action of a network operator
`A : ROp (Fin W) m κ` on these streams.

* `roleOpF_mul`, `roleOpF_one`: the operator calculus.
* `roleOpF_blockOp`, `roleOpF_wireOp`, `roleOpF_roleDiag`: block operators act per role.
* The steps of `P8Lift.RStep`: `roleOpF_gate` (a linear combination of whole roles),
  `roleOpF_perm` (role permutation), `roleOpF_phase`, `roleOpF_rowAdd`, `roleOpF_child`
  (`selOpF` on one role; a child acts with the restricted strides of its slot).
* **`roleOpF_net`**: the network acts as the sequence of its steps, and for an interface
  (`Prop8IfaceR.correct`) **`roleOpF_iface`**: running the steps of `J.net` one after the other
  computes `selOpF st n (tens Cm)` on every role.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open AModel Matrix Finset P8Lift Butterfly

section Role

variable {W m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ] {M : Type*} [AddCommMonoid M]
  [Module ℂ M]

/-- The action of a network operator on role streams. -/
noncomputable def roleOpF (st : Fin m × κ → ℕ) (n : ℕ) (A : ROp (Fin W) m κ)
    (X : Fin W → List M) : Fin W → List M := fun w =>
  (List.range n).map fun j =>
    ∑ q : Fin W × Addr m κ, A (w, fbit st j) q • (X q.1).getD (fset st j q.2) 0

/-! ## The steps -/

/-! ## The network -/

/-- Run the steps one after the other. -/
noncomputable def runSteps (st : Fin m × κ → ℕ) (n : ℕ) (T : Butterfly.Op κ) :
    List (RStep (Fin W) m) → (Fin W → List M) → (Fin W → List M)
  | [], X => X
  | s :: l, X => runSteps st n T l (roleOpF st n (s.den T) X)

end Role

end NLogN.KFN

end

section
-- Source module: Solutions.KFNPhase
/-!
# KFN: a linear-time phase pass on a blocked stream

`phL R0 t B u aE` multiplies the `B · (R0 · 2^t)` Gaussian entries on top of the data stack by
`i^{E}`, where the exponent of entry `q` is read from the table
`tblG (replicate R0 (u mod 4)) aE t` (built by `t` doublings, `KFG.dblP`) at position
`q mod (R0 · 2^t)`:
`E = u + Σ_{p<t} aE p · ⌊q / R0 / 2^p⌋ mod 2` (`getD_tblR`).

Steps: build the table (`R0` pushes and `t` doublings, `O(R0 2^t)`), then for each of the `B`
blocks one multiply loop (`KFG.mulB`, table `xA → xB`, results to `qO`) and one move of the table
back (`xB → xA`); finally drop the table and move the results back to the data stack.
Cost `O(B · R0 2^t · (w + 1) + (t + B) (Z + 1))` (`phL_fsp`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG Complex

/-- The counter of the block loop. -/
abbrev cO : Fin HS := 46

/-! ## Tables -/

/-- The table after `t` doublings of `T0`. -/
def tblG (T0 : List ℤ) (aE : ℕ → ℕ) : ℕ → List ℤ
  | 0 => T0
  | t + 1 => dblL ((aE t % 4 : ℕ) : ℤ) (tblG T0 aE t)

/-! ## Building the table -/

/-- `R0` copies of `v` on `xA`. -/
def initP (R0 : ℕ) (v : ℤ) : RS := .rep NLogN.KFB.c1 (.lit (R0 : ℤ)) (.push xA (.lit v))

/-- `t` doublings. -/
def cTblG (M t Z : ℕ) : ℕ := 120000 * M * 2 ^ t + t * (10000 * (Z + 2) + 20001) + 1

/-! ## The multiply loop with a nonempty result stack -/

/-! ## One block -/

/-- One block: multiply by the table, then move the table back. -/
def blkP : RS := .seq (.rep NLogN.KFB.c1 (.reg xN) mulB) (NLogN.mvK xB xA NLogN.KFB.c1 (.reg xN))

/-- Cost of one block of `M` entries. -/
def cBlk (w M Z : ℕ) : ℕ := M * (cMulB w + 5000) + 1000 * (Z + 2) + 5000

/-! ## All blocks -/

/-- `B` copies of the table. -/
def repT (B : ℕ) (T : List ℤ) : List ℤ := (List.replicate B T).flatten

/-! ## The phase pass -/

/-- **The phase pass** on the `B · R0 · 2^t` entries on top of the data stack. -/
def phL (R0 t B u : ℕ) (aE : ℕ → ℕ) : RS :=
  .seq (initP R0 ((u % 4 : ℕ) : ℤ))
  (.seq (.push xN (.lit (R0 : ℤ)))
  (.seq (tblP aE t)
  (.seq (.rep NLogN.KFN.cO (.lit (B : ℤ)) blkP)
  (.seq (NLogN.mvK xA xB NLogN.KFB.c1 (.reg xN))
  (.seq (.rep NLogN.KFB.c1 (.reg xN) (.pop xB))
  (.seq (.asg xN (.lit ((2 * (B * (R0 * 2 ^ t)) : ℕ) : ℤ)))
  (.seq (NLogN.mvK qO 0 NLogN.KFB.c1 (.reg xN)) (.pop xN))))))))

/-- Cost of the phase pass on `V = B · R0 · 2^t` entries. -/
def cPhL (w V t B Z : ℕ) : ℕ := V * (cMulB w + 300000 + 500 * w) + (t + B + 10) * (20000 * (Z + 2) + 30000)

/-- The rotated stream. -/
noncomputable def phOut (R0 t B u : ℕ) (aE : ℕ → ℕ) (zs : List GI) : List GI :=
  KFG.mulL (repT B (tblG (List.replicate R0 ((u % 4 : ℕ) : ℤ)) aE t)) zs

end NLogN.KFN

end

section
-- Source module: Solutions.KFNImpl
/-!
# KFN: implementations of selected-bit operators and of network steps

* `ClN f`: the clean scratch used by all KFN programs (`KFB.Cl`, `KFG.ClP`, `KFG.ClT`, and the
  block counter `cO`).
* `SImpl r st n T code e dw cost`: on a clean managed state whose data stack starts with a record
  array `zs` (scale `F`, width `w`) representing `xs` (`n` records of `r` entries), `code` leaves
  an array representing `selOpF st n T xs` at scale `F + e`, width `w + dw`, at fine cost
  `cost w`.  This is `KFH.LImpl` with the layout given by record strides (`lImpl_of_sImpl`).
* `RImpl r n st A code e dw cost`: the same for a network operator `A` acting on `W` role streams
  stored one after the other (`flatR X`), through `roleOpF`.
* `RImpl.comp`, `rImpl_list` (a list of steps implements `denListR`), and
  `sImpl_of_rImpl_roleDiag`: an implementation of `roleDiag T` is one of `T` on the
  concatenated streams (strides inside the role blocks, `selOpF_flatR`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

/-! ## Clean scratch -/

/-- The clean scratch of the KFN programs. -/
structure ClN (f : Fin HS → List ℤ) : Prop where
  cl : KFB.Cl f
  p : KFG.ClP f
  t : KFG.ClT f
  o : f NLogN.KFN.cO = []

/-! ## Role streams stored one after the other -/

/-- The streams one after the other. -/
def flatR {α : Type*} {W : ℕ} (X : Fin W → List α) : List α := (List.finRange W).flatMap X

/-! ## Implementations -/

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- `code` implements the selected-bit operator `T` (record strides `st`, `n` records). -/
def SImpl {ι : Type*} [Fintype ι] [DecidableEq ι] (r : ℕ) (st : ι → ℕ) (n : ℕ)
    (T : Matrix (ι → Fin 2) (ι → Fin 2) ℂ) (code : RS) (e dw : ℕ) (cost : ℕ → ℕ) : Prop :=
  ∀ (s : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ) (xs : List (Fin r → ℂ))
    (F w : ℕ), ClN f → f 0 = gflat zs ++ l → RecRep r F zs xs → xs.length = n → NLogN.KFB.GB w zs →
    ∃ zs', RecRep r (F + e) zs' (selOpF st n T xs) ∧ NLogN.KFB.GB (w + dw) zs' ∧
      FSp code (encK s f) (encK s (Function.update f 0 (gflat zs' ++ l))) (cost w)

section SImplLemmas

variable {ι : Type*} [Fintype ι] [DecidableEq ι] {r : ℕ} {st : ι → ℕ} {n : ℕ}

end SImplLemmas

section RImplDefs

variable {W m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ]

/-- `code` implements the network operator `A` on `W` role streams of `n` records each. -/
def RImpl (r n : ℕ) (st : Fin m × κ → ℕ) (A : ROp (Fin W) m κ) (code : RS) (e dw : ℕ)
    (cost : ℕ → ℕ) : Prop :=
  ∀ (s : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ)
    (X : Fin W → List (Fin r → ℂ)) (F w : ℕ), ClN f → f 0 = gflat zs ++ l →
    RecRep r F zs (flatR X) → (∀ v, (X v).length = n) → NLogN.KFB.GB w zs →
    ∃ zs', RecRep r (F + e) zs' (flatR (roleOpF st n A X)) ∧ NLogN.KFB.GB (w + dw) zs' ∧
      FSp code (encK s f) (encK s (Function.update f 0 (gflat zs' ++ l))) (cost w)

variable {r n : ℕ} {st : Fin m × κ → ℕ}

/-- Sequence of programs. -/
def seqL : List RS → RS
  | [] => .skip
  | c :: cs => .seq c (seqL cs)

/-- Cost of a list of steps. -/
def costL {α : Type*} (cost : α → ℕ → ℕ) (dw : α → ℕ) : List α → ℕ → ℕ
  | [], _ => 1
  | a :: l, w => cost a w + costL cost dw l (w + dw a)

end RImplDefs

/-! ## Operators acting inside the role blocks -/

section Blocks

variable {ι : Type*} [Fintype ι] [DecidableEq ι] {M : Type*} [AddCommMonoid M] [Module ℂ M]

end Blocks

section Diag

variable {W m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ] {r n : ℕ} {st : Fin m × κ → ℕ}

end Diag

end NLogN.KFN

end

section
-- Source module: Solutions.KFNGStep
/-!
# KFN: gates and role permutations on role streams of records

* `enc_of_rec` / `rec_of_enc`: role streams of records (`RecRep r F zs (flatR X)`) are the
  KFG stream encoding `Enc W (n r) (const F) zs (yR r X)`.
* Scale shifts: the gate programs depend only on scale differences (`gateP_shE`, …), so one
  program serves every scale `F`.
* Scale bookkeeping (`gateE_bnd`, `gateE_top`, `gateE_zero`, `gateE_lev`) and the width / cost
  bounds `gateW_le`, `cGate_le` (shifts of at most one bit per atom).
* **`rImpl_gate`**: a gate `gateMat W as` (atoms with `h ≤ 1`, targets never sources) followed by
  a leveling pass (`c = 0` atoms bringing every role to the common scale `F + gδ as`).
* **`rImpl_perm`**: a role permutation.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

/-! ## Streams of records as the KFG encoding -/

section EncBridge

variable {W r n F : ℕ}

/-- Entry `p` of role `v` (record `p / r`, coefficient `p mod r`). -/
noncomputable def yR (r : ℕ) (X : Fin W → List (Fin r → ℂ)) : ℕ → ℕ → ℂ := fun v p =>
  if h : v < W ∧ 0 < r then ((X ⟨v, h.1⟩).getD (p / r) 0) ⟨p % r, Nat.mod_lt _ h.2⟩ else 0

end EncBridge

/-! ## Gate programs depend only on scale differences -/

section Shift

/-- Shift all scales by `F`. -/
def shE (F : ℕ) (e : ℕ → ℕ) : ℕ → ℕ := fun t => e t + F

end Shift

/-! ## Scale bookkeeping -/

section Scales

/-- The leveling atoms: every role `t ≠ t0` gets `0 ·` role `t0`. -/
def levA (W t0 : ℕ) : List KFG.Atom :=
  ((List.range W).filter fun t => t ≠ t0).map fun t => ⟨t, t0, (0, 0), 0⟩

end Scales

/-! ## Width and cost of a gate pass -/

section Bounds

end Bounds

/-! ## The gate step -/

section GateStep

/-- The first atom with `h = 1`. -/
def h1Atom (as : List KFG.Atom) : Option KFG.Atom := as.find? fun a => a.h == 1

/-- The scale added by a gate step. -/
def gδ (as : List KFG.Atom) : ℕ := if (h1Atom as).isSome then 1 else 0

/-- The leveling atoms after a gate. -/
def levOf (W : ℕ) (as : List KFG.Atom) : List KFG.Atom :=
  match h1Atom as with
  | some a0 => levA W a0.t
  | none => []

/-- The program of a gate step: the atoms, then the leveling. -/
def gateC (W L : ℕ) (as : List KFG.Atom) : RS :=
  .seq (gateP as (fun _ => 0) L) (gateP (levOf W as) (gateE as (fun _ => 0)) L)

/-- Cost of a gate step. -/
def cGateC (W L Z A w : ℕ) : ℕ :=
  A * cAtom (w + 3 * A) 1 L W Z + W * cAtom (w + 3 * A + 3 * W) 1 L W Z + 2

variable {W m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ]

end GateStep

end NLogN.KFN

end

section
-- Source module: Solutions.KFNWrap
/-!
# KFN: running a stream program on one role

`wrapC W L w e code`: swap role `w` to the front (one pair pass), run `code` on the first
segment (the other roles stay below it on the data stack), bring the other roles to the new
scale `F + e` (leveling atoms `0 · role 0`), and swap back.

**`rImpl_onRole`**: if `code` implements a selected-bit operator `Tw` on single streams
(`SImpl`), the wrapper implements every network operator `A` that acts as `Tw` on role `w` and
as the identity on the other roles (`roleOpF_phase`, `roleOpF_rowAdd`, `roleOpF_child`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

/-! ## Widths and costs with a general scale window -/

section Bounds2

end Bounds2

/-! ## Segments of the encoding -/

section Seg

end Seg

/-! ## The wrapper -/

section Wrap

/-- The swap bringing role `w` to the front (none for `w = 0`). -/
def swL (w : ℕ) : List (ℕ × ℕ) := if w = 0 then [] else [(0, w)]

/-- Run `code` (scale `+e`) on role `w`. -/
def wrapC (W L w e : ℕ) (code : RS) : RS :=
  .seq (swapsP L (NLogN.KFN.swL w)) (.seq code
    (.seq (gateP (levA W 0) (Function.update (fun _ => 0) 0 e) L) (swapsP L (NLogN.KFN.swL w))))

/-- Cost of the wrapper around a program of cost `cost`. -/
def cWrapR (W L Z e dw : ℕ) (cost : ℕ → ℕ) (w : ℕ) : ℕ :=
  cost w + (W + 2) * cAtom (w + dw + (e + 2) * W) e L W Z + 3

variable {W m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ]

end Wrap

end NLogN.KFN

end

section
-- Source module: Solutions.KFNIL
/-!
# KFN: implementations by item-language programs

The row additions of Lemma 13 are item-language programs (`IStmt HS 1` calling the Lemma 12
routines under a fixed `iprog`), whose code depends on the input width `w`.  This file gives the
`IStmt` versions of the implementation predicates, with **code families `ℕ → IStmt HS Pp`**
indexed by the input width, run under a fixed `I`, `iprog`:

* `SImplI I iprog r st n T code e dw cost`, `RImplI ...` (as `SImpl`, `RImpl`, with
  `IRuns I iprog (code w)` in place of `FSp code`);
* lifts of the register-program versions (`SImpl.toI`, `RImpl.toI`, via `FSp.irun`);
* `SImplI.comp`, `RImplI.comp`, `rImplI_list`, `sImplI_of_rImplI_roleDiag`, `rImplI_onRole`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- `code w` implements the selected-bit operator `T` on inputs of width `w`. -/
def SImplI (I : List Sym) (iprog : Fin Pp → IStmt HS Pp) {ι : Type*} [Fintype ι] [DecidableEq ι]
    (r : ℕ) (st : ι → ℕ) (n : ℕ) (T : Matrix (ι → Fin 2) (ι → Fin 2) ℂ) (code : ℕ → IStmt HS Pp)
    (e dw : ℕ) (cost : ℕ → ℕ) : Prop :=
  ∀ (s : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ) (xs : List (Fin r → ℂ))
    (F w : ℕ), ClN f → f 0 = gflat zs ++ l → RecRep r F zs xs → xs.length = n → NLogN.KFB.GB w zs →
    ∃ zs', RecRep r (F + e) zs' (selOpF st n T xs) ∧ NLogN.KFB.GB (w + dw) zs' ∧
      IRuns I iprog (code w) (encK s f) (encK s (Function.update f 0 (gflat zs' ++ l))) (cost w)

section SImplILemmas

variable {ι : Type*} [Fintype ι] [DecidableEq ι] {r : ℕ} {st : ι → ℕ} {n : ℕ}

/-- Sequential composition of code families. -/
def seqI (c1 c2 : ℕ → IStmt HS Pp) (d1 : ℕ) : ℕ → IStmt HS Pp :=
  fun w => .seq (c1 w) (c2 (w + d1))

end SImplILemmas

section RImplIDefs

variable {W m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ]

/-- `code w` implements the network operator `A` on `W` role streams of `n` records each. -/
def RImplI (I : List Sym) (iprog : Fin Pp → IStmt HS Pp) (r n : ℕ) (st : Fin m × κ → ℕ)
    (A : ROp (Fin W) m κ) (code : ℕ → IStmt HS Pp) (e dw : ℕ) (cost : ℕ → ℕ) : Prop :=
  ∀ (s : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ)
    (X : Fin W → List (Fin r → ℂ)) (F w : ℕ), ClN f → f 0 = gflat zs ++ l →
    RecRep r F zs (flatR X) → (∀ v, (X v).length = n) → NLogN.KFB.GB w zs →
    ∃ zs', RecRep r (F + e) zs' (flatR (roleOpF st n A X)) ∧ NLogN.KFB.GB (w + dw) zs' ∧
      IRuns I iprog (code w) (encK s f) (encK s (Function.update f 0 (gflat zs' ++ l))) (cost w)

variable {r n : ℕ} {st : Fin m × κ → ℕ}

/-- Sequence of code families (the `a`-th one starts at the width left by the previous ones). -/
def seqLI {α : Type*} (code : α → ℕ → IStmt HS Pp) (dw : α → ℕ) : List α → ℕ → IStmt HS Pp
  | [] => fun _ => .skip
  | a :: l => seqI (code a) (seqLI code dw l) (dw a)

end RImplIDefs

section DiagI

variable {W m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ] {r n : ℕ} {st : Fin m × κ → ℕ}

end DiagI

section WrapI

/-- Run the code family `code` (scale `+e`) on role `w`. -/
def wrapCI (W L w e : ℕ) (code : ℕ → IStmt HS Pp) : ℕ → IStmt HS Pp := fun x =>
  .seq (fcode (swapsP L (NLogN.KFN.swL w))) (.seq (code x)
    (.seq (fcode (gateP (levA W 0) (Function.update (fun _ => 0) 0 e) L))
      (fcode (swapsP L (NLogN.KFN.swL w)))))

variable {W m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ]

end WrapI

end NLogN.KFN

end

section
-- Source module: Solutions.KFNPhSem
/-!
# KFN: the phase pass as a selected-bit operator

For selected bits `i : ι` at record strides `2^{pos i} · S` (`pos i < t`) on `B · 2^t · S`
records of `r` entries, the phase pass `phL (S r) t B u (aEof pos a')` implements the diagonal
`x ↦ i^{u + Σ_i a'_i x_i}` (**`sImpl_phase`**), exactly (no scale, no width change), at cost
`cPhL`, linear in the data.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

/-- The doubling exponents: `aE p = Σ_{pos i = p} a'_i`. -/
def aEof {ι : Type*} [Fintype ι] (pos : ι → ℕ) (a' : ι → ℕ) (p : ℕ) : ℕ :=
  ∑ i ∈ Finset.univ.filter (fun i => pos i = p), a' i

end NLogN.KFN

end

section
-- Source module: Solutions.KFNGate
/-!
# KFN: classification of the scalar gates of the complex network (task 1)

Every `RStep.gate G` of `P8Lift.complexIface.net` is one of the two kinds of fine IL gates of
KFG:

* a **dyadic atom gate** `G = gateMat W_c as` (`KFG.gateMat`: `1 + Σ_a (c_a/2^{h_a}) e_{t_a,s_a}`)
  for a list of atoms with `c_a ∈ {±1}`, `h_a ≤ 1`, every atom `Ok` and no target of the list a
  source of the list (`KFG.NoTS`), so `KFG.gate_iruns` applies; these are the side injections,
  scatters, copies and gathers of the schedule (2) of `NetComplexLemma5` and their inverses;
* the **sign gate** `fixG`, a diagonal `±1` matrix (one unit scaling `KFG.UAtom` with `u = -1`
  per physical `X` wire).

`GOK G`: `G = 1 + Σ_{x ∈ l} (ε_x / 2^{h_x}) e_{t_x, s_x}` with `ε_x = ±1`, `h_x ≤ 1`, targets
never sources.  `cxNamed_GOK`: every gate of `cxNamed` is `GOK`.  `GOK_atoms`: a `GOK` matrix
transported along `W ≃ Fin n` is `gateMat n as` for an `Ok`, `NoTS` atom list.
**`complexIface_gate_class`**: the classification of all gate steps of `complexIface`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option synthInstance.maxSize 512

namespace NLogN.KFN

open Matrix Finset Complex

section Generic

variable {W : Type*} [Fintype W] [DecidableEq W]

/-- The matrix of an elementary dyadic list. -/
noncomputable def eMat (l : List (W × W × ℤ × ℕ)) : Matrix W W ℂ :=
  (l.map fun x => single x.1 x.2.1 ((x.2.2.1 : ℂ) / 2 ^ x.2.2.2)).sum

/-- A dyadic elementary gate whose targets are never sources. -/
def GOK (G : Matrix W W ℂ) : Prop :=
  ∃ l : List (W × W × ℤ × ℕ), G = 1 + eMat l ∧ (∀ x ∈ l, ∀ y ∈ l, x.1 ≠ y.2.1) ∧
    ∀ x ∈ l, (x.2.2.1 = 1 ∨ x.2.2.1 = -1) ∧ x.2.2.2 ≤ 1

/-- The atoms of a `GOK` list, transported along `e : W ≃ Fin n`. -/
def atomsOf {n : ℕ} (e : W ≃ Fin n) (l : List (W × W × ℤ × ℕ)) : List KFG.Atom :=
  l.map fun x => ⟨(e x.1).val, (e x.2.1).val, (x.2.2.1, 0), x.2.2.2⟩

end Generic

/-! ## The gates of the complex construction -/

section Complex

variable {h : ℕ} {T : Type*} [Fintype T] [DecidableEq T] {ind : T → Finset (Fin h)}

variable {Wt : Type*} [Fintype Wt] [DecidableEq Wt]

end Complex

/-! ## The step lists -/

section Steps

open P8Lift

variable {R : Type} [Fintype R] [DecidableEq R]

end Steps

/-! ## The classification for `complexIface` -/

end NLogN.KFN

end

section
-- Source module: Solutions.KFNNode
/-!
# KFN: one node of the recursive `C`-layer

A node runs the steps of an interface `J : Prop8IfaceR W m s` on `W` role streams of `n`
records; its bits are `Fin m × κ` (rows × the bits of a child), at record strides `st`.

* **`sImpl_node`**: implementations of the steps compose to an implementation of `tens Cm`
  on the `W · n` records (`rImpl_list`, `J.correct`, `sImpl_of_rImpl_roleDiag`).
* Step builders on power-of-two strides `st p = 2^{pos p} · S₀` (`n = B · 2^t · S₀`):
  `rImpl_phaseStep` (phase pass on one role), `rImpl_fixG` (a `±1` role diagonal as a list of
  phase steps), `rImpl_rowStep` (a row addition from any single-stream implementation),
  `rImpl_childStep` (forward child from an implementation of the child, the inverse child via
  `tens Cinv = (-i)^{|κ|} ZD · tens Cm · ZD`, `sImpl_cinv`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

section Node

variable {Pp : ℕ} {Isym : List Sym} {iprog : Fin Pp → IStmt HS Pp}
variable {W m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ]

/-! ## Phase steps -/

/-- Code of a phase pass with exponents `u + Σ a'_i x_i` on strides `2^{pos i} S₀`. -/
def phC {ι : Type*} [Fintype ι] (pos : ι → ℕ) (S0 r t B u : ℕ) (a' : ι → ℕ) : RS :=
  phL (S0 * r) t B u (aEof pos a')

/-- The phase steps negating the roles of a list. -/
def negSteps (l : List (Fin W)) : List (RStep (Fin W) m) := l.map fun w => RStep.phase w 2 0 0

/-! ## Row additions and child calls -/

/-- The inverse child: `ZD`, then the forward child, then `(-i)^{|κ|} ZD`. -/
def cinvC (pos : κ → ℕ) (S0 r t B : ℕ) (code : ℕ → IStmt HS Pp) (dw : ℕ) : ℕ → IStmt HS Pp :=
  seqI (seqI (fun _ => fcode (phC pos S0 r t B 0 (fun _ => 2))) code 0)
    (fun _ => fcode (phC pos S0 r t B (3 * Fintype.card κ) (fun _ => 2))) (0 + dw)

end Node

end NLogN.KFN

end

section
-- Source module: Solutions.KFNRec
/-!
# KFN: the recursive `C`-layer — local step bounds

Parameters: an interface `J : Prop8IfaceR W m s` whose gates are atom lists (at most `A`
atoms) or `±1` role diagonals (`GClass`), records of `r` entries, a size bound `Z`, and a
monotone "log-size" cost profile `Ψ ≥ Z + 2`.  The row additions are an abstract hypothesis
(`RowHyp`, e.g. Lemma 13): cost `CR · n · (r (w+1) (e+1)^τ + Ψ w)` on `n` records whose rows have
`e` bits.

Unit of cost: `Xf r Ψ v = r (v + 1) + Ψ v` per record at width `v`.

This file: the recursion constants (`Eb` scale, `DWb` width, `Ob` offsets, `Whb` cumulative
widths, `cb` cost coefficients), the per-step cost bounds, and **`step_ok`**: every step of
`J.net` at level `k+1` has an implementation with the bounds, given the level-`k` claim `Lv`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

/-! ## Constants of the recursion -/

/-- Scale bound: the leaf kernel has scale `1`; a node adds `≤ 1` per local step. -/
def Eb (L s : ℕ) : ℕ → ℕ
  | 0 => 1
  | k + 1 => L + s * Eb L s k

/-- Width increase of a local step. -/
def dwloc (W A : ℕ) : ℕ := 3 * A + 3 * W + 2 * W * W + 2 * W

/-- Width bound of a level. -/
def DWb (L s W A : ℕ) : ℕ → ℕ
  | 0 => 2
  | k + 1 => s * (DWb L s W A k + (Eb L s k + 2) * W) + L * dwloc W A

/-- Width offset of the cost arguments inside the steps of level `k+1`. -/
def Ob (L s W A k : ℕ) : ℕ :=
  DWb L s W A k + (Eb L s k + 2) * W + Eb L s k + dwloc W A + 10

/-- Cumulative width offset of a level. -/
def Whb (L s W A : ℕ) : ℕ → ℕ
  | 0 => 2
  | k + 1 => DWb L s W A (k + 1) + Ob L s W A k + Whb L s W A k

/-- The local cost constant. -/
def alN (W A : ℕ) : ℕ := 10000000 * (W + A + 10) ^ 3

/-- Cost coefficients (per record, per unit `Xf`). -/
noncomputable def cb (L s W A m : ℕ) (CR τ : ℝ) : ℕ → ℝ
  | 0 => 100000
  | k + 1 => (s / W : ℝ) * cb L s W A m CR τ k +
      ((s + L : ℕ) : ℝ) * ((alN W A : ℝ) + CR * ((m : ℝ) ^ k + 1) ^ τ) + 1

/-- Cost unit per record. -/
def Xf (r : ℕ) (Ψ : ℕ → ℕ) (v : ℕ) : ℕ := r * (v + 1) + Ψ v

/-! ## Classes and hypotheses -/

/-- The gates of an interface: atom lists (at most `A` atoms) or `±1` role diagonals. -/
def GClass {W m s : ℕ} (J : Prop8IfaceR W m s) (A : ℕ) : Prop :=
  ∀ G, RStep.gate G ∈ J.net →
    (∃ as : List KFG.Atom, as.length ≤ A ∧ (∀ a ∈ as, a.Ok W ∧ a.h ≤ 1) ∧ KFG.NoTS as ∧
        G = KFG.gateMat W as) ∨
      (∃ U : Finset (Fin W), G = Matrix.diagonal fun w => if w ∈ U then -1 else 1)

/-- Positions of the bits of a row family at stride `K` (descending): `pos i + K i = e`, with
the top chunk `[e - e % K, e - e % K + K)` below `t`. -/
def AffK (K t : ℕ) {N : ℕ} (pos : Fin N → ℕ) : Prop :=
  ∃ e, (∀ i, pos i + K * i.val = e) ∧ e - e % K + K ≤ t

/-- The same for `Fin m × Fin (m^k)` (row `a`, bit `b` at index `b + m^k a`). -/
def AffK2 (K t m k : ℕ) (pos : Fin m × Fin (m ^ k) → ℕ) : Prop :=
  ∃ e, (∀ p, pos p + K * (p.2.val + m ^ k * p.1.val) = e) ∧ e - e % K + K ≤ t

/-- **Row additions** (abstract; e.g. Lemma 13): on the stride-`K` layouts of the bits
`Fin m × Fin (m^k)` (`AffK2`), as produced by the layer layout of stride `K`. -/
def RowHyp {Pp : ℕ} (Isym : List Sym) (iprog : Fin Pp → IStmt HS Pp) (K m r Z : ℕ)
    (Ψ : ℕ → ℕ) (CR τ : ℝ) : Prop :=
  ∀ (k : ℕ) (pos : Fin m × Fin (m ^ k) → ℕ) (t S0 B : ℕ), Function.Injective pos →
    (∀ p, pos p < t) → AffK2 K t m k pos → 1 ≤ S0 → 1 ≤ B →
    zsz ((16 * (B * 2 ^ t * S0 * r) : ℕ) : ℤ) ≤ Z →
    ∀ (i j : Fin m) (hij : i ≠ j), ∃ (code : ℕ → IStmt HS Pp) (cost : ℕ → ℕ),
      SImplI Isym iprog r (fun p => 2 ^ pos p * S0) (B * 2 ^ t * S0)
        (Pm i j : Butterfly.Op (Fin m × Fin (m ^ k))) code 0 0 cost ∧
      ∀ w, (cost w : ℝ) ≤ CR * ((B * 2 ^ t * S0 : ℕ) : ℝ) *
        ((r : ℝ) * (w + 1) * ((m : ℝ) ^ k + 1) ^ τ + Ψ w)

/-- The level claim: `tens Cm` on `m^k` bits at power-of-two strides. -/
def Lv {Pp : ℕ} (Isym : List Sym) (iprog : Fin Pp → IStmt HS Pp) (K W m r Z : ℕ) (Ψ : ℕ → ℕ)
    (k Ek DWk Wk : ℕ) (ck : ℝ) : Prop :=
  ∀ (pos : Fin (m ^ k) → ℕ) (t S0 B : ℕ), Function.Injective pos → (∀ i, pos i < t) →
    AffK K t pos →
    1 ≤ S0 → 1 ≤ B → W ^ k ∣ B → zsz ((16 * (B * 2 ^ t * S0 * r) : ℕ) : ℤ) ≤ Z →
    ∃ (code : ℕ → IStmt HS Pp) (e dw : ℕ) (cost : ℕ → ℕ),
      SImplI Isym iprog r (fun i => 2 ^ pos i * S0) (B * 2 ^ t * S0) (Butterfly.tens Cm : Butterfly.Op (Fin (m ^ k)))
        code e dw cost ∧ e ≤ Ek ∧ dw ≤ DWk ∧
      ∀ w, (cost w : ℝ) ≤ ck * ((B * 2 ^ t * S0 : ℕ) : ℝ) * (Xf r Ψ (w + Wk) : ℝ)

/-! ## Local cost bounds -/

section LocalBounds

variable {r Z : ℕ} {Ψ : ℕ → ℕ}

end LocalBounds

end NLogN.KFN

end

section
-- Source module: Solutions.KFNStep
/-!
# KFN: the steps of a node at level `k+1`, with bounds

Node bits `Fin m × Fin (m^k)` at strides `2^{pos p} S₀`, `W` roles of `n = B · 2^t · S₀`
records each.  **`step_ok`**: every step of `J.net` has an implementation (`RImpl`) with

* scale `≤ 1` (local) or `≤ E_k` (child),
* width `≤ dwloc W A` (local) or `≤ DW_k + (E_k + 2) W` (child),
* cost `≤ [c_k n Xf(w + Ŵ_k)]_{child} + (alN + CR Γ) (W n) Xf(w + O)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

section Steps

variable {Pp : ℕ} {Isym : List Sym} {iprog : Fin Pp → IStmt HS Pp} {K : ℕ}
variable {W m s : ℕ} (J : Prop8IfaceR W m s) {A r Z : ℕ} {Ψ : ℕ → ℕ}

/-- The number of records of a role. -/
abbrev nR (t S0 B : ℕ) : ℕ := B * 2 ^ t * S0

/-- Code of a phase step (other steps: `skip`). -/
def phStepC (W n r : ℕ) {k : ℕ} (pos : Fin m × Fin (m ^ k) → ℕ) (S0 t B : ℕ) :
    RStep (Fin W) m → RS
  | .phase w0 k₀ k₁ a => wrapC W (n * r) w0.val 0
      (phC pos S0 r t B (k₀ + k₁ * Fintype.card (Fin (m ^ k))) (fun p => a p.1))
  | _ => .skip

end Steps

end NLogN.KFN

end

section
-- Source module: Solutions.KFNPad
/-!
# KFN: running a program on a zero-padded stream

`padC L0 Q code`: park the `L0` entries of the stream (stack `qP`), push `Q` zeros, unpark;
run `code`; park the first `L0` entries, drop the next `Q`, unpark.

**`sImpl_pad`**: if `code` implements `T` on `np ≥ n` records (strides `st`), the padded program
implements `T` on `n` records (the padding records are zero, and the selected-bit structure on
`n` records never reads them: `selOpF_pad`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

section Lists

variable {ι : Type*} [Fintype ι] [DecidableEq ι] {M : Type*} [AddCommMonoid M] [Module ℂ M]

end Lists

/-! ## The padded program -/

/-- Pad by `Q` zero entries (the register part before the program). -/
def padPre (L0 Q : ℕ) : RS :=
  .seq (NLogN.mvK 0 qP NLogN.KFB.c1 (.lit (L0 : ℤ))) (.seq (zeroK 0 NLogN.KFB.c1 Q) (NLogN.mvK qP 0 NLogN.KFB.c1 (.lit (L0 : ℤ))))

/-- Unpad (the register part after the program). -/
def padPost (L0 Q : ℕ) : RS :=
  .seq (NLogN.mvK 0 qP NLogN.KFB.c1 (.lit (L0 : ℤ))) (.seq (dropK 0 NLogN.KFB.c1 Q) (NLogN.mvK qP 0 NLogN.KFB.c1 (.lit (L0 : ℤ))))

/-- Pad by `Q` zero entries, run `code`, unpad. -/
def padC {Pp : ℕ} (L0 Q : ℕ) (code : ℕ → IStmt HS Pp) : ℕ → IStmt HS Pp := fun w =>
  .seq (fcode (NLogN.KFN.padPre L0 Q)) (.seq (code w) (fcode (padPost L0 Q)))

/-- The cost of a move of `L0` entries of width `w`. -/
def cMv (L0 w : ℕ) : ℕ := zsz (L0 : ℤ) + 2 + cRep L0 + L0 * (250 * (w + 2))

/-- Cost of the padded program. -/
def cPadC (L0 Q Z w dw : ℕ) (cost : ℕ → ℕ) : ℕ :=
  2 * NLogN.KFN.cMv L0 w + cLp Q 8000 Z + cost w + 2 * NLogN.KFN.cMv L0 (w + dw) + cLp Q (2 * (w + dw) + 2) Z

end NLogN.KFN

end

section
-- Source module: Solutions.KFNPieces
/-!
# KFN: digit pieces

`pcs m k base R`: the pieces `(level j, offset o)` of the base-`m` digit decomposition of the
bit range `[base, base + R)`: `R / m^k` pieces of level `k`, then the remainder at the lower
levels.  Every bit of the range is covered by exactly one piece (`countP_pcs`); there are at most
`R / m^k + k m` pieces (`length_pcs`).

Operators: `subOp g (tens A) = tensF (indicator)` (`subOp_tens`); implementations of
`tensF Fₐ` compose to `tensF (∏ₐ Fₐ)` (`sImpl_tensF_list`); a family of indicators of a
partition multiplies to `A` (`prod_ind`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

/-! ## The pieces -/

/-- Pieces `(level, offset)` of the range `[base, base + R)`. -/
def pcs (m : ℕ) : ℕ → ℕ → ℕ → List (ℕ × ℕ)
  | 0, base, R => (List.range R).map fun c => (0, base + c)
  | k + 1, base, R => (List.range (R / m ^ (k + 1))).map (fun c => (k + 1, base + c * m ^ (k + 1))) ++
      pcs m k (base + R / m ^ (k + 1) * m ^ (k + 1)) (R % m ^ (k + 1))

/-- Piece `p` covers bit `i`. -/
def covers (m : ℕ) (i : ℕ) (p : ℕ × ℕ) : Bool := decide (p.2 ≤ i ∧ i < p.2 + m ^ p.1)

/-! ## Operators -/

section Ops

variable {ι κ : Type*} [Fintype ι] [DecidableEq ι] [Fintype κ] [DecidableEq κ]

end Ops

section Impl

variable {Pp : ℕ} {Isym : List Sym} {iprog : Fin Pp → IStmt HS Pp}
variable {ι : Type*} [Fintype ι] [DecidableEq ι] {r n : ℕ} {st : ι → ℕ}

end Impl

end NLogN.KFN

end

section
-- Source module: Solutions.KFNTop
/-!
# KFN: `tens Cm` on any number of bits at power-of-two strides

**`cLow`**: for `D'` bits at strides `2^{pos i} S₀` (`pos` injective, `< t`) on `B · 2^t · S₀`
records with `W^k ∣ B` and `D' < m^{k+1}`, the digit pieces (`pcs m k 0 D'`), each a node of
level `≤ k` (`lv_all`), implement `tens Cm`; at most `(k+1) m` pieces, so scale
`≤ (k+1) m E_k`, width `≤ (k+1) m DW_k`, cost `≤ ((k+1) m c_k + 1) n Xf(w + Ŵ)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

/-! ## Monotonicity of the constants -/

section Mono

variable {L s W A m : ℕ} {CR τ : ℝ}

end Mono

/-! ## The low bits -/

section Low

variable {Pp : ℕ} {Isym : List Sym} {iprog : Fin Pp → IStmt HS Pp} {K : ℕ}
variable {W m s : ℕ} (J : Prop8IfaceR W m s) {A r Z : ℕ} {Ψ : ℕ → ℕ} {CR τ : ℝ}

/-- The embedding of a piece. -/
def pcE {D' j o : ℕ} (h : o + m ^ j ≤ D') : Fin (m ^ j) → Fin D' := fun i => ⟨o + i, by omega⟩

end Low

end NLogN.KFN

end

section
-- Source module: Solutions.KFNLayer
/-!
# KFN: the recursive `C`-layer on a layer layout

**`cLayer`**: on the layer layout `(P, K, D, S, ρ)` (bits `Fin D` at strides `stT`), with
`k = log_m D` and `q₀ = k · L₂` (`W ≤ 2^{L₂}`):

* if `D ≤ q₀`: `D` kernels `C`;
* otherwise: kernels on the top `q₀` bits (so that the `q₀ K` top address bits join the outer
  count `B = P 2^{K q₀} ≥ W^k`), then the low `D - q₀` bits by digit pieces (`cLow`) on the
  stream zero-padded to a multiple `B' ≤ 2B` of `W^k` blocks (`sImpl_pad`).

Scale `E ≤ D + (k+1) m E_k`, width `≤ 2D + (k+1) m DW_k`, cost
`≤ (2 ((k+1) m c_k + 1) + 10^6 (D + 2)) · N · Xf(w + 2D + (k+1) m DW_k + Ŵ_k)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

section Bounds

variable {r : ℕ} {Ψ : ℕ → ℕ} {Z : ℕ}

end Bounds

section Layer

variable {W m s : ℕ} (J : Prop8IfaceR W m s) {A r Z : ℕ} {Ψ : ℕ → ℕ} {CR τ : ℝ}

/-- The top `q₀` bits. -/
def topL (D q0 : ℕ) (h : q0 ≤ D) : List (Fin D) := (List.finRange q0).map (Fin.castLE h)

/-- The low bits. -/
def lowE {D q0 : ℕ} (h : q0 ≤ D) : Fin (D - q0) → Fin D := fun i => ⟨q0 + i, by omega⟩

end Layer

end NLogN.KFN

end

section
-- Source module: Solutions.KFMErr
/-!
# KFM 14: the error chain of the computed pipeline (§9.7.1), abstract form

Arrays are functions `ι → ℂ` with the sup norm; `hp = 2^{-p}`. Every computed procedure is
an arbitrary function with an error bound on grid inputs (`Grid p f`: `Q_p f = f`), so the IL
denotations (and `P18IL`) plug in later.

* `QA p f`: `Q_p` on every entry; `QA_grid`, `norm_QA_le`, `norm_QA_sub_le` (`≤ √2 hp`).
* **`phase_err`**: `‖Q_p(ẽ·x) - e·x‖ ≤ 4hp`, `‖Q_p(ẽ·x)‖ ≤ ‖x‖` (unit phases, `|ẽ-e| ≤ 2hp`).
* `BilC M`: bilinear contraction; **`bil_err`**.
* **`mc_err`**: the computed complex convolution `Q_p(θ̃⁻ · M̃r(Q_p(θ̃⁺f), Q_p(θ̃⁺g)))` is within
  `(Er + 12) hp` of `θ⁻ · Mr(θ⁺ f, θ⁺ g)` (one operand of norm `≤ .7`).
* **`chirp_err`**: `Q_p(ã · M̃C(ã, Q_p(ã · x)))` is within `(Ec + 10) hp` of `a · MC(a, a · x)`.
* **`source_err`**: `2^γ B̃₀(C̃(Ã u))` is within `2^γ (εA + Ec + εB) hp` of `2^γ B₀(QF(A u))`.
* **`pipeline_err`**: with `z = Q_p(F̃u · F̃v)` and `F̃⁺ = conj ∘ F̃ ∘ conj`,
  `‖F̃⁺ z - G(Fu · Fv)‖ ≤ (3Es + 2) hp`.
-/

namespace NLogN.KFM

open Complex

noncomputable section

/-- The grid step `2^{-p}`. -/
abbrev hp (p : ℕ) : ℝ := (2 : ℝ)⁻¹ ^ p

section Arr

variable {ι : Type*} [Fintype ι]

/-- `Q_p` on every entry. -/
def QA (p : ℕ) (f : ι → ℂ) : ι → ℂ := fun i => Butterfly.Qp p (f i)

/-- Grid arrays: fixed by `Q_p`. -/
def Grid (p : ℕ) (f : ι → ℂ) : Prop := QA p f = f

/-- Entrywise conjugation. -/
def cnj (f : ι → ℂ) : ι → ℂ := fun i => starRingEnd ℂ (f i)

end Arr

/-! ## Bilinear contractions -/

section Bil

variable {ι κ : Type*} [Fintype ι] [Fintype κ]

/-- A bilinear contraction on arrays. -/
structure BilC (M : (ι → ℂ) → (ι → ℂ) → (κ → ℂ)) : Prop where
  sub : ∀ a a' b b', M a' b' - M a b = M (a' - a) b' + M a (b' - b)
  le : ∀ a b, ‖M a b‖ ≤ ‖a‖ * ‖b‖

/-- A linear contraction on arrays. -/
structure LinC (L : (ι → ℂ) → (κ → ℂ)) : Prop where
  sub : ∀ a b, L a - L b = L (a - b)
  le : ∀ a, ‖L a‖ ≤ ‖a‖

end Bil

/-! ## The computed complex convolution `M̃_C` (§9.4.1) -/

section MC

variable {ι : Type*} [Fintype ι]

/-- `M̃_C(f, g) = Q_p(θ̃⁻ · M̃r(Q_p(θ̃⁺ f), Q_p(θ̃⁺ g)))`. -/
def mcT (p : ℕ) (Mrt : (ι → ℂ) → (ι → ℂ) → (ι → ℂ)) (θt θit : ι → ℂ) (f g : ι → ℂ) : ι → ℂ :=
  QA p (θit * Mrt (QA p (θt * f)) (QA p (θt * g)))

/-- Phase data: exact unit phases `θ`, `θ⁻`, grid approximations within `2hp`, all of modulus `≤ 1`. -/
structure Phases (p : ℕ) (θ θt : ι → ℂ) : Prop where
  e : ∀ i, ‖θ i‖ ≤ 1
  t : ∀ i, ‖θt i - θ i‖ ≤ 2 * hp p
  t1 : ∀ i, ‖θt i‖ ≤ 1

end MC

/-! ## The computed chirp transform (§9.5) -/

section Chirp

variable {ι : Type*} [Fintype ι]

/-- `C̃(x) = Q_p(ã · M̃C(b̃, Q_p(ã · x)))` (kernel `b̃`, endpoint phases `ã`). -/
def chirpT (p : ℕ) (MCt : (ι → ℂ) → (ι → ℂ) → (ι → ℂ)) (bt at' x : ι → ℂ) : ι → ℂ :=
  QA p (at' * MCt bt (QA p (at' * x)))

end Chirp

/-! ## The source transform (§9.5.1) -/

section Source

variable {ι κ : Type*} [Fintype ι] [Fintype κ]

end Source

/-! ## The pipeline (§9.7.2) -/

section Pipe

variable {κ : Type*} [Fintype κ]

/-- `F̃⁺(z) = conj F̃(conj z)`. -/
def plusT (Ft : (κ → ℂ) → (κ → ℂ)) (z : κ → ℂ) : κ → ℂ := cnj (Ft (cnj z))

end Pipe

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFMErr3
/-!
# KFM 14: the exact side of the error chain

The exact maps of §9.4–9.7 as `LinC` / `BilC` instances, and the identities `chain_close` needs.

* `avg_le`: a normalized sum of terms of modulus `≤ C` has modulus `≤ C`.
* **`linC_clm`** (CLMs of norm `≤ 1`), **`linC_RF`** (`R F_s`), **`linC_QF`** (`Q F_t`).
* `mrX t f g`: the normalized ring convolution on the box `∏_{i ≤ D} ℤ/t_i`: cyclic on the first
  `D` axes, negacyclic on the last (coefficient) axis, divided by `T = ∏ t_i`; **`bilC_mrX`**.
* **`mc_exact`**: twisting the last coordinate by `ζ^k` (`ζ^{t_D} = -1`) turns `mrX` into the
  normalized cyclic convolution `(f ∗ g)/T` (§9.4, pointwise form of `RTw.MC_eq`).
* **`chirp_exact`**: `ā · ((a ∗ (ā · x))/T) = Q F_t x` (Lemma 21, (36)).
* **`final_exact`**: `2^{2b+4} S · (S · conj RF_s(conj(RF_s u · RF_s v))) = ū ∗ v̄` (`pipeline_exact`).
* **`chain_concrete`**: `chain_close` with all exact maps instantiated (Lemma 19/20 witnesses
  `tensCLM (kerA M.A)`, `tensCLM (kerB M.B)`); only the computed procedures remain as hypotheses.
-/

namespace NLogN.KFM

open Complex IntMul.OpenAI.RT IntMul.OpenAI.RC

noncomputable section

/-! ## Averages -/

variable {d : ℕ}

/-! ## The twisted ring convolution -/

variable {D : ℕ}

/-- The normalized ring convolution on `∏_{i ≤ D} ℤ/t_i`: cyclic on the first `D` axes and
negacyclic on the last axis, divided by `T = ∏ t_i`. -/
def mrX (t : Fin (D + 1) → ℕ) [∀ i, NeZero (t i)] (f g : ((i : Fin (D + 1)) → ZMod (t i)) → ℂ) :
    ((i : Fin (D + 1)) → ZMod (t i)) → ℂ := fun j =>
  (∏ i, (t i : ℂ))⁻¹ * ∑ k, f k * g (j - k) *
    (if (k (Fin.last D)).val ≤ (j (Fin.last D)).val then 1 else -1)

/-- The twist `θ_j = ζ^{j_D}` of the last coordinate. -/
def twX (t : Fin (D + 1) → ℕ) (ζ : ℂ) : ((i : Fin (D + 1)) → ZMod (t i)) → ℂ :=
  fun j => ζ ^ (j (Fin.last D)).val

/-! ## The chain with the exact maps instantiated -/

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFNP14
/-!
# KFN: Proposition 14 in IL around the recursive `C`-layer

* `p14WithI P K D S ρ r E Cc`: the Proposition 14 wrapper `S_D ; C ; S_D ; b^D ; Q_p` around a
  `C`-layer **code family** `Cc : ℕ → IStmt HS Pp` (an item-language program, e.g. calling the
  Lemma 12/13 routines under a fixed `iprog`) of scale `E` (the truncation removes `E + ⌈D/2⌉`
  bits instead of `D + ⌈D/2⌉`).
* **`p14WithI_iruns`**: for every such implementation (`SImplI` on the layer layout), the
  output represents `prop14Proc … (tens Cm) z = Q_p(H0^{⊗D} z)`, disk, error `< √2·2^{-p}`;
  cost `kC (w + 2D) + cWrapE …`.
* **`prop14_il`**: with the recursive `C`-layer `cLayer`, conditional only on the row-addition
  hypothesis `RowHyp` (an `IStmt` code family under the fixed `iprog`) and the gate class.
* `gclass_of`, **`prop14_il_complex`**: the interface `P8Lift.complexIface` (`complexIface_gate_class`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

/-- Cost of the wrapper (without the `C`-layer). -/
def cWrapE (D E N Z w dC : ℕ) : ℕ :=
  (D * ck N Z (w + 2 * D) + 1) + (D * ck N Z (w + (2 * D + dC) + 2 * D) + 1) +
    ck N Z (w + (2 * D + dC + 2 * D)) +
    cTrunc (w + 4 * D + dC + 2) (E + (D / 2 + D % 2)) (2 * N) Z

/-- Proposition 14 around a `C`-layer code family `Cc` of scale `E`. -/
def p14WithI {Pp : ℕ} (P K D S ρ r E : ℕ) (Cc : ℕ → IStmt HS Pp) : ℕ → IStmt HS Pp := fun w =>
  .seq (.seq (.seq (.seq (fcode (kersP K D S ρ P r MS ((List.finRange D).map Fin.val)))
    (Cc (w + 2 * D)))
    (fcode (kersP K D S ρ P r MS ((List.finRange D).map Fin.val))))
    (fcode (kerP K D S ρ P r 0 (MD D))))
    (fcode (truncP (E + (D / 2 + D % 2)) (2 * (vol (Butterfly.layerShape P K D S) * r))))

section Wrap

variable {Pp : ℕ} {Isym : List Sym} {iprog : Fin Pp → IStmt HS Pp} {P K D S ρ : ℕ}

end Wrap

/-! ## Gate classes of finite networks -/

/-! ## Proposition 14 in IL with the recursive `C`-layer -/

section Main

variable {Pp : ℕ} {Isym : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- The exponent of Lemma 13: `τ = 1 - 2^{-50}`. -/
noncomputable def τ50 : ℝ := 1 - (2 : ℝ)⁻¹ ^ 50

end Main

end NLogN.KFN

end

section
-- Source module: Solutions.KFNClosed
/-!
# KFN: closed bounds for the recursion constants

With `G = s + L + W + A + 10` and `R = max (s/W, m^τ, 1)`:

* `Eb_le`: `E_k ≤ (k+1) G^k`; `DWb_le`: `DW_k ≤ 4 (k+1)^2 G^{k+2}`;
  `Whb_le`: `Ŵ_k ≤ 20 (k+1)^3 G^{k+2}`;
* `cb_le`: `c_k ≤ (10^5 + k (a + 2b)) R^k` with `a = (s+L) alN + 1`, `b = (s+L) CR`;
* `pow_log_le`: `R^{log_m D} ≤ D^{log_m R}`;
* for `complexIface` (`m = 10^6`, `s_c / W_c < 10^6`, `τ = 1 - 2^{-50}`): `λ = log_m R < 1`
  (`lam_complex_lt_one`), and **`prop14_il_complex_closed`**: Proposition 14 in IL with the
  `C`-layer cost `≤ (2 ((k+1) m (10^5 + k(a+2b)) D^λ + 1) + 10^6 (D+2)) · vol ·
  Xf(w + 2D + (k+1) m · 4(k+1)^2 G^{k+2} + 20 (k+1)^3 G^{k+2} + 2)`, `k = ⌊log_m D⌋`,
  conditional only on `RowHyp` (row additions as `IStmt` code families under a fixed `iprog`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFN

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly

/-- The base of the width bounds. -/
def Gb (L s W A : ℕ) : ℕ := s + L + W + A + 10

section Nat

variable (L s W A : ℕ)

end Nat

/-! ## The cost coefficients -/

section Real

variable (L s W A m : ℕ) (CR τ : ℝ)

/-- The growth ratio of the cost coefficients. -/
noncomputable def Rb : ℝ := max (max ((s : ℝ) / W) ((m : ℝ) ^ τ)) 1

end Real

/-! ## `complexIface` -/

section Closed

variable {Pp : ℕ} {Isym : List Sym} {iprog : Fin Pp → IStmt HS Pp}

end Closed

end NLogN.KFN

end

section
-- Source module: Solutions.KFMPhase
/-!
# KFM 2: Proposition 14 with the phase diagonals `S_D` in one pass each

`KFN.p14WithI` implements `S_D = diag i^{x_1+⋯+x_D}` as `D` kernel passes (`kersP … MS`), so one
simultaneous butterfly round costs `Θ(D · N · w)` even before the `C`-layer.  Lemma 16 runs `≈ ℓ`
such rounds with `D ≈ d`, giving `ℓ d V ≈ V log T`: a full logarithmic factor, which `Ncost`
(`ℓ d^{λ'}` for this row) cannot absorb (MACHINE-PLAN.md, item 4, issue (A)).

Here `S_D` is the KFN phase pass `phL` (`KFN.sImpl_phase`): one pass with an exponent table built
by doubling, `O(N (w+1))` plus `O((K D + P) Z)` descriptor work, independent of `D`.

* `posL K D ρ i = ρ + K (D-1-i)`: the bit position of the selected bit of chunk `i`;
  `stT_posL`: the record stride is `2^{posL i} · S`.
* **`sImpl_SD`**: `sdP` (one `phL` pass) implements `SD` on the layer layout, scale and width
  unchanged, at cost `cSD … w = cPhL w N (K D) P Z`.
* `p14Ph`: the wrapper `S_D ; C ; S_D ; b^D ; Q_p` with one-pass phases.
* **`p14Ph_iruns`**: same conclusions as `KFN.p14WithI_iruns` (output `prop14Proc … (tens Cm)`,
  `GB (p+1)`, error `< √2·2^{-p}` against `layerH`), at cost `kC w + cWrapPh …`.
* **`cWrapPh_le`**: `cWrapPh ≤ 10^6 (w + dC + E + D + 10) N + 10^6 (K D + P + 10)(Z + 2)`:
  the data term carries **no factor `D`** (only the additive width `E + D`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFM

open IntMul RS AModel KIL KFD KFB KFG KFH KFN Complex P8Lift Butterfly

/-! ## 1. `S_D` as one phase pass -/

/-- Position (in the binary record index below `P`) of the selected bit of chunk `i`. -/
def posL (K D ρ : ℕ) (i : Fin D) : ℕ := ρ + K * (D - 1 - i)

/-- The `S_D` pass: `phL` with unit exponent on every selected bit. -/
def sdP (P K D S ρ r : ℕ) : RS := phL (S * r) (K * D) P 0 (aEof (posL K D ρ) (fun _ => 1))

/-- Cost of one `S_D` pass on `N = vol · r` Gaussian entries of width `w`. -/
def cSD (P K D S r Z w : ℕ) : ℕ := cPhL w (P * (S * r * 2 ^ (K * D))) (K * D) P Z

/-! ## 2. The wrapper -/

/-- Proposition 14 with one-pass phases: `S_D ; C ; S_D ; b^D ; Q_p`. -/
def p14Ph {Pp : ℕ} (P K D S ρ r E dC : ℕ) (Cc : ℕ → IStmt HS Pp) : ℕ → IStmt HS Pp := fun w =>
  .seq (seqI (seqI (seqI (fun _ => fcode (sdP P K D S ρ r)) Cc 0)
      (fun _ => fcode (sdP P K D S ρ r)) (0 + dC))
    (fun _ => fcode (kerP K D S ρ P r 0 (MD D))) (0 + dC + 0) w)
  (fcode (truncP (E + (D / 2 + D % 2)) (2 * (vol (Butterfly.layerShape P K D S) * r))))

/-- Wrapper cost (everything except the `C`-layer). -/
def cWrapPh (P K D S r E Z w dC : ℕ) : ℕ :=
  cSD P K D S r Z w + cSD P K D S r Z (w + dC) +
    ck (vol (Butterfly.layerShape P K D S) * r) Z (w + dC) +
    cTrunc (w + dC + 2) (E + (D / 2 + D % 2)) (2 * (vol (Butterfly.layerShape P K D S) * r)) Z

section Wrap

variable {Pp : ℕ} {Isym : List Sym} {iprog : Fin Pp → IStmt HS Pp} {P K D S ρ : ℕ}

end Wrap

/-! ## 3. The wrapper cost has no factor `D` on the data -/

end NLogN.KFM

end

section
-- Source module: Solutions.KFMSetup
/-!
# KFM 14: the setup arithmetic (§9.2), in the integer form the machine computes

* **`floor_root_eq`/`floor_root_spec`**: `⌊N^{1/m}⌋ = x` iff `x^m ≤ N < (x+1)^m`, so
  **`dP_spec`** (`m = 2^75`, of `b`) and **`KP_spec`** (`m = 2^56`, of `d`) are exact integer
  comparisons (computed by saturating powers and binary search).
* **`ceil_root_eq`**, **`αP_spec`**: `α = αP n` iff `(α-1)^4 < 12 d² b ≤ α^4`.
* **`LP_spec`/`LP_eq`**: `L = LP n` iff `4n ≤ 2^L b` and (`L = 0` or `4n > 2^{L-1} b`).
* **`srcS d d' r i`**: the explicit source lengths of `RCop.coprime_params`
  (`1 + Q m_i`, `Q = 2F(⌊r/(2FD)⌋+1)`, `F = D!`, `D = 256 d²`), and **`srcS_ok`**: odd, pairwise
  coprime, in the windows `((1-2η)t_i, (1-η)t_i]`; **`srcS_ok_n`** for the parameters of `n ≥ n₀`.
* **`mu_spec`**: `μ = (P mod s)⁻¹` from `Nat.gcdA` (extended Euclid) satisfies `μ P ≡ 1 (mod s)`.
-/

namespace NLogN.KFM

open IntMul.OpenAI.Glue IntMul.OpenAI.RCost

noncomputable section

/-! ## Integer roots -/

/-! ## Explicit source lengths -/

/-- `Q = 2F(⌊r/(2FD)⌋+1)`, `F = D!`, `D = 256 d²`. -/
def srcQ (d r : ℕ) : ℕ :=
  2 * (256 * d ^ 2).factorial * (r / (2 * (256 * d ^ 2).factorial * (256 * d ^ 2)) + 1)

/-- The top index `⌊((4d-1) t - 4d)/(4dQ)⌋` of a group. -/
def srcM (d r t : ℕ) : ℕ := ((4 * d - 1) * t - 4 * d) / (4 * d * srcQ d r)

/-- The source length of axis `i` (`t_i = r/2` for `i < d'`, `r` otherwise). -/
def srcS (d d' r i : ℕ) : ℕ :=
  1 + srcQ d r * (if i < d' then srcM d r (r / 2) - i else srcM d r r - (i - d'))

/-! ## Modular inverses -/

/-- The inverse `μ = gcdA(P, s) mod s` of `P` modulo `s` (extended Euclid). -/
def muInv (P s : ℕ) : ℕ := ((Nat.gcdA P s) % (s : ℤ)).toNat

end

end NLogN.KFM

end