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::Server::end_capture) records a graph,
55/// stores it in the backend's own registry, and returns this lightweight id;
56/// [`replay`](crate::server::Server::replay) and
57/// [`graph_destroy`](crate::server::Server::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    /// The kernel's type name, trimmed to its last path segment — the short
147    /// name a report can print without drowning the reader in type paths.
148    pub fn short_name(&self) -> &'static str {
149        let name = self.type_name.split('<').next().unwrap_or(self.type_name);
150        name.rsplit("::").next().unwrap_or(name)
151    }
152
153    /// The kernel's type, in full: what names it before its comptime
154    /// arguments tell instances apart.
155    pub fn type_name(&self) -> &'static str {
156        self.type_name
157    }
158
159    /// Create a new [kernel id](KernelId) for a type.
160    pub fn new<T: 'static>() -> Self {
161        Self {
162            type_id: core::any::TypeId::of::<T>(),
163            type_name: core::any::type_name::<T>(),
164            info: None,
165            cube_dim: Dim3::new_single(),
166            mode: ExecutionMode::Checked,
167            address_type: Default::default(),
168        }
169    }
170
171    /// Render the key in a standard format that can be used between runs.
172    ///
173    /// Can be used as a persistent kernel cache key.
174    pub fn stable_format(&self) -> String {
175        format!(
176            "{}-{}-{:?}-{:?}-{:?}",
177            self.type_name, self.address_type, self.cube_dim, self.mode, self.info
178        )
179    }
180
181    /// Hash the key in a stable way that can be used between runs.
182    ///
183    /// Can be used as a persistent kernel cache key.
184    pub fn stable_hash(&self) -> StableHash {
185        let mut hasher = StableHasher::new();
186        self.type_name.hash(&mut hasher);
187        self.address_type.hash(&mut hasher);
188        self.cube_dim.hash(&mut hasher);
189        self.mode.hash(&mut hasher);
190        self.info.hash(&mut hasher);
191
192        hasher.finalize()
193    }
194
195    /// Return the entrypoint name disambiguated with a stable hash discriminator.
196    pub fn entrypoint_name(&self, base: &str) -> String {
197        format!("{base}_{:08x}", self.stable_hash() as u32)
198    }
199
200    /// Add information to the [kernel id](KernelId).
201    ///
202    /// The information is used to differentiate kernels of the same kind but with different
203    /// configurations, which affect the generated code.
204    pub fn info<I: 'static + PartialEq + Eq + Hash + core::fmt::Debug + Send + Sync>(
205        mut self,
206        info: I,
207    ) -> Self {
208        self.info = Some(Info::new(info));
209        self
210    }
211
212    /// Set the [execution mode](ExecutionMode).
213    pub fn mode(mut self, mode: ExecutionMode) -> Self {
214        self.mode = mode;
215        self
216    }
217
218    /// Set the [cube dim](CubeDim).
219    pub fn cube_dim(mut self, cube_dim: Dim3) -> Self {
220        self.cube_dim = cube_dim;
221        self
222    }
223
224    /// Set the [`AddressType`].
225    pub fn address_type(mut self, addr_ty: AddressType) -> Self {
226        self.address_type = addr_ty;
227        self
228    }
229}
230
231impl core::fmt::Debug for Info {
232    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
233        self.value.fmt(f)
234    }
235}
236
237impl Info {
238    fn new<T: 'static + PartialEq + Eq + Hash + core::fmt::Debug + Send + Sync>(id: T) -> Self {
239        Self {
240            value: Arc::new(id),
241        }
242    }
243}
244
245/// This trait allows various types to be used as keys within a single data structure.
246///
247/// The downside is that the hashing method is hardcoded and cannot be configured using the
248/// [`core::hash::Hash`] function. The provided [Hasher] will be modified, but only based on the
249/// result of the hash from the [`DefaultHasher`].
250trait DynKey: core::fmt::Debug + Send + Sync {
251    fn dyn_type_id(&self) -> TypeId;
252    fn dyn_eq(&self, other: &dyn DynKey) -> bool;
253    fn dyn_hash(&self, state: &mut dyn Hasher);
254    fn dyn_hash_one(&self) -> StableHash;
255    fn as_any(&self) -> &dyn Any;
256}
257
258impl PartialEq for Info {
259    fn eq(&self, other: &Self) -> bool {
260        self.value.dyn_eq(other.value.as_ref())
261    }
262}
263
264/// Extra information
265#[derive(Clone)]
266pub(crate) struct Info {
267    value: Arc<dyn DynKey>,
268}
269impl Eq for Info {}
270
271impl Hash for Info {
272    fn hash<H: Hasher>(&self, state: &mut H) {
273        self.value.dyn_type_id().hash(state);
274        self.value.dyn_hash(state)
275    }
276}
277
278impl<T: 'static + PartialEq + Eq + Hash + core::fmt::Debug + Send + Sync> DynKey for T {
279    fn dyn_eq(&self, other: &dyn DynKey) -> bool {
280        if let Some(other) = other.as_any().downcast_ref::<T>() {
281            self == other
282        } else {
283            false
284        }
285    }
286
287    fn dyn_type_id(&self) -> TypeId {
288        TypeId::of::<T>()
289    }
290
291    fn dyn_hash(&self, state: &mut dyn Hasher) {
292        let hash = self.dyn_hash_one();
293        state.write_u128(hash);
294    }
295
296    fn dyn_hash_one(&self) -> StableHash {
297        let mut hasher = StableHasher::new();
298        self.hash(&mut hasher);
299        hasher.finalize()
300    }
301
302    fn as_any(&self) -> &dyn Any {
303        self
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310    use std::collections::HashSet;
311
312    #[test_log::test]
313    pub fn kernel_id_hash() {
314        let value_1 = KernelId::new::<()>().info("1");
315        let value_2 = KernelId::new::<()>().info("2");
316
317        let mut set = HashSet::new();
318
319        set.insert(value_1.clone());
320
321        assert!(set.contains(&value_1));
322        assert!(!set.contains(&value_2));
323    }
324}
cubecl_runtime/id.rs

cubecl_runtime/
id.rs
