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

rlx_driver/
device.rs
