cubecl_ir/
settings.rs

1use alloc::string::{String, ToString};
2use pliron::derive::format;
3
4use crate::AddressType;
5
6#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, serde::Serialize, serde::Deserialize)]
7#[allow(missing_docs)]
8#[format("`(` $x `, ` $y `, ` $z `)`")]
9/// The number of units across all 3 axis totalling to the number of working units in a cube.
10pub struct Dim3 {
11    /// The number of units in the x axis.
12    pub x: u32,
13    /// The number of units in the y axis.
14    pub y: u32,
15    /// The number of units in the z axis.
16    pub z: u32,
17}
18
19impl Dim3 {
20    /// Create a new `Dim3` with x = y = z = 1.
21    pub const fn new_single() -> Self {
22        Self { x: 1, y: 1, z: 1 }
23    }
24
25    /// Create a new `Dim3` with the given x, and y = z = 1.
26    pub const fn new_1d(x: u32) -> Self {
27        Self { x, y: 1, z: 1 }
28    }
29
30    /// Create a new dim3 with the given x and y, and z = 1.
31    pub const fn new_2d(x: u32, y: u32) -> Self {
32        Self { x, y, z: 1 }
33    }
34
35    /// Create a new `Dim3` with the given x, y and z.
36    /// This is equivalent to the [new](Dim3::new) function.
37    pub const fn new_3d(x: u32, y: u32, z: u32) -> Self {
38        Self { x, y, z }
39    }
40
41    /// Total numbers of units per cube
42    pub const fn num_elems(&self) -> u32 {
43        self.x * self.y * self.z
44    }
45
46    /// Whether this `Dim3` can fully contain `other`
47    pub const fn can_contain(&self, other: Dim3) -> bool {
48        self.x >= other.x && self.y >= other.y && self.z >= other.z
49    }
50}
51
52impl From<(u32, u32, u32)> for Dim3 {
53    fn from(value: (u32, u32, u32)) -> Self {
54        Dim3::new_3d(value.0, value.1, value.2)
55    }
56}
57
58impl From<Dim3> for (u32, u32, u32) {
59    fn from(val: Dim3) -> Self {
60        (val.x, val.y, val.z)
61    }
62}
63
64/// The kind of execution to be performed.
65#[derive(
66    Default, Hash, PartialEq, Eq, Clone, Debug, Copy, serde::Serialize, serde::Deserialize,
67)]
68pub enum ExecutionMode {
69    /// Checked kernels are safe.
70    #[default]
71    Checked,
72    /// Validate OOB and alert if OOB access occurs
73    Validate,
74    /// Unchecked kernels are unsafe.
75    Unchecked,
76}
77
78#[derive(Clone, Debug, PartialEq, Eq, Hash)]
79pub struct KernelSettings {
80    /// The cube dim of the kernel
81    pub cube_dim: Dim3,
82    /// The address type of the kernel
83    pub address_type: AddressType,
84    /// The name of the kernel
85    pub kernel_name: String,
86    /// Whether to include debug symbols
87    pub debug_symbols: bool,
88    /// CUDA Cluster dim, if any
89    pub cluster_dim: Option<Dim3>,
90    /// Execution mode
91    pub execution_mode: ExecutionMode,
92}
93
94impl KernelSettings {
95    pub fn new(cube_dim: Dim3, execution_mode: ExecutionMode, address_type: AddressType) -> Self {
96        Self {
97            cube_dim,
98            address_type,
99            kernel_name: String::new(),
100            debug_symbols: false,
101            cluster_dim: None,
102            execution_mode,
103        }
104    }
105}
106
107impl KernelSettings {
108    /// Set cube dimension.
109    pub fn cube_dim(mut self, cube_dim: Dim3) -> Self {
110        self.cube_dim = cube_dim;
111        self
112    }
113
114    /// Set address type.
115    pub fn address_type(mut self, ty: AddressType) -> Self {
116        self.address_type = ty;
117        self
118    }
119
120    /// Set kernel name.
121    pub fn kernel_name<S: AsRef<str>>(mut self, name: S) -> Self {
122        self.kernel_name = name.as_ref().to_string();
123        self
124    }
125
126    /// Activate debug symbols
127    pub fn debug_symbols(mut self) -> Self {
128        self.debug_symbols = true;
129        self
130    }
131
132    /// Set cluster dim
133    pub fn cluster_dim(mut self, cluster_dim: Dim3) -> Self {
134        self.cluster_dim = Some(cluster_dim);
135        self
136    }
137}
cubecl_ir/settings.rs

cubecl_ir/
settings.rs
