cubecl_core/frontend/
topology.rs

1//! In this file we use a trick where the constant has the same name as the module containing
2//! the expand function, so that a user implicitly imports the expand function when importing the constant.
3
4use cubecl_ir::{ExpandElement, Scope};
5
6use super::ExpandElementTyped;
7
8macro_rules! constant {
9    ($ident:ident, $var:expr, $doc:expr) => {
10        #[doc = $doc]
11        pub const $ident: u32 = 2;
12
13        #[allow(non_snake_case)]
14        #[doc = $doc]
15        pub mod $ident {
16            use super::*;
17
18            /// Expansion of the constant variable.
19            pub fn expand(_scope: &mut Scope) -> ExpandElementTyped<u32> {
20                ExpandElementTyped::new(ExpandElement::Plain(crate::ir::Variable::builtin($var)))
21            }
22        }
23    };
24}
25
26constant!(
27    PLANE_DIM,
28    crate::ir::Builtin::PlaneDim,
29    r"
30The total amount of working units in a plane.
31"
32);
33
34constant!(
35    UNIT_POS_PLANE,
36    crate::ir::Builtin::UnitPosPlane,
37    r"
38The relative position of the working unit inside the plane, without regards to cube dimensions.
39"
40);
41
42constant!(
43    UNIT_POS,
44    crate::ir::Builtin::UnitPos,
45    r"
46The position of the working unit inside the cube, without regards to axis.
47"
48);
49
50constant!(
51    UNIT_POS_X,
52    crate::ir::Builtin::UnitPosX,
53    r"
54The position of the working unit inside the cube along the X axis.
55"
56);
57
58constant!(
59    UNIT_POS_Y,
60    crate::ir::Builtin::UnitPosY,
61    r"
62The position of the working unit inside the cube along the Y axis.
63"
64);
65
66constant!(
67    UNIT_POS_Z,
68    crate::ir::Builtin::UnitPosZ,
69    r"
70The position of the working unit inside the cube along the Z axis.
71"
72);
73
74constant!(
75    CUBE_CLUSTER_DIM,
76    crate::ir::Builtin::CubeClusterDim,
77    r"
78The total amount of cubes in a cluster.
79"
80);
81
82constant!(
83    CUBE_CLUSTER_DIM_X,
84    crate::ir::Builtin::CubeClusterDimX,
85    r"
86The dimension of the cluster along the X axis.
87"
88);
89
90constant!(
91    CUBE_CLUSTER_DIM_Y,
92    crate::ir::Builtin::CubeClusterDimY,
93    r"
94The dimension of the cluster along the Y axis.
95"
96);
97
98constant!(
99    CUBE_CLUSTER_DIM_Z,
100    crate::ir::Builtin::CubeClusterDimZ,
101    r"
102The dimension of the cluster along the Z axis.
103"
104);
105
106constant!(
107    CUBE_DIM,
108    crate::ir::Builtin::CubeDim,
109    r"
110The total amount of working units in a cube.
111"
112);
113
114constant!(
115    CUBE_DIM_X,
116    crate::ir::Builtin::CubeDimX,
117    r"
118The dimension of the cube along the X axis.
119"
120);
121
122constant!(
123    CUBE_DIM_Y,
124    crate::ir::Builtin::CubeDimY,
125    r"
126The dimension of the cube along the Y axis.
127"
128);
129
130constant!(
131    CUBE_DIM_Z,
132    crate::ir::Builtin::CubeDimZ,
133    r"
134The dimension of the cube along the Z axis.
135"
136);
137
138constant!(
139    CUBE_POS,
140    crate::ir::Builtin::CubePos,
141    r"
142The cube position, without regards to axis.
143"
144);
145
146constant!(
147    CUBE_POS_X,
148    crate::ir::Builtin::CubePosX,
149    r"
150The cube position along the X axis.
151"
152);
153
154constant!(
155    CUBE_POS_Y,
156    crate::ir::Builtin::CubePosY,
157    r"
158The cube position along the Y axis.
159"
160);
161
162constant!(
163    CUBE_POS_Z,
164    crate::ir::Builtin::CubePosZ,
165    r"
166The cube position along the Z axis.
167"
168);
169
170constant!(
171    CUBE_POS_CLUSTER,
172    crate::ir::Builtin::CubePosCluster,
173    r"
174The cube position within the cluster.
175"
176);
177
178constant!(
179    CUBE_POS_CLUSTER_X,
180    crate::ir::Builtin::CubePosClusterX,
181    r"
182The cube position in the cluster along the X axis.
183"
184);
185
186constant!(
187    CUBE_POS_CLUSTER_Y,
188    crate::ir::Builtin::CubePosClusterY,
189    r"
190The cube position in the cluster along the Y axis.
191"
192);
193
194constant!(
195    CUBE_POS_CLUSTER_Z,
196    crate::ir::Builtin::CubePosClusterZ,
197    r"
198The cube position in the cluster along the Z axis.
199"
200);
201
202constant!(
203    CUBE_COUNT,
204    crate::ir::Builtin::CubeCount,
205    r"
206The number of cubes launched.
207"
208);
209
210constant!(
211    CUBE_COUNT_X,
212    crate::ir::Builtin::CubeCountX,
213    r"
214The number of cubes launched along the X axis.
215"
216);
217
218constant!(
219    CUBE_COUNT_Y,
220    crate::ir::Builtin::CubeCountY,
221    r"
222The number of cubes launched along the Y axis.
223"
224);
225
226constant!(
227    CUBE_COUNT_Z,
228    crate::ir::Builtin::CubeCountZ,
229    r"
230The number of cubes launched along the Z axis.
231"
232);
233
234constant!(
235    ABSOLUTE_POS,
236    crate::ir::Builtin::AbsolutePos,
237    r"
238The position of the working unit in the whole cube kernel, without regards to cubes and axis.
239"
240);
241
242constant!(
243    ABSOLUTE_POS_X,
244    crate::ir::Builtin::AbsolutePosX,
245    r"
246The index of the working unit in the whole cube kernel along the X axis, without regards to cubes.
247"
248);
249
250constant!(
251    ABSOLUTE_POS_Y,
252    crate::ir::Builtin::AbsolutePosY,
253    r"
254The index of the working unit in the whole cube kernel along the Y axis, without regards to cubes.
255"
256);
257
258constant!(
259    ABSOLUTE_POS_Z,
260    crate::ir::Builtin::AbsolutePosZ,
261    r"
262The index of the working unit in the whole cube kernel along the Z axis, without regards to cubes.
263"
264);
cubecl_core/frontend/topology.rs

cubecl_core/frontend/
topology.rs
