/home/noah/src/trueno/src/blis/backend_selection.rs
Line | Count | Source |
1 | | //! Backend Selection and Cost Model |
2 | | //! |
3 | | //! Automatic selection between CPU (SIMD), CUDA (PTX), and wgpu (WGSL) backends |
4 | | //! based on the 5× PCIe rule and roofline analysis. |
5 | | //! |
6 | | //! # Philosophy |
7 | | //! |
8 | | //! Uses Gregg & Hazelwood (2011) "5× PCIe rule": GPU worthwhile when |
9 | | //! compute time exceeds 5× data transfer time. |
10 | | //! |
11 | | //! # References |
12 | | //! |
13 | | //! - Gregg, C., & Hazelwood, K. (2011). Where is the Data? Why You Cannot |
14 | | //! Debate CPU vs. GPU Performance Without the Answer. IEEE ISPASS. |
15 | | //! - Volkov, V. (2010). Better Performance at Lower Occupancy. |
16 | | |
17 | | #[cfg(target_arch = "x86_64")] |
18 | | use std::arch::is_x86_feature_detected; |
19 | | |
20 | | use super::profiler::BlisProfiler; |
21 | | use super::{gemm_blis, TruenoError}; |
22 | | |
23 | | /// |
24 | | /// Maps to different ISA targets: |
25 | | /// - Cpu: x86 asm (AVX2/AVX-512), ARM asm (NEON) |
26 | | /// - Gpu: PTX (CUDA), wgpu compute shaders |
27 | | /// - Wgpu: WGSL for cross-platform GPU (Vulkan/Metal/DX12/WebGPU) |
28 | | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] |
29 | | pub enum ComputeBackend { |
30 | | /// CPU SIMD backend (AVX2, AVX-512, NEON, SSE2) |
31 | | Cpu, |
32 | | /// NVIDIA GPU backend (PTX) |
33 | | #[allow(dead_code)] |
34 | | Gpu, |
35 | | /// Cross-platform GPU backend (wgpu/WGSL) |
36 | | #[allow(dead_code)] |
37 | | Wgpu, |
38 | | /// Scalar fallback (no SIMD) |
39 | | Scalar, |
40 | | } |
41 | | |
42 | | /// ComputeBrick hierarchy level |
43 | | /// |
44 | | /// Maps BLIS loop structure to brick abstraction: |
45 | | /// - Nano: Microkernel (MR×NR×K) - register file |
46 | | /// - Micro: Midi loop (MC×NC×KC) - L1/L2 cache |
47 | | /// - Meso: Macro loop (full M×N×K) - L3/DRAM |
48 | | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] |
49 | | pub enum BrickLevel { |
50 | | /// Register-level compute (MR×NR tile) |
51 | | Nano, |
52 | | /// Cache-level compute (MC×NC block) |
53 | | Micro, |
54 | | /// Memory-level compute (full matrix) |
55 | | Meso, |
56 | | } |
57 | | |
58 | | /// Cost model for backend selection |
59 | | /// |
60 | | /// Based on Gregg & Hazelwood (2011): GPU worthwhile when compute > 5× transfer |
61 | | #[derive(Debug, Clone)] |
62 | | pub struct BackendCostModel { |
63 | | /// PCIe bandwidth in GB/s (e.g., 15.75 for PCIe 3.0 x16) |
64 | | pub pcie_bandwidth_gbps: f64, |
65 | | /// GPU peak TFLOP/s |
66 | | pub gpu_peak_tflops: f64, |
67 | | /// CPU peak GFLOP/s |
68 | | pub cpu_peak_gflops: f64, |
69 | | /// Minimum problem size for GPU (elements) |
70 | | pub gpu_min_elements: usize, |
71 | | } |
72 | | |
73 | | impl Default for BackendCostModel { |
74 | 0 | fn default() -> Self { |
75 | 0 | Self { |
76 | 0 | pcie_bandwidth_gbps: 15.75, // PCIe 3.0 x16 |
77 | 0 | gpu_peak_tflops: 10.0, // Mid-range GPU |
78 | 0 | cpu_peak_gflops: 400.0, // Modern AVX2 CPU |
79 | 0 | gpu_min_elements: 1_000_000, // ~1M elements |
80 | 0 | } |
81 | 0 | } |
82 | | } |
83 | | |
84 | | impl BackendCostModel { |
85 | | /// Select optimal backend based on 5× PCIe rule |
86 | | /// |
87 | | /// # References |
88 | | /// |
89 | | /// Gregg, C., & Hazelwood, K. (2011). Where is the Data? Why You Cannot |
90 | | /// Debate CPU vs. GPU Performance Without the Answer. IEEE ISPASS. |
91 | 0 | pub fn select_backend(&self, m: usize, n: usize, k: usize) -> ComputeBackend { |
92 | 0 | let flops = 2 * m * n * k; |
93 | 0 | let bytes = 4 * (m * k + k * n + m * n); // f32 = 4 bytes |
94 | 0 | let arithmetic_intensity = flops as f64 / bytes as f64; |
95 | | |
96 | | // Ridge point: where compute = memory bandwidth |
97 | 0 | let ridge_point = self.gpu_peak_tflops * 1000.0 / self.pcie_bandwidth_gbps; |
98 | | |
99 | | // GPU worthwhile if: |
100 | | // 1. High arithmetic intensity (compute-bound) |
101 | | // 2. Problem size exceeds minimum threshold |
102 | | // 3. Transfer time is amortized (5× rule) |
103 | 0 | let elements = m * n * k; |
104 | 0 | if arithmetic_intensity > ridge_point && elements > self.gpu_min_elements { |
105 | | // Check if wgpu available at runtime |
106 | | #[cfg(feature = "wgpu")] |
107 | 0 | return ComputeBackend::Wgpu; |
108 | | |
109 | | #[cfg(all(not(feature = "wgpu"), feature = "cuda"))] |
110 | | return ComputeBackend::Gpu; |
111 | | |
112 | | #[allow(unreachable_code)] |
113 | | ComputeBackend::Cpu |
114 | | } else { |
115 | | // CPU is better for small problems or memory-bound workloads |
116 | | #[cfg(target_arch = "x86_64")] |
117 | | { |
118 | 0 | if is_x86_feature_detected!("avx2") { |
119 | 0 | return ComputeBackend::Cpu; |
120 | 0 | } |
121 | | } |
122 | | #[cfg(target_arch = "aarch64")] |
123 | | { |
124 | | return ComputeBackend::Cpu; |
125 | | } |
126 | 0 | ComputeBackend::Scalar |
127 | | } |
128 | 0 | } |
129 | | |
130 | | /// Estimate execution time in microseconds |
131 | 0 | pub fn estimate_time_us(&self, m: usize, n: usize, k: usize, backend: ComputeBackend) -> f64 { |
132 | 0 | let flops = 2.0 * m as f64 * n as f64 * k as f64; |
133 | 0 | let bytes = 4.0 * (m * k + k * n + m * n) as f64; |
134 | | |
135 | 0 | match backend { |
136 | | ComputeBackend::Gpu | ComputeBackend::Wgpu => { |
137 | | // Transfer time + compute time |
138 | 0 | let transfer_us = bytes / (self.pcie_bandwidth_gbps * 1e3); |
139 | 0 | let compute_us = flops / (self.gpu_peak_tflops * 1e6); |
140 | 0 | transfer_us + compute_us |
141 | | } |
142 | | ComputeBackend::Cpu => { |
143 | 0 | flops / (self.cpu_peak_gflops * 1e3) |
144 | | } |
145 | | ComputeBackend::Scalar => { |
146 | | // Assume 1 GFLOP/s for scalar |
147 | 0 | flops / 1e3 |
148 | | } |
149 | | } |
150 | 0 | } |
151 | | } |
152 | | |
153 | | /// Unified profiler for all backends |
154 | | /// |
155 | | /// Collects metrics across CPU (RDTSC), GPU (CUDA events), and wgpu (timestamp queries) |
156 | | #[derive(Debug, Clone, Default)] |
157 | | pub struct UnifiedBrickProfiler { |
158 | | /// CPU profiling stats |
159 | | pub cpu_stats: BlisProfiler, |
160 | | /// Selected backend for this run |
161 | | pub backend: Option<ComputeBackend>, |
162 | | /// Total elements processed |
163 | | pub total_elements: u64, |
164 | | /// Backend selection decisions |
165 | | pub selection_history: Vec<(usize, usize, usize, ComputeBackend)>, |
166 | | } |
167 | | |
168 | | impl UnifiedBrickProfiler { |
169 | | /// Create a new unified profiler |
170 | 0 | pub fn new() -> Self { |
171 | 0 | Self { |
172 | 0 | cpu_stats: BlisProfiler::enabled(), |
173 | 0 | backend: None, |
174 | 0 | total_elements: 0, |
175 | 0 | selection_history: Vec::new(), |
176 | 0 | } |
177 | 0 | } |
178 | | |
179 | | /// Record backend selection |
180 | 0 | pub fn record_selection(&mut self, m: usize, n: usize, k: usize, backend: ComputeBackend) { |
181 | 0 | self.backend = Some(backend); |
182 | 0 | self.total_elements += (m * n) as u64; |
183 | 0 | self.selection_history.push((m, n, k, backend)); |
184 | 0 | } |
185 | | |
186 | | /// Get roofline analysis for current backend |
187 | 0 | pub fn roofline_analysis(&self, m: usize, n: usize, k: usize) -> RooflineResult { |
188 | 0 | let cost = BackendCostModel::default(); |
189 | 0 | let flops = 2.0 * m as f64 * n as f64 * k as f64; |
190 | 0 | let bytes = 4.0 * (m * k + k * n + m * n) as f64; |
191 | 0 | let ai = flops / bytes; |
192 | | |
193 | 0 | let ridge_point = match self.backend.unwrap_or(ComputeBackend::Cpu) { |
194 | | ComputeBackend::Gpu | ComputeBackend::Wgpu => { |
195 | 0 | cost.gpu_peak_tflops * 1000.0 / cost.pcie_bandwidth_gbps |
196 | | } |
197 | | ComputeBackend::Cpu | ComputeBackend::Scalar => { |
198 | 0 | cost.cpu_peak_gflops / 50.0 // ~50 GB/s memory bandwidth |
199 | | } |
200 | | }; |
201 | | |
202 | 0 | if ai < ridge_point { |
203 | 0 | RooflineResult::MemoryBound { ai, ridge_point } |
204 | | } else { |
205 | 0 | RooflineResult::ComputeBound { ai, ridge_point } |
206 | | } |
207 | 0 | } |
208 | | |
209 | | /// Generate summary report |
210 | 0 | pub fn summary(&self) -> String { |
211 | 0 | let mut s = String::new(); |
212 | 0 | s.push_str("Unified Brick Profiler Summary\n"); |
213 | 0 | s.push_str("==============================\n"); |
214 | 0 | s.push_str(&format!( |
215 | 0 | "Backend: {:?}\n", |
216 | 0 | self.backend.unwrap_or(ComputeBackend::Scalar) |
217 | 0 | )); |
218 | 0 | s.push_str(&format!("Total elements: {}\n", self.total_elements)); |
219 | 0 | s.push_str(&format!( |
220 | 0 | "Selections: {} decisions\n", |
221 | 0 | self.selection_history.len() |
222 | 0 | )); |
223 | 0 | s.push_str("\nCPU Stats:\n"); |
224 | 0 | s.push_str(&self.cpu_stats.summary()); |
225 | 0 | s |
226 | 0 | } |
227 | | } |
228 | | |
229 | | /// Roofline model result |
230 | | #[derive(Debug, Clone, Copy)] |
231 | | pub enum RooflineResult { |
232 | | /// Workload is memory-bound (AI < ridge point) |
233 | | MemoryBound { |
234 | | /// Arithmetic intensity (FLOP/byte) |
235 | | ai: f64, |
236 | | /// Ridge point where compute = memory |
237 | | ridge_point: f64, |
238 | | }, |
239 | | /// Workload is compute-bound (AI > ridge point) |
240 | | ComputeBound { |
241 | | /// Arithmetic intensity (FLOP/byte) |
242 | | ai: f64, |
243 | | /// Ridge point where compute = memory |
244 | | ridge_point: f64, |
245 | | }, |
246 | | } |
247 | | |
248 | | impl RooflineResult { |
249 | | /// Get arithmetic intensity |
250 | 0 | pub fn arithmetic_intensity(&self) -> f64 { |
251 | 0 | match self { |
252 | 0 | RooflineResult::MemoryBound { ai, .. } => *ai, |
253 | 0 | RooflineResult::ComputeBound { ai, .. } => *ai, |
254 | | } |
255 | 0 | } |
256 | | |
257 | | /// Check if compute-bound |
258 | 0 | pub fn is_compute_bound(&self) -> bool { |
259 | 0 | matches!(self, RooflineResult::ComputeBound { .. }) |
260 | 0 | } |
261 | | } |
262 | | |
263 | | /// PTX microkernel definition (for documentation and future CUDA support) |
264 | | /// |
265 | | /// This is a specification for the GPU microkernel. Actual PTX code generation |
266 | | /// would be done by the trueno-ptx crate. |
267 | | /// |
268 | | /// # References |
269 | | /// |
270 | | /// - NVIDIA PTX ISA Reference Manual |
271 | | /// - Volkov, V. (2010). Better Performance at Lower Occupancy. |
272 | | #[derive(Debug, Clone)] |
273 | | pub struct PtxMicrokernelSpec { |
274 | | /// PTX version (e.g., "8.0") |
275 | | pub ptx_version: &'static str, |
276 | | /// Target SM architecture (e.g., "sm_80") |
277 | | pub sm_target: &'static str, |
278 | | /// Register count per thread |
279 | | pub registers_per_thread: u32, |
280 | | /// Shared memory bytes per block |
281 | | pub smem_bytes: usize, |
282 | | /// Thread block dimensions |
283 | | pub block_dim: (u32, u32, u32), |
284 | | /// Tile dimensions (MR, NR) |
285 | | pub tile_dim: (usize, usize), |
286 | | } |
287 | | |
288 | | impl Default for PtxMicrokernelSpec { |
289 | 0 | fn default() -> Self { |
290 | 0 | Self { |
291 | 0 | ptx_version: "8.0", |
292 | 0 | sm_target: "sm_80", |
293 | 0 | registers_per_thread: 64, |
294 | 0 | smem_bytes: 48 * 1024, // 48KB shared memory |
295 | 0 | block_dim: (16, 16, 1), |
296 | 0 | tile_dim: (16, 16), // 16x16 output tile per warp |
297 | 0 | } |
298 | 0 | } |
299 | | } |
300 | | |
301 | | /// WGSL microkernel specification (for wgpu backend) |
302 | | /// |
303 | | /// Defines the compute shader for matrix multiplication. |
304 | | #[derive(Debug, Clone)] |
305 | | pub struct WgslMicrokernelSpec { |
306 | | /// Workgroup size (x, y, z) |
307 | | pub workgroup_size: (u32, u32, u32), |
308 | | /// Tile dimensions (MR, NR) |
309 | | pub tile_dim: (usize, usize), |
310 | | /// Use shared memory for tiling |
311 | | pub use_shared_memory: bool, |
312 | | } |
313 | | |
314 | | impl Default for WgslMicrokernelSpec { |
315 | 0 | fn default() -> Self { |
316 | 0 | Self { |
317 | 0 | workgroup_size: (8, 8, 1), |
318 | 0 | tile_dim: (8, 8), |
319 | 0 | use_shared_memory: true, |
320 | 0 | } |
321 | 0 | } |
322 | | } |
323 | | |
324 | | impl WgslMicrokernelSpec { |
325 | | /// Generate WGSL shader source |
326 | | /// |
327 | | /// This generates a basic tiled GEMM shader. For production use, |
328 | | /// this would be optimized with coalesced memory access and bank conflict avoidance. |
329 | 0 | pub fn generate_wgsl(&self) -> String { |
330 | 0 | format!( |
331 | 0 | r#"// WGSL GEMM Microkernel |
332 | 0 | // Generated by trueno BLIS module |
333 | 0 | // Tile: {}x{}, Workgroup: {}x{}x{} |
334 | 0 |
|
335 | 0 | struct GemmParams {{ |
336 | 0 | m: u32, |
337 | 0 | n: u32, |
338 | 0 | k: u32, |
339 | 0 | alpha: f32, |
340 | 0 | beta: f32, |
341 | 0 | }} |
342 | 0 |
|
343 | 0 | @group(0) @binding(0) var<uniform> params: GemmParams; |
344 | 0 | @group(0) @binding(1) var<storage, read> a: array<f32>; |
345 | 0 | @group(0) @binding(2) var<storage, read> b: array<f32>; |
346 | 0 | @group(0) @binding(3) var<storage, read_write> c: array<f32>; |
347 | 0 |
|
348 | 0 | var<workgroup> tile_a: array<f32, {tile_a_size}>; |
349 | 0 | var<workgroup> tile_b: array<f32, {tile_b_size}>; |
350 | 0 |
|
351 | 0 | @compute @workgroup_size({wx}, {wy}, {wz}) |
352 | 0 | fn main( |
353 | 0 | @builtin(global_invocation_id) global_id: vec3<u32>, |
354 | 0 | @builtin(local_invocation_id) local_id: vec3<u32>, |
355 | 0 | @builtin(workgroup_id) group_id: vec3<u32>, |
356 | 0 | ) {{ |
357 | 0 | let row = global_id.y; |
358 | 0 | let col = global_id.x; |
359 | 0 |
|
360 | 0 | if (row >= params.m || col >= params.n) {{ |
361 | 0 | return; |
362 | 0 | }} |
363 | 0 |
|
364 | 0 | var sum: f32 = 0.0; |
365 | 0 |
|
366 | 0 | // Tile over K dimension |
367 | 0 | let num_tiles = (params.k + {tile_k}u - 1u) / {tile_k}u; |
368 | 0 |
|
369 | 0 | for (var t: u32 = 0u; t < num_tiles; t++) {{ |
370 | 0 | let k_base = t * {tile_k}u; |
371 | 0 |
|
372 | 0 | // Load tile_a and tile_b into shared memory |
373 | 0 | // (simplified - production code would have proper coalescing) |
374 | 0 | let k_idx = k_base + local_id.x; |
375 | 0 | if (row < params.m && k_idx < params.k) {{ |
376 | 0 | tile_a[local_id.y * {tile_k}u + local_id.x] = a[row * params.k + k_idx]; |
377 | 0 | }} |
378 | 0 | if (k_idx < params.k && col < params.n) {{ |
379 | 0 | tile_b[local_id.y * {tile_k}u + local_id.x] = b[k_idx * params.n + col]; |
380 | 0 | }} |
381 | 0 |
|
382 | 0 | workgroupBarrier(); |
383 | 0 |
|
384 | 0 | // Compute partial sum |
385 | 0 | for (var kk: u32 = 0u; kk < {tile_k}u; kk++) {{ |
386 | 0 | if (k_base + kk < params.k) {{ |
387 | 0 | sum += tile_a[local_id.y * {tile_k}u + kk] * tile_b[kk * {tile_k}u + local_id.x]; |
388 | 0 | }} |
389 | 0 | }} |
390 | 0 |
|
391 | 0 | workgroupBarrier(); |
392 | 0 | }} |
393 | 0 |
|
394 | 0 | // Store result |
395 | 0 | let c_idx = row * params.n + col; |
396 | 0 | c[c_idx] = params.alpha * sum + params.beta * c[c_idx]; |
397 | 0 | }} |
398 | 0 | "#, |
399 | | self.tile_dim.0, |
400 | | self.tile_dim.1, |
401 | | self.workgroup_size.0, |
402 | | self.workgroup_size.1, |
403 | | self.workgroup_size.2, |
404 | 0 | tile_a_size = self.tile_dim.0 * self.tile_dim.0, |
405 | 0 | tile_b_size = self.tile_dim.0 * self.tile_dim.1, |
406 | | wx = self.workgroup_size.0, |
407 | | wy = self.workgroup_size.1, |
408 | | wz = self.workgroup_size.2, |
409 | | tile_k = self.tile_dim.0, |
410 | | ) |
411 | 0 | } |
412 | | } |
413 | | |
414 | | /// GEMM with automatic backend selection |
415 | | /// |
416 | | /// Uses the 5× PCIe rule to select between CPU (asm) and GPU (PTX/WGSL) backends. |
417 | 0 | pub fn gemm_auto( |
418 | 0 | m: usize, |
419 | 0 | n: usize, |
420 | 0 | k: usize, |
421 | 0 | a: &[f32], |
422 | 0 | b: &[f32], |
423 | 0 | c: &mut [f32], |
424 | 0 | profiler: Option<&mut UnifiedBrickProfiler>, |
425 | 0 | ) -> Result<(), TruenoError> { |
426 | 0 | let cost_model = BackendCostModel::default(); |
427 | 0 | let backend = cost_model.select_backend(m, n, k); |
428 | | |
429 | 0 | if let Some(prof) = profiler { |
430 | 0 | prof.record_selection(m, n, k, backend); |
431 | 0 | } |
432 | | |
433 | 0 | match backend { |
434 | | ComputeBackend::Cpu | ComputeBackend::Scalar => { |
435 | | // Use BLIS CPU implementation |
436 | 0 | gemm_blis(m, n, k, a, b, c, None) |
437 | | } |
438 | | ComputeBackend::Gpu => { |
439 | | // PTX backend (stub - requires CUDA support) |
440 | | // For now, fall back to CPU |
441 | 0 | gemm_blis(m, n, k, a, b, c, None) |
442 | | } |
443 | | ComputeBackend::Wgpu => { |
444 | | // WGSL backend (stub - requires wgpu support) |
445 | | // For now, fall back to CPU |
446 | 0 | gemm_blis(m, n, k, a, b, c, None) |
447 | | } |
448 | | } |
449 | 0 | } |