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 super::ExpandElementTyped;
5
6macro_rules! constant {
7    ($ident:ident, $var:expr, $doc:expr) => {
8        #[doc = $doc]
9        pub const $ident: u32 = 2;
10
11        #[allow(non_snake_case)]
12        #[doc = $doc]
13        pub mod $ident {
14            use super::*;
15            use crate::frontend::{CubeContext, ExpandElement};
16
17            /// Expansion of the constant variable.
18            pub fn expand(_context: &mut CubeContext) -> ExpandElementTyped<u32> {
19                ExpandElementTyped::new(ExpandElement::Plain(crate::ir::Variable::builtin($var)))
20            }
21        }
22    };
23}
24
25constant!(
26    PLANE_DIM,
27    crate::ir::Builtin::PlaneDim,
28    r"
29The total amount of working units in a plane.
30"
31);
32
33constant!(
34    UNIT_POS_PLANE,
35    crate::ir::Builtin::UnitPosPlane,
36    r"
37The relative position of the working unit inside the plane, without regards to cube dimensions.
38"
39);
40
41constant!(
42    UNIT_POS,
43    crate::ir::Builtin::UnitPos,
44    r"
45The position of the working unit inside the cube, without regards to axis.
46"
47);
48
49constant!(
50    UNIT_POS_X,
51    crate::ir::Builtin::UnitPosX,
52    r"
53The position of the working unit inside the cube along the X axis.
54"
55);
56
57constant!(
58    UNIT_POS_Y,
59    crate::ir::Builtin::UnitPosY,
60    r"
61The position of the working unit inside the cube along the Y axis.
62"
63);
64
65constant!(
66    UNIT_POS_Z,
67    crate::ir::Builtin::UnitPosZ,
68    r"
69The position of the working unit inside the cube along the Z axis.
70"
71);
72
73constant!(
74    CUBE_DIM,
75    crate::ir::Builtin::CubeDim,
76    r"
77The total amount of working units in a cube.
78"
79);
80
81constant!(
82    CUBE_DIM_X,
83    crate::ir::Builtin::CubeDimX,
84    r"
85The dimension of the cube along the X axis.
86"
87);
88
89constant!(
90    CUBE_DIM_Y,
91    crate::ir::Builtin::CubeDimY,
92    r"
93The dimension of the cube along the Y axis.
94"
95);
96
97constant!(
98    CUBE_DIM_Z,
99    crate::ir::Builtin::CubeDimZ,
100    r"
101The dimension of the cube along the Z axis.
102"
103);
104
105constant!(
106    CUBE_POS,
107    crate::ir::Builtin::CubePos,
108    r"
109The cube position, without regards to axis.
110"
111);
112
113constant!(
114    CUBE_POS_X,
115    crate::ir::Builtin::CubePosX,
116    r"
117The cube position along the X axis.
118"
119);
120
121constant!(
122    CUBE_POS_Y,
123    crate::ir::Builtin::CubePosY,
124    r"
125The cube position along the Y axis.
126"
127);
128
129constant!(
130    CUBE_POS_Z,
131    crate::ir::Builtin::CubePosZ,
132    r"
133The cube position along the Z axis.
134"
135);
136constant!(
137    CUBE_COUNT,
138    crate::ir::Builtin::CubeCount,
139    r"
140The number of cubes launched.
141"
142);
143
144constant!(
145    CUBE_COUNT_X,
146    crate::ir::Builtin::CubeCountX,
147    r"
148The number of cubes launched along the X axis.
149"
150);
151
152constant!(
153    CUBE_COUNT_Y,
154    crate::ir::Builtin::CubeCountY,
155    r"
156The number of cubes launched along the Y axis.
157"
158);
159
160constant!(
161    CUBE_COUNT_Z,
162    crate::ir::Builtin::CubeCountZ,
163    r"
164The number of cubes launched along the Z axis.
165"
166);
167
168constant!(
169    ABSOLUTE_POS,
170    crate::ir::Builtin::AbsolutePos,
171    r"
172The position of the working unit in the whole cube kernel, without regards to cubes and axis.
173"
174);
175
176constant!(
177    ABSOLUTE_POS_X,
178    crate::ir::Builtin::AbsolutePosX,
179    r"
180The index of the working unit in the whole cube kernel along the X axis, without regards to cubes.
181"
182);
183
184constant!(
185    ABSOLUTE_POS_Y,
186    crate::ir::Builtin::AbsolutePosY,
187    r"
188The index of the working unit in the whole cube kernel along the Y axis, without regards to cubes.
189"
190);
191
192constant!(
193    ABSOLUTE_POS_Z,
194    crate::ir::Builtin::AbsolutePosZ,
195    r"
196The index of the working unit in the whole cube kernel along the Z axis, without regards to cubes.
197"
198);
cubecl_core/frontend/topology.rs

cubecl_core/frontend/
topology.rs
