rlx_runtime/
device_parse.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//! String identifiers for [`rlx_driver::Device`] — config files, CLI, env vars.
17
18use rlx_driver::Device;
19
20/// Failed to parse a device name.
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct ParseDeviceError {
23    pub input: String,
24    pub message: String,
25}
26
27impl std::fmt::Display for ParseDeviceError {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        write!(f, "{}", self.message)
30    }
31}
32
33impl std::error::Error for ParseDeviceError {}
34
35/// Lower-case Cargo feature names and common aliases → [`Device`].
36pub fn parse_device(s: &str) -> Result<Device, ParseDeviceError> {
37    let key = s.trim().to_ascii_lowercase();
38    match key.as_str() {
39        "cpu" => Ok(Device::Cpu),
40        "metal" | "mtl" => Ok(Device::Metal),
41        "mlx" => Ok(Device::Mlx),
42        "ane" | "coreml" | "neural-engine" => Ok(Device::Ane),
43        "cuda" | "nvidia" => Ok(Device::Cuda),
44        "rocm" | "hip" | "amd" => Ok(Device::Rocm),
45        "oneapi" | "levelzero" | "level-zero" | "l0" | "intel" | "sycl" => Ok(Device::OneApi),
46        "gpu" | "wgpu" => Ok(Device::Gpu),
47        "vulkan" | "vk" => Ok(Device::Vulkan),
48        "opengl" | "gl" => Ok(Device::OpenGl),
49        "directx" | "dx12" | "d3d12" => Ok(Device::DirectX),
50        "webgpu" => Ok(Device::WebGpu),
51        "tpu" => Ok(Device::Tpu),
52        "hexagon" | "htp" | "qnn" => Ok(Device::Hexagon),
53        "" => Err(ParseDeviceError {
54            input: s.to_string(),
55            message: "empty device name".into(),
56        }),
57        other => Err(ParseDeviceError {
58            input: s.to_string(),
59            message: format!(
60                "unknown device '{other}' (try: cpu, metal, mlx, ane, coreml, cuda, rocm, gpu, vulkan, tpu)"
61            ),
62        }),
63    }
64}
65
66/// Stable lower-case label for `device` (matches Cargo feature names).
67pub fn device_label(device: Device) -> &'static str {
68    match device {
69        Device::Cpu => "cpu",
70        Device::Metal => "metal",
71        Device::Mlx => "mlx",
72        Device::Ane => "ane",
73        Device::Cuda => "cuda",
74        Device::Rocm => "rocm",
75        Device::OneApi => "oneapi",
76        Device::Gpu => "gpu",
77        Device::Vulkan => "vulkan",
78        Device::OpenGl => "opengl",
79        Device::DirectX => "directx",
80        Device::WebGpu => "webgpu",
81        Device::Tpu => "tpu",
82        Device::Hexagon => "hexagon",
83    }
84}
85
86/// Parse comma/semicolon/whitespace-separated device lists (`RLX_DEVICES=cpu,metal`).
87pub fn parse_device_list(s: &str) -> Result<Vec<Device>, ParseDeviceError> {
88    let mut out = Vec::new();
89    for part in s.split([',', ';', ' ']) {
90        let part = part.trim();
91        if part.is_empty() {
92            continue;
93        }
94        let dev = parse_device(part)?;
95        if !out.contains(&dev) {
96            out.push(dev);
97        }
98    }
99    if out.is_empty() {
100        return Err(ParseDeviceError {
101            input: s.to_string(),
102            message: "device list is empty".into(),
103        });
104    }
105    Ok(out)
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn parse_aliases() {
114        assert_eq!(parse_device("CUDA").unwrap(), Device::Cuda);
115        assert_eq!(parse_device("wgpu").unwrap(), Device::Gpu);
116        assert_eq!(parse_device("coreml").unwrap(), Device::Ane);
117        assert_eq!(
118            parse_device_list("cpu, metal;mlx").unwrap(),
119            vec![Device::Cpu, Device::Metal, Device::Mlx,]
120        );
121    }
122}
rlx_runtime/device_parse.rs

rlx_runtime/
device_parse.rs
