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    /// Whether tensors on this device share the host address space and can be
23    /// accessed by CPU code without an explicit copy.
24    pub fn is_host_accessible(self) -> bool {
25        // MLX targets Apple unified memory; CPU is trivially host-accessible.
26        matches!(self, DeviceType::Cpu | DeviceType::Mlx)
27    }
28}
29
30/// A specific device instance: a [`DeviceType`] plus an ordinal index.
31#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
32pub struct DeviceId {
33    pub device_type: DeviceType,
34    pub index: u32,
35}
36
37impl DeviceId {
38    /// Construct a device id.
39    pub fn new(device_type: DeviceType, index: u32) -> Self {
40        Self { device_type, index }
41    }
42
43    /// The default host device (`CPU:0`).
44    pub fn cpu() -> Self {
45        Self::new(DeviceType::Cpu, 0)
46    }
47
48    /// A CUDA device by ordinal.
49    pub fn cuda(index: u32) -> Self {
50        Self::new(DeviceType::Cuda, index)
51    }
52
53    /// Whether this device is host-accessible (see [`DeviceType::is_host_accessible`]).
54    pub fn is_host_accessible(self) -> bool {
55        self.device_type.is_host_accessible()
56    }
57}
58
59impl Default for DeviceId {
60    fn default() -> Self {
61        Self::cpu()
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn default_is_cpu0() {
71        assert_eq!(DeviceId::default(), DeviceId::cpu());
72        assert_eq!(DeviceId::default().index, 0);
73    }
74
75    #[test]
76    fn host_accessibility() {
77        assert!(DeviceId::cpu().is_host_accessible());
78        assert!(DeviceId::new(DeviceType::Mlx, 0).is_host_accessible());
79        assert!(!DeviceId::cuda(0).is_host_accessible());
80    }
81}
onnx_runtime_ir/device.rs

onnx_runtime_ir/
device.rs
