cubecl_ir/properties.rs
1use alloc::string::String;
2use core::hash::{BuildHasher, Hash, Hasher};
3
4use crate::{
5 AddressType, OpaqueType, SemanticType, StorageType, Type, TypeHash, VectorSize,
6 features::{AtomicUsage, Features, TypeUsage},
7};
8use cubecl_common::profile::TimingMethod;
9use enumset::EnumSet;
10
11/// Properties of the device related to the accelerator hardware.
12///
13/// # Plane size min/max
14///
15/// This is a range of possible values for the plane size.
16///
17/// For Nvidia GPUs and HIP, this is a single fixed value.
18///
19/// For wgpu with AMD GPUs this is a range of possible values, but the actual configured value
20/// is undefined and can only be queried at runtime. Should usually be 32, but not guaranteed.
21///
22/// For Intel GPUs, this is variable based on the number of registers used in the kernel. No way to
23/// query this at compile time is currently available. As a result, the minimum value should usually
24/// be assumed.
25#[derive(Debug, Clone, PartialEq, Eq, Hash)]
26pub struct HardwareProperties {
27 /// The maximum size of a single load instruction, in bits. Used for optimized vector sizes.
28 pub load_width: u32,
29 /// The minimum size of a plane on this device
30 pub plane_size_min: u32,
31 /// The maximum size of a plane on this device
32 pub plane_size_max: u32,
33 /// minimum number of bindings for a kernel that can be used at once.
34 pub max_bindings: u32,
35 /// Maximum amount of shared memory, in bytes
36 pub max_shared_memory_size: usize,
37 /// Maximum `CubeCount` in x, y and z dimensions
38 pub max_cube_count: (u32, u32, u32),
39 /// Maximum number of total units in a cube
40 pub max_units_per_cube: u32,
41 /// Maximum `CubeDim` in x, y, and z dimensions
42 pub max_cube_dim: (u32, u32, u32),
43 /// Number of streaming multiprocessors (SM), if available
44 pub num_streaming_multiprocessors: Option<u32>,
45 /// Number of available parallel cpu units, if the runtime is CPU.
46 pub num_cpu_cores: Option<u32>,
47 /// Number of tensor cores per SM, if any
48 pub num_tensor_cores: Option<u32>,
49 /// The minimum tiling dimension for a single axis in tensor cores.
50 ///
51 /// For a backend that only supports 16x16x16, the value would be 16.
52 /// For a backend that also supports 32x8x16, the value would be 8.
53 pub min_tensor_cores_dim: Option<u32>,
54 /// Maximum vector size supported by the device
55 pub max_vector_size: VectorSize,
56 /// Memory reserved for the driver when using cube-scoped matrices
57 pub cube_mma_reserved_shared_memory: usize,
58}
59
60/// Properties of the device related to allocation.
61#[derive(Debug, Clone, PartialEq, Eq, Hash)]
62pub struct MemoryDeviceProperties {
63 /// The maximum nr. of bytes that can be allocated in one go.
64 pub max_page_size: u64,
65 /// The required memory offset alignment in bytes.
66 pub alignment: u64,
67}
68
69/// Who a device is, and what its compiled code is keyed to.
70///
71/// The two fields answer different questions and must not be confused. `name`
72/// is for people: it names the physical part, and two machines holding the same
73/// part report the same name. `fingerprint` is for correctness: it is verbatim
74/// the string this runtime passes to
75/// [`compilation_store`](../../cubecl_runtime/compiler/fn.compilation_store.html),
76/// which is what puts a compiled artifact out of reach of a machine that cannot
77/// run it.
78///
79/// Reporting the fingerprint here rather than recomputing it is the whole
80/// point: a backend derives it once and hands it to both consumers, so the
81/// identity a bundle is stamped with and the namespace its kernels live under
82/// cannot drift apart.
83#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
84pub struct DeviceIdentity {
85 /// The device as it names itself — `AMD Radeon 8060S Graphics`,
86 /// `NVIDIA H100 PCIe`. Display only: distinct parts may share a name, so
87 /// nothing may gate on it.
88 pub name: String,
89 /// What this runtime compiles *for* — `hip-kernel_gfx1151`, `ptx_sm90`.
90 /// Verbatim the `compilation_store` fingerprint, so a namespace read back
91 /// out of a bundle compares against it directly.
92 pub fingerprint: String,
93}
94
95/// Properties of what the device can do, like what `Feature` are
96/// supported by it and what its memory properties are.
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct DeviceProperties {
99 /// The features supported by the runtime.
100 pub features: Features,
101 /// The memory properties of this client.
102 pub memory: MemoryDeviceProperties,
103 /// The topology properties of this client.
104 pub hardware: HardwareProperties,
105 /// The method used for profiling on the device.
106 pub timing_method: TimingMethod,
107 /// Who the device is, and what its kernels are keyed to.
108 pub identity: DeviceIdentity,
109}
110
111impl TypeHash for DeviceProperties {
112 fn write_hash(_hasher: &mut impl core::hash::Hasher) {
113 // ignored.
114 }
115}
116
117impl DeviceProperties {
118 /// Create a new feature set with the given features and memory properties.
119 ///
120 /// `identity` is a required argument rather than something a backend may
121 /// fill in afterwards, so a runtime cannot ship reporting an anonymous
122 /// device — the failure mode that leaves a bundle unable to say what it was
123 /// built for.
124 pub fn new(
125 features: Features,
126 memory_props: MemoryDeviceProperties,
127 hardware: HardwareProperties,
128 timing_method: TimingMethod,
129 identity: DeviceIdentity,
130 ) -> Self {
131 DeviceProperties {
132 features,
133 memory: memory_props,
134 hardware,
135 timing_method,
136 identity,
137 }
138 }
139
140 /// Get the usages for a type
141 pub fn type_usage(&self, ty: StorageType) -> EnumSet<TypeUsage> {
142 self.features.type_usage(ty)
143 }
144
145 /// Get the usages for an atomic type
146 pub fn atomic_type_usage(&self, ty: Type) -> EnumSet<AtomicUsage> {
147 self.features.atomic_type_usage(ty)
148 }
149
150 /// Whether the type is supported in any way
151 pub fn supports_type(&self, ty: impl Into<Type>) -> bool {
152 self.features.supports_type(ty)
153 }
154
155 /// Whether the address type is supported in any way
156 pub fn supports_address(&self, ty: impl Into<AddressType>) -> bool {
157 self.features.supports_address(ty)
158 }
159
160 /// Register an address type to the features
161 pub fn register_address_type(&mut self, ty: impl Into<AddressType>) {
162 self.features.types.address.insert(ty.into());
163 }
164
165 /// Register an address type to the features
166 pub fn register_atomic_type_usage(&mut self, ty: Type, uses: impl Into<EnumSet<AtomicUsage>>) {
167 *self.features.types.atomic.entry(ty).or_default() |= uses.into();
168 }
169
170 /// Register a storage type to the features
171 pub fn register_type_usage(
172 &mut self,
173 ty: impl Into<StorageType>,
174 uses: impl Into<EnumSet<TypeUsage>>,
175 ) {
176 *self.features.types.storage.entry(ty.into()).or_default() |= uses.into();
177 }
178
179 /// Register a semantic type to the features
180 pub fn register_semantic_type(&mut self, ty: SemanticType) {
181 self.features.types.semantic.insert(ty);
182 }
183
184 /// Register an opaque type to the features
185 pub fn register_opaque_type(&mut self, ty: OpaqueType) {
186 self.features.types.opaque.insert(ty);
187 }
188
189 /// Create a stable hash of all device properties relevant to kernel compilation. Can be used
190 /// as a stable checksum for a compilation cache.
191 pub fn checksum(&self) -> u64 {
192 let state = foldhash::fast::FixedState::default();
193 let mut hasher = state.build_hasher();
194 self.features.hash(&mut hasher);
195 self.hardware.hash(&mut hasher);
196 hasher.finish()
197 }
198}