ultralytics_inference/
device.rs

1// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
3//! Hardware device support and abstraction.
4use std::fmt;
5use std::str::FromStr;
6
7/// Hardware device for inference.
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub enum Device {
10    /// CPU (Central Processing Unit).
11    Cpu,
12    /// CUDA (Compute Unified Device Architecture) for NVIDIA GPUs.
13    /// The argument specifies the device index (e.g., 0 for the first GPU).
14    Cuda(usize),
15    /// `CoreML` execution provider for Apple Silicon / macOS.
16    CoreMl,
17    /// `DirectML` (Direct Machine Learning) for Windows.
18    /// The argument specifies the device index.
19    DirectMl(usize),
20    /// Intel CPU via the `OpenVINO` execution provider (`intel:cpu`).
21    IntelCpu,
22    /// Intel GPU via the `OpenVINO` execution provider (`intel:gpu`).
23    IntelGpu,
24    /// Intel NPU via the `OpenVINO` execution provider (`intel:npu`).
25    IntelNpu,
26    /// XNNPACK (optimized floating-point neural network inference operators) for CPU.
27    Xnnpack,
28    /// `TensorRT` (NVIDIA `TensorRT`) for high-performance deep learning inference.
29    /// The argument specifies the device index.
30    TensorRt(usize),
31    /// `ROCm` (Radeon Open Compute) for AMD GPUs.
32    /// The argument specifies the device index.
33    Rocm(usize),
34}
35
36impl fmt::Display for Device {
37    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38        match self {
39            Self::Cpu => write!(f, "cpu"),
40            Self::Cuda(i) => write!(f, "cuda:{i}"),
41            Self::CoreMl => write!(f, "coreml"),
42            Self::DirectMl(i) => write!(f, "directml:{i}"),
43            Self::IntelCpu => write!(f, "intel:cpu"),
44            Self::IntelGpu => write!(f, "intel:gpu"),
45            Self::IntelNpu => write!(f, "intel:npu"),
46            Self::Xnnpack => write!(f, "xnnpack"),
47            Self::TensorRt(i) => write!(f, "tensorrt:{i}"),
48            Self::Rocm(i) => write!(f, "rocm:{i}"),
49        }
50    }
51}
52
53impl FromStr for Device {
54    type Err = String;
55
56    fn from_str(s: &str) -> Result<Self, Self::Err> {
57        let s = s.to_lowercase();
58        if let Some(rest) = s.strip_prefix("cuda") {
59            return Ok(Self::Cuda(parse_device_index(rest, &s)?));
60        }
61        if let Some(rest) = s.strip_prefix("directml") {
62            return Ok(Self::DirectMl(parse_device_index(rest, &s)?));
63        }
64        if let Some(rest) = s.strip_prefix("tensorrt") {
65            return Ok(Self::TensorRt(parse_device_index(rest, &s)?));
66        }
67        if let Some(rest) = s.strip_prefix("rocm") {
68            return Ok(Self::Rocm(parse_device_index(rest, &s)?));
69        }
70        match s.as_str() {
71            "cpu" => Ok(Self::Cpu),
72            "coreml" => Ok(Self::CoreMl),
73            "xnnpack" => Ok(Self::Xnnpack),
74            // Ultralytics OpenVINO naming: `intel:cpu`, `intel:gpu`, `intel:npu`.
75            "intel:cpu" => Ok(Self::IntelCpu),
76            "intel:gpu" => Ok(Self::IntelGpu),
77            "intel:npu" => Ok(Self::IntelNpu),
78            _ => Err(format!("Unknown device: {s}")),
79        }
80    }
81}
82
83/// Parse a trailing device index like `":0"`, defaulting to `0` when absent.
84///
85/// Anything else is an error rather than a silent fallback to device 0: a typo such as
86/// `cuda:abc`, `cuda:-1`, or `cudax` would otherwise run the whole job on the wrong device
87/// without telling anyone. `full` is the complete device string, for the error message.
88fn parse_device_index(s: &str, full: &str) -> Result<usize, String> {
89    match s.strip_prefix(':') {
90        None if s.is_empty() => Ok(0),
91        Some(index) => index
92            .parse()
93            .map_err(|_| format!("Invalid device index in '{full}': expected an integer")),
94        None => Err(format!("Unknown device: {full}")),
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn test_parse_device() {
104        assert_eq!(Device::from_str("cpu").unwrap(), Device::Cpu);
105        assert_eq!(Device::from_str("cuda").unwrap(), Device::Cuda(0));
106        assert_eq!(Device::from_str("cuda:0").unwrap(), Device::Cuda(0));
107        assert_eq!(Device::from_str("cuda:1").unwrap(), Device::Cuda(1));
108        assert_eq!(Device::from_str("coreml").unwrap(), Device::CoreMl);
109        assert_eq!(Device::from_str("directml").unwrap(), Device::DirectMl(0));
110        assert_eq!(Device::from_str("directml:1").unwrap(), Device::DirectMl(1));
111        // OpenVINO uses the Ultralytics `intel:<type>` naming.
112        assert_eq!(Device::from_str("intel:cpu").unwrap(), Device::IntelCpu);
113        assert_eq!(Device::from_str("intel:gpu").unwrap(), Device::IntelGpu);
114        assert_eq!(Device::from_str("intel:npu").unwrap(), Device::IntelNpu);
115        assert!(Device::from_str("intel").is_err());
116        assert!(Device::from_str("intel:tpu").is_err());
117        assert!(Device::from_str("openvino").is_err());
118    }
119
120    /// A malformed index used to silently resolve to device 0.
121    #[test]
122    fn test_parse_device_rejects_bad_index() {
123        for s in [
124            "cuda:abc",
125            "cuda:-1",
126            "cuda:",
127            "cuda:1.5",
128            "cudax",
129            "cuda0",
130            "tensorrt:x",
131            "rocm:-2",
132            "directmlfoo",
133            "cuda:99999999999999999999",
134        ] {
135            assert!(Device::from_str(s).is_err(), "{s} should not parse");
136        }
137    }
138
139    #[test]
140    fn test_device_display_roundtrip() {
141        for s in [
142            "cpu",
143            "cuda:0",
144            "cuda:1",
145            "coreml",
146            "directml:0",
147            "tensorrt:2",
148            "rocm:3",
149            "intel:cpu",
150            "intel:gpu",
151            "intel:npu",
152            "xnnpack",
153        ] {
154            assert_eq!(Device::from_str(s).unwrap().to_string(), s);
155        }
156    }
157}
ultralytics_inference/device.rs

ultralytics_inference/
device.rs
