cubecl_ir/
properties.rs

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

cubecl_ir/
properties.rs
