1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
//! Device types and placement identifiers (see `docs/architecture/ORT2.md` ยง4.2).
//!
//! Device placement is a first-class annotation on every [`crate::Value`] and
//! [`crate::Node`], enabling multi-device partitioning without side tables.
/// A class of compute device / execution backend.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub enum DeviceType {
Cpu,
Cuda,
Rocm,
CoreMl,
Mlx,
WebGpu,
Qnn,
OpenVino,
/// Vendor / experimental backend keyed by an opaque id.
Custom(u32),
}
impl DeviceType {
/// The canonical lower-case name used in traces and diagnostics.
///
/// One owner for these spellings so a trace produced by any execution
/// provider labels its device the same way. `Custom` keeps its opaque id
/// rather than collapsing every vendor backend to one indistinguishable
/// name.
pub fn trace_name(self) -> std::borrow::Cow<'static, str> {
match self {
DeviceType::Cpu => "cpu".into(),
DeviceType::Cuda => "cuda".into(),
DeviceType::Rocm => "rocm".into(),
DeviceType::CoreMl => "coreml".into(),
DeviceType::Mlx => "mlx".into(),
DeviceType::WebGpu => "webgpu".into(),
DeviceType::Qnn => "qnn".into(),
DeviceType::OpenVino => "openvino".into(),
DeviceType::Custom(id) => format!("custom:{id}").into(),
}
}
/// The device a canonical name refers to, or `None` if it names none.
///
/// The inverse of [`trace_name`](DeviceType::trace_name), kept beside it so
/// the two cannot drift. A plugin execution provider is configured with a
/// device name from package metadata and has to report which device it
/// actually runs on; without this it could only guess, and guessing `Cpu`
/// makes a trace claim that Metal work happened on the host.
pub fn from_trace_name(name: &str) -> Option<Self> {
let name = name.trim().to_ascii_lowercase();
Some(match name.as_str() {
"cpu" => DeviceType::Cpu,
"cuda" => DeviceType::Cuda,
"rocm" => DeviceType::Rocm,
"coreml" => DeviceType::CoreMl,
// The Metal plugin is named for the API it targets; the device
// class it runs on is MLX.
"mlx" | "metal" => DeviceType::Mlx,
"webgpu" => DeviceType::WebGpu,
"qnn" => DeviceType::Qnn,
"openvino" => DeviceType::OpenVino,
other => {
let id = other.strip_prefix("custom:")?.parse().ok()?;
DeviceType::Custom(id)
}
})
}
/// Whether tensors on this device share the host address space and can be
/// accessed by CPU code without an explicit copy.
pub fn is_host_accessible(self) -> bool {
// MLX targets Apple unified memory; CPU is trivially host-accessible.
matches!(self, DeviceType::Cpu | DeviceType::Mlx)
}
}
/// A specific device instance: a [`DeviceType`] plus an ordinal index.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct DeviceId {
pub device_type: DeviceType,
pub index: u32,
}
impl DeviceId {
/// Construct a device id.
pub fn new(device_type: DeviceType, index: u32) -> Self {
Self { device_type, index }
}
/// The default host device (`CPU:0`).
pub fn cpu() -> Self {
Self::new(DeviceType::Cpu, 0)
}
/// A CUDA device by ordinal.
pub fn cuda(index: u32) -> Self {
Self::new(DeviceType::Cuda, index)
}
/// Whether this device is host-accessible (see [`DeviceType::is_host_accessible`]).
pub fn is_host_accessible(self) -> bool {
self.device_type.is_host_accessible()
}
}
impl Default for DeviceId {
fn default() -> Self {
Self::cpu()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_is_cpu0() {
assert_eq!(DeviceId::default(), DeviceId::cpu());
assert_eq!(DeviceId::default().index, 0);
}
#[test]
fn host_accessibility() {
assert!(DeviceId::cpu().is_host_accessible());
assert!(DeviceId::new(DeviceType::Mlx, 0).is_host_accessible());
assert!(!DeviceId::cuda(0).is_host_accessible());
}
}