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/// How much source-level debug data a kernel carries. Profilers and debuggers read it to map
79/// kernel code back to the `#[cube]` source.
80#[derive(
81    Default,
82    Hash,
83    PartialEq,
84    Eq,
85    PartialOrd,
86    Ord,
87    Clone,
88    Copy,
89    Debug,
90    serde::Serialize,
91    serde::Deserialize,
92)]
93pub enum DebugInfo {
94    /// No debug data.
95    #[default]
96    #[serde(rename = "none")]
97    None,
98    /// Source lines, with each inlined `#[cube]` function as its own frame.
99    #[serde(rename = "line-tables")]
100    LineTables,
101    /// Line tables and the embedded source text.
102    #[serde(rename = "full")]
103    Full,
104}
105
106#[derive(Clone, Debug, PartialEq, Eq, Hash)]
107pub struct KernelSettings {
108    /// The cube dim of the kernel
109    pub cube_dim: Dim3,
110    /// The address type of the kernel
111    pub address_type: AddressType,
112    /// The name of the kernel
113    pub kernel_name: String,
114    /// The source-level debug data to include
115    pub debug_info: DebugInfo,
116    /// CUDA Cluster dim, if any
117    pub cluster_dim: Option<Dim3>,
118    /// Execution mode
119    pub execution_mode: ExecutionMode,
120}
121
122impl KernelSettings {
123    pub fn new(cube_dim: Dim3, execution_mode: ExecutionMode, address_type: AddressType) -> Self {
124        Self {
125            cube_dim,
126            address_type,
127            kernel_name: String::new(),
128            debug_info: DebugInfo::None,
129            cluster_dim: None,
130            execution_mode,
131        }
132    }
133}
134
135impl KernelSettings {
136    /// Set cube dimension.
137    pub fn cube_dim(mut self, cube_dim: Dim3) -> Self {
138        self.cube_dim = cube_dim;
139        self
140    }
141
142    /// Set address type.
143    pub fn address_type(mut self, ty: AddressType) -> Self {
144        self.address_type = ty;
145        self
146    }
147
148    /// Set kernel name.
149    pub fn kernel_name<S: AsRef<str>>(mut self, name: S) -> Self {
150        self.kernel_name = name.as_ref().to_string();
151        self
152    }
153
154    /// Include full debug data, whatever the cargo profile asks for.
155    #[must_use]
156    pub fn debug_symbols(self) -> Self {
157        self.debug_info(DebugInfo::Full)
158    }
159
160    /// Set the source-level debug data to include.
161    #[must_use]
162    pub fn debug_info(mut self, debug_info: DebugInfo) -> Self {
163        self.debug_info = debug_info;
164        self
165    }
166
167    /// Set cluster dim
168    pub fn cluster_dim(mut self, cluster_dim: Dim3) -> Self {
169        self.cluster_dim = Some(cluster_dim);
170        self
171    }
172}
cubecl_ir/settings.rs

cubecl_ir/
settings.rs
