/home/noah/src/trueno/src/tiling.rs
Line | Count | Source |
1 | | //! Tiling Compute Blocks (TCB) - Work Partitioning for High-Performance Kernels |
2 | | //! |
3 | | //! TCBs represent the fundamental unit of work partitioning within `ComputeBrick` kernels. |
4 | | //! While a `ComputeBrick` defines a logical operation (e.g., Q4_K MatMul), a TCB defines |
5 | | //! the physical execution strategy—how data is partitioned across the memory hierarchy. |
6 | | //! |
7 | | //! # Architecture |
8 | | //! |
9 | | //! Tiling occurs at three levels: |
10 | | //! 1. **Macro-Tile (L3/Global Memory)**: Partitioning across CPU sockets or GPU SMs |
11 | | //! 2. **Midi-Tile (L2/Shared Memory)**: Partitioning within a thread block or Rayon task |
12 | | //! 3. **Micro-Tile (Registers)**: Smallest unit processed by SIMD or CUDA warps |
13 | | //! |
14 | | //! # Scientific Basis |
15 | | //! |
16 | | //! Per Lam et al. (1991), optimal tile sizes balance: |
17 | | //! - Cache capacity (tile must fit in target cache level) |
18 | | //! - Reuse factor (maximize arithmetic intensity) |
19 | | //! - Alignment (match hardware vector width) |
20 | | //! |
21 | | //! # Example |
22 | | //! |
23 | | //! ```rust |
24 | | //! use trueno::tiling::{TcbGeometry, TcbLevel, TilingConfig}; |
25 | | //! |
26 | | //! // GPU Q4_K MatVec tiling |
27 | | //! let config = TilingConfig::gpu_q4k_matvec(); |
28 | | //! assert_eq!(config.macro_tile.m, 1); |
29 | | //! assert_eq!(config.macro_tile.k, 256); // Q4_K superblock alignment |
30 | | //! ``` |
31 | | |
32 | | use serde::{Deserialize, Serialize}; |
33 | | use std::fmt; |
34 | | |
35 | | // ============================================================================ |
36 | | // TILE-001: TcbGeometry Struct |
37 | | // ============================================================================ |
38 | | |
39 | | /// Dimensions for a Tiling Compute Block |
40 | | /// |
41 | | /// Represents the (M, N, K) dimensions of a tile in matrix operations: |
42 | | /// - M: Output rows |
43 | | /// - N: Output columns |
44 | | /// - K: Reduction dimension (inner product) |
45 | | /// |
46 | | /// # Alignment Constraints |
47 | | /// |
48 | | /// Per the TCB-03 pattern (Tile Quantization Alignment), K must align with |
49 | | /// the quantization superblock size: |
50 | | /// - Q4_0: K % 32 == 0 |
51 | | /// - Q4_K: K % 256 == 0 |
52 | | /// - Q8_0: K % 32 == 0 |
53 | | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
54 | | pub struct TcbGeometry { |
55 | | /// Items processed in M dimension (rows) |
56 | | pub m: u32, |
57 | | /// Items processed in N dimension (columns) |
58 | | pub n: u32, |
59 | | /// Reduction dimension (inner product) |
60 | | pub k: u32, |
61 | | /// Alignment requirement in bytes (typically 16 for SIMD, 32 for AVX2, 64 for AVX-512) |
62 | | pub alignment: u32, |
63 | | } |
64 | | |
65 | | impl TcbGeometry { |
66 | | /// Create a new TCB geometry |
67 | | /// |
68 | | /// # Panics |
69 | | /// Panics if any dimension is zero. |
70 | | #[must_use] |
71 | 0 | pub fn new(m: u32, n: u32, k: u32) -> Self { |
72 | 0 | assert!(m > 0 && n > 0 && k > 0, "TCB dimensions must be non-zero"); |
73 | 0 | Self { |
74 | 0 | m, |
75 | 0 | n, |
76 | 0 | k, |
77 | 0 | alignment: 16, // Default to SSE/NEON alignment |
78 | 0 | } |
79 | 0 | } |
80 | | |
81 | | /// Create geometry with explicit alignment |
82 | | #[must_use] |
83 | 0 | pub fn with_alignment(m: u32, n: u32, k: u32, alignment: u32) -> Self { |
84 | 0 | assert!(m > 0 && n > 0 && k > 0, "TCB dimensions must be non-zero"); |
85 | 0 | assert!( |
86 | 0 | alignment.is_power_of_two(), |
87 | 0 | "Alignment must be power of 2" |
88 | | ); |
89 | 0 | Self { m, n, k, alignment } |
90 | 0 | } |
91 | | |
92 | | /// Calculate arithmetic intensity (FLOPS per byte loaded) |
93 | | /// |
94 | | /// For GEMM: AI = (2 * M * N * K) / (M*K + K*N) * sizeof(f32) |
95 | | /// |
96 | | /// Higher AI means compute-bound; lower means memory-bound. |
97 | | #[must_use] |
98 | 0 | pub fn arithmetic_intensity(&self) -> f32 { |
99 | 0 | let flops = 2.0 * self.m as f64 * self.n as f64 * self.k as f64; |
100 | 0 | let bytes = (self.m as f64 * self.k as f64 + self.k as f64 * self.n as f64) * 4.0; |
101 | 0 | (flops / bytes) as f32 |
102 | 0 | } |
103 | | |
104 | | /// Calculate total elements in the tile |
105 | | #[must_use] |
106 | 0 | pub fn total_elements(&self) -> u64 { |
107 | 0 | self.m as u64 * self.n as u64 |
108 | 0 | } |
109 | | |
110 | | /// Calculate total FLOPs for this tile |
111 | | #[must_use] |
112 | 0 | pub fn total_flops(&self) -> u64 { |
113 | 0 | 2 * self.m as u64 * self.n as u64 * self.k as u64 |
114 | 0 | } |
115 | | |
116 | | /// Check if K dimension aligns with Q4_K superblock (256) |
117 | | #[must_use] |
118 | 0 | pub fn is_q4k_aligned(&self) -> bool { |
119 | 0 | self.k % 256 == 0 |
120 | 0 | } |
121 | | |
122 | | /// Check if K dimension aligns with Q4_0/Q8_0 block (32) |
123 | | #[must_use] |
124 | 0 | pub fn is_q4_0_aligned(&self) -> bool { |
125 | 0 | self.k % 32 == 0 |
126 | 0 | } |
127 | | |
128 | | /// Calculate bytes needed for A tile (M × K × sizeof(f32)) |
129 | | #[must_use] |
130 | 0 | pub fn a_tile_bytes(&self) -> usize { |
131 | 0 | self.m as usize * self.k as usize * 4 |
132 | 0 | } |
133 | | |
134 | | /// Calculate bytes needed for B tile (K × N × sizeof(f32)) |
135 | | #[must_use] |
136 | 0 | pub fn b_tile_bytes(&self) -> usize { |
137 | 0 | self.k as usize * self.n as usize * 4 |
138 | 0 | } |
139 | | |
140 | | /// Calculate bytes needed for C tile (M × N × sizeof(f32)) |
141 | | #[must_use] |
142 | 0 | pub fn c_tile_bytes(&self) -> usize { |
143 | 0 | self.m as usize * self.n as usize * 4 |
144 | 0 | } |
145 | | |
146 | | /// Check if tile fits in given cache size (bytes) |
147 | | #[must_use] |
148 | 0 | pub fn fits_in_cache(&self, cache_bytes: usize) -> bool { |
149 | 0 | self.a_tile_bytes() + self.b_tile_bytes() <= cache_bytes |
150 | 0 | } |
151 | | } |
152 | | |
153 | | impl Default for TcbGeometry { |
154 | 0 | fn default() -> Self { |
155 | | // Sensible default: 4×4 micro-tile for SIMD |
156 | 0 | Self { |
157 | 0 | m: 4, |
158 | 0 | n: 4, |
159 | 0 | k: 4, |
160 | 0 | alignment: 16, |
161 | 0 | } |
162 | 0 | } |
163 | | } |
164 | | |
165 | | impl fmt::Display for TcbGeometry { |
166 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
167 | 0 | write!( |
168 | 0 | f, |
169 | 0 | "TCB({}×{}×{}, align={}, AI={:.2})", |
170 | | self.m, |
171 | | self.n, |
172 | | self.k, |
173 | | self.alignment, |
174 | 0 | self.arithmetic_intensity() |
175 | | ) |
176 | 0 | } |
177 | | } |
178 | | |
179 | | // ============================================================================ |
180 | | // TILE-001: Tiling Levels |
181 | | // ============================================================================ |
182 | | |
183 | | /// Tiling hierarchy level |
184 | | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
185 | | pub enum TcbLevel { |
186 | | /// Macro-tile: L3 cache / GPU global memory partitioning |
187 | | Macro, |
188 | | /// Midi-tile: L2 cache / GPU shared memory |
189 | | Midi, |
190 | | /// Micro-tile: Registers / SIMD lanes |
191 | | Micro, |
192 | | } |
193 | | |
194 | | impl TcbLevel { |
195 | | /// Get typical cache size for this level (x86_64) |
196 | | #[must_use] |
197 | 0 | pub fn typical_cache_bytes(&self) -> usize { |
198 | 0 | match self { |
199 | 0 | TcbLevel::Macro => 32 * 1024 * 1024, // 32 MB L3 |
200 | 0 | TcbLevel::Midi => 256 * 1024, // 256 KB L2 |
201 | 0 | TcbLevel::Micro => 32 * 1024, // 32 KB L1 |
202 | | } |
203 | 0 | } |
204 | | } |
205 | | |
206 | | // ============================================================================ |
207 | | // TILE-001: Complete Tiling Configuration |
208 | | // ============================================================================ |
209 | | |
210 | | /// Complete tiling configuration for a kernel |
211 | | /// |
212 | | /// Contains geometry for all three tiling levels, enabling hierarchical |
213 | | /// cache-aware execution. |
214 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
215 | | pub struct TilingConfig { |
216 | | /// Kernel name for identification |
217 | | pub name: String, |
218 | | /// Macro-tile geometry (L3/Global) |
219 | | pub macro_tile: TcbGeometry, |
220 | | /// Midi-tile geometry (L2/Shared) |
221 | | pub midi_tile: TcbGeometry, |
222 | | /// Micro-tile geometry (Registers) |
223 | | pub micro_tile: TcbGeometry, |
224 | | /// Target backend |
225 | | pub backend: TilingBackend, |
226 | | } |
227 | | |
228 | | /// Backend target for tiling configuration |
229 | | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
230 | | pub enum TilingBackend { |
231 | | /// CPU with AVX2 (256-bit SIMD) |
232 | | CpuAvx2, |
233 | | /// CPU with AVX-512 (512-bit SIMD) |
234 | | CpuAvx512, |
235 | | /// CPU with NEON (128-bit SIMD) |
236 | | CpuNeon, |
237 | | /// GPU (CUDA/wgpu) |
238 | | Gpu, |
239 | | /// Scalar fallback |
240 | | Scalar, |
241 | | } |
242 | | |
243 | | impl TilingConfig { |
244 | | /// Create configuration for GPU Q4_K MatVec |
245 | | /// |
246 | | /// Optimized for single-token generation where M=1. |
247 | | #[must_use] |
248 | 0 | pub fn gpu_q4k_matvec() -> Self { |
249 | 0 | Self { |
250 | 0 | name: "Q4K_MatVec_GPU".into(), |
251 | 0 | macro_tile: TcbGeometry::with_alignment(1, 4096, 256, 64), |
252 | 0 | midi_tile: TcbGeometry::with_alignment(1, 256, 256, 64), |
253 | 0 | micro_tile: TcbGeometry::with_alignment(1, 32, 256, 64), |
254 | 0 | backend: TilingBackend::Gpu, |
255 | 0 | } |
256 | 0 | } |
257 | | |
258 | | /// Create configuration for GPU Q4_K MatMul (batched) |
259 | | /// |
260 | | /// Optimized for prefill where M > 1. |
261 | | #[must_use] |
262 | 0 | pub fn gpu_q4k_matmul() -> Self { |
263 | 0 | Self { |
264 | 0 | name: "Q4K_MatMul_GPU".into(), |
265 | 0 | macro_tile: TcbGeometry::with_alignment(128, 128, 256, 64), |
266 | 0 | midi_tile: TcbGeometry::with_alignment(32, 32, 256, 64), |
267 | 0 | micro_tile: TcbGeometry::with_alignment(8, 8, 256, 64), |
268 | 0 | backend: TilingBackend::Gpu, |
269 | 0 | } |
270 | 0 | } |
271 | | |
272 | | /// Create configuration for GPU Softmax |
273 | | #[must_use] |
274 | 0 | pub fn gpu_softmax() -> Self { |
275 | 0 | Self { |
276 | 0 | name: "Softmax_GPU".into(), |
277 | 0 | macro_tile: TcbGeometry::with_alignment(1, 32000, 1, 64), |
278 | 0 | midi_tile: TcbGeometry::with_alignment(1, 1024, 1, 64), |
279 | 0 | micro_tile: TcbGeometry::with_alignment(1, 32, 1, 64), |
280 | 0 | backend: TilingBackend::Gpu, |
281 | 0 | } |
282 | 0 | } |
283 | | |
284 | | /// Create configuration for CPU AVX-512 MatMul |
285 | | /// |
286 | | /// Optimized for 512-bit wide SIMD: |
287 | | /// - 16 floats per ZMM register |
288 | | /// - 32 ZMM registers available |
289 | | /// - 4×16 micro-kernel uses 8 registers (4 accumulators + 4 scratch) |
290 | | #[must_use] |
291 | 0 | pub fn cpu_avx512_matmul() -> Self { |
292 | 0 | Self { |
293 | 0 | name: "MatMul_AVX512".into(), |
294 | 0 | macro_tile: TcbGeometry::with_alignment(512, 512, 512, 64), |
295 | 0 | midi_tile: TcbGeometry::with_alignment(128, 128, 128, 64), |
296 | 0 | // 16 floats wide × 4 rows = 64 elements in registers |
297 | 0 | micro_tile: TcbGeometry::with_alignment(4, 16, 128, 64), |
298 | 0 | backend: TilingBackend::CpuAvx512, |
299 | 0 | } |
300 | 0 | } |
301 | | |
302 | | /// Create configuration for CPU AVX-512 Q4K MatVec |
303 | | /// |
304 | | /// Optimized for Q4_K quantized inference with 512-bit SIMD. |
305 | | /// Key differences from AVX2: |
306 | | /// - 64-byte aligned for cache line optimization |
307 | | /// - 4×1 micro-kernel processes 4 rows simultaneously |
308 | | /// - K=256 aligned to Q4_K superblock |
309 | | #[must_use] |
310 | 0 | pub fn cpu_avx512_q4k_matvec() -> Self { |
311 | 0 | Self { |
312 | 0 | name: "Q4K_MatVec_AVX512".into(), |
313 | 0 | // Large macro-tile to amortize L3 access |
314 | 0 | macro_tile: TcbGeometry::with_alignment(4096, 1, 4096, 64), |
315 | 0 | // Midi-tile fits in L2 (256KB) |
316 | 0 | // 64 rows × 256 K × 0.5625 bytes/element ≈ 9KB weights |
317 | 0 | midi_tile: TcbGeometry::with_alignment(64, 1, 256, 64), |
318 | 0 | // 4 rows × 1 output, K=256 (Q4_K superblock) |
319 | 0 | micro_tile: TcbGeometry::with_alignment(4, 1, 256, 64), |
320 | 0 | backend: TilingBackend::CpuAvx512, |
321 | 0 | } |
322 | 0 | } |
323 | | |
324 | | /// Create configuration for AVX-512 VNNI Q4K×Q8K integer dot product |
325 | | /// |
326 | | /// AVX-512 VNNI (Vector Neural Network Instructions) provides: |
327 | | /// - VPDPBUSD: 8-bit unsigned × 8-bit signed multiply-add to i32 |
328 | | /// - VPDPWSSD: 16-bit signed × 16-bit signed multiply-add to i32 |
329 | | /// |
330 | | /// This enables pure integer Q4K×Q8K without intermediate f32 conversion. |
331 | | #[must_use] |
332 | 0 | pub fn cpu_avx512_vnni_q4k_q8k() -> Self { |
333 | 0 | Self { |
334 | 0 | name: "Q4K_Q8K_VNNI".into(), |
335 | 0 | macro_tile: TcbGeometry::with_alignment(4096, 1, 4096, 64), |
336 | 0 | midi_tile: TcbGeometry::with_alignment(64, 1, 256, 64), |
337 | 0 | // VNNI processes 64 i8 values per ZMM register |
338 | 0 | micro_tile: TcbGeometry::with_alignment(4, 1, 256, 64), |
339 | 0 | backend: TilingBackend::CpuAvx512, |
340 | 0 | } |
341 | 0 | } |
342 | | |
343 | | /// Create configuration for CPU AVX2 MatMul |
344 | | #[must_use] |
345 | 0 | pub fn cpu_avx2_matmul() -> Self { |
346 | 0 | Self { |
347 | 0 | name: "MatMul_AVX2".into(), |
348 | 0 | macro_tile: TcbGeometry::with_alignment(256, 256, 256, 32), |
349 | 0 | midi_tile: TcbGeometry::with_alignment(64, 64, 64, 32), |
350 | 0 | // 8 floats wide × 4 rows = 32 elements in registers |
351 | 0 | micro_tile: TcbGeometry::with_alignment(4, 8, 64, 32), |
352 | 0 | backend: TilingBackend::CpuAvx2, |
353 | 0 | } |
354 | 0 | } |
355 | | |
356 | | /// Create configuration for CPU Q4_K MatVec (AVX2) |
357 | | #[must_use] |
358 | 0 | pub fn cpu_avx2_q4k_matvec() -> Self { |
359 | 0 | Self { |
360 | 0 | name: "Q4K_MatVec_AVX2".into(), |
361 | 0 | // Process 4 rows at a time (4×1 micro-kernel) |
362 | 0 | macro_tile: TcbGeometry::with_alignment(4096, 1, 4096, 32), |
363 | 0 | midi_tile: TcbGeometry::with_alignment(64, 1, 256, 32), |
364 | 0 | // 4 rows × 1 output, K=256 (Q4_K superblock) |
365 | 0 | micro_tile: TcbGeometry::with_alignment(4, 1, 256, 32), |
366 | 0 | backend: TilingBackend::CpuAvx2, |
367 | 0 | } |
368 | 0 | } |
369 | | |
370 | | /// Create configuration for RMSNorm (CPU) |
371 | | #[must_use] |
372 | 0 | pub fn cpu_rmsnorm() -> Self { |
373 | 0 | Self { |
374 | 0 | name: "RMSNorm_CPU".into(), |
375 | 0 | macro_tile: TcbGeometry::with_alignment(1, 4096, 1, 32), |
376 | 0 | midi_tile: TcbGeometry::with_alignment(1, 256, 1, 32), |
377 | 0 | micro_tile: TcbGeometry::with_alignment(1, 16, 1, 32), |
378 | 0 | backend: TilingBackend::CpuAvx512, |
379 | 0 | } |
380 | 0 | } |
381 | | |
382 | | /// Validate that tiling configuration is internally consistent |
383 | 0 | pub fn validate(&self) -> Result<(), TilingError> { |
384 | | // Macro must be >= Midi >= Micro |
385 | 0 | if self.midi_tile.m > self.macro_tile.m |
386 | 0 | || self.midi_tile.n > self.macro_tile.n |
387 | 0 | || self.midi_tile.k > self.macro_tile.k |
388 | | { |
389 | 0 | return Err(TilingError::InvalidHierarchy { |
390 | 0 | reason: "Midi-tile larger than macro-tile".into(), |
391 | 0 | }); |
392 | 0 | } |
393 | | |
394 | 0 | if self.micro_tile.m > self.midi_tile.m |
395 | 0 | || self.micro_tile.n > self.midi_tile.n |
396 | 0 | || self.micro_tile.k > self.midi_tile.k |
397 | | { |
398 | 0 | return Err(TilingError::InvalidHierarchy { |
399 | 0 | reason: "Micro-tile larger than midi-tile".into(), |
400 | 0 | }); |
401 | 0 | } |
402 | | |
403 | | // Check divisibility |
404 | 0 | if self.macro_tile.m % self.midi_tile.m != 0 { |
405 | 0 | return Err(TilingError::DivisibilityError { |
406 | 0 | level: "macro/midi", |
407 | 0 | dimension: "M", |
408 | 0 | larger: self.macro_tile.m, |
409 | 0 | smaller: self.midi_tile.m, |
410 | 0 | }); |
411 | 0 | } |
412 | | |
413 | 0 | if self.midi_tile.m % self.micro_tile.m != 0 { |
414 | 0 | return Err(TilingError::DivisibilityError { |
415 | 0 | level: "midi/micro", |
416 | 0 | dimension: "M", |
417 | 0 | larger: self.midi_tile.m, |
418 | 0 | smaller: self.micro_tile.m, |
419 | 0 | }); |
420 | 0 | } |
421 | | |
422 | 0 | Ok(()) |
423 | 0 | } |
424 | | |
425 | | /// Calculate total number of macro-tiles for given problem size |
426 | | #[must_use] |
427 | 0 | pub fn num_macro_tiles(&self, m: u32, n: u32) -> u32 { |
428 | 0 | let m_tiles = (m + self.macro_tile.m - 1) / self.macro_tile.m; |
429 | 0 | let n_tiles = (n + self.macro_tile.n - 1) / self.macro_tile.n; |
430 | 0 | m_tiles * n_tiles |
431 | 0 | } |
432 | | |
433 | | /// Calculate total number of midi-tiles within a macro-tile |
434 | | #[must_use] |
435 | 0 | pub fn midi_tiles_per_macro(&self) -> u32 { |
436 | 0 | let m_tiles = self.macro_tile.m / self.midi_tile.m; |
437 | 0 | let n_tiles = self.macro_tile.n / self.midi_tile.n; |
438 | 0 | m_tiles * n_tiles |
439 | 0 | } |
440 | | |
441 | | /// Calculate total number of micro-tiles within a midi-tile |
442 | | #[must_use] |
443 | 0 | pub fn micro_tiles_per_midi(&self) -> u32 { |
444 | 0 | let m_tiles = self.midi_tile.m / self.micro_tile.m; |
445 | 0 | let n_tiles = self.midi_tile.n / self.micro_tile.n; |
446 | 0 | m_tiles * n_tiles |
447 | 0 | } |
448 | | } |
449 | | |
450 | | // ============================================================================ |
451 | | // TILE-002: Hierarchical Index Calculator |
452 | | // ============================================================================ |
453 | | |
454 | | /// Index calculator for hierarchical tiling |
455 | | /// |
456 | | /// Converts between linear indices and (row, col) coordinates at each tiling level. |
457 | | #[derive(Debug, Clone)] |
458 | | pub struct TcbIndexCalculator { |
459 | | /// Tiling configuration |
460 | | config: TilingConfig, |
461 | | /// Problem dimensions |
462 | | problem_m: u32, |
463 | | problem_n: u32, |
464 | | problem_k: u32, |
465 | | } |
466 | | |
467 | | impl TcbIndexCalculator { |
468 | | /// Create a new index calculator for the given problem size |
469 | | #[must_use] |
470 | 0 | pub fn new(config: TilingConfig, m: u32, n: u32, k: u32) -> Self { |
471 | 0 | Self { |
472 | 0 | config, |
473 | 0 | problem_m: m, |
474 | 0 | problem_n: n, |
475 | 0 | problem_k: k, |
476 | 0 | } |
477 | 0 | } |
478 | | |
479 | | /// Get macro-tile offset for a given block index |
480 | | /// |
481 | | /// Returns (row_offset, col_offset) in the output matrix. |
482 | | #[must_use] |
483 | 0 | pub fn macro_tile_offset(&self, block_idx: u32) -> (u32, u32) { |
484 | 0 | let tiles_per_row = (self.problem_n + self.config.macro_tile.n - 1) / self.config.macro_tile.n; |
485 | 0 | let row = (block_idx / tiles_per_row) * self.config.macro_tile.m; |
486 | 0 | let col = (block_idx % tiles_per_row) * self.config.macro_tile.n; |
487 | 0 | (row, col) |
488 | 0 | } |
489 | | |
490 | | /// Get midi-tile offset within a macro-tile |
491 | | #[must_use] |
492 | 0 | pub fn midi_tile_offset(&self, midi_idx: u32) -> (u32, u32) { |
493 | 0 | let tiles_per_row = self.config.macro_tile.n / self.config.midi_tile.n; |
494 | 0 | let row = (midi_idx / tiles_per_row) * self.config.midi_tile.m; |
495 | 0 | let col = (midi_idx % tiles_per_row) * self.config.midi_tile.n; |
496 | 0 | (row, col) |
497 | 0 | } |
498 | | |
499 | | /// Get micro-tile offset within a midi-tile |
500 | | #[must_use] |
501 | 0 | pub fn micro_tile_offset(&self, micro_idx: u32) -> (u32, u32) { |
502 | 0 | let tiles_per_row = self.config.midi_tile.n / self.config.micro_tile.n; |
503 | 0 | let row = (micro_idx / tiles_per_row) * self.config.micro_tile.m; |
504 | 0 | let col = (micro_idx % tiles_per_row) * self.config.micro_tile.n; |
505 | 0 | (row, col) |
506 | 0 | } |
507 | | |
508 | | /// Convert block index to linear memory offset |
509 | | /// |
510 | | /// For row-major C matrix with given stride. |
511 | | #[must_use] |
512 | | #[inline] |
513 | 0 | pub fn block_to_linear_offset(&self, block_idx: u32, stride: u32) -> usize { |
514 | 0 | let (row, col) = self.macro_tile_offset(block_idx); |
515 | 0 | (row * stride + col) as usize |
516 | 0 | } |
517 | | |
518 | | /// Calculate A matrix offset for K-dimension blocking |
519 | | #[must_use] |
520 | | #[inline] |
521 | 0 | pub fn a_offset(&self, macro_row: u32, k_block: u32) -> usize { |
522 | 0 | let row = macro_row * self.config.macro_tile.m; |
523 | 0 | let col = k_block * self.config.macro_tile.k; |
524 | 0 | (row * self.problem_k + col) as usize |
525 | 0 | } |
526 | | |
527 | | /// Calculate B matrix offset for K-dimension blocking |
528 | | #[must_use] |
529 | | #[inline] |
530 | 0 | pub fn b_offset(&self, k_block: u32, macro_col: u32) -> usize { |
531 | 0 | let row = k_block * self.config.macro_tile.k; |
532 | 0 | let col = macro_col * self.config.macro_tile.n; |
533 | 0 | (row * self.problem_n + col) as usize |
534 | 0 | } |
535 | | |
536 | | /// Get number of K blocks needed |
537 | | #[must_use] |
538 | 0 | pub fn num_k_blocks(&self) -> u32 { |
539 | 0 | (self.problem_k + self.config.macro_tile.k - 1) / self.config.macro_tile.k |
540 | 0 | } |
541 | | |
542 | | /// Check if this is a boundary tile (may need masking) |
543 | | #[must_use] |
544 | 0 | pub fn is_boundary_tile(&self, block_idx: u32) -> bool { |
545 | 0 | let (row, col) = self.macro_tile_offset(block_idx); |
546 | 0 | row + self.config.macro_tile.m > self.problem_m |
547 | 0 | || col + self.config.macro_tile.n > self.problem_n |
548 | 0 | } |
549 | | |
550 | | /// Get actual tile dimensions (may be smaller at boundaries) |
551 | | #[must_use] |
552 | 0 | pub fn actual_tile_dims(&self, block_idx: u32) -> (u32, u32) { |
553 | 0 | let (row, col) = self.macro_tile_offset(block_idx); |
554 | 0 | let actual_m = (self.problem_m - row).min(self.config.macro_tile.m); |
555 | 0 | let actual_n = (self.problem_n - col).min(self.config.macro_tile.n); |
556 | 0 | (actual_m, actual_n) |
557 | 0 | } |
558 | | } |
559 | | |
560 | | // ============================================================================ |
561 | | // TILE-002: Memory Layout Helpers |
562 | | // ============================================================================ |
563 | | |
564 | | /// Memory layout for packed matrices |
565 | | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
566 | | pub enum PackingLayout { |
567 | | /// Row-major (C-style) |
568 | | RowMajor, |
569 | | /// Column-major (Fortran-style) |
570 | | ColumnMajor, |
571 | | /// Panel-major for A (Goto algorithm) |
572 | | PanelMajorA, |
573 | | /// Panel-major for B (Goto algorithm) |
574 | | PanelMajorB, |
575 | | } |
576 | | |
577 | | /// Calculate packed index for panel-major A layout |
578 | | /// |
579 | | /// Panel-major stores micro-panels contiguously for sequential access. |
580 | | #[must_use] |
581 | | #[inline] |
582 | 0 | pub fn pack_a_index(row: usize, col: usize, mr: usize, kc: usize, _mc: usize) -> usize { |
583 | 0 | let panel = row / mr; |
584 | 0 | let row_in_panel = row % mr; |
585 | 0 | panel * mr * kc + col * mr + row_in_panel |
586 | 0 | } |
587 | | |
588 | | /// Calculate packed index for panel-major B layout |
589 | | #[must_use] |
590 | | #[inline] |
591 | 0 | pub fn pack_b_index(row: usize, col: usize, nr: usize, kc: usize, _nc: usize) -> usize { |
592 | 0 | let panel = col / nr; |
593 | 0 | let col_in_panel = col % nr; |
594 | 0 | panel * kc * nr + row * nr + col_in_panel |
595 | 0 | } |
596 | | |
597 | | /// Apply XOR swizzling for shared memory bank conflict avoidance |
598 | | /// |
599 | | /// Pattern: idx_swizzled = idx ^ (idx >> 5) for 32-bank architectures. |
600 | | #[must_use] |
601 | | #[inline] |
602 | 0 | pub fn swizzle_index(idx: usize) -> usize { |
603 | 0 | idx ^ (idx >> 5) |
604 | 0 | } |
605 | | |
606 | | // ============================================================================ |
607 | | // TILE-003: Prefetch Helpers |
608 | | // ============================================================================ |
609 | | |
610 | | /// Prefetch locality hint |
611 | | #[derive(Debug, Clone, Copy)] |
612 | | pub enum PrefetchLocality { |
613 | | /// Non-temporal (streaming, evict soon) |
614 | | NonTemporal, |
615 | | /// L3 cache |
616 | | T2, |
617 | | /// L2 cache |
618 | | T1, |
619 | | /// L1 cache (highest priority) |
620 | | T0, |
621 | | } |
622 | | |
623 | | /// Calculate optimal prefetch distance based on tile geometry and cache level |
624 | | /// |
625 | | /// Per Ding & Kennedy (2004): distance = memory_latency / compute_time_per_iter |
626 | | #[must_use] |
627 | 0 | pub fn optimal_prefetch_distance(geometry: &TcbGeometry, level: TcbLevel) -> usize { |
628 | | // Approximate cycles per micro-tile |
629 | 0 | let compute_cycles = geometry.m as usize * geometry.n as usize * geometry.k as usize / 8; |
630 | | |
631 | | // Memory latency in cycles (approximate for modern x86) |
632 | 0 | let mem_latency = match level { |
633 | 0 | TcbLevel::Micro => 4, // L1: ~4 cycles |
634 | 0 | TcbLevel::Midi => 12, // L2: ~12 cycles |
635 | 0 | TcbLevel::Macro => 40, // L3: ~40 cycles |
636 | | }; |
637 | | |
638 | | // Distance = latency / compute_time, minimum 1 |
639 | 0 | (mem_latency / compute_cycles.max(1)).max(1) |
640 | 0 | } |
641 | | |
642 | | // ============================================================================ |
643 | | // TILE-003: Tiled Q4_K MatVec (TCB-01 Pattern) |
644 | | // ============================================================================ |
645 | | |
646 | | /// Q4_K superblock constants (per GGML specification) |
647 | | pub const Q4K_SUPERBLOCK_SIZE: usize = 256; |
648 | | pub const Q4K_SUPERBLOCK_BYTES: usize = 144; |
649 | | |
650 | | /// Tiled Q4_K MatVec executor |
651 | | /// |
652 | | /// Implements TCB-01 pattern: Cache-blocked matvec with 4×1 micro-kernel. |
653 | | /// |
654 | | /// # Memory Layout |
655 | | /// |
656 | | /// Weights are stored in Q4_K superblock format (144 bytes per 256 elements): |
657 | | /// - d: f16 (2 bytes) - block scale |
658 | | /// - dmin: f16 (2 bytes) - block minimum |
659 | | /// - scales: 12 bytes - 8 sub-block scales (6-bit packed) |
660 | | /// - qs: 128 bytes - 256 quantized values (4-bit packed) |
661 | | /// |
662 | | /// # Performance Characteristics |
663 | | /// |
664 | | /// - L2-resident: Process midi_tile.m rows at a time |
665 | | /// - Vectorized: 4×1 micro-kernel processes 4 output rows simultaneously |
666 | | /// - Aligned: K dimension aligned to Q4_K superblock (256) |
667 | | #[derive(Debug, Clone)] |
668 | | pub struct TiledQ4KMatvec { |
669 | | /// Tiling configuration |
670 | | pub config: TilingConfig, |
671 | | /// Number of rows (M dimension) |
672 | | pub m: usize, |
673 | | /// Number of columns (K dimension) |
674 | | pub k: usize, |
675 | | } |
676 | | |
677 | | impl TiledQ4KMatvec { |
678 | | /// Create a new tiled Q4K matvec executor |
679 | | /// |
680 | | /// # Panics |
681 | | /// Panics if K is not aligned to Q4_K superblock size (256). |
682 | | #[must_use] |
683 | 0 | pub fn new(m: usize, k: usize) -> Self { |
684 | 0 | assert!( |
685 | 0 | k % Q4K_SUPERBLOCK_SIZE == 0, |
686 | 0 | "K dimension ({}) must be aligned to Q4_K superblock size ({})", |
687 | | k, |
688 | | Q4K_SUPERBLOCK_SIZE |
689 | | ); |
690 | | |
691 | 0 | Self { |
692 | 0 | config: TilingConfig::cpu_avx2_q4k_matvec(), |
693 | 0 | m, |
694 | 0 | k, |
695 | 0 | } |
696 | 0 | } |
697 | | |
698 | | /// Get number of superblocks per row |
699 | | #[must_use] |
700 | 0 | pub fn superblocks_per_row(&self) -> usize { |
701 | 0 | self.k / Q4K_SUPERBLOCK_SIZE |
702 | 0 | } |
703 | | |
704 | | /// Get total number of superblocks |
705 | | #[must_use] |
706 | 0 | pub fn total_superblocks(&self) -> usize { |
707 | 0 | self.m * self.superblocks_per_row() |
708 | 0 | } |
709 | | |
710 | | /// Get weight bytes offset for a given row |
711 | | #[must_use] |
712 | | #[inline] |
713 | 0 | pub fn weight_row_offset(&self, row: usize) -> usize { |
714 | 0 | row * self.superblocks_per_row() * Q4K_SUPERBLOCK_BYTES |
715 | 0 | } |
716 | | |
717 | | /// Calculate optimal number of parallel rows based on L2 cache |
718 | | /// |
719 | | /// Goal: Keep working set in L2 (256KB typical) |
720 | | /// Working set = midi_tile.m rows × K × sizeof(Q4K) + K × sizeof(f32) |
721 | | #[must_use] |
722 | 0 | pub fn optimal_parallel_rows(&self, l2_bytes: usize) -> usize { |
723 | | // Q4K: 144 bytes per 256 elements = 0.5625 bytes/element |
724 | 0 | let row_bytes = (self.k as f32 * 0.5625) as usize; |
725 | | // Input vector: K × 4 bytes |
726 | 0 | let input_bytes = self.k * 4; |
727 | | // Available for rows |
728 | 0 | let available = l2_bytes.saturating_sub(input_bytes); |
729 | | // Rows that fit (minimum 4 for micro-kernel) |
730 | 0 | (available / row_bytes).max(4) |
731 | 0 | } |
732 | | |
733 | | /// Execute tiled matvec (reference scalar implementation) |
734 | | /// |
735 | | /// This is the reference implementation for correctness testing. |
736 | | /// Actual SIMD implementation would be in the backends. |
737 | | /// |
738 | | /// For parallel execution, use [`execute_parallel`] when the `parallel` feature is enabled. |
739 | 0 | pub fn execute_scalar(&self, weights: &[u8], input: &[f32], output: &mut [f32]) { |
740 | 0 | assert_eq!(weights.len(), self.total_superblocks() * Q4K_SUPERBLOCK_BYTES); |
741 | 0 | assert_eq!(input.len(), self.k); |
742 | 0 | assert_eq!(output.len(), self.m); |
743 | | |
744 | 0 | let superblocks_per_row = self.superblocks_per_row(); |
745 | | |
746 | 0 | for row in 0..self.m { |
747 | 0 | let mut sum = 0.0f32; |
748 | 0 | let row_offset = row * superblocks_per_row * Q4K_SUPERBLOCK_BYTES; |
749 | | |
750 | 0 | for sb in 0..superblocks_per_row { |
751 | 0 | let sb_offset = row_offset + sb * Q4K_SUPERBLOCK_BYTES; |
752 | 0 | let sb_data = &weights[sb_offset..sb_offset + Q4K_SUPERBLOCK_BYTES]; |
753 | 0 |
|
754 | 0 | // Dequantize and dot product for this superblock |
755 | 0 | let input_offset = sb * Q4K_SUPERBLOCK_SIZE; |
756 | 0 | sum += self.scalar_superblock_dot(sb_data, &input[input_offset..input_offset + Q4K_SUPERBLOCK_SIZE]); |
757 | 0 | } |
758 | | |
759 | 0 | output[row] = sum; |
760 | | } |
761 | 0 | } |
762 | | |
763 | | /// Execute tiled matvec with parallel row processing |
764 | | /// |
765 | | /// Uses Rayon to parallelize across rows for multi-core speedup. |
766 | | /// Falls back to scalar execution if the `parallel` feature is not enabled. |
767 | | /// |
768 | | /// # Performance |
769 | | /// |
770 | | /// Achieves near-linear speedup with core count for large matrices. |
771 | | /// For small matrices (< 256 rows), scalar may be faster due to overhead. |
772 | | #[cfg(feature = "parallel")] |
773 | | pub fn execute_parallel(&self, weights: &[u8], input: &[f32], output: &mut [f32]) { |
774 | | use rayon::prelude::*; |
775 | | |
776 | | assert_eq!(weights.len(), self.total_superblocks() * Q4K_SUPERBLOCK_BYTES); |
777 | | assert_eq!(input.len(), self.k); |
778 | | assert_eq!(output.len(), self.m); |
779 | | |
780 | | let superblocks_per_row = self.superblocks_per_row(); |
781 | | let row_stride = superblocks_per_row * Q4K_SUPERBLOCK_BYTES; |
782 | | |
783 | | output.par_iter_mut().enumerate().for_each(|(row, out)| { |
784 | | let mut sum = 0.0f32; |
785 | | let row_offset = row * row_stride; |
786 | | |
787 | | for sb in 0..superblocks_per_row { |
788 | | let sb_offset = row_offset + sb * Q4K_SUPERBLOCK_BYTES; |
789 | | let sb_data = &weights[sb_offset..sb_offset + Q4K_SUPERBLOCK_BYTES]; |
790 | | |
791 | | let input_offset = sb * Q4K_SUPERBLOCK_SIZE; |
792 | | sum += self.scalar_superblock_dot(sb_data, &input[input_offset..input_offset + Q4K_SUPERBLOCK_SIZE]); |
793 | | } |
794 | | |
795 | | *out = sum; |
796 | | }); |
797 | | } |
798 | | |
799 | | /// Execute tiled matvec with parallel row processing (fallback) |
800 | | /// |
801 | | /// When `parallel` feature is not enabled, this is equivalent to `execute_scalar`. |
802 | | #[cfg(not(feature = "parallel"))] |
803 | 0 | pub fn execute_parallel(&self, weights: &[u8], input: &[f32], output: &mut [f32]) { |
804 | 0 | self.execute_scalar(weights, input, output); |
805 | 0 | } |
806 | | |
807 | | /// Scalar dot product for a single Q4_K superblock |
808 | | /// |
809 | | /// # Performance |
810 | | /// |
811 | | /// Optimized version with: |
812 | | /// - Precomputed scale/min pairs |
813 | | /// - Loop unrolling hints |
814 | | /// - Minimized branching in inner loop |
815 | | #[inline] |
816 | 0 | fn scalar_superblock_dot(&self, sb_data: &[u8], input: &[f32]) -> f32 { |
817 | | // Read header (hot path optimized) |
818 | 0 | let d = f16_to_f32(&sb_data[0..2]); |
819 | 0 | let dmin = f16_to_f32(&sb_data[2..4]); |
820 | 0 | let scales = &sb_data[4..16]; |
821 | 0 | let qs = &sb_data[16..144]; |
822 | | |
823 | | // Precompute all scale/min pairs upfront |
824 | 0 | let scale_mins = precompute_scales_mins(scales); |
825 | | |
826 | 0 | let mut sum = 0.0f32; |
827 | | |
828 | | // Process 256 values in 8 chunks of 32 |
829 | 0 | for chunk in 0..8 { |
830 | 0 | let (sc, m) = scale_mins[chunk]; |
831 | 0 | let d_scale = d * sc; |
832 | 0 | let dm = dmin * m; |
833 | | |
834 | 0 | let q_offset = chunk * 16; // 32 nibbles = 16 bytes |
835 | 0 | let input_offset = chunk * 32; |
836 | | |
837 | | // Process 32 values: low nibbles then high nibbles |
838 | | // Manually unroll inner loop for better optimization |
839 | 0 | let mut chunk_sum = 0.0f32; |
840 | | |
841 | | // Process 16 byte pairs (32 nibbles) |
842 | 0 | for i in 0..16 { |
843 | 0 | let byte = qs[q_offset + i]; |
844 | 0 |
|
845 | 0 | // Extract nibbles |
846 | 0 | let q_lo = (byte & 0x0F) as f32; |
847 | 0 | let q_hi = (byte >> 4) as f32; |
848 | 0 |
|
849 | 0 | // Dequantize: val = d * scale * q - dmin * min |
850 | 0 | let val_lo = d_scale * q_lo - dm; |
851 | 0 | let val_hi = d_scale * q_hi - dm; |
852 | 0 |
|
853 | 0 | // Accumulate dot product |
854 | 0 | chunk_sum += val_lo * input[input_offset + i]; |
855 | 0 | chunk_sum += val_hi * input[input_offset + 16 + i]; |
856 | 0 | } |
857 | | |
858 | 0 | sum += chunk_sum; |
859 | | } |
860 | | |
861 | 0 | sum |
862 | 0 | } |
863 | | |
864 | | /// Get tiling statistics for profiling |
865 | | #[must_use] |
866 | 0 | pub fn stats(&self) -> TilingStats { |
867 | 0 | let bytes_per_row = self.superblocks_per_row() * Q4K_SUPERBLOCK_BYTES; |
868 | 0 | let total_weight_bytes = self.m * bytes_per_row; |
869 | 0 | let input_bytes = self.k * 4; |
870 | 0 | let output_bytes = self.m * 4; |
871 | | |
872 | 0 | TilingStats { |
873 | 0 | total_weight_bytes, |
874 | 0 | input_bytes, |
875 | 0 | output_bytes, |
876 | 0 | superblocks: self.total_superblocks(), |
877 | 0 | arithmetic_ops: self.m * self.k * 2, // 2 ops per element (mul + add) |
878 | 0 | arithmetic_intensity: (self.m * self.k * 2) as f32 / (total_weight_bytes + input_bytes) as f32, |
879 | 0 | } |
880 | 0 | } |
881 | | } |
882 | | |
883 | | /// Statistics for a tiled operation |
884 | | #[derive(Debug, Clone)] |
885 | | pub struct TilingStats { |
886 | | /// Total weight bytes |
887 | | pub total_weight_bytes: usize, |
888 | | /// Input vector bytes |
889 | | pub input_bytes: usize, |
890 | | /// Output vector bytes |
891 | | pub output_bytes: usize, |
892 | | /// Number of superblocks |
893 | | pub superblocks: usize, |
894 | | /// Total arithmetic operations |
895 | | pub arithmetic_ops: usize, |
896 | | /// Arithmetic intensity (FLOPS/byte) |
897 | | pub arithmetic_intensity: f32, |
898 | | } |
899 | | |
900 | | /// Convert 2 bytes (f16 IEEE 754) to f32 |
901 | | /// |
902 | | /// Manual implementation to avoid half crate dependency. |
903 | | /// Format: 1 sign bit, 5 exponent bits, 10 mantissa bits. |
904 | | /// |
905 | | /// # Performance |
906 | | /// |
907 | | /// Optimized for the common case (normal numbers). Special cases (zero, |
908 | | /// subnormal, inf, nan) use branches but are rare in practice for model weights. |
909 | | #[inline] |
910 | 0 | fn f16_to_f32(bytes: &[u8]) -> f32 { |
911 | 0 | let bits = u16::from_le_bytes([bytes[0], bytes[1]]); |
912 | 0 | f16_bits_to_f32(bits) |
913 | 0 | } |
914 | | |
915 | | /// Fast path f16 to f32 conversion from raw bits |
916 | | /// |
917 | | /// Optimized version that handles the common case (normal numbers) with |
918 | | /// minimal branching. Uses branchless bit manipulation for the hot path. |
919 | | #[inline(always)] |
920 | 0 | fn f16_bits_to_f32(bits: u16) -> f32 { |
921 | 0 | let sign = (bits >> 15) & 0x1; |
922 | 0 | let exponent = (bits >> 10) & 0x1F; |
923 | 0 | let mantissa = bits & 0x3FF; |
924 | | |
925 | | // Fast path: normal numbers (exponent != 0 && exponent != 31) |
926 | | // This is the common case for model weights |
927 | 0 | if exponent != 0 && exponent != 31 { |
928 | | // Branchless conversion for normal numbers |
929 | | // f16 bias = 15, f32 bias = 127 |
930 | 0 | let f32_exp = (exponent as u32 + 112) as u32; // 127 - 15 = 112 |
931 | 0 | let f32_mant = (mantissa as u32) << 13; // 10 bits -> 23 bits |
932 | 0 | let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | f32_mant; |
933 | 0 | return f32::from_bits(f32_bits); |
934 | 0 | } |
935 | | |
936 | | // Slow path: special cases |
937 | 0 | f16_special_to_f32(sign, exponent, mantissa) |
938 | 0 | } |
939 | | |
940 | | /// Handle f16 special cases (zero, subnormal, inf, nan) |
941 | | /// |
942 | | /// Cold path - marked to help branch prediction |
943 | | #[cold] |
944 | | #[inline(never)] |
945 | 0 | fn f16_special_to_f32(sign: u16, exponent: u16, mantissa: u16) -> f32 { |
946 | 0 | if exponent == 0 { |
947 | 0 | if mantissa == 0 { |
948 | | // Zero (positive or negative) |
949 | 0 | return if sign == 1 { -0.0 } else { 0.0 }; |
950 | 0 | } |
951 | | // Subnormal f16 -> normalized f32 |
952 | | // 2^-14 as constant to avoid powi() call |
953 | | const TWO_POW_NEG_14: f32 = 6.103_515_625e-5; // 2^-14 |
954 | 0 | let m = mantissa as f32 * (1.0 / 1024.0); |
955 | 0 | let result = m * TWO_POW_NEG_14; |
956 | 0 | return if sign == 1 { -result } else { result }; |
957 | 0 | } |
958 | | |
959 | | // exponent == 31: Inf or NaN |
960 | 0 | if mantissa == 0 { |
961 | 0 | if sign == 1 { f32::NEG_INFINITY } else { f32::INFINITY } |
962 | | } else { |
963 | 0 | f32::NAN |
964 | | } |
965 | 0 | } |
966 | | |
967 | | /// Extract 6-bit scale and min values from packed scales array |
968 | | /// |
969 | | /// Q4_K uses 6-bit packed scales: 12 bytes encode 8 (scale, min) pairs. |
970 | | /// |
971 | | /// # Performance |
972 | | /// |
973 | | /// Uses bitwise operations to avoid branches and bounds checks in the hot path. |
974 | | /// The scales array is always 12 bytes, so we use unchecked access after |
975 | | /// validating at the entry point. |
976 | | #[inline(always)] |
977 | 0 | fn extract_scale_min_6bit(scales: &[u8], idx: usize) -> (f32, f32) { |
978 | 0 | debug_assert!(scales.len() >= 12, "scales array must be at least 12 bytes"); |
979 | 0 | debug_assert!(idx < 8, "idx must be < 8"); |
980 | | |
981 | | // Precomputed base offsets: idx * 3 / 2 for idx 0..8 |
982 | | // [0, 1, 3, 4, 6, 7, 9, 10] |
983 | | // Using bitwise: base = idx + (idx >> 1) |
984 | 0 | let base = idx + (idx >> 1); |
985 | | |
986 | | // Branchless extraction using bitwise selection |
987 | | // Even indices: scale = byte[base] & 0x3F |
988 | | // Odd indices: scale = (byte[base] >> 6) | ((byte[base+1] & 0x0F) << 2) |
989 | 0 | let is_odd = idx & 1; |
990 | | |
991 | | // Safety: base is always < 11 for idx < 8, and scales.len() >= 12 |
992 | 0 | let b0 = scales[base]; |
993 | 0 | let b1 = scales[base + 1]; |
994 | | |
995 | | // Extract scale: branchless using masking |
996 | 0 | let scale_even = (b0 & 0x3F) as u32; |
997 | 0 | let scale_odd = ((b0 >> 6) | ((b1 & 0x0F) << 2)) as u32; |
998 | 0 | let scale = if is_odd == 0 { scale_even } else { scale_odd }; |
999 | | |
1000 | | // Extract min: branchless using masking |
1001 | 0 | let min_even = ((b0 >> 6) | ((b1 & 0x0F) << 2)) as u32; |
1002 | | // For odd indices, we need byte at base+2, but use 0 if at boundary |
1003 | 0 | let b2 = if base + 2 < scales.len() { scales[base + 2] } else { 0 }; |
1004 | 0 | let min_odd = ((b1 >> 4) | ((b2 & 0x03) << 4)) as u32; |
1005 | 0 | let min = if is_odd == 0 { min_even } else { min_odd }; |
1006 | | |
1007 | 0 | (scale as f32, min as f32) |
1008 | 0 | } |
1009 | | |
1010 | | /// Precompute all 8 scale/min pairs for a Q4_K superblock |
1011 | | /// |
1012 | | /// More efficient than calling extract_scale_min_6bit 8 times when |
1013 | | /// we need all values (which is the common case). |
1014 | | #[inline] |
1015 | 0 | fn precompute_scales_mins(scales: &[u8]) -> [(f32, f32); 8] { |
1016 | 0 | debug_assert!(scales.len() >= 12); |
1017 | | |
1018 | | // Unroll the extraction for all 8 chunks |
1019 | 0 | [ |
1020 | 0 | extract_scale_min_6bit(scales, 0), |
1021 | 0 | extract_scale_min_6bit(scales, 1), |
1022 | 0 | extract_scale_min_6bit(scales, 2), |
1023 | 0 | extract_scale_min_6bit(scales, 3), |
1024 | 0 | extract_scale_min_6bit(scales, 4), |
1025 | 0 | extract_scale_min_6bit(scales, 5), |
1026 | 0 | extract_scale_min_6bit(scales, 6), |
1027 | 0 | extract_scale_min_6bit(scales, 7), |
1028 | 0 | ] |
1029 | 0 | } |
1030 | | |
1031 | | // ============================================================================ |
1032 | | // Error Types |
1033 | | // ============================================================================ |
1034 | | |
1035 | | /// Tiling configuration errors |
1036 | | #[derive(Debug, Clone)] |
1037 | | pub enum TilingError { |
1038 | | /// Tile hierarchy is invalid (e.g., micro > midi) |
1039 | | InvalidHierarchy { reason: String }, |
1040 | | /// Tile dimensions not divisible |
1041 | | DivisibilityError { |
1042 | | level: &'static str, |
1043 | | dimension: &'static str, |
1044 | | larger: u32, |
1045 | | smaller: u32, |
1046 | | }, |
1047 | | /// Tile doesn't fit in cache |
1048 | | CacheOverflow { |
1049 | | level: TcbLevel, |
1050 | | required_bytes: usize, |
1051 | | available_bytes: usize, |
1052 | | }, |
1053 | | /// Alignment violation |
1054 | | AlignmentError { |
1055 | | required: u32, |
1056 | | actual: u32, |
1057 | | }, |
1058 | | /// Quantization alignment violated |
1059 | | QuantAlignmentError { |
1060 | | format: &'static str, |
1061 | | required_k: u32, |
1062 | | actual_k: u32, |
1063 | | }, |
1064 | | } |
1065 | | |
1066 | | impl fmt::Display for TilingError { |
1067 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
1068 | 0 | match self { |
1069 | 0 | TilingError::InvalidHierarchy { reason } => { |
1070 | 0 | write!(f, "Invalid tiling hierarchy: {}", reason) |
1071 | | } |
1072 | | TilingError::DivisibilityError { |
1073 | 0 | level, |
1074 | 0 | dimension, |
1075 | 0 | larger, |
1076 | 0 | smaller, |
1077 | | } => { |
1078 | 0 | write!( |
1079 | 0 | f, |
1080 | 0 | "Tiling divisibility error at {}: {} ({}) not divisible by {}", |
1081 | | level, dimension, larger, smaller |
1082 | | ) |
1083 | | } |
1084 | | TilingError::CacheOverflow { |
1085 | 0 | level, |
1086 | 0 | required_bytes, |
1087 | 0 | available_bytes, |
1088 | | } => { |
1089 | 0 | write!( |
1090 | 0 | f, |
1091 | 0 | "Tile exceeds {:?} cache: {} bytes required, {} available", |
1092 | | level, required_bytes, available_bytes |
1093 | | ) |
1094 | | } |
1095 | 0 | TilingError::AlignmentError { required, actual } => { |
1096 | 0 | write!( |
1097 | 0 | f, |
1098 | 0 | "Alignment error: required {} bytes, actual {} bytes", |
1099 | | required, actual |
1100 | | ) |
1101 | | } |
1102 | | TilingError::QuantAlignmentError { |
1103 | 0 | format, |
1104 | 0 | required_k, |
1105 | 0 | actual_k, |
1106 | | } => { |
1107 | 0 | write!( |
1108 | 0 | f, |
1109 | 0 | "Quantization alignment error for {}: K must be multiple of {}, got {}", |
1110 | | format, required_k, actual_k |
1111 | | ) |
1112 | | } |
1113 | | } |
1114 | 0 | } |
1115 | | } |
1116 | | |
1117 | | impl std::error::Error for TilingError {} |
1118 | | |
1119 | | // ============================================================================ |
1120 | | // Tests |
1121 | | // ============================================================================ |
1122 | | |
1123 | | #[cfg(test)] |
1124 | | mod tests { |
1125 | | use super::*; |
1126 | | |
1127 | | // F301: TCB Output Equivalence - tested via property tests |
1128 | | // F302: Tile Size Power of 2 - static analysis |
1129 | | |
1130 | | #[test] |
1131 | | fn test_tcb_geometry_creation() { |
1132 | | let geom = TcbGeometry::new(4, 8, 256); |
1133 | | assert_eq!(geom.m, 4); |
1134 | | assert_eq!(geom.n, 8); |
1135 | | assert_eq!(geom.k, 256); |
1136 | | assert_eq!(geom.alignment, 16); |
1137 | | } |
1138 | | |
1139 | | #[test] |
1140 | | fn test_tcb_geometry_alignment() { |
1141 | | let geom = TcbGeometry::with_alignment(4, 16, 128, 64); |
1142 | | assert_eq!(geom.alignment, 64); |
1143 | | } |
1144 | | |
1145 | | #[test] |
1146 | | #[should_panic(expected = "TCB dimensions must be non-zero")] |
1147 | | fn test_tcb_geometry_zero_dimension() { |
1148 | | let _ = TcbGeometry::new(0, 8, 256); |
1149 | | } |
1150 | | |
1151 | | #[test] |
1152 | | #[should_panic(expected = "Alignment must be power of 2")] |
1153 | | fn test_tcb_geometry_invalid_alignment() { |
1154 | | let _ = TcbGeometry::with_alignment(4, 8, 256, 17); |
1155 | | } |
1156 | | |
1157 | | #[test] |
1158 | | fn test_arithmetic_intensity() { |
1159 | | // 4×8×256 tile |
1160 | | let geom = TcbGeometry::new(4, 8, 256); |
1161 | | let ai = geom.arithmetic_intensity(); |
1162 | | // AI = 2*4*8*256 / ((4*256 + 256*8) * 4) = 16384 / 12288 ≈ 1.33 |
1163 | | assert!((ai - 1.33).abs() < 0.1); |
1164 | | } |
1165 | | |
1166 | | #[test] |
1167 | | fn test_q4k_alignment() { |
1168 | | let aligned = TcbGeometry::new(4, 8, 256); |
1169 | | assert!(aligned.is_q4k_aligned()); |
1170 | | |
1171 | | let unaligned = TcbGeometry::new(4, 8, 128); |
1172 | | assert!(!unaligned.is_q4k_aligned()); |
1173 | | } |
1174 | | |
1175 | | #[test] |
1176 | | fn test_cache_fitting() { |
1177 | | let geom = TcbGeometry::new(64, 64, 64); |
1178 | | // A: 64*64*4 = 16KB, B: 64*64*4 = 16KB, total = 32KB |
1179 | | assert!(geom.fits_in_cache(64 * 1024)); // 64KB cache |
1180 | | assert!(!geom.fits_in_cache(16 * 1024)); // 16KB cache |
1181 | | } |
1182 | | |
1183 | | #[test] |
1184 | | fn test_tiling_config_gpu_q4k_matvec() { |
1185 | | let config = TilingConfig::gpu_q4k_matvec(); |
1186 | | assert_eq!(config.macro_tile.m, 1); |
1187 | | assert_eq!(config.macro_tile.k, 256); |
1188 | | assert!(config.macro_tile.is_q4k_aligned()); |
1189 | | assert!(config.validate().is_ok()); |
1190 | | } |
1191 | | |
1192 | | #[test] |
1193 | | fn test_tiling_config_cpu_avx2() { |
1194 | | let config = TilingConfig::cpu_avx2_matmul(); |
1195 | | assert_eq!(config.micro_tile.n, 8); // AVX2 = 8 floats |
1196 | | assert!(config.validate().is_ok()); |
1197 | | } |
1198 | | |
1199 | | #[test] |
1200 | | fn test_tiling_config_validation_failure() { |
1201 | | let mut config = TilingConfig::cpu_avx2_matmul(); |
1202 | | // Make midi larger than macro (invalid) |
1203 | | config.midi_tile.m = config.macro_tile.m + 1; |
1204 | | assert!(config.validate().is_err()); |
1205 | | } |
1206 | | |
1207 | | #[test] |
1208 | | fn test_index_calculator_macro_offset() { |
1209 | | let config = TilingConfig::cpu_avx2_matmul(); |
1210 | | let calc = TcbIndexCalculator::new(config.clone(), 1024, 1024, 1024); |
1211 | | |
1212 | | let (row, col) = calc.macro_tile_offset(0); |
1213 | | assert_eq!((row, col), (0, 0)); |
1214 | | |
1215 | | let (_row, col) = calc.macro_tile_offset(1); |
1216 | | assert_eq!(col, config.macro_tile.n); |
1217 | | } |
1218 | | |
1219 | | #[test] |
1220 | | fn test_index_calculator_boundary() { |
1221 | | let config = TilingConfig::cpu_avx2_matmul(); |
1222 | | |
1223 | | // With 512×512 problem and 256×256 tiles, first tile is NOT a boundary |
1224 | | let calc_large = TcbIndexCalculator::new(config.clone(), 512, 512, 256); |
1225 | | assert!(!calc_large.is_boundary_tile(0)); |
1226 | | |
1227 | | // With 100×100 problem and 256×256 tiles, first (only) tile IS a boundary |
1228 | | let calc_small = TcbIndexCalculator::new(config, 100, 100, 256); |
1229 | | assert!(calc_small.is_boundary_tile(0)); |
1230 | | |
1231 | | // Actual dimensions should be clamped to problem size |
1232 | | let (actual_m, actual_n) = calc_small.actual_tile_dims(0); |
1233 | | assert_eq!(actual_m, 100); |
1234 | | assert_eq!(actual_n, 100); |
1235 | | } |
1236 | | |
1237 | | #[test] |
1238 | | fn test_pack_a_index() { |
1239 | | // mr=4, kc=256, panel 0 |
1240 | | let idx = pack_a_index(0, 0, 4, 256, 64); |
1241 | | assert_eq!(idx, 0); |
1242 | | |
1243 | | // Second element in first panel |
1244 | | let idx = pack_a_index(1, 0, 4, 256, 64); |
1245 | | assert_eq!(idx, 1); |
1246 | | |
1247 | | // First element, second k |
1248 | | let idx = pack_a_index(0, 1, 4, 256, 64); |
1249 | | assert_eq!(idx, 4); |
1250 | | } |
1251 | | |
1252 | | #[test] |
1253 | | fn test_swizzle_index() { |
1254 | | // XOR swizzling should avoid bank conflicts |
1255 | | let idx0 = swizzle_index(0); |
1256 | | let idx32 = swizzle_index(32); |
1257 | | // These would conflict without swizzling (both bank 0) |
1258 | | // With swizzling: 0 ^ 0 = 0, 32 ^ 1 = 33 |
1259 | | assert_ne!(idx0 % 32, idx32 % 32); |
1260 | | } |
1261 | | |
1262 | | #[test] |
1263 | | fn test_optimal_prefetch_distance() { |
1264 | | let geom = TcbGeometry::new(4, 8, 64); |
1265 | | let dist = optimal_prefetch_distance(&geom, TcbLevel::Midi); |
1266 | | assert!(dist >= 1); |
1267 | | } |
1268 | | |
1269 | | // F321: Odd-Sized Matrix Handling |
1270 | | #[test] |
1271 | | fn test_odd_sized_matrices() { |
1272 | | let config = TilingConfig::cpu_avx2_matmul(); |
1273 | | |
1274 | | // Test various odd sizes |
1275 | | for (m, n, k) in [(127, 255, 513), (1, 1, 1), (7, 13, 31)] { |
1276 | | let calc = TcbIndexCalculator::new(config.clone(), m, n, k); |
1277 | | let num_tiles = calc.num_k_blocks(); |
1278 | | assert!(num_tiles >= 1); |
1279 | | } |
1280 | | } |
1281 | | |
1282 | | // F322: Zero-Padding Efficiency |
1283 | | #[test] |
1284 | | fn test_tile_count_calculation() { |
1285 | | let config = TilingConfig::cpu_avx2_matmul(); |
1286 | | let calc = TcbIndexCalculator::new(config.clone(), 1024, 1024, 1024); |
1287 | | |
1288 | | let num_macro = calc.config.num_macro_tiles(1024, 1024); |
1289 | | let num_midi = calc.config.midi_tiles_per_macro(); |
1290 | | let num_micro = calc.config.micro_tiles_per_midi(); |
1291 | | |
1292 | | assert!(num_macro > 0); |
1293 | | assert!(num_midi > 0); |
1294 | | assert!(num_micro > 0); |
1295 | | } |
1296 | | |
1297 | | // TILE-003: Q4K MatVec Tests |
1298 | | #[test] |
1299 | | fn test_tiled_q4k_matvec_creation() { |
1300 | | let matvec = TiledQ4KMatvec::new(4096, 4096); |
1301 | | assert_eq!(matvec.m, 4096); |
1302 | | assert_eq!(matvec.k, 4096); |
1303 | | assert_eq!(matvec.superblocks_per_row(), 16); // 4096 / 256 |
1304 | | assert_eq!(matvec.total_superblocks(), 4096 * 16); |
1305 | | } |
1306 | | |
1307 | | #[test] |
1308 | | #[should_panic(expected = "K dimension")] |
1309 | | fn test_tiled_q4k_matvec_unaligned_k() { |
1310 | | let _ = TiledQ4KMatvec::new(4096, 100); // Not aligned to 256 |
1311 | | } |
1312 | | |
1313 | | #[test] |
1314 | | fn test_tiled_q4k_matvec_weight_offset() { |
1315 | | let matvec = TiledQ4KMatvec::new(100, 512); |
1316 | | // Row 0: offset 0 |
1317 | | assert_eq!(matvec.weight_row_offset(0), 0); |
1318 | | // Row 1: offset = 2 superblocks * 144 bytes = 288 |
1319 | | assert_eq!(matvec.weight_row_offset(1), 2 * Q4K_SUPERBLOCK_BYTES); |
1320 | | } |
1321 | | |
1322 | | #[test] |
1323 | | fn test_tiled_q4k_matvec_optimal_rows() { |
1324 | | let matvec = TiledQ4KMatvec::new(4096, 4096); |
1325 | | // With 256KB L2, should fit many rows |
1326 | | let rows = matvec.optimal_parallel_rows(256 * 1024); |
1327 | | assert!(rows >= 4); // At least micro-kernel size |
1328 | | assert!(rows <= 4096); // At most all rows |
1329 | | } |
1330 | | |
1331 | | #[test] |
1332 | | fn test_tiled_q4k_matvec_stats() { |
1333 | | let matvec = TiledQ4KMatvec::new(4096, 4096); |
1334 | | let stats = matvec.stats(); |
1335 | | |
1336 | | // Weight bytes: 4096 * 16 * 144 = 9,437,184 bytes |
1337 | | assert_eq!(stats.superblocks, 4096 * 16); |
1338 | | // Arithmetic ops: 4096 * 4096 * 2 = 33,554,432 |
1339 | | assert_eq!(stats.arithmetic_ops, 4096 * 4096 * 2); |
1340 | | // AI should be reasonable for Q4K |
1341 | | assert!(stats.arithmetic_intensity > 1.0); |
1342 | | } |
1343 | | |
1344 | | #[test] |
1345 | | fn test_q4k_constants() { |
1346 | | assert_eq!(Q4K_SUPERBLOCK_SIZE, 256); |
1347 | | assert_eq!(Q4K_SUPERBLOCK_BYTES, 144); |
1348 | | } |
1349 | | |
1350 | | // TILE-004: AVX-512 Register Tiling Tests |
1351 | | #[test] |
1352 | | fn test_tiling_config_avx512_matmul() { |
1353 | | let config = TilingConfig::cpu_avx512_matmul(); |
1354 | | assert_eq!(config.micro_tile.n, 16); // AVX-512 = 16 floats |
1355 | | assert_eq!(config.micro_tile.alignment, 64); // 64-byte alignment |
1356 | | assert!(config.validate().is_ok()); |
1357 | | } |
1358 | | |
1359 | | #[test] |
1360 | | fn test_tiling_config_avx512_q4k_matvec() { |
1361 | | let config = TilingConfig::cpu_avx512_q4k_matvec(); |
1362 | | assert!(config.micro_tile.is_q4k_aligned()); |
1363 | | assert_eq!(config.micro_tile.m, 4); // 4×1 micro-kernel |
1364 | | assert_eq!(config.micro_tile.n, 1); // Single output column (matvec) |
1365 | | assert!(config.validate().is_ok()); |
1366 | | } |
1367 | | |
1368 | | #[test] |
1369 | | fn test_tiling_config_avx512_vnni() { |
1370 | | let config = TilingConfig::cpu_avx512_vnni_q4k_q8k(); |
1371 | | assert!(config.micro_tile.is_q4k_aligned()); |
1372 | | assert_eq!(config.backend, TilingBackend::CpuAvx512); |
1373 | | assert!(config.validate().is_ok()); |
1374 | | } |
1375 | | |
1376 | | #[test] |
1377 | | fn test_avx512_vs_avx2_tile_sizes() { |
1378 | | let avx2 = TilingConfig::cpu_avx2_matmul(); |
1379 | | let avx512 = TilingConfig::cpu_avx512_matmul(); |
1380 | | |
1381 | | // AVX-512 should have 2x wider micro-tiles |
1382 | | assert_eq!(avx512.micro_tile.n, avx2.micro_tile.n * 2); |
1383 | | |
1384 | | // AVX-512 should have stricter alignment |
1385 | | assert!(avx512.micro_tile.alignment >= avx2.micro_tile.alignment); |
1386 | | } |
1387 | | |
1388 | | // TILE-005: F321-F340 Boundary Handling Tests |
1389 | | // F321: Odd-Sized Matrix Handling (already exists above) |
1390 | | |
1391 | | // F323: Single-element matrices |
1392 | | #[test] |
1393 | | fn test_single_element_matrix() { |
1394 | | let config = TilingConfig::cpu_avx2_matmul(); |
1395 | | let calc = TcbIndexCalculator::new(config, 1, 1, 256); |
1396 | | |
1397 | | assert!(calc.is_boundary_tile(0)); |
1398 | | let (actual_m, actual_n) = calc.actual_tile_dims(0); |
1399 | | assert_eq!(actual_m, 1); |
1400 | | assert_eq!(actual_n, 1); |
1401 | | } |
1402 | | |
1403 | | // F324: Prime-sized matrices (no clean tiling) |
1404 | | #[test] |
1405 | | fn test_prime_sized_matrices() { |
1406 | | let config = TilingConfig::cpu_avx2_matmul(); |
1407 | | |
1408 | | // Prime sizes: 127, 251, 509 (all < macro_tile.m which is 256) |
1409 | | for size in [127, 251] { |
1410 | | let calc = TcbIndexCalculator::new(config.clone(), size, size, 256); |
1411 | | let num_tiles = config.num_macro_tiles(size, size); |
1412 | | assert!(num_tiles >= 1); |
1413 | | |
1414 | | // Tiles smaller than macro size are boundary tiles |
1415 | | assert!(calc.is_boundary_tile(0)); |
1416 | | } |
1417 | | |
1418 | | // 509 > 256, so first tile is NOT a boundary, but second tile IS |
1419 | | let calc = TcbIndexCalculator::new(config.clone(), 509, 509, 256); |
1420 | | // First tile (0,0 to 255,255) is not boundary for 509×509 |
1421 | | assert!(!calc.is_boundary_tile(0)); |
1422 | | // Second tile (0,256 to 255,508) IS boundary (509-256=253 < 256) |
1423 | | assert!(calc.is_boundary_tile(1)); |
1424 | | } |
1425 | | |
1426 | | // F325: K dimension exactly equals superblock |
1427 | | #[test] |
1428 | | fn test_k_equals_superblock() { |
1429 | | let matvec = TiledQ4KMatvec::new(100, 256); |
1430 | | assert_eq!(matvec.superblocks_per_row(), 1); |
1431 | | assert_eq!(matvec.total_superblocks(), 100); |
1432 | | } |
1433 | | |
1434 | | // F326: Very large M dimension |
1435 | | #[test] |
1436 | | fn test_large_m_dimension() { |
1437 | | let matvec = TiledQ4KMatvec::new(100_000, 256); |
1438 | | assert_eq!(matvec.superblocks_per_row(), 1); |
1439 | | assert_eq!(matvec.total_superblocks(), 100_000); |
1440 | | // Should still compute optimal rows |
1441 | | let rows = matvec.optimal_parallel_rows(256 * 1024); |
1442 | | assert!(rows >= 4); |
1443 | | } |
1444 | | |
1445 | | // F327: Very large K dimension |
1446 | | #[test] |
1447 | | fn test_large_k_dimension() { |
1448 | | let matvec = TiledQ4KMatvec::new(10, 32768); // 32K hidden dim |
1449 | | assert_eq!(matvec.superblocks_per_row(), 128); |
1450 | | let stats = matvec.stats(); |
1451 | | assert!(stats.arithmetic_intensity > 0.0); |
1452 | | } |
1453 | | |
1454 | | // F328: Tile offset at boundaries |
1455 | | #[test] |
1456 | | fn test_tile_offset_boundaries() { |
1457 | | let config = TilingConfig::cpu_avx2_matmul(); |
1458 | | let calc = TcbIndexCalculator::new(config.clone(), 1000, 1000, 256); |
1459 | | |
1460 | | // Last tile index |
1461 | | let num_tiles = config.num_macro_tiles(1000, 1000); |
1462 | | let last_idx = num_tiles - 1; |
1463 | | |
1464 | | let (row, col) = calc.macro_tile_offset(last_idx); |
1465 | | // Should be within bounds |
1466 | | assert!(row < 1000 + config.macro_tile.m); |
1467 | | assert!(col < 1000 + config.macro_tile.n); |
1468 | | } |
1469 | | |
1470 | | // F329: Index calculator consistency |
1471 | | #[test] |
1472 | | fn test_index_calculator_consistency() { |
1473 | | let config = TilingConfig::cpu_avx2_matmul(); |
1474 | | let calc = TcbIndexCalculator::new(config.clone(), 512, 512, 256); |
1475 | | |
1476 | | // Macro offset for tile 0 should be (0, 0) |
1477 | | let (r0, c0) = calc.macro_tile_offset(0); |
1478 | | assert_eq!((r0, c0), (0, 0)); |
1479 | | |
1480 | | // Linear offset should match |
1481 | | let linear = calc.block_to_linear_offset(0, 512); |
1482 | | assert_eq!(linear, 0); |
1483 | | |
1484 | | // A and B offsets at k_block=0 should also be 0 |
1485 | | let a_off = calc.a_offset(0, 0); |
1486 | | let b_off = calc.b_offset(0, 0); |
1487 | | assert_eq!(a_off, 0); |
1488 | | assert_eq!(b_off, 0); |
1489 | | } |
1490 | | |
1491 | | // F330: Midi/micro tile divisibility |
1492 | | #[test] |
1493 | | fn test_tile_divisibility() { |
1494 | | let config = TilingConfig::cpu_avx512_matmul(); |
1495 | | |
1496 | | // Macro should be divisible by midi |
1497 | | assert_eq!(config.macro_tile.m % config.midi_tile.m, 0); |
1498 | | assert_eq!(config.macro_tile.n % config.midi_tile.n, 0); |
1499 | | |
1500 | | // Midi should be divisible by micro |
1501 | | assert_eq!(config.midi_tile.m % config.micro_tile.m, 0); |
1502 | | assert_eq!(config.midi_tile.n % config.micro_tile.n, 0); |
1503 | | } |
1504 | | |
1505 | | // F331: f16 to f32 conversion |
1506 | | #[test] |
1507 | | fn test_f16_conversion() { |
1508 | | // Zero |
1509 | | assert_eq!(f16_to_f32(&[0x00, 0x00]), 0.0); |
1510 | | |
1511 | | // One (0x3C00 in f16) |
1512 | | let one = f16_to_f32(&[0x00, 0x3C]); |
1513 | | assert!((one - 1.0).abs() < 0.001); |
1514 | | |
1515 | | // Negative one (0xBC00) |
1516 | | let neg_one = f16_to_f32(&[0x00, 0xBC]); |
1517 | | assert!((neg_one - (-1.0)).abs() < 0.001); |
1518 | | |
1519 | | // Infinity (0x7C00) |
1520 | | assert!(f16_to_f32(&[0x00, 0x7C]).is_infinite()); |
1521 | | |
1522 | | // NaN (0x7C01) |
1523 | | assert!(f16_to_f32(&[0x01, 0x7C]).is_nan()); |
1524 | | } |
1525 | | |
1526 | | // F332: f16 subnormal conversion |
1527 | | #[test] |
1528 | | fn test_f16_subnormal() { |
1529 | | // Smallest positive subnormal: 0x0001 |
1530 | | let subnormal = f16_to_f32(&[0x01, 0x00]); |
1531 | | assert!(subnormal > 0.0); |
1532 | | assert!(subnormal < 0.001); // Very small |
1533 | | |
1534 | | // Negative zero: 0x8000 |
1535 | | let neg_zero = f16_to_f32(&[0x00, 0x80]); |
1536 | | assert_eq!(neg_zero, -0.0); |
1537 | | assert!(neg_zero.is_sign_negative()); |
1538 | | |
1539 | | // Negative infinity: 0xFC00 |
1540 | | let neg_inf = f16_to_f32(&[0x00, 0xFC]); |
1541 | | assert!(neg_inf.is_infinite()); |
1542 | | assert!(neg_inf.is_sign_negative()); |
1543 | | } |
1544 | | |
1545 | | // F333: Execute scalar implementation |
1546 | | #[test] |
1547 | | fn test_execute_scalar() { |
1548 | | let matvec = TiledQ4KMatvec::new(2, 256); |
1549 | | |
1550 | | // Create minimal valid Q4K weights (2 rows × 1 superblock each) |
1551 | | let mut weights = vec![0u8; 2 * Q4K_SUPERBLOCK_BYTES]; |
1552 | | |
1553 | | // Set up first row: d=1.0, dmin=0.0, all qs=0 |
1554 | | // f16 for 1.0 is 0x3C00 |
1555 | | weights[0] = 0x00; |
1556 | | weights[1] = 0x3C; |
1557 | | // dmin = 0 |
1558 | | weights[2] = 0x00; |
1559 | | weights[3] = 0x00; |
1560 | | // scales all zero (simplified) |
1561 | | // qs all zero -> dequantized values will be 0 |
1562 | | |
1563 | | // Second row: same setup |
1564 | | let offset = Q4K_SUPERBLOCK_BYTES; |
1565 | | weights[offset] = 0x00; |
1566 | | weights[offset + 1] = 0x3C; |
1567 | | |
1568 | | let input = vec![1.0f32; 256]; |
1569 | | let mut output = vec![0.0f32; 2]; |
1570 | | |
1571 | | matvec.execute_scalar(&weights, &input, &mut output); |
1572 | | |
1573 | | // With zero quantized values, output should be 0 or near 0 |
1574 | | // (The exact value depends on the scale/min extraction) |
1575 | | assert!(output[0].is_finite()); |
1576 | | assert!(output[1].is_finite()); |
1577 | | } |
1578 | | |
1579 | | // F334: TcbGeometry helper methods |
1580 | | #[test] |
1581 | | fn test_tcb_geometry_helpers() { |
1582 | | let geom = TcbGeometry::new(8, 16, 256); |
1583 | | |
1584 | | // total_elements = m * n |
1585 | | assert_eq!(geom.total_elements(), 8 * 16); |
1586 | | |
1587 | | // total_flops = 2 * m * n * k |
1588 | | assert_eq!(geom.total_flops(), 2 * 8 * 16 * 256); |
1589 | | |
1590 | | // Tile bytes |
1591 | | assert_eq!(geom.a_tile_bytes(), 8 * 256 * 4); |
1592 | | assert_eq!(geom.b_tile_bytes(), 256 * 16 * 4); |
1593 | | assert_eq!(geom.c_tile_bytes(), 8 * 16 * 4); |
1594 | | |
1595 | | // Q4_0 alignment |
1596 | | assert!(geom.is_q4_0_aligned()); // 256 % 32 == 0 |
1597 | | let unaligned = TcbGeometry::new(4, 4, 17); |
1598 | | assert!(!unaligned.is_q4_0_aligned()); |
1599 | | } |
1600 | | |
1601 | | // F335: TcbGeometry Display |
1602 | | #[test] |
1603 | | fn test_tcb_geometry_display() { |
1604 | | let geom = TcbGeometry::new(4, 8, 256); |
1605 | | let display = format!("{}", geom); |
1606 | | assert!(display.contains("TCB")); |
1607 | | assert!(display.contains("4×8×256")); |
1608 | | assert!(display.contains("align=16")); |
1609 | | assert!(display.contains("AI=")); |
1610 | | } |
1611 | | |
1612 | | // F336: TcbGeometry Default |
1613 | | #[test] |
1614 | | fn test_tcb_geometry_default() { |
1615 | | let geom = TcbGeometry::default(); |
1616 | | assert_eq!(geom.m, 4); |
1617 | | assert_eq!(geom.n, 4); |
1618 | | assert_eq!(geom.k, 4); |
1619 | | assert_eq!(geom.alignment, 16); |
1620 | | } |
1621 | | |
1622 | | // F337: TcbLevel typical cache bytes |
1623 | | #[test] |
1624 | | fn test_tcb_level_cache_bytes() { |
1625 | | assert_eq!(TcbLevel::Macro.typical_cache_bytes(), 32 * 1024 * 1024); |
1626 | | assert_eq!(TcbLevel::Midi.typical_cache_bytes(), 256 * 1024); |
1627 | | assert_eq!(TcbLevel::Micro.typical_cache_bytes(), 32 * 1024); |
1628 | | } |
1629 | | |
1630 | | // F338: TilingConfig factory methods |
1631 | | #[test] |
1632 | | fn test_tiling_config_gpu_softmax() { |
1633 | | let config = TilingConfig::gpu_softmax(); |
1634 | | assert_eq!(config.name, "Softmax_GPU"); |
1635 | | assert_eq!(config.macro_tile.m, 1); |
1636 | | assert_eq!(config.macro_tile.n, 32000); // Vocab size |
1637 | | assert_eq!(config.backend, TilingBackend::Gpu); |
1638 | | assert!(config.validate().is_ok()); |
1639 | | } |
1640 | | |
1641 | | #[test] |
1642 | | fn test_tiling_config_cpu_rmsnorm() { |
1643 | | let config = TilingConfig::cpu_rmsnorm(); |
1644 | | assert_eq!(config.name, "RMSNorm_CPU"); |
1645 | | assert_eq!(config.macro_tile.m, 1); |
1646 | | assert_eq!(config.backend, TilingBackend::CpuAvx512); |
1647 | | assert!(config.validate().is_ok()); |
1648 | | } |
1649 | | |
1650 | | #[test] |
1651 | | fn test_tiling_config_gpu_q4k_matmul() { |
1652 | | let config = TilingConfig::gpu_q4k_matmul(); |
1653 | | assert_eq!(config.name, "Q4K_MatMul_GPU"); |
1654 | | assert_eq!(config.macro_tile.m, 128); |
1655 | | assert!(config.macro_tile.is_q4k_aligned()); |
1656 | | assert!(config.validate().is_ok()); |
1657 | | } |
1658 | | |
1659 | | #[test] |
1660 | | fn test_tiling_config_cpu_avx2_q4k_matvec() { |
1661 | | let config = TilingConfig::cpu_avx2_q4k_matvec(); |
1662 | | assert_eq!(config.name, "Q4K_MatVec_AVX2"); |
1663 | | assert!(config.micro_tile.is_q4k_aligned()); |
1664 | | assert_eq!(config.backend, TilingBackend::CpuAvx2); |
1665 | | assert!(config.validate().is_ok()); |
1666 | | } |
1667 | | |
1668 | | // F339: TcbIndexCalculator midi/micro offsets |
1669 | | #[test] |
1670 | | fn test_index_calculator_midi_offset() { |
1671 | | let config = TilingConfig::cpu_avx2_matmul(); |
1672 | | let calc = TcbIndexCalculator::new(config, 1024, 1024, 1024); |
1673 | | |
1674 | | let (row, col) = calc.midi_tile_offset(0); |
1675 | | assert_eq!((row, col), (0, 0)); |
1676 | | |
1677 | | let (row1, col1) = calc.midi_tile_offset(1); |
1678 | | // Second midi tile should be one midi_tile.n to the right |
1679 | | assert_eq!(row1, 0); |
1680 | | assert!(col1 > 0); |
1681 | | } |
1682 | | |
1683 | | #[test] |
1684 | | fn test_index_calculator_micro_offset() { |
1685 | | let config = TilingConfig::cpu_avx2_matmul(); |
1686 | | let calc = TcbIndexCalculator::new(config, 1024, 1024, 1024); |
1687 | | |
1688 | | let (row, col) = calc.micro_tile_offset(0); |
1689 | | assert_eq!((row, col), (0, 0)); |
1690 | | |
1691 | | let (row1, col1) = calc.micro_tile_offset(1); |
1692 | | assert_eq!(row1, 0); |
1693 | | assert!(col1 > 0); |
1694 | | } |
1695 | | |
1696 | | // F340: pack_b_index |
1697 | | #[test] |
1698 | | fn test_pack_b_index() { |
1699 | | // nr=8, kc=64, panel 0 |
1700 | | let idx = pack_b_index(0, 0, 8, 64, 64); |
1701 | | assert_eq!(idx, 0); |
1702 | | |
1703 | | // Second element in first panel (col 1) |
1704 | | let idx = pack_b_index(0, 1, 8, 64, 64); |
1705 | | assert_eq!(idx, 1); |
1706 | | |
1707 | | // First element, second row |
1708 | | let idx = pack_b_index(1, 0, 8, 64, 64); |
1709 | | assert_eq!(idx, 8); |
1710 | | |
1711 | | // Second panel (col 8) |
1712 | | let idx = pack_b_index(0, 8, 8, 64, 64); |
1713 | | // panel 1 * 64 * 8 + 0 * 8 + 0 = 512 |
1714 | | assert_eq!(idx, 512); |
1715 | | } |
1716 | | |
1717 | | // F341: TilingError Display implementations |
1718 | | #[test] |
1719 | | fn test_tiling_error_display() { |
1720 | | let err1 = TilingError::InvalidHierarchy { |
1721 | | reason: "test".into(), |
1722 | | }; |
1723 | | assert!(format!("{}", err1).contains("Invalid tiling hierarchy")); |
1724 | | assert!(format!("{}", err1).contains("test")); |
1725 | | |
1726 | | let err2 = TilingError::DivisibilityError { |
1727 | | level: "macro/midi", |
1728 | | dimension: "M", |
1729 | | larger: 256, |
1730 | | smaller: 17, |
1731 | | }; |
1732 | | assert!(format!("{}", err2).contains("Tiling divisibility error")); |
1733 | | assert!(format!("{}", err2).contains("256")); |
1734 | | |
1735 | | let err3 = TilingError::CacheOverflow { |
1736 | | level: TcbLevel::Midi, |
1737 | | required_bytes: 1000, |
1738 | | available_bytes: 500, |
1739 | | }; |
1740 | | assert!(format!("{}", err3).contains("exceeds")); |
1741 | | assert!(format!("{}", err3).contains("Midi")); |
1742 | | |
1743 | | let err4 = TilingError::AlignmentError { |
1744 | | required: 64, |
1745 | | actual: 32, |
1746 | | }; |
1747 | | assert!(format!("{}", err4).contains("Alignment error")); |
1748 | | |
1749 | | let err5 = TilingError::QuantAlignmentError { |
1750 | | format: "Q4_K", |
1751 | | required_k: 256, |
1752 | | actual_k: 100, |
1753 | | }; |
1754 | | assert!(format!("{}", err5).contains("Quantization alignment")); |
1755 | | assert!(format!("{}", err5).contains("Q4_K")); |
1756 | | } |
1757 | | |
1758 | | // F342: TilingError as std::error::Error |
1759 | | #[test] |
1760 | | fn test_tiling_error_trait() { |
1761 | | let err = TilingError::InvalidHierarchy { |
1762 | | reason: "test".into(), |
1763 | | }; |
1764 | | // Ensure it implements Error trait |
1765 | | let _: &dyn std::error::Error = &err; |
1766 | | } |
1767 | | |
1768 | | // F343: TilingConfig validation - divisibility errors |
1769 | | #[test] |
1770 | | fn test_tiling_config_divisibility_error() { |
1771 | | let mut config = TilingConfig::cpu_avx2_matmul(); |
1772 | | // Make midi not divide evenly into macro |
1773 | | config.midi_tile.m = 17; // 256 % 17 != 0 |
1774 | | let result = config.validate(); |
1775 | | assert!(result.is_err()); |
1776 | | if let Err(TilingError::DivisibilityError { level, .. }) = result { |
1777 | | assert_eq!(level, "macro/midi"); |
1778 | | } else { |
1779 | | panic!("Expected DivisibilityError"); |
1780 | | } |
1781 | | } |
1782 | | |
1783 | | #[test] |
1784 | | fn test_tiling_config_micro_divisibility_error() { |
1785 | | let mut config = TilingConfig::cpu_avx2_matmul(); |
1786 | | // Make micro not divide evenly into midi |
1787 | | config.micro_tile.m = 17; // midi.m % 17 != 0 |
1788 | | let result = config.validate(); |
1789 | | assert!(result.is_err()); |
1790 | | } |
1791 | | |
1792 | | // F344: TilingStats fields |
1793 | | #[test] |
1794 | | fn test_tiling_stats_complete() { |
1795 | | let matvec = TiledQ4KMatvec::new(100, 512); |
1796 | | let stats = matvec.stats(); |
1797 | | |
1798 | | assert_eq!(stats.input_bytes, 512 * 4); |
1799 | | assert_eq!(stats.output_bytes, 100 * 4); |
1800 | | assert_eq!(stats.superblocks, 100 * 2); // 512/256 = 2 per row |
1801 | | assert!(stats.total_weight_bytes > 0); |
1802 | | } |
1803 | | |
1804 | | // F345: extract_scale_min_6bit function |
1805 | | #[test] |
1806 | | fn test_extract_scale_min_6bit() { |
1807 | | // Test with known byte patterns |
1808 | | let scales = [0x3F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]; |
1809 | | |
1810 | | // idx 0: scale from bits 0-5 of byte 0 = 0x3F = 63 |
1811 | | let (sc, _m) = extract_scale_min_6bit(&scales, 0); |
1812 | | assert_eq!(sc, 63.0); |
1813 | | |
1814 | | // Test odd index |
1815 | | let scales2 = [0xC0, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]; |
1816 | | let (sc1, _m1) = extract_scale_min_6bit(&scales2, 1); |
1817 | | assert!(sc1 >= 0.0); // Just ensure it doesn't panic |
1818 | | } |
1819 | | |
1820 | | // F346: TilingBackend enum |
1821 | | #[test] |
1822 | | fn test_tiling_backend_equality() { |
1823 | | assert_eq!(TilingBackend::CpuAvx2, TilingBackend::CpuAvx2); |
1824 | | assert_ne!(TilingBackend::CpuAvx2, TilingBackend::CpuAvx512); |
1825 | | assert_ne!(TilingBackend::Gpu, TilingBackend::Scalar); |
1826 | | assert_eq!(TilingBackend::CpuNeon, TilingBackend::CpuNeon); |
1827 | | } |
1828 | | |
1829 | | // F347: PackingLayout enum |
1830 | | #[test] |
1831 | | fn test_packing_layout_equality() { |
1832 | | assert_eq!(PackingLayout::RowMajor, PackingLayout::RowMajor); |
1833 | | assert_ne!(PackingLayout::RowMajor, PackingLayout::ColumnMajor); |
1834 | | assert_ne!(PackingLayout::PanelMajorA, PackingLayout::PanelMajorB); |
1835 | | } |
1836 | | |
1837 | | // F348: TcbLevel enum |
1838 | | #[test] |
1839 | | fn test_tcb_level_equality() { |
1840 | | assert_eq!(TcbLevel::Macro, TcbLevel::Macro); |
1841 | | assert_ne!(TcbLevel::Macro, TcbLevel::Midi); |
1842 | | assert_ne!(TcbLevel::Midi, TcbLevel::Micro); |
1843 | | } |
1844 | | |
1845 | | // F349: Serialization round-trip for TcbGeometry |
1846 | | #[test] |
1847 | | fn test_tcb_geometry_serde() { |
1848 | | let geom = TcbGeometry::with_alignment(4, 8, 256, 64); |
1849 | | let json = serde_json::to_string(&geom).unwrap(); |
1850 | | let decoded: TcbGeometry = serde_json::from_str(&json).unwrap(); |
1851 | | assert_eq!(geom, decoded); |
1852 | | } |
1853 | | |
1854 | | // F350: Serialization round-trip for TilingConfig |
1855 | | #[test] |
1856 | | fn test_tiling_config_serde() { |
1857 | | let config = TilingConfig::cpu_avx512_matmul(); |
1858 | | let json = serde_json::to_string(&config).unwrap(); |
1859 | | let decoded: TilingConfig = serde_json::from_str(&json).unwrap(); |
1860 | | assert_eq!(config.name, decoded.name); |
1861 | | assert_eq!(config.backend, decoded.backend); |
1862 | | } |
1863 | | |
1864 | | // F351: Index calculator k_blocks |
1865 | | #[test] |
1866 | | fn test_index_calculator_k_blocks() { |
1867 | | let config = TilingConfig::cpu_avx2_matmul(); |
1868 | | let calc = TcbIndexCalculator::new(config.clone(), 512, 512, 1024); |
1869 | | |
1870 | | // 1024 / 256 = 4 K blocks |
1871 | | assert_eq!(calc.num_k_blocks(), 4); |
1872 | | |
1873 | | // Non-divisible case |
1874 | | let calc2 = TcbIndexCalculator::new(config, 512, 512, 300); |
1875 | | // ceil(300 / 256) = 2 |
1876 | | assert_eq!(calc2.num_k_blocks(), 2); |
1877 | | } |
1878 | | |
1879 | | // F352: A and B offset calculations |
1880 | | #[test] |
1881 | | fn test_ab_offset_calculations() { |
1882 | | let config = TilingConfig::cpu_avx2_matmul(); |
1883 | | let calc = TcbIndexCalculator::new(config.clone(), 512, 512, 512); |
1884 | | |
1885 | | // A offset: row * problem_k + col |
1886 | | let a_off = calc.a_offset(1, 0); // macro_row=1, k_block=0 |
1887 | | assert_eq!(a_off, (config.macro_tile.m * 512) as usize); |
1888 | | |
1889 | | // B offset: row * problem_n + col |
1890 | | let b_off = calc.b_offset(0, 1); // k_block=0, macro_col=1 |
1891 | | assert_eq!(b_off, config.macro_tile.n as usize); |
1892 | | } |
1893 | | |
1894 | | // F353: Large tile calculations |
1895 | | #[test] |
1896 | | fn test_large_tile_arithmetic() { |
1897 | | // Very large geometry to test u64 arithmetic |
1898 | | let geom = TcbGeometry::new(10000, 10000, 1000); |
1899 | | let total = geom.total_elements(); |
1900 | | assert_eq!(total, 100_000_000); |
1901 | | |
1902 | | let flops = geom.total_flops(); |
1903 | | assert_eq!(flops, 200_000_000_000); |
1904 | | } |
1905 | | |
1906 | | // F354: Prefetch with different levels |
1907 | | #[test] |
1908 | | fn test_prefetch_all_levels() { |
1909 | | let geom = TcbGeometry::new(4, 8, 64); |
1910 | | |
1911 | | let dist_micro = optimal_prefetch_distance(&geom, TcbLevel::Micro); |
1912 | | let dist_midi = optimal_prefetch_distance(&geom, TcbLevel::Midi); |
1913 | | let dist_macro = optimal_prefetch_distance(&geom, TcbLevel::Macro); |
1914 | | |
1915 | | // Macro should have larger distance (higher latency) |
1916 | | assert!(dist_macro >= dist_midi); |
1917 | | assert!(dist_midi >= dist_micro); |
1918 | | // All should be at least 1 |
1919 | | assert!(dist_micro >= 1); |
1920 | | } |
1921 | | |
1922 | | // F355: PrefetchLocality Debug |
1923 | | #[test] |
1924 | | fn test_prefetch_locality_debug() { |
1925 | | let loc = PrefetchLocality::T0; |
1926 | | let debug = format!("{:?}", loc); |
1927 | | assert!(debug.contains("T0")); |
1928 | | |
1929 | | let loc2 = PrefetchLocality::NonTemporal; |
1930 | | let debug2 = format!("{:?}", loc2); |
1931 | | assert!(debug2.contains("NonTemporal")); |
1932 | | } |
1933 | | } |