Skip to main content

cubecl_runtime/device/
metal.rs

1use cubecl_common::device::{Device, DeviceId};
2
3/// A device of the native Metal runtime.
4#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
5pub enum MetalDevice {
6    #[default]
7    /// Default Metal device (usually the first GPU)
8    DefaultDevice,
9    /// Discrete GPU by index
10    DiscreteGpu(usize),
11    /// Integrated GPU by index
12    IntegratedGpu(usize),
13    /// Existing device with unique ID, registered through the Metal runtime.
14    Existing(u32),
15}
16
17impl Device for MetalDevice {
18    fn from_id(device_id: DeviceId) -> Self {
19        match device_id.type_id {
20            0 => Self::DefaultDevice,
21            1 => Self::DiscreteGpu(device_id.index_id as usize),
22            2 => Self::IntegratedGpu(device_id.index_id as usize),
23            3 => Self::Existing(device_id.index_id as u32),
24            _ => panic!("Invalid Metal device ID: {:?}", device_id),
25        }
26    }
27
28    fn to_id(&self) -> DeviceId {
29        match self {
30            Self::DefaultDevice => DeviceId {
31                type_id: 0,
32                index_id: 0,
33            },
34            Self::DiscreteGpu(idx) => DeviceId {
35                type_id: 1,
36                index_id: super::index_id(*idx),
37            },
38            Self::IntegratedGpu(idx) => DeviceId {
39                type_id: 2,
40                index_id: super::index_id(*idx),
41            },
42            Self::Existing(id) => DeviceId {
43                type_id: 3,
44                index_id: *id as u16,
45            },
46        }
47    }
48}
49
50impl core::fmt::Display for MetalDevice {
51    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
52        match self {
53            Self::DefaultDevice => write!(f, "Metal (default)"),
54            Self::DiscreteGpu(idx) => write!(f, "Metal DiscreteGpu {}", idx),
55            Self::IntegratedGpu(idx) => write!(f, "Metal IntegratedGpu {}", idx),
56            Self::Existing(id) => write!(f, "Metal Device {}", id),
57        }
58    }
59}