oxicuda_backend/
registry.rs

1//! Backend registry, capability-based selection, and fallback chains.
2//!
3//! This is the host-side "control plane" of the abstraction layer. It does
4//! **not** execute any GPU work; it decides *which* backend should run a
5//! given task, given:
6//!
7//! * which backends a build registered ([`BackendRegistry::register`]),
8//! * whether each is actually present at runtime (its `available` flag,
9//!   set from a probe in the concrete backend crate), and
10//! * what each one can do ([`Capabilities`]).
11//!
12//! The selection logic is fully deterministic and testable without any GPU:
13//! it always prefers the highest-priority *available* backend that satisfies
14//! a [`SelectionRequest`], and degrades down a [`fallback_chain`] that ends
15//! at the CPU reference backend whenever one is registered.
16
17use crate::backend_kind::BackendKind;
18use crate::capabilities::Capabilities;
19use crate::error::{BackendError, BackendResult};
20
21/// A backend known to the registry, plus the runtime facts needed to pick it.
22///
23/// `available` is supplied by the caller (typically the result of a cheap
24/// driver probe such as "did `cuInit` succeed?"). The registry never sets it
25/// itself, so this crate stays free of any device dependency.
26#[derive(Debug, Clone)]
27pub struct BackendEntry {
28    /// Which concrete backend this describes.
29    pub kind: BackendKind,
30    /// Whether the backend is usable on this machine right now.
31    pub available: bool,
32    /// Selection priority — higher wins. Defaults to
33    /// [`BackendKind::default_priority`].
34    pub priority: u32,
35    /// What the backend can do (filled from a driver query, or a sensible
36    /// default for the CPU/portable backends).
37    pub capabilities: Capabilities,
38}
39
40impl BackendEntry {
41    /// Create an entry using the kind's default priority and CPU-profile
42    /// capabilities. Callers override `capabilities` after construction for
43    /// real GPU backends.
44    #[must_use]
45    pub fn new(kind: BackendKind, available: bool) -> Self {
46        Self {
47            kind,
48            available,
49            priority: kind.default_priority(),
50            capabilities: if kind == BackendKind::Cpu {
51                Capabilities::cpu()
52            } else {
53                Capabilities::default()
54            },
55        }
56    }
57
58    /// Builder-style override of the priority.
59    #[must_use]
60    pub fn with_priority(mut self, priority: u32) -> Self {
61        self.priority = priority;
62        self
63    }
64
65    /// Builder-style override of the capability report.
66    #[must_use]
67    pub fn with_capabilities(mut self, capabilities: Capabilities) -> Self {
68        self.capabilities = capabilities;
69        self
70    }
71}
72
73/// Required-feature predicate used to filter backends during selection.
74///
75/// A backend is eligible only if it is available **and** its capabilities
76/// satisfy every requested flag. All fields default to "don't care".
77#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
78pub struct SelectionRequest {
79    /// Require a GPU backend (exclude the CPU fallback from the primary pick).
80    pub require_gpu: bool,
81    /// Require native FP16 compute.
82    pub require_fp16: bool,
83    /// Require native BF16 compute.
84    pub require_bf16: bool,
85    /// Require FP8 compute.
86    pub require_fp8: bool,
87    /// Require Tensor-Core / matrix units.
88    pub require_tensor_cores: bool,
89    /// Require unified / managed memory.
90    pub require_unified_memory: bool,
91    /// Require peer-to-peer device access.
92    pub require_peer_access: bool,
93    /// If `Some`, only this exact backend kind is acceptable.
94    pub pin: Option<BackendKind>,
95}
96
97impl SelectionRequest {
98    /// A request with no constraints (any available backend qualifies).
99    #[must_use]
100    pub const fn any() -> Self {
101        Self {
102            require_gpu: false,
103            require_fp16: false,
104            require_bf16: false,
105            require_fp8: false,
106            require_tensor_cores: false,
107            require_unified_memory: false,
108            require_peer_access: false,
109            pin: None,
110        }
111    }
112
113    /// A request that demands a GPU (used by callers that must not silently
114    /// fall back to the CPU reference path).
115    #[must_use]
116    pub const fn require_gpu() -> Self {
117        let mut r = Self::any();
118        r.require_gpu = true;
119        r
120    }
121
122    /// Pin selection to one specific backend kind.
123    #[must_use]
124    pub const fn pinned(kind: BackendKind) -> Self {
125        let mut r = Self::any();
126        r.pin = Some(kind);
127        r
128    }
129
130    /// Narrow this request to the CPU reference backend when the workload is
131    /// too small to pay for a GPU dispatch.
132    ///
133    /// A GPU dispatch costs a host→device copy, a command submission and a
134    /// synchronisation; below some byte count that overhead dominates the
135    /// kernel and the host path finishes first. This helper expresses that
136    /// policy *at selection time* — i.e. before any memory is allocated — so
137    /// that the whole workload (allocations included) lands on one backend.
138    /// Per-operation routing is deliberately **not** offered: a device pointer
139    /// belongs to the backend that allocated it, so a per-op switch would
140    /// require duplicating every buffer.
141    ///
142    /// The request is returned unchanged when the caller has already expressed
143    /// an explicit preference ([`require_gpu`](Self::require_gpu) or a
144    /// [`pin`](Self::pin)), or when `workload_bytes >= gpu_threshold_bytes`.
145    ///
146    /// # Example
147    ///
148    /// ```
149    /// use oxicuda_backend::{BackendKind, SelectionRequest};
150    ///
151    /// // 1 KiB with a 64 KiB threshold → pinned to the host backend.
152    /// let small = SelectionRequest::any().for_workload(1024, 64 * 1024);
153    /// assert_eq!(small.pin, Some(BackendKind::Cpu));
154    ///
155    /// // 1 MiB → unchanged, so the GPU can win the selection.
156    /// let large = SelectionRequest::any().for_workload(1024 * 1024, 64 * 1024);
157    /// assert_eq!(large, SelectionRequest::any());
158    /// ```
159    #[must_use]
160    pub const fn for_workload(mut self, workload_bytes: usize, gpu_threshold_bytes: usize) -> Self {
161        if !self.require_gpu && self.pin.is_none() && workload_bytes < gpu_threshold_bytes {
162            self.pin = Some(BackendKind::Cpu);
163        }
164        self
165    }
166
167    /// Returns `true` if `entry` satisfies every constraint in this request.
168    #[must_use]
169    pub fn is_satisfied_by(&self, entry: &BackendEntry) -> bool {
170        if !entry.available {
171            return false;
172        }
173        if let Some(pin) = self.pin {
174            if entry.kind != pin {
175                return false;
176            }
177        }
178        let caps = &entry.capabilities;
179        if self.require_gpu && !entry.kind.is_gpu() {
180            return false;
181        }
182        if self.require_fp16 && !caps.supports_fp16 {
183            return false;
184        }
185        if self.require_bf16 && !caps.supports_bf16 {
186            return false;
187        }
188        if self.require_fp8 && !caps.supports_fp8 {
189            return false;
190        }
191        if self.require_tensor_cores && !caps.tensor_cores {
192            return false;
193        }
194        if self.require_unified_memory && !caps.unified_memory {
195            return false;
196        }
197        if self.require_peer_access && !caps.peer_access {
198            return false;
199        }
200        true
201    }
202}
203
204/// Routable operation classes used by [`BackendRegistry::route`].
205///
206/// This is a coarse routing key, not the full op set: it lets a consumer ask
207/// "which backend should run my GEMMs?" so that, for example, dense linear
208/// algebra can be pinned to a Tensor-Core backend while element-wise work
209/// goes to whatever is cheapest.
210#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
211pub enum OpClass {
212    /// Dense matrix multiply (`gemm` / `batched_gemm`).
213    MatMul,
214    /// Convolution (`conv2d_forward`).
215    Convolution,
216    /// Scaled dot-product attention.
217    Attention,
218    /// Axis reductions.
219    Reduction,
220    /// Element-wise unary / binary / softmax.
221    Elementwise,
222    /// Host/device memory transfers.
223    Memory,
224}
225
226impl OpClass {
227    /// Every routable op class.
228    pub const ALL: [OpClass; 6] = [
229        OpClass::MatMul,
230        OpClass::Convolution,
231        OpClass::Attention,
232        OpClass::Reduction,
233        OpClass::Elementwise,
234        OpClass::Memory,
235    ];
236
237    /// `true` if this op class benefits from Tensor-Core units, so the
238    /// router prefers a matrix-capable backend when one is available.
239    #[must_use]
240    pub const fn prefers_tensor_cores(self) -> bool {
241        matches!(self, Self::MatMul | Self::Convolution | Self::Attention)
242    }
243}
244
245/// A registry of compute backends with capability-aware selection.
246///
247/// Construction registers nothing; call [`register`](Self::register) (or
248/// [`with_defaults`](Self::with_defaults)) to populate it. Selection is pure
249/// and side-effect-free, so it is trivially unit-testable.
250#[derive(Debug, Clone, Default)]
251pub struct BackendRegistry {
252    entries: Vec<BackendEntry>,
253}
254
255impl BackendRegistry {
256    /// An empty registry.
257    #[must_use]
258    pub fn new() -> Self {
259        Self {
260            entries: Vec::new(),
261        }
262    }
263
264    /// A registry pre-populated with one entry per [`BackendKind`], all
265    /// flagged **unavailable** except the CPU reference backend.
266    ///
267    /// Concrete crates then flip the `available` flag (and refine the
268    /// capabilities) for backends they detect via
269    /// [`set_available`](Self::set_available) /
270    /// [`set_capabilities`](Self::set_capabilities). This guarantees there
271    /// is always at least one usable backend (the CPU fallback).
272    #[must_use]
273    pub fn with_defaults() -> Self {
274        let mut reg = Self::new();
275        for kind in BackendKind::ALL {
276            reg.register(BackendEntry::new(kind, kind == BackendKind::Cpu));
277        }
278        reg
279    }
280
281    /// Register (or replace) an entry. If a backend of the same kind is
282    /// already present, it is overwritten so a later, more-informed probe
283    /// wins.
284    pub fn register(&mut self, entry: BackendEntry) {
285        if let Some(slot) = self.entries.iter_mut().find(|e| e.kind == entry.kind) {
286            *slot = entry;
287        } else {
288            self.entries.push(entry);
289        }
290    }
291
292    /// Number of registered backends.
293    #[must_use]
294    pub fn len(&self) -> usize {
295        self.entries.len()
296    }
297
298    /// `true` if no backend is registered.
299    #[must_use]
300    pub fn is_empty(&self) -> bool {
301        self.entries.is_empty()
302    }
303
304    /// Borrow the entry for `kind`, if registered.
305    #[must_use]
306    pub fn get(&self, kind: BackendKind) -> Option<&BackendEntry> {
307        self.entries.iter().find(|e| e.kind == kind)
308    }
309
310    /// Mark `kind`'s availability, returning `true` if the kind was present.
311    pub fn set_available(&mut self, kind: BackendKind, available: bool) -> bool {
312        if let Some(e) = self.entries.iter_mut().find(|e| e.kind == kind) {
313            e.available = available;
314            true
315        } else {
316            false
317        }
318    }
319
320    /// Update `kind`'s capability report, returning `true` if present.
321    pub fn set_capabilities(&mut self, kind: BackendKind, caps: Capabilities) -> bool {
322        if let Some(e) = self.entries.iter_mut().find(|e| e.kind == kind) {
323            e.capabilities = caps;
324            true
325        } else {
326            false
327        }
328    }
329
330    /// All currently-available backend kinds.
331    #[must_use]
332    pub fn available_kinds(&self) -> Vec<BackendKind> {
333        self.entries
334            .iter()
335            .filter(|e| e.available)
336            .map(|e| e.kind)
337            .collect()
338    }
339
340    /// Select the single best backend satisfying `req`, or
341    /// [`BackendError::Unsupported`] if none qualifies.
342    ///
343    /// "Best" = highest `priority` among satisfying entries; ties break by
344    /// [`BackendKind::default_priority`] and then declaration order, so the
345    /// result is deterministic.
346    pub fn select(&self, req: &SelectionRequest) -> BackendResult<BackendKind> {
347        self.entries
348            .iter()
349            .filter(|e| req.is_satisfied_by(e))
350            .max_by(|a, b| {
351                a.priority
352                    .cmp(&b.priority)
353                    .then(a.kind.default_priority().cmp(&b.kind.default_priority()))
354            })
355            .map(|e| e.kind)
356            .ok_or_else(|| {
357                BackendError::Unsupported(format!(
358                    "no registered backend satisfies the request {req:?}"
359                ))
360            })
361    }
362
363    /// Convenience: select the best available backend with no constraints,
364    /// preferring a GPU but accepting the CPU fallback.
365    pub fn select_best(&self) -> BackendResult<BackendKind> {
366        self.select(&SelectionRequest::any())
367    }
368
369    /// Select the best backend for a workload of `workload_bytes`, sending
370    /// workloads smaller than `gpu_threshold_bytes` to the CPU reference
371    /// backend (see [`SelectionRequest::for_workload`] for why this is a
372    /// selection-time and not a per-operation decision).
373    ///
374    /// If the narrowed request finds nothing — e.g. a small workload on a
375    /// registry with no CPU entry — the unnarrowed `req` is retried, so this
376    /// never fails where [`select`](Self::select) would have succeeded.
377    pub fn select_for_workload(
378        &self,
379        req: &SelectionRequest,
380        workload_bytes: usize,
381        gpu_threshold_bytes: usize,
382    ) -> BackendResult<BackendKind> {
383        let narrowed = req.for_workload(workload_bytes, gpu_threshold_bytes);
384        match self.select(&narrowed) {
385            Ok(kind) => Ok(kind),
386            Err(_) if narrowed != *req => self.select(req),
387            Err(e) => Err(e),
388        }
389    }
390
391    /// The ordered fallback chain for `req`: every satisfying backend, most-
392    /// to least-preferred, with the CPU reference backend forced to the end
393    /// if it is available (even when it would also satisfy `req` earlier).
394    ///
395    /// Consumers walk this chain, trying each backend until one initializes
396    /// and runs the op, so a transient GPU failure degrades gracefully to
397    /// the host.
398    #[must_use]
399    pub fn fallback_chain(&self, req: &SelectionRequest) -> Vec<BackendKind> {
400        // Collect satisfying entries together with their configured priority,
401        // then sort most-preferred first: by explicit priority desc, then by
402        // the kind's default priority desc as a deterministic tie-break.
403        let mut ranked: Vec<(u32, BackendKind)> = self
404            .entries
405            .iter()
406            .filter(|e| req.is_satisfied_by(e))
407            .map(|e| (e.priority, e.kind))
408            .collect();
409        ranked.sort_by(|(pa, ka), (pb, kb)| {
410            pb.cmp(pa)
411                .then(kb.default_priority().cmp(&ka.default_priority()))
412        });
413        let mut chain: Vec<BackendKind> = ranked.into_iter().map(|(_, k)| k).collect();
414        // Force the CPU reference backend to the very end if it is present,
415        // so the host path is always the last resort.
416        if let Some(pos) = chain.iter().position(|&k| k == BackendKind::Cpu) {
417            let cpu = chain.remove(pos);
418            chain.push(cpu);
419        }
420        chain
421    }
422
423    /// Route an [`OpClass`] to the best backend.
424    ///
425    /// For Tensor-Core-friendly classes the router first tries to find an
426    /// available matrix-capable backend; if none exists it falls back to the
427    /// plain best-available selection (so routing never fails when *any*
428    /// backend is available).
429    pub fn route(&self, op: OpClass) -> BackendResult<BackendKind> {
430        if op.prefers_tensor_cores() {
431            let tc_req = SelectionRequest {
432                require_tensor_cores: true,
433                ..SelectionRequest::any()
434            };
435            if let Ok(kind) = self.select(&tc_req) {
436                return Ok(kind);
437            }
438        }
439        self.select_best()
440    }
441}
442
443#[cfg(test)]
444mod tests {
445    use super::*;
446
447    /// Build a registry with CUDA (no tensor cores), ROCm (tensor cores),
448    /// and CPU, all available.
449    fn three_backend_registry() -> BackendRegistry {
450        let mut reg = BackendRegistry::new();
451        reg.register(BackendEntry::new(BackendKind::Cuda, true));
452        let rocm_caps = Capabilities {
453            tensor_cores: true,
454            supports_fp16: true,
455            ..Capabilities::default()
456        };
457        reg.register(BackendEntry::new(BackendKind::Rocm, true).with_capabilities(rocm_caps));
458        reg.register(BackendEntry::new(BackendKind::Cpu, true));
459        reg
460    }
461
462    #[test]
463    fn defaults_have_cpu_available_and_gpus_absent() {
464        let reg = BackendRegistry::with_defaults();
465        assert_eq!(reg.len(), 7);
466        assert!(reg.get(BackendKind::Cpu).unwrap().available);
467        assert!(!reg.get(BackendKind::Cuda).unwrap().available);
468        // With only the CPU available, best selection is CPU.
469        assert_eq!(reg.select_best().unwrap(), BackendKind::Cpu);
470    }
471
472    #[test]
473    fn select_picks_highest_priority_available() {
474        let reg = three_backend_registry();
475        // CUDA has the highest default priority and is available.
476        assert_eq!(reg.select_best().unwrap(), BackendKind::Cuda);
477    }
478
479    #[test]
480    fn select_skips_unavailable_highest_priority() {
481        let mut reg = three_backend_registry();
482        reg.set_available(BackendKind::Cuda, false);
483        // Now ROCm (next highest) wins.
484        assert_eq!(reg.select_best().unwrap(), BackendKind::Rocm);
485        reg.set_available(BackendKind::Rocm, false);
486        // Only CPU left.
487        assert_eq!(reg.select_best().unwrap(), BackendKind::Cpu);
488    }
489
490    #[test]
491    fn select_honors_capability_requirements() {
492        let reg = three_backend_registry();
493        // Require tensor cores → only ROCm qualifies (CUDA entry has none).
494        let req = SelectionRequest {
495            require_tensor_cores: true,
496            ..SelectionRequest::any()
497        };
498        assert_eq!(reg.select(&req).unwrap(), BackendKind::Rocm);
499    }
500
501    #[test]
502    fn select_require_gpu_excludes_cpu() {
503        let mut reg = BackendRegistry::new();
504        reg.register(BackendEntry::new(BackendKind::Cpu, true));
505        // Only CPU available; require_gpu must fail rather than pick CPU.
506        assert!(reg.select(&SelectionRequest::require_gpu()).is_err());
507        // ... but a constraint-free request happily returns CPU.
508        assert_eq!(reg.select_best().unwrap(), BackendKind::Cpu);
509    }
510
511    #[test]
512    fn select_pinned_backend() {
513        let reg = three_backend_registry();
514        let req = SelectionRequest::pinned(BackendKind::Cpu);
515        assert_eq!(reg.select(&req).unwrap(), BackendKind::Cpu);
516        // Pinning an unavailable backend fails.
517        let mut reg2 = reg.clone();
518        reg2.set_available(BackendKind::Rocm, false);
519        assert!(
520            reg2.select(&SelectionRequest::pinned(BackendKind::Rocm))
521                .is_err()
522        );
523    }
524
525    #[test]
526    fn higher_explicit_priority_overrides_default() {
527        let mut reg = BackendRegistry::new();
528        // Give CPU an absurd priority so it beats CUDA's default 100.
529        reg.register(BackendEntry::new(BackendKind::Cuda, true));
530        reg.register(BackendEntry::new(BackendKind::Cpu, true).with_priority(1000));
531        assert_eq!(reg.select_best().unwrap(), BackendKind::Cpu);
532    }
533
534    #[test]
535    fn fallback_chain_orders_by_priority_with_cpu_last() {
536        let reg = three_backend_registry();
537        let chain = reg.fallback_chain(&SelectionRequest::any());
538        assert_eq!(
539            chain,
540            vec![BackendKind::Cuda, BackendKind::Rocm, BackendKind::Cpu]
541        );
542        // CPU is always the final element when present.
543        assert_eq!(*chain.last().unwrap(), BackendKind::Cpu);
544    }
545
546    #[test]
547    fn fallback_chain_respects_constraints() {
548        let reg = three_backend_registry();
549        let req = SelectionRequest {
550            require_tensor_cores: true,
551            ..SelectionRequest::any()
552        };
553        // Only ROCm has tensor cores; CPU/CUDA excluded.
554        assert_eq!(reg.fallback_chain(&req), vec![BackendKind::Rocm]);
555    }
556
557    #[test]
558    fn fallback_chain_empty_when_nothing_qualifies() {
559        let mut reg = three_backend_registry();
560        reg.set_available(BackendKind::Cuda, false);
561        reg.set_available(BackendKind::Rocm, false);
562        reg.set_available(BackendKind::Cpu, false);
563        assert!(reg.fallback_chain(&SelectionRequest::any()).is_empty());
564        assert!(reg.select_best().is_err());
565    }
566
567    #[test]
568    fn route_matmul_prefers_tensor_core_backend() {
569        let reg = three_backend_registry();
570        // CUDA has higher priority but no tensor cores; ROCm has them.
571        // MatMul prefers tensor cores → ROCm.
572        assert_eq!(reg.route(OpClass::MatMul).unwrap(), BackendKind::Rocm);
573        // Elementwise does not prefer tensor cores → highest priority (CUDA).
574        assert_eq!(reg.route(OpClass::Elementwise).unwrap(), BackendKind::Cuda);
575    }
576
577    #[test]
578    fn route_falls_back_when_no_tensor_cores_anywhere() {
579        let mut reg = BackendRegistry::new();
580        reg.register(BackendEntry::new(BackendKind::Cuda, true)); // no TC
581        reg.register(BackendEntry::new(BackendKind::Cpu, true));
582        // No tensor-core backend exists → MatMul routes to best available.
583        assert_eq!(reg.route(OpClass::MatMul).unwrap(), BackendKind::Cuda);
584    }
585
586    #[test]
587    fn register_replaces_same_kind() {
588        let mut reg = BackendRegistry::new();
589        reg.register(BackendEntry::new(BackendKind::Cuda, false));
590        assert!(!reg.get(BackendKind::Cuda).unwrap().available);
591        reg.register(BackendEntry::new(BackendKind::Cuda, true));
592        assert_eq!(reg.len(), 1);
593        assert!(reg.get(BackendKind::Cuda).unwrap().available);
594    }
595
596    #[test]
597    fn available_kinds_lists_only_available() {
598        let mut reg = three_backend_registry();
599        reg.set_available(BackendKind::Rocm, false);
600        let avail = reg.available_kinds();
601        assert!(avail.contains(&BackendKind::Cuda));
602        assert!(avail.contains(&BackendKind::Cpu));
603        assert!(!avail.contains(&BackendKind::Rocm));
604    }
605
606    #[test]
607    fn set_methods_report_presence() {
608        let mut reg = BackendRegistry::new();
609        assert!(!reg.set_available(BackendKind::Cuda, true));
610        reg.register(BackendEntry::new(BackendKind::Cuda, false));
611        assert!(reg.set_available(BackendKind::Cuda, true));
612        assert!(reg.set_capabilities(BackendKind::Cuda, Capabilities::default()));
613        assert!(!reg.set_capabilities(BackendKind::Metal, Capabilities::default()));
614    }
615
616    #[test]
617    fn op_class_tensor_core_preference() {
618        assert!(OpClass::MatMul.prefers_tensor_cores());
619        assert!(OpClass::Attention.prefers_tensor_cores());
620        assert!(!OpClass::Elementwise.prefers_tensor_cores());
621        assert!(!OpClass::Memory.prefers_tensor_cores());
622        assert_eq!(OpClass::ALL.len(), 6);
623    }
624
625    // ── Workload-size-aware selection ────────────────────────────────────────
626
627    /// 64 KiB, mirroring `oxicuda::AUTO_SELECT_THRESHOLD_BYTES`.
628    const THRESHOLD: usize = 64 * 1024;
629
630    #[test]
631    fn for_workload_pins_cpu_below_threshold_only() {
632        let below = SelectionRequest::any().for_workload(THRESHOLD - 1, THRESHOLD);
633        assert_eq!(below.pin, Some(BackendKind::Cpu));
634        // Exactly at the threshold is *not* below it → GPU still eligible.
635        let at = SelectionRequest::any().for_workload(THRESHOLD, THRESHOLD);
636        assert_eq!(at, SelectionRequest::any());
637        let above = SelectionRequest::any().for_workload(THRESHOLD + 1, THRESHOLD);
638        assert_eq!(above, SelectionRequest::any());
639    }
640
641    #[test]
642    fn for_workload_respects_explicit_preferences() {
643        // An explicit require_gpu is never downgraded to the host backend.
644        let gpu = SelectionRequest::require_gpu().for_workload(16, THRESHOLD);
645        assert_eq!(gpu.pin, None);
646        assert!(gpu.require_gpu);
647        // An explicit pin is never overwritten.
648        let pinned = SelectionRequest::pinned(BackendKind::Metal).for_workload(16, THRESHOLD);
649        assert_eq!(pinned.pin, Some(BackendKind::Metal));
650    }
651
652    #[test]
653    fn select_for_workload_routes_small_to_cpu_and_large_to_gpu() {
654        let reg = three_backend_registry();
655        let any = SelectionRequest::any();
656        assert_eq!(
657            reg.select_for_workload(&any, 1024, THRESHOLD)
658                .expect("small workload must select a backend"),
659            BackendKind::Cpu,
660            "a 1 KiB workload must stay on the host"
661        );
662        assert_eq!(
663            reg.select_for_workload(&any, 1024 * 1024, THRESHOLD)
664                .expect("large workload must select a backend"),
665            BackendKind::Cuda,
666            "a 1 MiB workload must reach the highest-priority GPU"
667        );
668    }
669
670    #[test]
671    fn select_for_workload_falls_back_when_no_cpu_entry() {
672        let mut reg = BackendRegistry::new();
673        reg.register(BackendEntry::new(BackendKind::Cuda, true));
674        // No CPU entry at all: the narrowed (CPU-pinned) request cannot be
675        // satisfied, so the original request must still be honoured.
676        assert_eq!(
677            reg.select_for_workload(&SelectionRequest::any(), 16, THRESHOLD)
678                .expect("must fall back to the unnarrowed request"),
679            BackendKind::Cuda
680        );
681    }
682
683    #[test]
684    fn select_for_workload_still_fails_when_nothing_qualifies() {
685        let mut reg = three_backend_registry();
686        for kind in BackendKind::ALL {
687            reg.set_available(kind, false);
688        }
689        assert!(
690            reg.select_for_workload(&SelectionRequest::any(), 16, THRESHOLD)
691                .is_err()
692        );
693    }
694
695    #[test]
696    fn selection_request_constructors() {
697        assert!(SelectionRequest::require_gpu().require_gpu);
698        assert_eq!(
699            SelectionRequest::pinned(BackendKind::Metal).pin,
700            Some(BackendKind::Metal)
701        );
702        assert_eq!(SelectionRequest::any(), SelectionRequest::default());
703    }
704}
oxicuda_backend/registry.rs

oxicuda_backend/
registry.rs
