/home/noah/src/trueno/src/tiling/config.rs

Source
//! Tiling configuration and backend selection.

use serde::{Deserialize, Serialize};
use super::geometry::TcbGeometry;
use super::error::TilingError;

/// Complete tiling configuration for a kernel
///
/// Contains geometry for all three tiling levels, enabling hierarchical
/// cache-aware execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TilingConfig {
    /// Kernel name for identification
    pub name: String,
    /// Macro-tile geometry (L3/Global)
    pub macro_tile: TcbGeometry,
    /// Midi-tile geometry (L2/Shared)
    pub midi_tile: TcbGeometry,
    /// Micro-tile geometry (Registers)
    pub micro_tile: TcbGeometry,
    /// Target backend
    pub backend: TilingBackend,
}

/// Backend target for tiling configuration
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TilingBackend {
    /// CPU with AVX2 (256-bit SIMD)
    CpuAvx2,
    /// CPU with AVX-512 (512-bit SIMD)
    CpuAvx512,
    /// CPU with NEON (128-bit SIMD)
    CpuNeon,
    /// GPU (CUDA/wgpu)
    Gpu,
    /// Scalar fallback
    Scalar,
}

impl TilingConfig {
    /// Create configuration for GPU Q4_K MatVec
    ///
    /// Optimized for single-token generation where M=1.
    #[must_use]
    pub fn gpu_q4k_matvec() -> Self {
        Self {
            name: "Q4K_MatVec_GPU".into(),
            macro_tile: TcbGeometry::with_alignment(1, 4096, 256, 64),
            midi_tile: TcbGeometry::with_alignment(1, 256, 256, 64),
            micro_tile: TcbGeometry::with_alignment(1, 32, 256, 64),
            backend: TilingBackend::Gpu,
        }
    }

    /// Create configuration for GPU Q4_K MatMul (batched)
    ///
    /// Optimized for prefill where M > 1.
    #[must_use]
    pub fn gpu_q4k_matmul() -> Self {
        Self {
            name: "Q4K_MatMul_GPU".into(),
            macro_tile: TcbGeometry::with_alignment(128, 128, 256, 64),
            midi_tile: TcbGeometry::with_alignment(32, 32, 256, 64),
            micro_tile: TcbGeometry::with_alignment(8, 8, 256, 64),
            backend: TilingBackend::Gpu,
        }
    }

    /// Create configuration for GPU Softmax
    #[must_use]
    pub fn gpu_softmax() -> Self {
        Self {
            name: "Softmax_GPU".into(),
            macro_tile: TcbGeometry::with_alignment(1, 32000, 1, 64),
            midi_tile: TcbGeometry::with_alignment(1, 1024, 1, 64),
            micro_tile: TcbGeometry::with_alignment(1, 32, 1, 64),
            backend: TilingBackend::Gpu,
        }
    }

    /// Create configuration for CPU AVX-512 MatMul
    ///
    /// Optimized for 512-bit wide SIMD:
    /// - 16 floats per ZMM register
    /// - 32 ZMM registers available
    /// - 4×16 micro-kernel uses 8 registers (4 accumulators + 4 scratch)
    #[must_use]
    pub fn cpu_avx512_matmul() -> Self {
        Self {
            name: "MatMul_AVX512".into(),
            macro_tile: TcbGeometry::with_alignment(512, 512, 512, 64),
            midi_tile: TcbGeometry::with_alignment(128, 128, 128, 64),
            // 16 floats wide × 4 rows = 64 elements in registers
            micro_tile: TcbGeometry::with_alignment(4, 16, 128, 64),
            backend: TilingBackend::CpuAvx512,
        }
    }

    /// Create configuration for CPU AVX-512 Q4K MatVec
    ///
    /// Optimized for Q4_K quantized inference with 512-bit SIMD.
    /// Key differences from AVX2:
    /// - 64-byte aligned for cache line optimization
    /// - 4×1 micro-kernel processes 4 rows simultaneously
    /// - K=256 aligned to Q4_K superblock
    #[must_use]
    pub fn cpu_avx512_q4k_matvec() -> Self {
        Self {
            name: "Q4K_MatVec_AVX512".into(),
            // Large macro-tile to amortize L3 access
            macro_tile: TcbGeometry::with_alignment(4096, 1, 4096, 64),
            // Midi-tile fits in L2 (256KB)
            // 64 rows × 256 K × 0.5625 bytes/element ≈ 9KB weights
            midi_tile: TcbGeometry::with_alignment(64, 1, 256, 64),
            // 4 rows × 1 output, K=256 (Q4_K superblock)
            micro_tile: TcbGeometry::with_alignment(4, 1, 256, 64),
            backend: TilingBackend::CpuAvx512,
        }
    }

    /// Create configuration for AVX-512 VNNI Q4K×Q8K integer dot product
    ///
    /// AVX-512 VNNI (Vector Neural Network Instructions) provides:
    /// - VPDPBUSD: 8-bit unsigned × 8-bit signed multiply-add to i32
    /// - VPDPWSSD: 16-bit signed × 16-bit signed multiply-add to i32
    ///
    /// This enables pure integer Q4K×Q8K without intermediate f32 conversion.
    #[must_use]
    pub fn cpu_avx512_vnni_q4k_q8k() -> Self {
        Self {
            name: "Q4K_Q8K_VNNI".into(),
            macro_tile: TcbGeometry::with_alignment(4096, 1, 4096, 64),
            midi_tile: TcbGeometry::with_alignment(64, 1, 256, 64),
            // VNNI processes 64 i8 values per ZMM register
            micro_tile: TcbGeometry::with_alignment(4, 1, 256, 64),
            backend: TilingBackend::CpuAvx512,
        }
    }

    /// Create configuration for CPU AVX2 MatMul
    #[must_use]
    pub fn cpu_avx2_matmul() -> Self {
        Self {
            name: "MatMul_AVX2".into(),
            macro_tile: TcbGeometry::with_alignment(256, 256, 256, 32),
            midi_tile: TcbGeometry::with_alignment(64, 64, 64, 32),
            // 8 floats wide × 4 rows = 32 elements in registers
            micro_tile: TcbGeometry::with_alignment(4, 8, 64, 32),
            backend: TilingBackend::CpuAvx2,
        }
    }

    /// Create configuration for CPU Q4_K MatVec (AVX2)
    #[must_use]
    pub fn cpu_avx2_q4k_matvec() -> Self {
        Self {
            name: "Q4K_MatVec_AVX2".into(),
            // Process 4 rows at a time (4×1 micro-kernel)
            macro_tile: TcbGeometry::with_alignment(4096, 1, 4096, 32),
            midi_tile: TcbGeometry::with_alignment(64, 1, 256, 32),
            // 4 rows × 1 output, K=256 (Q4_K superblock)
            micro_tile: TcbGeometry::with_alignment(4, 1, 256, 32),
            backend: TilingBackend::CpuAvx2,
        }
    }

    /// Create configuration for RMSNorm (CPU)
    #[must_use]
    pub fn cpu_rmsnorm() -> Self {
        Self {
            name: "RMSNorm_CPU".into(),
            macro_tile: TcbGeometry::with_alignment(1, 4096, 1, 32),
            midi_tile: TcbGeometry::with_alignment(1, 256, 1, 32),
            micro_tile: TcbGeometry::with_alignment(1, 16, 1, 32),
            backend: TilingBackend::CpuAvx512,
        }
    }

    /// Validate that tiling configuration is internally consistent
    pub fn validate(&self) -> Result<(), TilingError> {
        // Macro must be >= Midi >= Micro
        if self.midi_tile.m > self.macro_tile.m
            || self.midi_tile.n > self.macro_tile.n
            || self.midi_tile.k > self.macro_tile.k
        {
            return Err(TilingError::InvalidHierarchy {
                reason: "Midi-tile larger than macro-tile".into(),
            });
        }

        if self.micro_tile.m > self.midi_tile.m
            || self.micro_tile.n > self.midi_tile.n
            || self.micro_tile.k > self.midi_tile.k
        {
            return Err(TilingError::InvalidHierarchy {
                reason: "Micro-tile larger than midi-tile".into(),
            });
        }

        // Check divisibility
        if self.macro_tile.m % self.midi_tile.m != 0 {
            return Err(TilingError::DivisibilityError {
                level: "macro/midi",
                dimension: "M",
                larger: self.macro_tile.m,
                smaller: self.midi_tile.m,
            });
        }

        if self.midi_tile.m % self.micro_tile.m != 0 {
            return Err(TilingError::DivisibilityError {
                level: "midi/micro",
                dimension: "M",
                larger: self.midi_tile.m,
                smaller: self.micro_tile.m,
            });
        }

        Ok(())
    }

    /// Calculate total number of macro-tiles for given problem size
    #[must_use]
    pub fn num_macro_tiles(&self, m: u32, n: u32) -> u32 {
        let m_tiles = (m + self.macro_tile.m - 1) / self.macro_tile.m;
        let n_tiles = (n + self.macro_tile.n - 1) / self.macro_tile.n;
        m_tiles * n_tiles
    }

    /// Calculate total number of midi-tiles within a macro-tile
    #[must_use]
    pub fn midi_tiles_per_macro(&self) -> u32 {
        let m_tiles = self.macro_tile.m / self.midi_tile.m;
        let n_tiles = self.macro_tile.n / self.midi_tile.n;
        m_tiles * n_tiles
    }

    /// Calculate total number of micro-tiles within a midi-tile
    #[must_use]
    pub fn micro_tiles_per_midi(&self) -> u32 {
        let m_tiles = self.midi_tile.m / self.micro_tile.m;
        let n_tiles = self.midi_tile.n / self.micro_tile.n;
        m_tiles * n_tiles
    }
}

Coverage Report

Created: 2026-01-25 15:05

Line	Count	Source
1		//! Tiling configuration and backend selection.
2
3		use serde::{Deserialize, Serialize};
4		use super::geometry::TcbGeometry;
5		use super::error::TilingError;
6
7		/// Complete tiling configuration for a kernel
8		///
9		/// Contains geometry for all three tiling levels, enabling hierarchical
10		/// cache-aware execution.
11		#[derive(Debug, Clone, Serialize, Deserialize)]
12		pub struct TilingConfig {
13		/// Kernel name for identification
14		pub name: String,
15		/// Macro-tile geometry (L3/Global)
16		pub macro_tile: TcbGeometry,
17		/// Midi-tile geometry (L2/Shared)
18		pub midi_tile: TcbGeometry,
19		/// Micro-tile geometry (Registers)
20		pub micro_tile: TcbGeometry,
21		/// Target backend
22		pub backend: TilingBackend,
23		}
24
25		/// Backend target for tiling configuration
26		#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
27		pub enum TilingBackend {
28		/// CPU with AVX2 (256-bit SIMD)
29		CpuAvx2,
30		/// CPU with AVX-512 (512-bit SIMD)
31		CpuAvx512,
32		/// CPU with NEON (128-bit SIMD)
33		CpuNeon,
34		/// GPU (CUDA/wgpu)
35		Gpu,
36		/// Scalar fallback
37		Scalar,
38		}
39
40		impl TilingConfig {
41		/// Create configuration for GPU Q4_K MatVec
42		///
43		/// Optimized for single-token generation where M=1.
44		#[must_use]
45	0	pub fn gpu_q4k_matvec() -> Self {
46	0	Self {
47	0	name: "Q4K_MatVec_GPU".into(),
48	0	macro_tile: TcbGeometry::with_alignment(1, 4096, 256, 64),
49	0	midi_tile: TcbGeometry::with_alignment(1, 256, 256, 64),
50	0	micro_tile: TcbGeometry::with_alignment(1, 32, 256, 64),
51	0	backend: TilingBackend::Gpu,
52	0	}
53	0	}
54
55		/// Create configuration for GPU Q4_K MatMul (batched)
56		///
57		/// Optimized for prefill where M > 1.
58		#[must_use]
59	0	pub fn gpu_q4k_matmul() -> Self {
60	0	Self {
61	0	name: "Q4K_MatMul_GPU".into(),
62	0	macro_tile: TcbGeometry::with_alignment(128, 128, 256, 64),
63	0	midi_tile: TcbGeometry::with_alignment(32, 32, 256, 64),
64	0	micro_tile: TcbGeometry::with_alignment(8, 8, 256, 64),
65	0	backend: TilingBackend::Gpu,
66	0	}
67	0	}
68
69		/// Create configuration for GPU Softmax
70		#[must_use]
71	0	pub fn gpu_softmax() -> Self {
72	0	Self {
73	0	name: "Softmax_GPU".into(),
74	0	macro_tile: TcbGeometry::with_alignment(1, 32000, 1, 64),
75	0	midi_tile: TcbGeometry::with_alignment(1, 1024, 1, 64),
76	0	micro_tile: TcbGeometry::with_alignment(1, 32, 1, 64),
77	0	backend: TilingBackend::Gpu,
78	0	}
79	0	}
80
81		/// Create configuration for CPU AVX-512 MatMul
82		///
83		/// Optimized for 512-bit wide SIMD:
84		/// - 16 floats per ZMM register
85		/// - 32 ZMM registers available
86		/// - 4×16 micro-kernel uses 8 registers (4 accumulators + 4 scratch)
87		#[must_use]
88	0	pub fn cpu_avx512_matmul() -> Self {
89	0	Self {
90	0	name: "MatMul_AVX512".into(),
91	0	macro_tile: TcbGeometry::with_alignment(512, 512, 512, 64),
92	0	midi_tile: TcbGeometry::with_alignment(128, 128, 128, 64),
93	0	// 16 floats wide × 4 rows = 64 elements in registers
94	0	micro_tile: TcbGeometry::with_alignment(4, 16, 128, 64),
95	0	backend: TilingBackend::CpuAvx512,
96	0	}
97	0	}
98
99		/// Create configuration for CPU AVX-512 Q4K MatVec
100		///
101		/// Optimized for Q4_K quantized inference with 512-bit SIMD.
102		/// Key differences from AVX2:
103		/// - 64-byte aligned for cache line optimization
104		/// - 4×1 micro-kernel processes 4 rows simultaneously
105		/// - K=256 aligned to Q4_K superblock
106		#[must_use]
107	0	pub fn cpu_avx512_q4k_matvec() -> Self {
108	0	Self {
109	0	name: "Q4K_MatVec_AVX512".into(),
110	0	// Large macro-tile to amortize L3 access
111	0	macro_tile: TcbGeometry::with_alignment(4096, 1, 4096, 64),
112	0	// Midi-tile fits in L2 (256KB)
113	0	// 64 rows × 256 K × 0.5625 bytes/element ≈ 9KB weights
114	0	midi_tile: TcbGeometry::with_alignment(64, 1, 256, 64),
115	0	// 4 rows × 1 output, K=256 (Q4_K superblock)
116	0	micro_tile: TcbGeometry::with_alignment(4, 1, 256, 64),
117	0	backend: TilingBackend::CpuAvx512,
118	0	}
119	0	}
120
121		/// Create configuration for AVX-512 VNNI Q4K×Q8K integer dot product
122		///
123		/// AVX-512 VNNI (Vector Neural Network Instructions) provides:
124		/// - VPDPBUSD: 8-bit unsigned × 8-bit signed multiply-add to i32
125		/// - VPDPWSSD: 16-bit signed × 16-bit signed multiply-add to i32
126		///
127		/// This enables pure integer Q4K×Q8K without intermediate f32 conversion.
128		#[must_use]
129	0	pub fn cpu_avx512_vnni_q4k_q8k() -> Self {
130	0	Self {
131	0	name: "Q4K_Q8K_VNNI".into(),
132	0	macro_tile: TcbGeometry::with_alignment(4096, 1, 4096, 64),
133	0	midi_tile: TcbGeometry::with_alignment(64, 1, 256, 64),
134	0	// VNNI processes 64 i8 values per ZMM register
135	0	micro_tile: TcbGeometry::with_alignment(4, 1, 256, 64),
136	0	backend: TilingBackend::CpuAvx512,
137	0	}
138	0	}
139
140		/// Create configuration for CPU AVX2 MatMul
141		#[must_use]
142	0	pub fn cpu_avx2_matmul() -> Self {
143	0	Self {
144	0	name: "MatMul_AVX2".into(),
145	0	macro_tile: TcbGeometry::with_alignment(256, 256, 256, 32),
146	0	midi_tile: TcbGeometry::with_alignment(64, 64, 64, 32),
147	0	// 8 floats wide × 4 rows = 32 elements in registers
148	0	micro_tile: TcbGeometry::with_alignment(4, 8, 64, 32),
149	0	backend: TilingBackend::CpuAvx2,
150	0	}
151	0	}
152
153		/// Create configuration for CPU Q4_K MatVec (AVX2)
154		#[must_use]
155	0	pub fn cpu_avx2_q4k_matvec() -> Self {
156	0	Self {
157	0	name: "Q4K_MatVec_AVX2".into(),
158	0	// Process 4 rows at a time (4×1 micro-kernel)
159	0	macro_tile: TcbGeometry::with_alignment(4096, 1, 4096, 32),
160	0	midi_tile: TcbGeometry::with_alignment(64, 1, 256, 32),
161	0	// 4 rows × 1 output, K=256 (Q4_K superblock)
162	0	micro_tile: TcbGeometry::with_alignment(4, 1, 256, 32),
163	0	backend: TilingBackend::CpuAvx2,
164	0	}
165	0	}
166
167		/// Create configuration for RMSNorm (CPU)
168		#[must_use]
169	0	pub fn cpu_rmsnorm() -> Self {
170	0	Self {
171	0	name: "RMSNorm_CPU".into(),
172	0	macro_tile: TcbGeometry::with_alignment(1, 4096, 1, 32),
173	0	midi_tile: TcbGeometry::with_alignment(1, 256, 1, 32),
174	0	micro_tile: TcbGeometry::with_alignment(1, 16, 1, 32),
175	0	backend: TilingBackend::CpuAvx512,
176	0	}
177	0	}
178
179		/// Validate that tiling configuration is internally consistent
180	0	pub fn validate(&self) -> Result<(), TilingError> {
181		// Macro must be >= Midi >= Micro
182	0	if self.midi_tile.m > self.macro_tile.m
183	0	\|\| self.midi_tile.n > self.macro_tile.n
184	0	\|\| self.midi_tile.k > self.macro_tile.k
185		{
186	0	return Err(TilingError::InvalidHierarchy {
187	0	reason: "Midi-tile larger than macro-tile".into(),
188	0	});
189	0	}
190
191	0	if self.micro_tile.m > self.midi_tile.m
192	0	\|\| self.micro_tile.n > self.midi_tile.n
193	0	\|\| self.micro_tile.k > self.midi_tile.k
194		{
195	0	return Err(TilingError::InvalidHierarchy {
196	0	reason: "Micro-tile larger than midi-tile".into(),
197	0	});
198	0	}
199
200		// Check divisibility
201	0	if self.macro_tile.m % self.midi_tile.m != 0 {
202	0	return Err(TilingError::DivisibilityError {
203	0	level: "macro/midi",
204	0	dimension: "M",
205	0	larger: self.macro_tile.m,
206	0	smaller: self.midi_tile.m,
207	0	});
208	0	}
209
210	0	if self.midi_tile.m % self.micro_tile.m != 0 {
211	0	return Err(TilingError::DivisibilityError {
212	0	level: "midi/micro",
213	0	dimension: "M",
214	0	larger: self.midi_tile.m,
215	0	smaller: self.micro_tile.m,
216	0	});
217	0	}
218
219	0	Ok(())
220	0	}
221
222		/// Calculate total number of macro-tiles for given problem size
223		#[must_use]
224	0	pub fn num_macro_tiles(&self, m: u32, n: u32) -> u32 {
225	0	let m_tiles = (m + self.macro_tile.m - 1) / self.macro_tile.m;
226	0	let n_tiles = (n + self.macro_tile.n - 1) / self.macro_tile.n;
227	0	m_tiles * n_tiles
228	0	}
229
230		/// Calculate total number of midi-tiles within a macro-tile
231		#[must_use]
232	0	pub fn midi_tiles_per_macro(&self) -> u32 {
233	0	let m_tiles = self.macro_tile.m / self.midi_tile.m;
234	0	let n_tiles = self.macro_tile.n / self.midi_tile.n;
235	0	m_tiles * n_tiles
236	0	}
237
238		/// Calculate total number of micro-tiles within a midi-tile
239		#[must_use]
240	0	pub fn micro_tiles_per_midi(&self) -> u32 {
241	0	let m_tiles = self.midi_tile.m / self.micro_tile.m;
242	0	let n_tiles = self.midi_tile.n / self.micro_tile.n;
243	0	m_tiles * n_tiles
244	0	}
245		}