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_map(|registration| {
114                    (registration.id == forced && !registration.reference_oracle)
115                        .then_some(registration)
116                });
117                let hit = match hit {
118                    Some(registration) if backend_dispatches(registration.id)? => {
119                        Some(registration)
120                    }
121                    _ => None,
122                };
123                return match hit {
124                    Some(reg) => Ok(RouterDecision {
125                        backend: reg.id,
126                        reason: Reason::EnvOverride,
127                    }),
128                    None => Err(BackendError::new(format!(
129                        "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."
130                    ))),
131                };
132            }
133        }
134
135        // V7-EXT-021: precedence comes from the BackendPrecedence inventory
136        // submitted by each backend crate, not a hardcoded driver-side table.
137        // Walk backends in precedence order and return the first hit.
138        for registration in registered_backends_by_precedence_slice()? {
139            if registered
140                .iter()
141                .any(|candidate| candidate.id == registration.id)
142                && backend_dispatches(registration.id)?
143                && !registration.reference_oracle
144            {
145                return Ok(RouterDecision {
146                    backend: registration.id,
147                    reason: Reason::Precedence,
148                });
149            }
150        }
151
152        Err(BackendError::new(
153            "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.",
154        ))
155    }
156
157    /// Enumerate every registered backend in precedence order. Inventory-driven
158    /// per V7-EXT-021  -  backends without a submitted `BackendPrecedence`
159    /// trail every backend that has one (rank `u32::MAX`).
160    /// # Errors
161    ///
162    /// Returns the validated registry startup error when providers conflict.
163    pub fn enumerate_by_precedence() -> Result<Vec<&'static BackendRegistration>, BackendError> {
164        Ok(registered_backends_by_precedence_slice()?.to_vec())
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171    use vyre_driver::backend::backend_precedence;
172
173    fn noop_program() -> Program {
174        // Programs built without any buffers / nodes are valid for
175        // the router's purposes  -  we don't dispatch, we just pick.
176        Program::wrapped(Vec::new(), [1, 1, 1], Vec::new())
177    }
178
179    #[test]
180    fn enumerate_by_precedence_puts_wgpu_before_reference() {
181        // V7-EXT-021: precedence is now inventory-driven. wgpu submits
182        // rank 30 in this crate's lib.rs; cpu-ref (when registered)
183        // must trail it.
184        let wgpu_rank = backend_precedence("wgpu").expect("valid backend registry");
185        let ref_rank = backend_precedence("cpu-ref").expect("valid backend registry");
186        assert!(
187            wgpu_rank < ref_rank || ref_rank == u32::MAX,
188            "wgpu (rank {wgpu_rank}) must take precedence over the CPU reference oracle (rank {ref_rank})"
189        );
190    }
191
192    #[test]
193    fn enumerate_by_precedence_is_inventory_driven() {
194        // Replaces the BACKEND_PRECEDENCE static-slice assertion.
195        let ranked = BackendRouter::enumerate_by_precedence().expect("valid backend registry");
196        // wgpu registers in this crate; it must appear with a finite rank.
197        let wgpu = ranked.iter().find(|r| r.id == "wgpu").expect(
198            "Fix: wgpu backend registered in this crate; restore this invariant before continuing.",
199        );
200        assert_eq!(
201            backend_precedence(wgpu.id).expect("valid backend registry"),
202            30
203        );
204    }
205
206    #[test]
207    fn explicit_override_with_unknown_backend_surfaces_error() {
208        let router = BackendRouter::new();
209        let err = router
210            .pick_with_override(
211                &noop_program(),
212                Override::Explicit("does-not-exist-backend"),
213            )
214            .expect_err("unknown backend must error");
215        let msg = format!("{err}");
216        assert!(msg.contains("does-not-exist-backend"));
217        assert!(msg.contains("Fix:"));
218    }
219
220    #[test]
221    fn explicit_override_picks_the_named_backend_when_registered() {
222        let router = BackendRouter::new();
223        // wgpu registers via inventory::submit! in lib.rs.
224        let decision = router
225            .pick_with_override(&noop_program(), Override::Explicit("wgpu"))
226            .expect("Fix: wgpu backend is registered in this crate");
227        assert_eq!(decision.backend, "wgpu");
228        assert_eq!(decision.reason, Reason::EnvOverride);
229    }
230
231    #[test]
232    fn precedence_picks_wgpu_when_registered() {
233        let router = BackendRouter::new();
234        let decision = router
235            .pick_with_override(&noop_program(), Override::None)
236            .expect("Fix: at least one backend must register");
237        assert_eq!(decision.reason, Reason::Precedence);
238        // The picked backend must have a registered precedence rank
239        // (V7-EXT-021: replaces the BACKEND_PRECEDENCE static-slice check).
240        assert!(
241            backend_precedence(decision.backend).expect("valid backend registry") < u32::MAX,
242            "picked backend {} did not submit a BackendPrecedence inventory entry",
243            decision.backend
244        );
245    }
246}