qualia_core_db/platform/compute_bridge/
execute.rs

1//! The dispatch entry the STEM substrate calls: `accelerated_gemm_f32` runs `C = A·B`
2//! on the GPU when the **measured** capability matrix says it wins and the job is big
3//! enough to be worth the dispatch, and on a `rayon` CPU path otherwise. The CPU path
4//! is always present and never hard-fails (§7).
5//!
6//! The machine benchmark ([`crate::device_benchmark`]) runs **once** here, lazily, to
7//! build the shared [`ComputePolicy`]; every subsequent call is the O(1) `select`.
8//!
9//! f64 vs f32: this accelerated path is **f32** (the GPU/WGSL reality and what the
10//! throughput-bound callers want). The exact-f64 scientific GEMM
11//! ([`crate::solvers::linear_algebra::gemm`]) is unchanged and stays on the CPU — the
12//! bridge never silently downcasts a caller that asked for f64.
13
14use std::sync::OnceLock;
15
16use super::gpu_gemm;
17use super::kernel_class::KernelClass;
18use super::policy::ComputePolicy;
19use crate::platform::hetero_dispatch::{HostCapabilities, PowerThermalBudget};
20
21/// Below this FLOP count a GEMM stays on the CPU regardless of the matrix — GPU upload
22/// + dispatch + readback overhead dominates a small job. (`m·k·n` multiply-adds.)
23const GPU_MIN_FLOPS: u64 = 1 << 20; // ~100³
24
25/// Which backend actually ran a dispatch (for observability and the correctness gate).
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum RanOn {
28    Gpu,
29    Cpu,
30}
31
32static POLICY: OnceLock<ComputePolicy> = OnceLock::new();
33
34/// The shared compute policy, built once by probing this machine (the benchmark). The
35/// heavy probe runs on first call; thereafter `select` is O(1).
36pub fn shared_policy() -> &'static ComputePolicy {
37    POLICY.get_or_init(|| {
38        let registry = super::default_registry();
39        let panel = super::backend::KernelPanel::default();
40        let budget = PowerThermalBudget {
41            vram_budget_bytes: 4u64 << 30,
42            power_budget_mw: 45_000,
43            thermal_headroom_c: 20.0,
44        };
45        let host = HostCapabilities {
46            gpu_available: gpu_gemm::shared().is_some(),
47            vram_available: 2u64 << 30,
48            npu_available: false,
49            cpu_threads: num_cpus::get() as u32,
50        };
51        ComputePolicy::probe(&registry, &panel, budget, host)
52    })
53}
54
55/// Multi-threaded CPU `C = A·B` (f32, row-major) — the always-present fallback.
56fn cpu_gemm_f32(_m: usize, k: usize, n: usize, a: &[f32], b: &[f32], c: &mut [f32]) {
57    use rayon::prelude::*;
58    c.par_chunks_mut(n).enumerate().for_each(|(i, row)| {
59        let a_row = &a[i * k..(i + 1) * k];
60        for (j, cell) in row.iter_mut().enumerate() {
61            let mut s = 0.0f32;
62            for l in 0..k {
63                s += a_row[l] * b[l * n + j];
64            }
65            *cell = s;
66        }
67    });
68}
69
70/// Accelerated `C = A·B` (f32). `a` is `m×k`, `b` is `k×n`, `c` is `m×n` (overwritten).
71/// Routes to the GPU when the measured matrix favours it for `DenseLinear`, the job
72/// clears [`GPU_MIN_FLOPS`], and it fits in GPU buffers; otherwise the CPU path. Returns
73/// which backend ran. On any GPU shortfall it falls back to CPU — never a hard fail.
74pub fn accelerated_gemm_f32(
75    m: usize,
76    k: usize,
77    n: usize,
78    a: &[f32],
79    b: &[f32],
80    c: &mut [f32],
81) -> RanOn {
82    debug_assert_eq!(a.len(), m * k);
83    debug_assert_eq!(b.len(), k * n);
84    debug_assert_eq!(c.len(), m * n);
85
86    let flops = (m as u64) * (k as u64) * (n as u64);
87    if flops >= GPU_MIN_FLOPS {
88        if let Some(ctx) = gpu_gemm::shared() {
89            if ctx.fits(m, k, n) {
90                let plan = shared_policy().select(KernelClass::DenseLinear, (m * n * 4) as u64);
91                if !plan.is_cpu() {
92                    if let Some(result) = ctx.gemm(m, k, n, a, b) {
93                        c.copy_from_slice(&result);
94                        return RanOn::Gpu;
95                    }
96                }
97            }
98        }
99    }
100    cpu_gemm_f32(m, k, n, a, b, c);
101    RanOn::Cpu
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    fn ref_gemm(m: usize, k: usize, n: usize, a: &[f32], b: &[f32]) -> Vec<f32> {
109        let mut c = vec![0.0f32; m * n];
110        for i in 0..m {
111            for j in 0..n {
112                let mut s = 0.0f32;
113                for l in 0..k {
114                    s += a[i * k + l] * b[l * n + j];
115                }
116                c[i * n + j] = s;
117            }
118        }
119        c
120    }
121
122    #[test]
123    fn small_job_runs_on_cpu_and_is_correct() {
124        let (m, k, n) = (4, 5, 3);
125        let a: Vec<f32> = (0..m * k).map(|i| i as f32 * 0.1).collect();
126        let b: Vec<f32> = (0..k * n).map(|i| i as f32 * 0.2 - 0.5).collect();
127        let mut c = vec![0.0f32; m * n];
128        let ran = accelerated_gemm_f32(m, k, n, &a, &b, &mut c);
129        assert_eq!(ran, RanOn::Cpu, "a tiny job must stay on CPU");
130        assert_eq!(c, ref_gemm(m, k, n, &a, &b));
131    }
132
133    #[test]
134    fn large_job_is_correct_on_whichever_backend_the_machine_chose() {
135        // 128³ clears the FLOP floor. Whether it runs on GPU or CPU depends on the
136        // measured matrix; either way the result must match the reference.
137        let (m, k, n) = (128, 128, 128);
138        let a: Vec<f32> = (0..m * k).map(|i| ((i % 17) as f32) * 0.05 - 0.3).collect();
139        let b: Vec<f32> = (0..k * n).map(|i| ((i % 11) as f32) * 0.07 - 0.2).collect();
140        let mut c = vec![0.0f32; m * n];
141        let ran = accelerated_gemm_f32(m, k, n, &a, &b, &mut c);
142        eprintln!("[accelerated_gemm_f32] 128³ ran on {ran:?}");
143        let reference = ref_gemm(m, k, n, &a, &b);
144        let max_err = c
145            .iter()
146            .zip(&reference)
147            .map(|(x, y)| (x - y).abs())
148            .fold(0.0f32, f32::max);
149        assert!(max_err < 1e-2, "max abs err {max_err} on {ran:?}");
150    }
151}
qualia_core_db/platform/compute_bridge/execute.rs

qualia_core_db/platform/compute_bridge/
execute.rs
