qualia_core_db/platform/compute_bridge/
policy.rs

1//! `ComputePolicy` — the one shared dispatch surface the whole engine calls
2//! (HARDWARE_BACKEND_AUTOSELECT_PLAN.md §4).
3//!
4//! A STEM call site asks `policy.select(class, problem_bytes)` and runs the returned
5//! [`Plan`]. No module embeds a device decision; they all defer to the measured
6//! per-class matrix. `select` is **O(1) and zero-heap** (it reads an already-built
7//! matrix and returns a `Copy` plan) — the heavy probe ran once at boot. CPU is
8//! always a valid plan and never hard-fails (plan §7).
9//!
10//! This *wraps* the existing vendor-neutral `hetero_dispatch` helpers (precision,
11//! tiling, zero-copy) rather than duplicating them, and adds the part they cannot do
12//! alone: the measured per-class backend/circuit choice from the matrix.
13
14use crate::device_benchmark::CircuitKind;
15use crate::platform::hetero_dispatch::{
16    select_precision, HeterogeneousDispatcher, HostCapabilities, PowerThermalBudget, Precision,
17    ZeroCopyStrategy,
18};
19
20use super::backend::{BackendId, KernelPanel};
21use super::kernel_class::KernelClass;
22use super::matrix::{probe_class_matrix, ClassMatrix};
23
24/// The resolved execution plan for one kernel dispatch. `Copy`, zero-heap — it
25/// names the backend/circuit and the precision/tiling/transfer policy, all of which
26/// are small scalars. The human-readable circuit label stays in the matrix.
27#[derive(Debug, Clone, Copy, PartialEq)]
28pub struct Plan {
29    /// Which acceleration method to run on (`"cpu"`, `"wgpu"`, …).
30    pub backend: BackendId,
31    /// The circuit class chosen (discrete GPU / iGPU / CPU / NPU).
32    pub circuit_kind: CircuitKind,
33    /// Numeric precision to use (from the host's VRAM/power/thermal envelope).
34    pub precision: Precision,
35    /// Number of sequential tiles a GPU job is split into so each fits in VRAM
36    /// (graceful degradation, never an OOM hard-fail). 1 on CPU or when it fits.
37    pub tiles: u32,
38    /// How data reaches the device (mmap-direct on unified memory, else staging).
39    pub zero_copy: ZeroCopyStrategy,
40}
41
42impl Plan {
43    /// Is this plan running on the CPU fallback?
44    pub fn is_cpu(self) -> bool {
45        self.backend == BackendId::CPU || self.circuit_kind == CircuitKind::Cpu
46    }
47}
48
49/// The shared compute policy: a measured per-class matrix plus the host's
50/// precision/VRAM envelope. Built once at startup; `select` is the hot, O(1) call.
51pub struct ComputePolicy {
52    matrix: ClassMatrix,
53    budget: PowerThermalBudget,
54    host: HostCapabilities,
55}
56
57impl ComputePolicy {
58    /// Build from an already-measured per-class matrix (e.g. one loaded from the
59    /// passport, or a synthetic one in tests — no GPU required).
60    pub fn from_class_matrix(
61        matrix: ClassMatrix,
62        budget: PowerThermalBudget,
63        host: HostCapabilities,
64    ) -> Self {
65        Self {
66            matrix,
67            budget,
68            host,
69        }
70    }
71
72    /// Probe the registry once (heavy) and build the policy. Call at startup; cache
73    /// the matrix in the passport to avoid re-probing every boot.
74    pub fn probe(
75        registry: &super::backend::BackendRegistry,
76        panel: &KernelPanel,
77        budget: PowerThermalBudget,
78        host: HostCapabilities,
79    ) -> Self {
80        Self::from_class_matrix(probe_class_matrix(registry, panel), budget, host)
81    }
82
83    /// The measured per-class matrix (for inspection / passport caching).
84    pub fn matrix(&self) -> &ClassMatrix {
85        &self.matrix
86    }
87
88    /// Select the execution plan for a `class` kernel touching `problem_bytes` of
89    /// data. O(1), zero-heap, never fails: returns a CPU plan when no accelerator
90    /// wins, when no GPU was probed, or when the measured GPU win is within noise of
91    /// CPU for a class that is not typically GPU-amenable (the §13 tie-break — a
92    /// measured GPU win on an amenable class is always honoured).
93    pub fn select(&self, class: KernelClass, problem_bytes: u64) -> Plan {
94        // Precision from the host's VRAM/power/thermal envelope (reuse hetero_dispatch).
95        // Treat the problem as f32 elements for the param-count proxy.
96        let param_count = (problem_bytes / 4).max(1);
97        let precision = select_precision(param_count, &self.budget);
98
99        let best = self.matrix.best_for(class);
100        let cpu_ms = self
101            .matrix
102            .rows(class)
103            .iter()
104            .find(|r| r.kind == CircuitKind::Cpu)
105            .map(|r| r.ms_per_gemv);
106
107        let choose_gpu = match best {
108            None => false,                                  // class not probed at all → CPU
109            Some(b) if b.kind == CircuitKind::Cpu => false, // CPU already won
110            Some(b) => {
111                // A GPU/other circuit measured fastest. Honour it unless the class is
112                // not GPU-amenable AND the win over CPU is within ~5% (measurement
113                // noise) — then prefer the simpler CPU path.
114                if class.is_typically_gpu_amenable() {
115                    true
116                } else if let Some(cms) = cpu_ms {
117                    b.ms_per_gemv < cms * 0.95
118                } else {
119                    true
120                }
121            }
122        };
123
124        if let (true, Some(b)) = (choose_gpu, best) {
125            let dispatcher = HeterogeneousDispatcher::new(self.host);
126            let tiles = dispatcher.gpu_tiles(problem_bytes);
127            let zero_copy = match b.kind {
128                CircuitKind::DiscreteGpu => ZeroCopyStrategy::StagingUpload,
129                _ => ZeroCopyStrategy::MmapDirect, // integrated/unified or NPU
130            };
131            Plan {
132                backend: BackendId(backend_id_for(&b.backend)),
133                circuit_kind: b.kind,
134                precision,
135                tiles,
136                zero_copy,
137            }
138        } else {
139            // CPU fallback — always valid, single pass, in-pool.
140            Plan {
141                backend: BackendId::CPU,
142                circuit_kind: CircuitKind::Cpu,
143                precision: Precision::F32, // CPU reference runs in f32
144                tiles: 1,
145                zero_copy: ZeroCopyStrategy::MmapDirect,
146            }
147        }
148    }
149}
150
151/// Map a `CircuitBench.backend` string (`"Vulkan"`, `"Dx12"`, `"native"`, …) to a
152/// stable `BackendId`. wgpu circuits all carry the `"wgpu"` id (the adapter API is
153/// in `circuit_kind`/the matrix label); native is CPU.
154fn backend_id_for(backend: &str) -> &'static str {
155    match backend {
156        "native" => "cpu",
157        _ => "wgpu",
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use crate::device_benchmark::CircuitBench;
165
166    const GIB: u64 = 1 << 30;
167
168    fn roomy_host() -> (PowerThermalBudget, HostCapabilities) {
169        (
170            PowerThermalBudget {
171                vram_budget_bytes: 16 * GIB,
172                power_budget_mw: 60_000,
173                thermal_headroom_c: 30.0,
174            },
175            HostCapabilities {
176                gpu_available: true,
177                vram_available: 8 * GIB,
178                npu_available: false,
179                cpu_threads: 16,
180            },
181        )
182    }
183
184    fn row(kind: CircuitKind, backend: &str, ms: f64) -> CircuitBench {
185        CircuitBench {
186            label: format!("{backend} circuit"),
187            kind,
188            backend: backend.to_string(),
189            ms_per_gemv: ms,
190            gflops: 0.0,
191            upload_gbps: if kind == CircuitKind::Cpu {
192                f64::INFINITY
193            } else {
194                5.0
195            },
196            rel_score: 1.0,
197            decode_proxy_tok_s: None,
198        }
199    }
200
201    /// A synthetic matrix: GPU clearly wins DenseLinear; CPU only for Divergent.
202    fn synth_matrix() -> ClassMatrix {
203        let mut per_class = Vec::new();
204        for class in KernelClass::ALL {
205            let rows = match class {
206                KernelClass::DenseLinear => vec![
207                    row(CircuitKind::DiscreteGpu, "Vulkan", 0.4),
208                    row(CircuitKind::Cpu, "native", 20.0),
209                ],
210                KernelClass::Divergent => vec![row(CircuitKind::Cpu, "native", 3.0)],
211                _ => vec![row(CircuitKind::Cpu, "native", 5.0)],
212            };
213            per_class.push((class, rows));
214        }
215        ClassMatrix::from_per_class(per_class)
216    }
217
218    #[test]
219    fn selects_measured_gpu_winner_for_dense_linear() {
220        let (budget, host) = roomy_host();
221        let policy = ComputePolicy::from_class_matrix(synth_matrix(), budget, host);
222        let plan = policy.select(KernelClass::DenseLinear, 64 * 1024 * 1024);
223        assert_eq!(plan.backend, BackendId::WGPU);
224        assert_eq!(plan.circuit_kind, CircuitKind::DiscreteGpu);
225        assert_eq!(plan.zero_copy, ZeroCopyStrategy::StagingUpload); // discrete → staging
226        assert!(!plan.is_cpu());
227    }
228
229    #[test]
230    fn falls_back_to_cpu_when_only_cpu_probed() {
231        let (budget, host) = roomy_host();
232        let policy = ComputePolicy::from_class_matrix(synth_matrix(), budget, host);
233        let plan = policy.select(KernelClass::Divergent, 1 << 20);
234        assert!(plan.is_cpu());
235        assert_eq!(plan.backend, BackendId::CPU);
236        assert_eq!(plan.tiles, 1);
237    }
238
239    #[test]
240    fn unprobed_class_is_cpu_not_a_panic() {
241        // A matrix with no rows for a class → CPU plan, never a panic (plan §7).
242        let (budget, host) = roomy_host();
243        let empty = ClassMatrix::from_per_class(vec![(KernelClass::Fft, Vec::new())]);
244        let policy = ComputePolicy::from_class_matrix(empty, budget, host);
245        assert!(policy.select(KernelClass::Fft, 1 << 20).is_cpu());
246        assert!(policy.select(KernelClass::AllPairs, 1 << 20).is_cpu()); // class absent entirely
247    }
248
249    #[test]
250    fn vram_pressure_tiles_instead_of_failing() {
251        let budget = PowerThermalBudget {
252            vram_budget_bytes: 16 * GIB,
253            power_budget_mw: 60_000,
254            thermal_headroom_c: 30.0,
255        };
256        let host = HostCapabilities {
257            gpu_available: true,
258            vram_available: 2 * GIB,
259            npu_available: false,
260            cpu_threads: 8,
261        };
262        let policy = ComputePolicy::from_class_matrix(synth_matrix(), budget, host);
263        // A 5 GiB DenseLinear job on a 2 GiB GPU → 3 tiles, not an OOM.
264        let plan = policy.select(KernelClass::DenseLinear, 5 * GIB);
265        assert!(!plan.is_cpu());
266        assert_eq!(plan.tiles, 3);
267    }
268}
qualia_core_db/platform/compute_bridge/policy.rs

qualia_core_db/platform/compute_bridge/
policy.rs
