PacketCostModels.lean

1 module
2 public import Mathlib.Data.Fintype.Card
3 public import Mathlib.Data.Fintype.Sigma
4 public import Mathlib.Analysis.Real.Sqrt
5 public import Mathlib.Tactic
6 
7 /-!
8 Author: Komorebi P.
9 Paper VIII 3.1, 3.2, 3.4, 3.6, 3.8, 3.9, 3.10, 4.2.
10 Counts come from explicit finite address/event types. A modeled multiply-accumulate
11 costs two scalar operations; this is not a measurement of GPU instructions/time.
12 The thin output product accumulates into the diagonal result.
13 -/
14 public section
15 set_option autoImplicit false
16 namespace KomorebiPaperVIII
17 open scoped BigOperators
18 
19 /-- Scalar storage addresses for a finite collection of rectangular blocks. -/
20 abbrev PacketEntry (E : Type*) (rows cols : E → ℕ) :=
21   (e : E) × (Fin (rows e) × Fin (cols e))
22 
23 /-- Two operations per scalar coefficient and right-hand side, by declared model. -/
24 abbrev PacketOperation (E : Type*) (rows cols : E → ℕ) (q : ℕ) :=
25   PacketEntry E rows cols × Fin q × Bool
26 
27 theorem packet_storage_count {E : Type*} [Fintype E] (rows cols : E → ℕ) :
28     Fintype.card (PacketEntry E rows cols) = ∑ e, rows e * cols e := by
29   simp [PacketEntry, Fintype.card_sigma, Fintype.card_prod]
30 
31 theorem packet_application_count {E : Type*} [Fintype E] (rows cols : E → ℕ) (q : ℕ) :
32     Fintype.card (PacketOperation E rows cols q) = 2 * q * ∑ e, rows e * cols e := by
33   simp [PacketOperation, Fintype.card_prod]
34   ring
35 
36 theorem dense_storage_count (n : ℕ) : Fintype.card (Fin n × Fin n) = n ^ 2 := by
37   simp [pow_two]
38 
39 theorem dense_application_count (n q : ℕ) :
40     Fintype.card ((Fin n × Fin n) × Fin q × Bool) = 2 * n ^ 2 * q := by
41   simp [Fintype.card_prod]
42   ring
43 
44 theorem equal_diagonal_storage {I : Type*} [Fintype I] (b : ℕ) :
45     Fintype.card (PacketEntry I (fun _ => b) (fun _ => b)) = (Fintype.card I * b) * b := by
46   simp [mul_assoc]
47 
48 theorem equal_diagonal_work {I : Type*} [Fintype I] (b q : ℕ) :
49     Fintype.card (PacketOperation I (fun _ => b) (fun _ => b) q) =
50       2 * q * (Fintype.card I * b) * b := by
51   simp
52   ring
53 
54 /-- The complete BDLR scalar representation, including both thin factors. -/
55 abbrev BDLREntry (I : Type*) (width : I → ℕ) (n r : ℕ) :=
56   PacketEntry I width width ⊕ ((Fin n × Fin r) ⊕ (Fin n × Fin r))
57 
58 theorem bdlr_storage_count {I : Type*} [Fintype I] (width : I → ℕ) (n r : ℕ) :
59     Fintype.card (BDLREntry I width n r) = (∑ i, width i ^ 2) + 2 * n * r := by
60   simp [BDLREntry, pow_two]
61   ring
62 
63 /-- The three stages are diagonal application and two thin products. -/
64 abbrev BDLROperation (I : Type*) (width : I → ℕ) (n r q : ℕ) :=
65   PacketOperation I width width q ⊕
66     (((Fin n × Fin r) × Fin q × Bool) ⊕ ((Fin n × Fin r) × Fin q × Bool))
67 
68 theorem bdlr_application_count {I : Type*} [Fintype I]
69     (width : I → ℕ) (n r q : ℕ) :
70     Fintype.card (BDLROperation I width n r q) =
71       2 * q * ((∑ i, width i ^ 2) + 2 * n * r) := by
72   simp [BDLROperation, pow_two]
73   ring
74 
75 theorem equal_bdlr_storage {I : Type*} [Fintype I] (b r : ℕ) :
76     Fintype.card (BDLREntry I (fun _ => b) (Fintype.card I * b) r) =
77       (Fintype.card I * b) * (b + 2 * r) := by
78   simp
79   ring
80 
81 theorem equal_bdlr_work {I : Type*} [Fintype I] (b r q : ℕ) :
82     Fintype.card (BDLROperation I (fun _ => b) (Fintype.card I * b) r q) =
83       2 * q * (Fintype.card I * b) * (b + 2 * r) := by
84   simp
85   ring
86 
87 theorem equal_block_reduction (n b : ℝ) (hn : n ≠ 0) (hb : b ≠ 0) :
88     n ^ 2 / (n * b) = n / b := by
89   field_simp
90 
91 theorem equal_bdlr_reduction (n b r q : ℝ) (hn : n ≠ 0) (hq : q ≠ 0)
92     (hbr : b + 2 * r ≠ 0) :
93     n ^ 2 / (n * (b + 2 * r)) = n / (b + 2 * r) ∧
94     (2 * q * n ^ 2) / (2 * q * n * (b + 2 * r)) = n / (b + 2 * r) := by
95   constructor <;> field_simp
96 
97 /-- One block read, one source read, and target read plus write for each edge. -/
98 abbrev StreamEvent (E : Type*) (rows cols : E → ℕ) (q : ℕ) :=
99   (e : E) × ((Fin (rows e) × Fin (cols e)) ⊕
100     ((Fin (cols e) × Fin q) ⊕ ((Fin (rows e) × Fin q) × Bool)))
101 
102 theorem streaming_value_count {E : Type*} [Fintype E] (rows cols : E → ℕ) (q : ℕ) :
103     Fintype.card (StreamEvent E rows cols q) =
104       ∑ e, (rows e * cols e + cols e * q + 2 * rows e * q) := by
105   simp only [StreamEvent, Fintype.card_sigma, Fintype.card_sum, Fintype.card_prod,
106     Fintype.card_fin, Fintype.card_bool]
107   apply Finset.sum_congr rfl
108   intro e _
109   ring
110 
111 theorem streaming_byte_count {E : Type*} [Fintype E]
112     (rows cols : E → ℕ) (q bytesPerScalar : ℕ) :
113     Fintype.card (Fin bytesPerScalar × StreamEvent E rows cols q) =
114       bytesPerScalar * ∑ e, (rows e * cols e + cols e * q + 2 * rows e * q) := by
115   rw [Fintype.card_prod, Fintype.card_fin, streaming_value_count]
116 
117 /-- Both physical capacity inequalities are hypotheses, not conclusions of the model. -/
118 theorem roofline_time_bound (F B P W T : ℝ) (hP : 0 < P) (hW : 0 < W)
119     (hF : F ≤ P * T) (hB : B ≤ W * T) : max (F / P) (B / W) ≤ T := by
120   apply max_le
121   · apply (div_le_iff₀ hP).2
122     nlinarith
123   · apply (div_le_iff₀ hW).2
124     nlinarith
125 
126 theorem roofline_work_rate (F B P W T : ℝ) (hF : 0 ≤ F) (hB : 0 < B)
127     (hT : 0 < T) (hp : F ≤ P * T) (hw : B ≤ W * T) :
128     F / T ≤ min P (W * (F / B)) := by
129   apply le_min
130   · apply (div_le_iff₀ hT).2
131     nlinarith
132   · apply (div_le_iff₀ hT).2
133     have hprod := mul_le_mul_of_nonneg_left hw hF
134     have hclear : F * B ≤ (W * F) * T := by nlinarith
135     calc
136       F ≤ (W * F * T) / B := (le_div_iff₀ hB).2 hclear
137       _ = W * (F / B) * T := by ring
138 
139 theorem dense_capacity_threshold (n M s : ℝ) (hn : 0 ≤ n) (hM : 0 ≤ M) (hs : 0 < s) :
140     n ^ 2 * s ≤ M ↔ n ≤ Real.sqrt (M / s) := by
141   rw [Real.le_sqrt hn (div_nonneg hM (le_of_lt hs)), le_div_iff₀ hs]
142 
143 theorem bdlr_capacity_threshold (n b r s M : ℕ) :
144     n * (b + 2 * r) * s ≤ M ↔ s * (n * b + 2 * n * r) ≤ M := by
145   rw [show n * (b + 2 * r) * s = s * (n * b + 2 * n * r) by ring]
146 
147 /-- An explicit scalar addition for each entry is a concrete O(n^2) remainder. -/
148 abbrev MaterializeOperation (n r : ℕ) :=
149   ((Fin n × Fin n) × Fin r × Bool) ⊕ (Fin n × Fin n)
150 
151 theorem materialization_work_count (n r : ℕ) :
152     Fintype.card (MaterializeOperation n r) = 2 * n ^ 2 * r + n ^ 2 := by
153   simp [MaterializeOperation]
154   ring
155 
156 theorem materialization_quadratic_remainder (n r : ℕ) :
157     Fintype.card (MaterializeOperation n r) - 2 * n ^ 2 * r ≤ n ^ 2 := by
158   rw [materialization_work_count]
159   omega
160 
161 /-- Sparse indices are excluded exactly as in the source's scalar-value model. -/
162 abbrev SparseLowRankEntry (z n r : ℕ) := Fin z ⊕ ((Fin n × Fin r) ⊕ (Fin n × Fin r))
163 
164 theorem sparse_low_rank_storage_count (z n r : ℕ) :
165     Fintype.card (SparseLowRankEntry z n r) = z + 2 * n * r := by
166   simp [SparseLowRankEntry]
167   ring
168 
169 theorem sparse_low_rank_work_count (z n r q : ℕ) :
170     Fintype.card (SparseLowRankEntry z n r × Fin q × Bool) = 2 * q * (z + 2 * n * r) := by
171   simp
172   ring
173 
174 end KomorebiPaperVIII
175 
176 -- FINAL AXIOM AUDITS
177 #print axioms KomorebiPaperVIII.packet_storage_count
178 #print axioms KomorebiPaperVIII.packet_application_count
179 #print axioms KomorebiPaperVIII.dense_storage_count
180 #print axioms KomorebiPaperVIII.dense_application_count
181 #print axioms KomorebiPaperVIII.equal_diagonal_storage
182 #print axioms KomorebiPaperVIII.equal_diagonal_work
183 #print axioms KomorebiPaperVIII.bdlr_storage_count
184 #print axioms KomorebiPaperVIII.bdlr_application_count
185 #print axioms KomorebiPaperVIII.equal_bdlr_storage
186 #print axioms KomorebiPaperVIII.equal_bdlr_work
187 #print axioms KomorebiPaperVIII.equal_block_reduction
188 #print axioms KomorebiPaperVIII.equal_bdlr_reduction
189 #print axioms KomorebiPaperVIII.streaming_value_count
190 #print axioms KomorebiPaperVIII.streaming_byte_count
191 #print axioms KomorebiPaperVIII.roofline_time_bound
192 #print axioms KomorebiPaperVIII.roofline_work_rate
193 #print axioms KomorebiPaperVIII.dense_capacity_threshold
194 #print axioms KomorebiPaperVIII.bdlr_capacity_threshold
195 #print axioms KomorebiPaperVIII.materialization_work_count
196 #print axioms KomorebiPaperVIII.materialization_quadratic_remainder
197 #print axioms KomorebiPaperVIII.sparse_low_rank_storage_count
198 #print axioms KomorebiPaperVIII.sparse_low_rank_work_count

Back to source controls