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        "oneapi" | "levelzero" | "level-zero" | "l0" | "intel" | "sycl" => Ok(Device::OneApi),
36        "gpu" | "wgpu" => Ok(Device::Gpu),
37        "vulkan" | "vk" => Ok(Device::Vulkan),
38        "opengl" | "gl" => Ok(Device::OpenGl),
39        "directx" | "dx12" | "d3d12" => Ok(Device::DirectX),
40        "webgpu" => Ok(Device::WebGpu),
41        "tpu" => Ok(Device::Tpu),
42        "hexagon" | "htp" | "qnn" => Ok(Device::Hexagon),
43        "" => Err(ParseDeviceError {
44            input: s.to_string(),
45            message: "empty device name".into(),
46        }),
47        other => Err(ParseDeviceError {
48            input: s.to_string(),
49            message: format!(
50                "unknown device '{other}' (try: cpu, metal, mlx, ane, coreml, cuda, rocm, gpu, vulkan, tpu)"
51            ),
52        }),
53    }
54}
55
56/// Stable lower-case label for `device` (matches Cargo feature names).
57pub fn device_label(device: Device) -> &'static str {
58    match device {
59        Device::Cpu => "cpu",
60        Device::Metal => "metal",
61        Device::Mlx => "mlx",
62        Device::Ane => "ane",
63        Device::Cuda => "cuda",
64        Device::Rocm => "rocm",
65        Device::Xdna => "xdna",
66        Device::OneApi => "oneapi",
67        Device::Gpu => "gpu",
68        Device::Vulkan => "vulkan",
69        Device::OpenGl => "opengl",
70        Device::DirectX => "directx",
71        Device::WebGpu => "webgpu",
72        Device::Tpu => "tpu",
73        Device::Hexagon => "hexagon",
74    }
75}
76
77/// Parse comma/semicolon/whitespace-separated device lists (`RLX_DEVICES=cpu,metal`).
78pub fn parse_device_list(s: &str) -> Result<Vec<Device>, ParseDeviceError> {
79    let mut out = Vec::new();
80    for part in s.split([',', ';', ' ']) {
81        let part = part.trim();
82        if part.is_empty() {
83            continue;
84        }
85        let dev = parse_device(part)?;
86        if !out.contains(&dev) {
87            out.push(dev);
88        }
89    }
90    if out.is_empty() {
91        return Err(ParseDeviceError {
92            input: s.to_string(),
93            message: "device list is empty".into(),
94        });
95    }
96    Ok(out)
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn parse_aliases() {
105        assert_eq!(parse_device("CUDA").unwrap(), Device::Cuda);
106        assert_eq!(parse_device("wgpu").unwrap(), Device::Gpu);
107        assert_eq!(parse_device("coreml").unwrap(), Device::Ane);
108        assert_eq!(
109            parse_device_list("cpu, metal;mlx").unwrap(),
110            vec![Device::Cpu, Device::Metal, Device::Mlx,]
111        );
112    }
113}
rlx_runtime/device_parse.rs

rlx_runtime/
device_parse.rs
