rlx_driver/
device.rs

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

rlx_driver/
device.rs
