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::{DebugInfo, 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    /// The debug data of this kernel, after the configuration lowered it.
101    pub debug_info: DebugInfo,
102    pub(crate) info: Option<Info>,
103}
104
105impl Hash for KernelId {
106    fn hash<H: Hasher>(&self, state: &mut H) {
107        self.type_id.hash(state);
108        self.address_type.hash(state);
109        self.cube_dim.hash(state);
110        self.mode.hash(state);
111        self.info.hash(state);
112        // A kernel without debug data keeps the hash it had before the level existed.
113        if self.debug_info != DebugInfo::None {
114            self.debug_info.hash(state);
115        }
116    }
117}
118
119impl core::fmt::Debug for KernelId {
120    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
121        let mut debug_str = f.debug_struct("KernelId");
122        debug_str
123            .field("type", &DebugRaw(self.type_name))
124            .field("address_type", &self.address_type);
125        debug_str.field("cube_dim", &self.cube_dim);
126        debug_str.field("mode", &self.mode);
127        debug_str.field("debug_info", &self.debug_info);
128        match &self.info {
129            Some(info) => debug_str.field("info", info),
130            None => debug_str.field("info", &self.info),
131        };
132        debug_str.finish()
133    }
134}
135
136impl Display for KernelId {
137    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
138        match &self.info {
139            Some(info) => f.write_str(
140                format_str(
141                    format!("{info:?}").as_str(),
142                    &[('(', ')'), ('[', ']'), ('{', '}')],
143                    true,
144                )
145                .as_str(),
146            ),
147            None => f.write_str("No info"),
148        }
149    }
150}
151
152impl KernelId {
153    /// The kernel's type name, trimmed to its last path segment — the short
154    /// name a report can print without drowning the reader in type paths.
155    pub fn short_name(&self) -> &'static str {
156        let name = self.type_name.split('<').next().unwrap_or(self.type_name);
157        name.rsplit("::").next().unwrap_or(name)
158    }
159
160    /// The kernel's type, in full: what names it before its comptime
161    /// arguments tell instances apart.
162    pub fn type_name(&self) -> &'static str {
163        self.type_name
164    }
165
166    /// Create a new [kernel id](KernelId) for a type.
167    pub fn new<T: 'static>() -> Self {
168        Self {
169            type_id: core::any::TypeId::of::<T>(),
170            type_name: core::any::type_name::<T>(),
171            info: None,
172            cube_dim: Dim3::new_single(),
173            mode: ExecutionMode::Checked,
174            address_type: Default::default(),
175            debug_info: DebugInfo::None,
176        }
177    }
178
179    /// Render the key in a standard format that can be used between runs.
180    ///
181    /// Can be used as a persistent kernel cache key.
182    pub fn stable_format(&self) -> String {
183        let format = format!(
184            "{}-{}-{:?}-{:?}-{:?}",
185            self.type_name, self.address_type, self.cube_dim, self.mode, self.info
186        );
187        match self.debug_info {
188            DebugInfo::None => format,
189            level => format!("{format}-{level:?}"),
190        }
191    }
192
193    /// Hash the key in a stable way that can be used between runs.
194    ///
195    /// Can be used as a persistent kernel cache key.
196    pub fn stable_hash(&self) -> StableHash {
197        let mut hasher = StableHasher::new();
198        self.type_name.hash(&mut hasher);
199        self.address_type.hash(&mut hasher);
200        self.cube_dim.hash(&mut hasher);
201        self.mode.hash(&mut hasher);
202        self.info.hash(&mut hasher);
203        // A kernel without debug data keeps its entry point name and cache key.
204        if self.debug_info != DebugInfo::None {
205            self.debug_info.hash(&mut hasher);
206        }
207
208        hasher.finalize()
209    }
210
211    /// Return the entrypoint name disambiguated with a stable hash discriminator.
212    pub fn entrypoint_name(&self, base: &str) -> String {
213        format!("{base}_{:08x}", self.stable_hash() as u32)
214    }
215
216    /// Add information to the [kernel id](KernelId).
217    ///
218    /// The information is used to differentiate kernels of the same kind but with different
219    /// configurations, which affect the generated code.
220    pub fn info<I: 'static + PartialEq + Eq + Hash + core::fmt::Debug + Send + Sync>(
221        mut self,
222        info: I,
223    ) -> Self {
224        self.info = Some(Info::new(info));
225        self
226    }
227
228    /// Set the [execution mode](ExecutionMode).
229    pub fn mode(mut self, mode: ExecutionMode) -> Self {
230        self.mode = mode;
231        self
232    }
233
234    /// Set the debug data that the kernel asks for. The id keeps the level after the
235    /// configuration lowers it, because that level changes the compiled kernel.
236    #[must_use]
237    pub fn debug_info(mut self, requested: DebugInfo) -> Self {
238        self.debug_info = crate::config::compilation::effective_debug_info(requested);
239        self
240    }
241
242    /// Set the [cube dim](CubeDim).
243    pub fn cube_dim(mut self, cube_dim: Dim3) -> Self {
244        self.cube_dim = cube_dim;
245        self
246    }
247
248    /// Set the [`AddressType`].
249    pub fn address_type(mut self, addr_ty: AddressType) -> Self {
250        self.address_type = addr_ty;
251        self
252    }
253}
254
255impl core::fmt::Debug for Info {
256    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
257        self.value.fmt(f)
258    }
259}
260
261impl Info {
262    fn new<T: 'static + PartialEq + Eq + Hash + core::fmt::Debug + Send + Sync>(id: T) -> Self {
263        Self {
264            value: Arc::new(id),
265        }
266    }
267}
268
269/// This trait allows various types to be used as keys within a single data structure.
270///
271/// The downside is that the hashing method is hardcoded and cannot be configured using the
272/// [`core::hash::Hash`] function. The provided [Hasher] will be modified, but only based on the
273/// result of the hash from the [`DefaultHasher`].
274trait DynKey: core::fmt::Debug + Send + Sync {
275    fn dyn_type_id(&self) -> TypeId;
276    fn dyn_eq(&self, other: &dyn DynKey) -> bool;
277    fn dyn_hash(&self, state: &mut dyn Hasher);
278    fn dyn_hash_one(&self) -> StableHash;
279    fn as_any(&self) -> &dyn Any;
280}
281
282impl PartialEq for Info {
283    fn eq(&self, other: &Self) -> bool {
284        self.value.dyn_eq(other.value.as_ref())
285    }
286}
287
288/// Extra information
289#[derive(Clone)]
290pub(crate) struct Info {
291    value: Arc<dyn DynKey>,
292}
293impl Eq for Info {}
294
295impl Hash for Info {
296    fn hash<H: Hasher>(&self, state: &mut H) {
297        self.value.dyn_type_id().hash(state);
298        self.value.dyn_hash(state)
299    }
300}
301
302impl<T: 'static + PartialEq + Eq + Hash + core::fmt::Debug + Send + Sync> DynKey for T {
303    fn dyn_eq(&self, other: &dyn DynKey) -> bool {
304        if let Some(other) = other.as_any().downcast_ref::<T>() {
305            self == other
306        } else {
307            false
308        }
309    }
310
311    fn dyn_type_id(&self) -> TypeId {
312        TypeId::of::<T>()
313    }
314
315    fn dyn_hash(&self, state: &mut dyn Hasher) {
316        let hash = self.dyn_hash_one();
317        state.write_u128(hash);
318    }
319
320    fn dyn_hash_one(&self) -> StableHash {
321        let mut hasher = StableHasher::new();
322        self.hash(&mut hasher);
323        hasher.finalize()
324    }
325
326    fn as_any(&self) -> &dyn Any {
327        self
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334    use std::collections::HashSet;
335
336    #[test_log::test]
337    pub fn kernel_id_hash() {
338        let value_1 = KernelId::new::<()>().info("1");
339        let value_2 = KernelId::new::<()>().info("2");
340
341        let mut set = HashSet::new();
342
343        set.insert(value_1.clone());
344
345        assert!(set.contains(&value_1));
346        assert!(!set.contains(&value_2));
347    }
348
349    #[test]
350    fn debug_info_changes_the_hash_and_the_cache_key() {
351        let with = |debug_info| KernelId {
352            debug_info,
353            ..KernelId::new::<()>().info("1")
354        };
355        let (none, lines, full) = (
356            with(DebugInfo::None),
357            with(DebugInfo::LineTables),
358            with(DebugInfo::Full),
359        );
360        assert_ne!(none.stable_hash(), lines.stable_hash());
361        assert_ne!(lines.stable_hash(), full.stable_hash());
362        assert_ne!(none.stable_format(), lines.stable_format());
363        assert_ne!(lines.stable_format(), full.stable_format());
364        assert_ne!(none, lines);
365        let state = std::collections::hash_map::RandomState::new();
366        assert_ne!(
367            core::hash::BuildHasher::hash_one(&state, &none),
368            core::hash::BuildHasher::hash_one(&state, &lines)
369        );
370        // Without debug data, the hash and the entry point name are the same as before.
371        let plain = KernelId::new::<()>().info("1");
372        assert_eq!(none.stable_hash(), plain.stable_hash());
373        assert_eq!(none.stable_format(), plain.stable_format());
374    }
375}
cubecl_runtime/id.rs

cubecl_runtime/
id.rs
