cubecl_ir/properties.rs
1use core::hash::{BuildHasher, Hash, Hasher};
2
3use crate::{
4 AddressType, OpaqueType, SemanticType, StorageType, Type, TypeHash, VectorSize,
5 features::{AtomicUsage, Features, TypeUsage},
6};
7use cubecl_common::profile::TimingMethod;
8use enumset::EnumSet;
9
10/// Properties of the device related to the accelerator hardware.
11///
12/// # Plane size min/max
13///
14/// This is a range of possible values for the plane size.
15///
16/// For Nvidia GPUs and HIP, this is a single fixed value.
17///
18/// For wgpu with AMD GPUs this is a range of possible values, but the actual configured value
19/// is undefined and can only be queried at runtime. Should usually be 32, but not guaranteed.
20///
21/// For Intel GPUs, this is variable based on the number of registers used in the kernel. No way to
22/// query this at compile time is currently available. As a result, the minimum value should usually
23/// be assumed.
24#[derive(Debug, Clone, PartialEq, Eq, Hash)]
25pub struct HardwareProperties {
26 /// The maximum size of a single load instruction, in bits. Used for optimized vector sizes.
27 pub load_width: u32,
28 /// The minimum size of a plane on this device
29 pub plane_size_min: u32,
30 /// The maximum size of a plane on this device
31 pub plane_size_max: u32,
32 /// minimum number of bindings for a kernel that can be used at once.
33 pub max_bindings: u32,
34 /// Maximum amount of shared memory, in bytes
35 pub max_shared_memory_size: usize,
36 /// Maximum `CubeCount` in x, y and z dimensions
37 pub max_cube_count: (u32, u32, u32),
38 /// Maximum number of total units in a cube
39 pub max_units_per_cube: u32,
40 /// Maximum `CubeDim` in x, y, and z dimensions
41 pub max_cube_dim: (u32, u32, u32),
42 /// Number of streaming multiprocessors (SM), if available
43 pub num_streaming_multiprocessors: Option<u32>,
44 /// Number of available parallel cpu units, if the runtime is CPU.
45 pub num_cpu_cores: Option<u32>,
46 /// Number of tensor cores per SM, if any
47 pub num_tensor_cores: Option<u32>,
48 /// The minimum tiling dimension for a single axis in tensor cores.
49 ///
50 /// For a backend that only supports 16x16x16, the value would be 16.
51 /// For a backend that also supports 32x8x16, the value would be 8.
52 pub min_tensor_cores_dim: Option<u32>,
53 /// Maximum vector size supported by the device
54 pub max_vector_size: VectorSize,
55 /// Memory reserved for the driver when using cube-scoped matrices
56 pub cube_mma_reserved_shared_memory: usize,
57}
58
59/// Properties of the device related to allocation.
60#[derive(Debug, Clone, PartialEq, Eq, Hash)]
61pub struct MemoryDeviceProperties {
62 /// The maximum nr. of bytes that can be allocated in one go.
63 pub max_page_size: u64,
64 /// The required memory offset alignment in bytes.
65 pub alignment: u64,
66}
67
68/// Properties of what the device can do, like what `Feature` are
69/// supported by it and what its memory properties are.
70#[derive(Debug, Clone, PartialEq, Eq)]
71pub struct DeviceProperties {
72 /// The features supported by the runtime.
73 pub features: Features,
74 /// The memory properties of this client.
75 pub memory: MemoryDeviceProperties,
76 /// The topology properties of this client.
77 pub hardware: HardwareProperties,
78 /// The method used for profiling on the device.
79 pub timing_method: TimingMethod,
80}
81
82impl TypeHash for DeviceProperties {
83 fn write_hash(_hasher: &mut impl core::hash::Hasher) {
84 // ignored.
85 }
86}
87
88impl DeviceProperties {
89 /// Create a new feature set with the given features and memory properties.
90 pub fn new(
91 features: Features,
92 memory_props: MemoryDeviceProperties,
93 hardware: HardwareProperties,
94 timing_method: TimingMethod,
95 ) -> Self {
96 DeviceProperties {
97 features,
98 memory: memory_props,
99 hardware,
100 timing_method,
101 }
102 }
103
104 /// Get the usages for a type
105 pub fn type_usage(&self, ty: StorageType) -> EnumSet<TypeUsage> {
106 self.features.type_usage(ty)
107 }
108
109 /// Get the usages for an atomic type
110 pub fn atomic_type_usage(&self, ty: Type) -> EnumSet<AtomicUsage> {
111 self.features.atomic_type_usage(ty)
112 }
113
114 /// Whether the type is supported in any way
115 pub fn supports_type(&self, ty: impl Into<Type>) -> bool {
116 self.features.supports_type(ty)
117 }
118
119 /// Whether the address type is supported in any way
120 pub fn supports_address(&self, ty: impl Into<AddressType>) -> bool {
121 self.features.supports_address(ty)
122 }
123
124 /// Register an address type to the features
125 pub fn register_address_type(&mut self, ty: impl Into<AddressType>) {
126 self.features.types.address.insert(ty.into());
127 }
128
129 /// Register an address type to the features
130 pub fn register_atomic_type_usage(&mut self, ty: Type, uses: impl Into<EnumSet<AtomicUsage>>) {
131 *self.features.types.atomic.entry(ty).or_default() |= uses.into();
132 }
133
134 /// Register a storage type to the features
135 pub fn register_type_usage(
136 &mut self,
137 ty: impl Into<StorageType>,
138 uses: impl Into<EnumSet<TypeUsage>>,
139 ) {
140 *self.features.types.storage.entry(ty.into()).or_default() |= uses.into();
141 }
142
143 /// Register a semantic type to the features
144 pub fn register_semantic_type(&mut self, ty: SemanticType) {
145 self.features.types.semantic.insert(ty);
146 }
147
148 /// Register an opaque type to the features
149 pub fn register_opaque_type(&mut self, ty: OpaqueType) {
150 self.features.types.opaque.insert(ty);
151 }
152
153 /// Create a stable hash of all device properties relevant to kernel compilation. Can be used
154 /// as a stable checksum for a compilation cache.
155 pub fn checksum(&self) -> u64 {
156 let state = foldhash::fast::FixedState::default();
157 let mut hasher = state.build_hasher();
158 self.features.hash(&mut hasher);
159 self.hardware.hash(&mut hasher);
160 hasher.finish()
161 }
162}