Skip to main content

vyre_driver_wgpu/runtime/
router.rs

1//! Backend auto-picker.
2//!
3//! `BackendRouter` walks `inventory::iter::<BackendRegistration>`,
4//! filters out registered backends that cannot dispatch or are CPU
5//! reference oracles, and picks the best executable GPU backend available
6//! by precedence. Override via `VYRE_BACKEND=<id>`. The router is
7//! intentionally stateless: backend precedence lives in inventory
8//! registrations and adapter-specific persistence belongs to the backend
9//! cache layer, not routing.
10//!
11//! Precedence (high → low):
12//!
13//! 1. `VYRE_BACKEND=<id>`  -  if set and the backend is registered,
14//!    wins only when the backend is registered, executable, and GPU-backed.
15//! 2. `cuda`  -  when an NVIDIA/CUDA backend is linked, registered, and executable.
16//! 3. `wgpu`  -  portable GPU backend after CUDA.
17//! 4. `spirv`  -  when the SPIR-V backend is registered.
18//!
19//! `BackendRouter::pick()` returns the selected backend id on success,
20//! or a structured `BackendError` when no executable backend is linked.
21
22use std::env;
23
24use vyre_driver::backend::{backend_dispatches, registered_backends_by_precedence_slice};
25use vyre_driver::{BackendError, BackendRegistration};
26use vyre_foundation::ir::Program;
27
28/// How to source the forced-backend override.
29#[derive(Debug, Clone, Copy)]
30#[non_exhaustive]
31pub enum Override<'a> {
32    /// Read `VYRE_BACKEND` from the process environment.
33    FromEnv,
34    /// Use the explicit override regardless of environment.
35    Explicit(&'a str),
36    /// No override  -  router runs on precedence alone.
37    None,
38}
39
40const OVERRIDE_ENV: &str = "VYRE_BACKEND";
41
42/// Routing decision produced by the backend auto-picker.
43#[derive(Debug, Clone)]
44pub struct RouterDecision {
45    /// The selected backend id.
46    pub backend: &'static str,
47    /// Reason the decision fell to this backend.
48    pub reason: Reason,
49}
50
51/// How the decision was reached.
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53#[non_exhaustive]
54pub enum Reason {
55    /// `VYRE_BACKEND=<id>` forced the selection.
56    EnvOverride,
57    /// Highest-precedence registered backend that covers the
58    /// Program's dialects.
59    Precedence,
60}
61
62/// Backend auto-picker.
63///
64/// Constructed with [`BackendRouter::new`]; queries the runtime
65/// inventory on demand so newly-registered backends participate
66/// without router rebuild.
67#[derive(Default)]
68pub struct BackendRouter;
69
70impl BackendRouter {
71    /// New router.
72    #[must_use]
73    pub fn new() -> Self {
74        Self
75    }
76
77    /// Pick the best-available backend for `_program`.
78    ///
79    /// # Errors
80    ///
81    /// Returns `BackendError` when:
82    ///
83    /// * `VYRE_BACKEND` is set to a backend id that is not
84    ///   registered.
85    /// * No executable registered GPU backend is found. Vyre reports this as
86    ///   a linkage or driver-visibility error instead of routing into
87    ///   reference evaluation.
88    pub fn pick(&self, program: &Program) -> Result<RouterDecision, BackendError> {
89        self.pick_with_override(program, Override::FromEnv)
90    }
91
92    /// Pick with an explicit override source  -  the testable form of
93    /// [`pick`](Self::pick).
94    ///
95    /// # Errors
96    ///
97    /// Same conditions as [`pick`](Self::pick).
98    pub fn pick_with_override(
99        &self,
100        _program: &Program,
101        source: Override<'_>,
102    ) -> Result<RouterDecision, BackendError> {
103        let registered = vyre_driver::backend::registered_backends();
104
105        let forced: Option<String> = match source {
106            Override::FromEnv => env::var(OVERRIDE_ENV).ok(),
107            Override::Explicit(s) => Some(s.to_owned()),
108            Override::None => None,
109        };
110        if let Some(forced) = forced {
111            let forced = forced.trim();
112            if !forced.is_empty() {
113                let hit = registered.iter().find(|r| {
114                    r.id == forced && backend_dispatches(r.id) && !is_reference_oracle_backend(r.id)
115                });
116                return match hit {
117                    Some(reg) => Ok(RouterDecision {
118                        backend: reg.id,
119                        reason: Reason::EnvOverride,
120                    }),
121                    None => Err(BackendError::new(format!(
122                        "VYRE_BACKEND={forced} is not an executable registered GPU backend. Fix: link CUDA/WGPU or unset VYRE_BACKEND; cpu-ref/reference are explicit conformance oracles, not runtime router targets."
123                    ))),
124                };
125            }
126        }
127
128        // V7-EXT-021: precedence comes from the BackendPrecedence inventory
129        // submitted by each backend crate, not a hardcoded driver-side table.
130        // Walk backends in precedence order and return the first hit.
131        for reg in registered_backends_by_precedence_slice() {
132            if registered.iter().any(|r| r.id == reg.id)
133                && backend_dispatches(reg.id)
134                && !is_reference_oracle_backend(reg.id)
135            {
136                return Ok(RouterDecision {
137                    backend: reg.id,
138                    reason: Reason::Precedence,
139                });
140            }
141        }
142
143        Err(BackendError::new(
144            "no executable GPU backend is registered. Fix: link vyre-driver-cuda or vyre-driver-wgpu into the binary and verify the GPU driver probe succeeds.",
145        ))
146    }
147
148    /// Enumerate every registered backend in precedence order. Inventory-driven
149    /// per V7-EXT-021  -  backends without a submitted `BackendPrecedence`
150    /// trail every backend that has one (rank `u32::MAX`).
151    #[must_use]
152    pub fn enumerate_by_precedence() -> Vec<&'static BackendRegistration> {
153        registered_backends_by_precedence_slice().to_vec()
154    }
155}
156
157fn is_reference_oracle_backend(id: &str) -> bool {
158    matches!(id, "cpu-ref" | "reference")
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use vyre_driver::backend::backend_precedence;
165
166    fn noop_program() -> Program {
167        // Programs built without any buffers / nodes are valid for
168        // the router's purposes  -  we don't dispatch, we just pick.
169        Program::wrapped(Vec::new(), [1, 1, 1], Vec::new())
170    }
171
172    #[test]
173    fn enumerate_by_precedence_puts_wgpu_before_reference() {
174        // V7-EXT-021: precedence is now inventory-driven. wgpu submits
175        // rank 30 in this crate's lib.rs; cpu-ref (when registered)
176        // must trail it.
177        let wgpu_rank = backend_precedence("wgpu");
178        let ref_rank = backend_precedence("cpu-ref");
179        assert!(
180            wgpu_rank < ref_rank || ref_rank == u32::MAX,
181            "wgpu (rank {wgpu_rank}) must take precedence over the CPU reference oracle (rank {ref_rank})"
182        );
183    }
184
185    #[test]
186    fn enumerate_by_precedence_is_inventory_driven() {
187        // Replaces the BACKEND_PRECEDENCE static-slice assertion.
188        let ranked = BackendRouter::enumerate_by_precedence();
189        // wgpu registers in this crate; it must appear with a finite rank.
190        let wgpu = ranked.iter().find(|r| r.id == "wgpu").expect(
191            "Fix: wgpu backend registered in this crate; restore this invariant before continuing.",
192        );
193        assert_eq!(backend_precedence(wgpu.id), 30);
194    }
195
196    #[test]
197    fn explicit_override_with_unknown_backend_surfaces_error() {
198        let router = BackendRouter::new();
199        let err = router
200            .pick_with_override(
201                &noop_program(),
202                Override::Explicit("does-not-exist-backend"),
203            )
204            .expect_err("unknown backend must error");
205        let msg = format!("{err}");
206        assert!(msg.contains("does-not-exist-backend"));
207        assert!(msg.contains("Fix:"));
208    }
209
210    #[test]
211    fn explicit_override_picks_the_named_backend_when_registered() {
212        let router = BackendRouter::new();
213        // wgpu registers via inventory::submit! in lib.rs.
214        let decision = router
215            .pick_with_override(&noop_program(), Override::Explicit("wgpu"))
216            .expect("Fix: wgpu backend is registered in this crate");
217        assert_eq!(decision.backend, "wgpu");
218        assert_eq!(decision.reason, Reason::EnvOverride);
219    }
220
221    #[test]
222    fn precedence_picks_wgpu_when_registered() {
223        let router = BackendRouter::new();
224        let decision = router
225            .pick_with_override(&noop_program(), Override::None)
226            .expect("Fix: at least one backend must register");
227        assert_eq!(decision.reason, Reason::Precedence);
228        // The picked backend must have a registered precedence rank
229        // (V7-EXT-021: replaces the BACKEND_PRECEDENCE static-slice check).
230        assert!(
231            backend_precedence(decision.backend) < u32::MAX,
232            "picked backend {} did not submit a BackendPrecedence inventory entry",
233            decision.backend
234        );
235    }
236}