cubecl_runtime/
id.rs

1use alloc::format;
2use alloc::string::String;
3use alloc::sync::Arc;
4use core::{
5    any::{Any, TypeId},
6    fmt::Display,
7    hash::{Hash, Hasher},
8};
9use cubecl_common::{
10    format::{DebugRaw, format_str},
11    hash::{StableHash, StableHasher},
12};
13use cubecl_ir::{
14    AddressType,
15    settings::{Dim3, ExecutionMode},
16};
17use derive_more::{Eq, PartialEq};
18
19#[macro_export(local_inner_macros)]
20/// Create a new storage ID type.
21macro_rules! storage_id_type {
22    ($name:ident) => {
23        /// Storage ID.
24        #[derive(Copy, Clone, Hash, PartialEq, Eq, Debug, PartialOrd, Ord)]
25        pub struct $name {
26            value: usize,
27        }
28
29        impl $name {
30            /// Create a new ID.
31            pub fn new() -> Self {
32                use core::sync::atomic::{AtomicUsize, Ordering};
33
34                static COUNTER: AtomicUsize = AtomicUsize::new(0);
35
36                let value = COUNTER.fetch_add(1, Ordering::Relaxed);
37                if value == usize::MAX {
38                    core::panic!("Memory ID overflowed");
39                }
40                Self { value }
41            }
42        }
43
44        impl Default for $name {
45            fn default() -> Self {
46                Self::new()
47            }
48        }
49    };
50}
51
52/// Identifies a backend-owned captured graph.
53///
54/// [`end_capture`](crate::server::ComputeServer::end_capture) records a graph,
55/// stores it in the backend's own registry, and returns this lightweight id;
56/// [`replay`](crate::server::ComputeServer::replay) and
57/// [`graph_destroy`](crate::server::ComputeServer::graph_destroy) take the id
58/// back to look the graph up. Referencing the graph by id keeps the raw
59/// executable inside the server — it never crosses the actor boundary in a box —
60/// exactly as memory is referenced by [`Handle`](crate::server::Handle) rather
61/// than by raw pointer.
62#[derive(Copy, Clone, Hash, PartialEq, Eq, Debug, PartialOrd, Ord)]
63pub struct GraphId {
64    value: u64,
65}
66
67impl GraphId {
68    /// Allocate a fresh, process-unique graph id.
69    pub fn new() -> Self {
70        use core::sync::atomic::{AtomicU64, Ordering};
71
72        static COUNTER: AtomicU64 = AtomicU64::new(0);
73
74        let value = COUNTER.fetch_add(1, Ordering::Relaxed);
75        if value == u64::MAX {
76            core::panic!("Graph ID overflowed");
77        }
78        Self { value }
79    }
80}
81
82impl Default for GraphId {
83    fn default() -> Self {
84        Self::new()
85    }
86}
87
88/// Kernel unique identifier.
89#[derive(Clone, PartialEq, Eq)]
90pub struct KernelId {
91    #[eq(skip)]
92    type_name: &'static str,
93    pub(crate) type_id: core::any::TypeId,
94    /// The [`CubeDim`] for this kernel
95    pub cube_dim: Dim3,
96    /// The address type for this kernel
97    pub address_type: AddressType,
98    /// The execution mode for this kernel
99    pub mode: ExecutionMode,
100    pub(crate) info: Option<Info>,
101}
102
103impl Hash for KernelId {
104    fn hash<H: Hasher>(&self, state: &mut H) {
105        self.type_id.hash(state);
106        self.address_type.hash(state);
107        self.cube_dim.hash(state);
108        self.mode.hash(state);
109        self.info.hash(state);
110    }
111}
112
113impl core::fmt::Debug for KernelId {
114    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
115        let mut debug_str = f.debug_struct("KernelId");
116        debug_str
117            .field("type", &DebugRaw(self.type_name))
118            .field("address_type", &self.address_type);
119        debug_str.field("cube_dim", &self.cube_dim);
120        debug_str.field("mode", &self.mode);
121        match &self.info {
122            Some(info) => debug_str.field("info", info),
123            None => debug_str.field("info", &self.info),
124        };
125        debug_str.finish()
126    }
127}
128
129impl Display for KernelId {
130    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
131        match &self.info {
132            Some(info) => f.write_str(
133                format_str(
134                    format!("{info:?}").as_str(),
135                    &[('(', ')'), ('[', ']'), ('{', '}')],
136                    true,
137                )
138                .as_str(),
139            ),
140            None => f.write_str("No info"),
141        }
142    }
143}
144
145impl KernelId {
146    /// Create a new [kernel id](KernelId) for a type.
147    pub fn new<T: 'static>() -> Self {
148        Self {
149            type_id: core::any::TypeId::of::<T>(),
150            type_name: core::any::type_name::<T>(),
151            info: None,
152            cube_dim: Dim3::new_single(),
153            mode: ExecutionMode::Checked,
154            address_type: Default::default(),
155        }
156    }
157
158    /// Render the key in a standard format that can be used between runs.
159    ///
160    /// Can be used as a persistent kernel cache key.
161    pub fn stable_format(&self) -> String {
162        format!(
163            "{}-{}-{:?}-{:?}-{:?}",
164            self.type_name, self.address_type, self.cube_dim, self.mode, self.info
165        )
166    }
167
168    /// Hash the key in a stable way that can be used between runs.
169    ///
170    /// Can be used as a persistent kernel cache key.
171    pub fn stable_hash(&self) -> StableHash {
172        let mut hasher = StableHasher::new();
173        self.type_name.hash(&mut hasher);
174        self.address_type.hash(&mut hasher);
175        self.cube_dim.hash(&mut hasher);
176        self.mode.hash(&mut hasher);
177        self.info.hash(&mut hasher);
178
179        hasher.finalize()
180    }
181
182    /// Add information to the [kernel id](KernelId).
183    ///
184    /// The information is used to differentiate kernels of the same kind but with different
185    /// configurations, which affect the generated code.
186    pub fn info<I: 'static + PartialEq + Eq + Hash + core::fmt::Debug + Send + Sync>(
187        mut self,
188        info: I,
189    ) -> Self {
190        self.info = Some(Info::new(info));
191        self
192    }
193
194    /// Set the [execution mode](ExecutionMode).
195    pub fn mode(mut self, mode: ExecutionMode) -> Self {
196        self.mode = mode;
197        self
198    }
199
200    /// Set the [cube dim](CubeDim).
201    pub fn cube_dim(mut self, cube_dim: Dim3) -> Self {
202        self.cube_dim = cube_dim;
203        self
204    }
205
206    /// Set the [`AddressType`].
207    pub fn address_type(mut self, addr_ty: AddressType) -> Self {
208        self.address_type = addr_ty;
209        self
210    }
211}
212
213impl core::fmt::Debug for Info {
214    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
215        self.value.fmt(f)
216    }
217}
218
219impl Info {
220    fn new<T: 'static + PartialEq + Eq + Hash + core::fmt::Debug + Send + Sync>(id: T) -> Self {
221        Self {
222            value: Arc::new(id),
223        }
224    }
225}
226
227/// This trait allows various types to be used as keys within a single data structure.
228///
229/// The downside is that the hashing method is hardcoded and cannot be configured using the
230/// [`core::hash::Hash`] function. The provided [Hasher] will be modified, but only based on the
231/// result of the hash from the [`DefaultHasher`].
232trait DynKey: core::fmt::Debug + Send + Sync {
233    fn dyn_type_id(&self) -> TypeId;
234    fn dyn_eq(&self, other: &dyn DynKey) -> bool;
235    fn dyn_hash(&self, state: &mut dyn Hasher);
236    fn dyn_hash_one(&self) -> StableHash;
237    fn as_any(&self) -> &dyn Any;
238}
239
240impl PartialEq for Info {
241    fn eq(&self, other: &Self) -> bool {
242        self.value.dyn_eq(other.value.as_ref())
243    }
244}
245
246/// Extra information
247#[derive(Clone)]
248pub(crate) struct Info {
249    value: Arc<dyn DynKey>,
250}
251impl Eq for Info {}
252
253impl Hash for Info {
254    fn hash<H: Hasher>(&self, state: &mut H) {
255        self.value.dyn_type_id().hash(state);
256        self.value.dyn_hash(state)
257    }
258}
259
260impl<T: 'static + PartialEq + Eq + Hash + core::fmt::Debug + Send + Sync> DynKey for T {
261    fn dyn_eq(&self, other: &dyn DynKey) -> bool {
262        if let Some(other) = other.as_any().downcast_ref::<T>() {
263            self == other
264        } else {
265            false
266        }
267    }
268
269    fn dyn_type_id(&self) -> TypeId {
270        TypeId::of::<T>()
271    }
272
273    fn dyn_hash(&self, state: &mut dyn Hasher) {
274        let hash = self.dyn_hash_one();
275        state.write_u128(hash);
276    }
277
278    fn dyn_hash_one(&self) -> StableHash {
279        let mut hasher = StableHasher::new();
280        self.hash(&mut hasher);
281        hasher.finalize()
282    }
283
284    fn as_any(&self) -> &dyn Any {
285        self
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use std::collections::HashSet;
293
294    #[test_log::test]
295    pub fn kernel_id_hash() {
296        let value_1 = KernelId::new::<()>().info("1");
297        let value_2 = KernelId::new::<()>().info("2");
298
299        let mut set = HashSet::new();
300
301        set.insert(value_1.clone());
302
303        assert!(set.contains(&value_1));
304        assert!(!set.contains(&value_2));
305    }
306}
cubecl_runtime/id.rs

cubecl_runtime/
id.rs
