cubecl_common/device/
base.rs

1use core::{
2    any::{Any, TypeId},
3    cmp::Ordering,
4};
5use cubecl_environment::sync::Arc;
6
7/// The device id.
8#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy, new)]
9pub struct DeviceId {
10    /// The type id identifies the type of the device.
11    pub type_id: u16,
12    /// The index id identifies the device number.
13    pub index_id: u16,
14}
15
16/// Device trait for all cubecl devices.
17pub trait Device: Default + Clone + core::fmt::Debug + Send + Sync + 'static {
18    /// Create a device from its [id](DeviceId).
19    fn from_id(device_id: DeviceId) -> Self;
20    /// Retrieve the [device id](DeviceId) from the device.
21    fn to_id(&self) -> DeviceId;
22}
23
24impl core::fmt::Display for DeviceId {
25    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
26        f.write_fmt(format_args!(
27            "DeviceId(type={}, index={})",
28            self.type_id, self.index_id
29        ))
30    }
31}
32
33impl Ord for DeviceId {
34    fn cmp(&self, other: &Self) -> Ordering {
35        match self.type_id.cmp(&other.type_id) {
36            Ordering::Equal => self.index_id.cmp(&other.index_id),
37            other => other,
38        }
39    }
40}
41
42impl PartialOrd for DeviceId {
43    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
44        Some(self.cmp(other))
45    }
46}
47
48/// An pointer to a service's server utilities.
49pub type ServerUtilitiesHandle = Arc<dyn Any + Send + Sync>;
50
51/// Represent a service that runs on a device.
52pub trait DeviceService: Send + 'static {
53    /// Initializes the service. It is only called once per device.
54    fn init(device_id: DeviceId) -> Self
55    where
56        Self: Sized;
57    /// Get the service utilities.
58    fn utilities(&self) -> ServerUtilitiesHandle;
59    /// Which pipeline stage this service runs on.
60    ///
61    /// Services on [`DeviceServiceStage::Upstream`] produce work ahead of time (e.g. autodiff graph
62    /// construction, kernel fusion) that is consumed by [`DeviceServiceStage::Downstream`] services,
63    /// which stream kernels to the device.
64    fn stage() -> DeviceServiceStage
65    where
66        Self: Sized,
67    {
68        DeviceServiceStage::Downstream
69    }
70}
71
72/// Pipeline stage a [`DeviceService`] runs on. Each stage gets its own runner thread per device,
73/// allowing upstream work to overlap with downstream kernel dispatch.
74#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
75pub enum DeviceServiceStage {
76    /// Produces work ahead of kernel dispatch (e.g. autodiff graph, fusion).
77    Upstream = 0,
78    /// Consumes upstream work and streams kernels to the device.
79    Downstream = 1,
80}
81
82/// One service instance: the device it runs on and the type of service running
83/// there. This pair is what the device registry keys on, so two instances never
84/// share one, whatever [`DeviceId`] two runtimes happen to hand out.
85#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
86pub struct ServiceId {
87    /// The device the service runs on.
88    pub device: DeviceId,
89    /// The service type, as the registry keys it.
90    pub service: TypeId,
91}
92
93impl ServiceId {
94    /// The service of type `S` on `device`.
95    pub fn of<S: 'static>(device: DeviceId) -> Self {
96        Self {
97            device,
98            service: TypeId::of::<S>(),
99        }
100    }
101}
102
103impl core::fmt::Display for ServiceId {
104    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
105        f.write_fmt(format_args!("{} ({:?})", self.device, self.service))
106    }
107}
cubecl_common/device/base.rs

cubecl_common/device/
base.rs
