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

rlx_runtime/
device_parse.rs
