rlx_driver/
device.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Device selection — which backend to use.
6
7/// Target device for graph execution.
8///
9/// Each variant maps to a backend crate gated by a Cargo feature.
10/// Use `Device::is_available()` to check if the feature is enabled.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub enum Device {
13    // ── CPU ─────────────────────────────────────────────────
14    /// CPU with SIMD (NEON/AVX) + optional BLAS.
15    Cpu,
16
17    // ── Apple ───────────────────────────────────────────────
18    /// GPU via Apple Metal (Metal Performance Shaders).
19    Metal,
20    /// Apple MLX framework (unified memory GPU).
21    Mlx,
22    /// Apple Neural Engine.
23    Ane,
24
25    // ── NVIDIA ──────────────────────────────────────────────
26    /// NVIDIA GPU via native CUDA (cuBLAS, cuDNN).
27    Cuda,
28
29    // ── AMD ─────────────────────────────────────────────────
30    /// AMD GPU via ROCm/HIP.
31    Rocm,
32    /// AMD XDNA / Ryzen AI NPU (AI Engine `aie2`) via the `amdxdna` driver.
33    Xdna,
34
35    // ── External (USB4 / Thunderbolt) ───────────────────────
36    /// Discrete GPU on a PCIe tunnel, driven from userspace. The path for
37    /// hosts with no in-tree driver for PCI class 0x03 (Apple Silicon).
38    Egpu,
39
40    // ── Intel ───────────────────────────────────────────────
41    /// Intel GPU (Arc / Data Center Max) via oneAPI Level Zero.
42    OneApi,
43
44    // ── Google ──────────────────────────────────────────────
45    /// Google TPU via libtpu's PJRT plugin (no Python).
46    Tpu,
47
48    // ── Qualcomm ────────────────────────────────────────────
49    /// Qualcomm Hexagon NPU via QNN (AI Engine Direct).
50    Hexagon,
51
52    // ── Cross-platform GPU ──────────────────────────────────
53    /// Portable GPU via wgpu (Metal/Vulkan/DX12/WebGPU).
54    Gpu,
55    /// Vulkan compute shaders.
56    Vulkan,
57    /// OpenGL compute shaders (legacy).
58    OpenGl,
59    /// DirectX 12 compute (Windows).
60    DirectX,
61    /// WebGPU (WASM target).
62    WebGpu,
63}
64
65impl Device {
66    /// Human-readable name (no engine-layer info).
67    /// `is_available` / `available` live in rlx-runtime since they
68    /// consult the engine's backend registry — keeping them out of
69    /// the driver layer preserves the one-way dep direction.
70    pub fn name(self) -> &'static str {
71        match self {
72            Device::Cpu => "CPU",
73            Device::Metal => "Metal",
74            Device::Mlx => "MLX",
75            Device::Ane => "ANE",
76            Device::Cuda => "CUDA",
77            Device::Rocm => "ROCm",
78            Device::Xdna => "XDNA NPU",
79            Device::Egpu => "eGPU (PCIe tunnel)",
80            Device::OneApi => "oneAPI (Level Zero)",
81            Device::Tpu => "TPU",
82            Device::Hexagon => "Hexagon NPU",
83            Device::Gpu => "GPU (wgpu)",
84            Device::Vulkan => "Vulkan",
85            Device::OpenGl => "OpenGL",
86            Device::DirectX => "DirectX 12",
87            Device::WebGpu => "WebGPU",
88        }
89    }
90
91    /// Canonical lowercase token that round-trips through `FromStr`.
92    /// Use this when forwarding a device across a CLI / `--device`
93    /// boundary; [`name`](Self::name) is human-facing and does not always
94    /// round-trip (e.g. `"GPU (wgpu)"`).
95    pub fn as_arg(self) -> &'static str {
96        match self {
97            Device::Cpu => "cpu",
98            Device::Metal => "metal",
99            Device::Mlx => "mlx",
100            Device::Ane => "ane",
101            Device::Cuda => "cuda",
102            Device::Rocm => "rocm",
103            Device::Xdna => "xdna",
104            Device::Egpu => "egpu",
105            Device::OneApi => "oneapi",
106            Device::Tpu => "tpu",
107            Device::Hexagon => "hexagon",
108            Device::Gpu => "gpu",
109            Device::Vulkan => "vulkan",
110            Device::OpenGl => "opengl",
111            Device::DirectX => "directx",
112            Device::WebGpu => "webgpu",
113        }
114    }
115
116    /// All variant labels — convenience for callers that want to
117    /// enumerate without listing every variant manually. Pair
118    /// with `rlx_runtime::available_devices()` to filter.
119    pub fn all() -> &'static [Device] {
120        &[
121            Device::Cpu,
122            Device::Metal,
123            Device::Mlx,
124            Device::Ane,
125            Device::Cuda,
126            Device::Rocm,
127            Device::Xdna,
128            Device::Egpu,
129            Device::OneApi,
130            Device::Tpu,
131            Device::Hexagon,
132            Device::Gpu,
133            Device::Vulkan,
134            Device::OpenGl,
135            Device::DirectX,
136            Device::WebGpu,
137        ]
138    }
139}
140
141impl std::fmt::Display for Device {
142    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
143        write!(f, "{}", self.name())
144    }
145}
146
147/// Error returned by `Device::from_str` when the input doesn't match
148/// any known device alias.
149#[derive(Debug, Clone, PartialEq, Eq)]
150pub struct DeviceFromStrError(pub String);
151
152impl std::fmt::Display for DeviceFromStrError {
153    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
154        write!(
155            f,
156            "unknown device '{}' (try: cpu, metal, mlx, ane, cuda, rocm, xdna, egpu, oneapi, gpu, vulkan, opengl, directx, webgpu, tpu)",
157            self.0
158        )
159    }
160}
161
162impl std::error::Error for DeviceFromStrError {}
163
164impl std::str::FromStr for Device {
165    type Err = DeviceFromStrError;
166
167    fn from_str(s: &str) -> Result<Self, Self::Err> {
168        let key = s.trim().to_ascii_lowercase();
169        Ok(match key.as_str() {
170            "cpu" => Device::Cpu,
171            "metal" | "mps" | "mtl" => Device::Metal,
172            "mlx" => Device::Mlx,
173            "ane" | "coreml" | "neural-engine" => Device::Ane,
174            "cuda" | "nvidia" => Device::Cuda,
175            "rocm" | "hip" | "amd" => Device::Rocm,
176            "xdna" | "aie" | "aie2" | "ryzenai" | "ryzen-ai" | "amdnpu" => Device::Xdna,
177            "egpu" | "tbgpu" | "thunderbolt" => Device::Egpu,
178            "oneapi" | "levelzero" | "level-zero" | "l0" | "intel" | "sycl" => Device::OneApi,
179            "gpu" | "wgpu" => Device::Gpu,
180            "vulkan" | "vk" => Device::Vulkan,
181            "opengl" | "gl" => Device::OpenGl,
182            "directx" | "dx12" | "d3d12" => Device::DirectX,
183            "webgpu" => Device::WebGpu,
184            "tpu" => Device::Tpu,
185            "hexagon" | "htp" | "qnn" => Device::Hexagon,
186            _ => return Err(DeviceFromStrError(s.to_string())),
187        })
188    }
189}
190
191/// Per-family backend support filter.
192///
193/// Each model family declares which devices it can execute on (e.g.,
194/// SAM adds TPU on top of the standard set; some VLM crates exclude
195/// MLX until the vision tower lands). A single shared
196/// [`BackendSupport`] impl per family lets [`validate_device`] return
197/// uniform error messages instead of every model crate hand-rolling
198/// the same `match` ladder.
199pub trait BackendSupport {
200    /// Short stable family identifier (`"qwen3"`, `"llama32"`, `"sam"`).
201    fn family(&self) -> &'static str;
202
203    /// `true` if this family can execute on `device` today.
204    fn supports(&self, device: Device) -> bool;
205}
206
207/// Workspace-wide standard backend set: CPU, Metal, MLX, CUDA, ROCm, GPU.
208///
209/// New families default to this set via [`StandardBackends`] until they
210/// need to opt in/out. Mirrors `rlx_core::STANDARD_DEVICES`.
211pub const STANDARD_DEVICES: &[Device] = &[
212    Device::Cpu,
213    Device::Metal,
214    Device::Mlx,
215    Device::Cuda,
216    Device::Rocm,
217    Device::Gpu,
218];
219
220/// Default [`BackendSupport`] for families on the standard backend set.
221#[derive(Debug, Clone, Copy)]
222pub struct StandardBackends(pub &'static str);
223
224impl BackendSupport for StandardBackends {
225    fn family(&self) -> &'static str {
226        self.0
227    }
228    fn supports(&self, device: Device) -> bool {
229        STANDARD_DEVICES.contains(&device)
230    }
231}
232
233/// Validate that `device` is supported by `family`. Returns the same device
234/// on success; on failure, formats a uniform error string. Callers that
235/// need a typed error should use [`BackendSupport::supports`] directly.
236pub fn validate_device<S: BackendSupport>(support: &S, device: Device) -> Result<Device, String> {
237    if support.supports(device) {
238        Ok(device)
239    } else {
240        Err(format!(
241            "device {} is not supported by family `{}`",
242            device.name(),
243            support.family()
244        ))
245    }
246}
247
248#[cfg(test)]
249mod from_str_tests {
250    use super::*;
251    use std::str::FromStr;
252
253    #[test]
254    fn parse_basics() {
255        assert_eq!(Device::from_str("cpu").unwrap(), Device::Cpu);
256        assert_eq!(Device::from_str("CUDA").unwrap(), Device::Cuda);
257        assert_eq!(Device::from_str("mps").unwrap(), Device::Metal);
258        assert_eq!(Device::from_str("wgpu").unwrap(), Device::Gpu);
259        assert!(Device::from_str("nothing").is_err());
260    }
261
262    #[test]
263    fn as_arg_round_trips_through_from_str() {
264        // The canonical CLI token must parse back to the same device for
265        // every variant (unlike `name`, e.g. "GPU (wgpu)").
266        for &d in Device::all() {
267            assert_eq!(
268                Device::from_str(d.as_arg()).unwrap(),
269                d,
270                "round-trip failed for {d:?}"
271            );
272        }
273    }
274
275    #[test]
276    fn standard_backends_set() {
277        let s = StandardBackends("qwen3");
278        assert!(s.supports(Device::Cpu));
279        assert!(s.supports(Device::Metal));
280        assert!(!s.supports(Device::Tpu));
281        assert!(validate_device(&s, Device::Tpu).is_err());
282    }
283}
rlx_driver/device.rs

rlx_driver/
device.rs
