onnx_runtime_ir/
device.rs

1//! Device types and placement identifiers (see `docs/ORT2.md` §4.2).
2//!
3//! Device placement is a first-class annotation on every [`crate::Value`] and
4//! [`crate::Node`], enabling multi-device partitioning without side tables.
5
6/// A class of compute device / execution backend.
7#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
8pub enum DeviceType {
9    Cpu,
10    Cuda,
11    Rocm,
12    CoreMl,
13    Mlx,
14    WebGpu,
15    Qnn,
16    OpenVino,
17    /// Vendor / experimental backend keyed by an opaque id.
18    Custom(u32),
19}
20
21impl DeviceType {
22    /// The canonical lower-case name used in traces and diagnostics.
23    ///
24    /// One owner for these spellings so a trace produced by any execution
25    /// provider labels its device the same way. `Custom` keeps its opaque id
26    /// rather than collapsing every vendor backend to one indistinguishable
27    /// name.
28    pub fn trace_name(self) -> std::borrow::Cow<'static, str> {
29        match self {
30            DeviceType::Cpu => "cpu".into(),
31            DeviceType::Cuda => "cuda".into(),
32            DeviceType::Rocm => "rocm".into(),
33            DeviceType::CoreMl => "coreml".into(),
34            DeviceType::Mlx => "mlx".into(),
35            DeviceType::WebGpu => "webgpu".into(),
36            DeviceType::Qnn => "qnn".into(),
37            DeviceType::OpenVino => "openvino".into(),
38            DeviceType::Custom(id) => format!("custom:{id}").into(),
39        }
40    }
41
42    /// Whether tensors on this device share the host address space and can be
43    /// accessed by CPU code without an explicit copy.
44    pub fn is_host_accessible(self) -> bool {
45        // MLX targets Apple unified memory; CPU is trivially host-accessible.
46        matches!(self, DeviceType::Cpu | DeviceType::Mlx)
47    }
48}
49
50/// A specific device instance: a [`DeviceType`] plus an ordinal index.
51#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
52pub struct DeviceId {
53    pub device_type: DeviceType,
54    pub index: u32,
55}
56
57impl DeviceId {
58    /// Construct a device id.
59    pub fn new(device_type: DeviceType, index: u32) -> Self {
60        Self { device_type, index }
61    }
62
63    /// The default host device (`CPU:0`).
64    pub fn cpu() -> Self {
65        Self::new(DeviceType::Cpu, 0)
66    }
67
68    /// A CUDA device by ordinal.
69    pub fn cuda(index: u32) -> Self {
70        Self::new(DeviceType::Cuda, index)
71    }
72
73    /// Whether this device is host-accessible (see [`DeviceType::is_host_accessible`]).
74    pub fn is_host_accessible(self) -> bool {
75        self.device_type.is_host_accessible()
76    }
77}
78
79impl Default for DeviceId {
80    fn default() -> Self {
81        Self::cpu()
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn default_is_cpu0() {
91        assert_eq!(DeviceId::default(), DeviceId::cpu());
92        assert_eq!(DeviceId::default().index, 0);
93    }
94
95    #[test]
96    fn host_accessibility() {
97        assert!(DeviceId::cpu().is_host_accessible());
98        assert!(DeviceId::new(DeviceType::Mlx, 0).is_host_accessible());
99        assert!(!DeviceId::cuda(0).is_host_accessible());
100    }
101}
onnx_runtime_ir/device.rs

onnx_runtime_ir/
device.rs
