cubecl_core/frontend/
plane.rs

1use cubecl_ir::ExpandValue;
2use half::{bf16, f16};
3
4use super::{CubePrimitive, Vector};
5use crate::prelude::*;
6use crate::{
7    ir::{Scope, attributes::IndexAttr, dialect::plane::*},
8    unexpanded,
9};
10
11pub trait PlaneNumeric {
12    fn __expand_native_sum(scope: &Scope, value: ExpandValue) -> ExpandValue;
13    fn __expand_native_inclusive_sum(scope: &Scope, value: ExpandValue) -> ExpandValue;
14    fn __expand_native_exclusive_sum(scope: &Scope, value: ExpandValue) -> ExpandValue;
15
16    fn __expand_native_prod(scope: &Scope, value: ExpandValue) -> ExpandValue;
17    fn __expand_native_inclusive_prod(scope: &Scope, value: ExpandValue) -> ExpandValue;
18    fn __expand_native_exclusive_prod(scope: &Scope, value: ExpandValue) -> ExpandValue;
19
20    fn __expand_native_plane_min(scope: &Scope, value: ExpandValue) -> ExpandValue;
21    fn __expand_native_plane_max(scope: &Scope, value: ExpandValue) -> ExpandValue;
22}
23
24macro_rules! plane_numeric {
25    ($($ty: ty),*; $sum: ty, $inc_sum: ty, $exc_sum: ty, $prod: ty, $inc_prod: ty, $exc_prod: ty, $min: ty, $max: ty) => {
26        $(impl PlaneNumeric for $ty {
27            fn __expand_native_sum(scope: &Scope, value: ExpandValue) -> ExpandValue {
28                unary_expand(scope, value, <$sum>::new)
29            }
30            fn __expand_native_inclusive_sum(scope: &Scope, value: ExpandValue) -> ExpandValue {
31                unary_expand(scope, value, <$inc_sum>::new)
32            }
33            fn __expand_native_exclusive_sum(scope: &Scope, value: ExpandValue) -> ExpandValue {
34                unary_expand(scope, value, <$exc_sum>::new)
35            }
36
37            fn __expand_native_prod(scope: &Scope, value: ExpandValue) -> ExpandValue {
38                unary_expand(scope, value, <$prod>::new)
39            }
40            fn __expand_native_inclusive_prod(scope: &Scope, value: ExpandValue) -> ExpandValue {
41                unary_expand(scope, value, <$inc_prod>::new)
42            }
43            fn __expand_native_exclusive_prod(scope: &Scope, value: ExpandValue) -> ExpandValue {
44                unary_expand(scope, value, <$exc_prod>::new)
45            }
46
47            fn __expand_native_plane_min(scope: &Scope, value: ExpandValue) -> ExpandValue {
48                unary_expand(scope, value, <$min>::new)
49            }
50            fn __expand_native_plane_max(scope: &Scope, value: ExpandValue) -> ExpandValue {
51                unary_expand(scope, value, <$max>::new)
52            }
53        })*
54    };
55}
56
57plane_numeric!(i8, i16, i32, i64, isize; ISumOp, InclusiveISumOp, ExclusiveISumOp, IProdOp, InclusiveIProdOp, ExclusiveIProdOp, SMinOp, SMaxOp);
58plane_numeric!(u8, u16, u32, u64, usize; ISumOp, InclusiveISumOp, ExclusiveISumOp, IProdOp, InclusiveIProdOp, ExclusiveIProdOp, UMinOp, UMaxOp);
59plane_numeric!(f16, bf16, f32, flex32, tf32, f64; FSumOp, InclusiveFSumOp, ExclusiveFSumOp, FProdOp, InclusiveFProdOp, ExclusiveFProdOp, FMinOp, FMaxOp);
60
61/// Returns true if the cube unit has the lowest `plane_unit_id` among active unit in the plane
62pub fn plane_elect() -> bool {
63    unexpanded!()
64}
65
66/// Module containing the expand function for [`plane_elect()`].
67pub mod plane_elect {
68    use super::*;
69
70    /// Expand method of [`plane_elect()`].
71    pub fn expand(scope: &Scope) -> NativeExpand<bool> {
72        let op = ElectOp::new(scope.ctx_mut());
73        scope.register_with_result(&op).into()
74    }
75}
76
77/// Broadcasts the value from the specified plane unit at the given index
78/// to all active units within that plane. Requires a constant index. For non-constant indices,
79/// use [`plane_shuffle()`].
80#[allow(unused_variables)]
81pub fn plane_broadcast<E: CubePrimitive>(value: E, index: u32) -> E {
82    unexpanded!()
83}
84
85/// Module containing the expand function for [`plane_broadcast()`].
86pub mod plane_broadcast {
87    use super::*;
88
89    /// Expand method of [`plane_broadcast()`].
90    pub fn expand<E: CubePrimitive>(
91        scope: &Scope,
92        value: NativeExpand<E>,
93        id: u32,
94    ) -> NativeExpand<E> {
95        let value = value.read_value(scope);
96        let op = BroadcastOp::new(scope.ctx_mut(), value, IndexAttr::new(id as usize));
97        scope.register_with_result(&op).into()
98    }
99}
100
101/// Perform an arbitrary lane shuffle operation across the plane.
102/// Each unit reads the value from the specified source lane.
103///
104/// # Example
105/// `plane_shuffle(value, 0)` - all lanes read from lane 0 (same as broadcast)
106/// `plane_shuffle(value, lane_id ^ 1)` - butterfly pattern (same as `shuffle_xor`)
107#[allow(unused_variables)]
108pub fn plane_shuffle<E: CubePrimitive>(value: E, src_lane: u32) -> E {
109    unexpanded!()
110}
111
112/// Module containing the expand function for [`plane_shuffle()`].
113pub mod plane_shuffle {
114    use super::*;
115
116    /// Expand method of [`plane_shuffle()`].
117    pub fn expand<E: CubePrimitive>(
118        scope: &Scope,
119        value: NativeExpand<E>,
120        src_lane: NativeExpand<u32>,
121    ) -> NativeExpand<E> {
122        let value = value.read_value(scope);
123        let src_lane = src_lane.read_value(scope);
124        let op = ShuffleOp::new(scope.ctx_mut(), value, src_lane);
125        scope.register_with_result(&op).into()
126    }
127}
128
129/// Perform a shuffle XOR operation across the plane.
130/// Each unit exchanges its value with another unit at an index determined by XOR with the mask.
131/// This is useful for butterfly reduction patterns.
132///
133/// # Example
134/// For a 32-lane warp with mask=1:
135/// - Lane 0 gets value from lane 1, lane 1 gets value from lane 0
136/// - Lane 2 gets value from lane 3, lane 3 gets value from lane 2
137/// - etc.
138#[allow(unused_variables)]
139pub fn plane_shuffle_xor<E: CubePrimitive>(value: E, mask: u32) -> E {
140    unexpanded!()
141}
142
143/// Module containing the expand function for [`plane_shuffle_xor()`].
144pub mod plane_shuffle_xor {
145    use super::*;
146
147    /// Expand method of [`plane_shuffle_xor()`].
148    pub fn expand<E: CubePrimitive>(
149        scope: &Scope,
150        value: NativeExpand<E>,
151        mask: NativeExpand<u32>,
152    ) -> NativeExpand<E> {
153        let value = value.read_value(scope);
154        let mask = mask.read_value(scope);
155        let op = ShuffleXorOp::new(scope.ctx_mut(), value, mask);
156        scope.register_with_result(&op).into()
157    }
158}
159
160/// Perform a shuffle up operation across the plane.
161/// Each unit reads the value from a unit with a lower lane ID (`current_id` - delta).
162/// A unit with `lane_id` < delta reads an indeterminate value that differs by backend, so
163/// mask those units at the call site.
164///
165/// # Example
166/// For delta=1, writing `?` for the indeterminate read: `[a, b, c, d] -> [?, a, b, c]`
167#[allow(unused_variables)]
168pub fn plane_shuffle_up<E: CubePrimitive>(value: E, delta: u32) -> E {
169    unexpanded!()
170}
171
172/// Module containing the expand function for [`plane_shuffle_up()`].
173pub mod plane_shuffle_up {
174    use super::*;
175
176    /// Expand method of [`plane_shuffle_up()`].
177    pub fn expand<E: CubePrimitive>(
178        scope: &Scope,
179        value: NativeExpand<E>,
180        delta: NativeExpand<u32>,
181    ) -> NativeExpand<E> {
182        let value = value.read_value(scope);
183        let delta = delta.read_value(scope);
184        let op = ShuffleUpOp::new(scope.ctx_mut(), value, delta);
185        scope.register_with_result(&op).into()
186    }
187}
188
189/// Perform a shuffle down operation across the plane.
190/// Each unit reads the value from a unit with a higher lane ID (`current_id` + delta).
191/// A unit whose `lane_id` + delta >= `plane_dim` reads an indeterminate value that differs
192/// by backend, so mask those units at the call site.
193///
194/// # Example
195/// For delta=1, writing `?` for the indeterminate read: `[a, b, c, d] -> [b, c, d, ?]`
196#[allow(unused_variables)]
197pub fn plane_shuffle_down<E: CubePrimitive>(value: E, delta: u32) -> E {
198    unexpanded!()
199}
200
201/// Module containing the expand function for [`plane_shuffle_down()`].
202pub mod plane_shuffle_down {
203    use super::*;
204
205    /// Expand method of [`plane_shuffle_down()`].
206    pub fn expand<E: CubePrimitive>(
207        scope: &Scope,
208        value: NativeExpand<E>,
209        delta: NativeExpand<u32>,
210    ) -> NativeExpand<E> {
211        let value = value.read_value(scope);
212        let delta = delta.read_value(scope);
213        let op = ShuffleDownOp::new(scope.ctx_mut(), value, delta);
214        scope.register_with_result(&op).into()
215    }
216}
217
218/// Perform a reduce sum operation across all units in a plane.
219#[allow(unused_variables)]
220pub fn plane_sum<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
221    unexpanded!()
222}
223
224/// Module containing the expand function for [`plane_sum()`].
225pub mod plane_sum {
226    use super::*;
227
228    /// Expand method of [`plane_sum()`].
229    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
230        scope: &Scope,
231        elem: NativeExpand<E>,
232    ) -> NativeExpand<E> {
233        E::Scalar::__expand_native_sum(scope, elem.into()).into()
234    }
235}
236
237/// Perform an inclusive sum operation across all units in a plane.
238/// This sums all values to the "left" of the unit, including this unit's value.
239/// Also known as "prefix sum" or "inclusive scan".
240///
241/// # Example
242/// `inclusive_sum([1, 2, 3, 4, 5]) == [1, 3, 6, 10, 15]`
243#[allow(unused_variables)]
244pub fn plane_inclusive_sum<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
245    unexpanded!()
246}
247
248/// Module containing the expand function for [`plane_inclusive_sum()`].
249pub mod plane_inclusive_sum {
250    use super::*;
251
252    /// Expand method of [`plane_inclusive_sum()`].
253    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
254        scope: &Scope,
255        elem: NativeExpand<E>,
256    ) -> NativeExpand<E> {
257        E::Scalar::__expand_native_inclusive_sum(scope, elem.into()).into()
258    }
259}
260
261/// Perform an exclusive sum operation across all units in a plane.
262/// This sums all values to the "left" of the unit, excluding this unit's value. The 0th unit will
263/// be set to `E::zero()`.
264/// Also known as "exclusive prefix sum" or "exclusive scan".
265///
266/// # Example
267/// `exclusive_sum([1, 2, 3, 4, 5]) == [0, 1, 3, 6, 10]`
268#[allow(unused_variables)]
269pub fn plane_exclusive_sum<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
270    unexpanded!()
271}
272
273/// Module containing the expand function for [`plane_exclusive_sum()`].
274pub mod plane_exclusive_sum {
275    use super::*;
276
277    /// Expand method of [`plane_exclusive_sum()`].
278    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
279        scope: &Scope,
280        elem: NativeExpand<E>,
281    ) -> NativeExpand<E> {
282        E::Scalar::__expand_native_exclusive_sum(scope, elem.into()).into()
283    }
284}
285
286/// Perform a reduce prod operation across all units in a plane.
287pub fn plane_prod<E: CubePrimitive<Scalar: PlaneNumeric>>(_elem: E) -> E {
288    unexpanded!()
289}
290
291/// Module containing the expand function for [`plane_prod()`].
292pub mod plane_prod {
293    use super::*;
294
295    /// Expand method of [`plane_prod()`].
296    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
297        scope: &Scope,
298        elem: NativeExpand<E>,
299    ) -> NativeExpand<E> {
300        E::Scalar::__expand_native_prod(scope, elem.into()).into()
301    }
302}
303
304/// Perform an inclusive product operation across all units in a plane.
305/// This multiplies all values to the "left" of the unit, including this unit's value.
306/// Also known as "prefix product" or "inclusive scan".
307///
308/// # Example
309/// `exclusive_prod([1, 2, 3, 4, 5]) == [1, 2, 6, 24, 120]`
310#[allow(unused_variables)]
311pub fn plane_inclusive_prod<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
312    unexpanded!()
313}
314
315/// Module containing the expand function for [`plane_inclusive_prod()`].
316pub mod plane_inclusive_prod {
317    use super::*;
318
319    /// Expand method of [`plane_inclusive_prod()`].
320    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
321        scope: &Scope,
322        elem: NativeExpand<E>,
323    ) -> NativeExpand<E> {
324        E::Scalar::__expand_native_inclusive_prod(scope, elem.into()).into()
325    }
326}
327
328/// Perform an exclusive product operation across all units in a plane.
329/// This multiplies all values to the "left" of the unit, excluding this unit's value. The 0th unit
330/// will be set to `E::one()`.
331/// Also known as "exclusive prefix product" or "exclusive scan".
332///
333/// # Example
334/// `exclusive_prod([1, 2, 3, 4, 5]) == [1, 1, 2, 6, 24]`
335#[allow(unused_variables)]
336pub fn plane_exclusive_prod<E: CubePrimitive<Scalar: PlaneNumeric>>(value: E) -> E {
337    unexpanded!()
338}
339
340/// Module containing the expand function for [`plane_exclusive_prod()`].
341pub mod plane_exclusive_prod {
342    use super::*;
343
344    /// Expand method of [`plane_exclusive_prod()`].
345    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
346        scope: &Scope,
347        elem: NativeExpand<E>,
348    ) -> NativeExpand<E> {
349        E::Scalar::__expand_native_exclusive_prod(scope, elem.into()).into()
350    }
351}
352
353/// Perform a reduce max operation across all units in a plane.
354pub fn plane_max<E: CubePrimitive<Scalar: PlaneNumeric>>(_elem: E) -> E {
355    unexpanded!()
356}
357
358/// Module containing the expand function for [`plane_max()`].
359pub mod plane_max {
360    use super::*;
361
362    /// Expand method of [`plane_max()`].
363    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
364        scope: &Scope,
365        elem: NativeExpand<E>,
366    ) -> NativeExpand<E> {
367        E::Scalar::__expand_native_plane_max(scope, elem.into()).into()
368    }
369}
370
371/// Perform a reduce min operation across all units in a plane.
372pub fn plane_min<E: CubePrimitive<Scalar: PlaneNumeric>>(_elem: E) -> E {
373    unexpanded!()
374}
375
376/// Module containing the expand function for [`plane_min()`].
377pub mod plane_min {
378    use super::*;
379
380    /// Expand method of [`plane_min()`].
381    pub fn expand<E: CubePrimitive<Scalar: PlaneNumeric>>(
382        scope: &Scope,
383        elem: NativeExpand<E>,
384    ) -> NativeExpand<E> {
385        E::Scalar::__expand_native_plane_min(scope, elem.into()).into()
386    }
387}
388
389/// Perform a reduce all operation across all units in a plane.
390pub fn plane_all(_elem: bool) -> bool {
391    unexpanded!()
392}
393
394/// Module containing the expand function for [`plane_all()`].
395pub mod plane_all {
396    use super::*;
397
398    /// Expand method of [`plane_all()`].
399    pub fn expand(scope: &Scope, elem: NativeExpand<bool>) -> NativeExpand<bool> {
400        let value = elem.read_value(scope);
401        let op = AllOp::new(scope.ctx_mut(), value);
402        scope.register_with_result(&op).into()
403    }
404}
405
406/// Perform a reduce any operation across all units in a plane.
407pub fn plane_any(_elem: bool) -> bool {
408    unexpanded!()
409}
410
411/// Module containing the expand function for [`plane_any()`].
412pub mod plane_any {
413    use super::*;
414
415    /// Expand method of [`plane_any()`].
416    pub fn expand(scope: &Scope, elem: NativeExpand<bool>) -> NativeExpand<bool> {
417        let value = elem.read_value(scope);
418        let op = AnyOp::new(scope.ctx_mut(), value);
419        scope.register_with_result(&op).into()
420    }
421}
422
423/// Perform a ballot operation across all units in a plane.
424/// Returns a set of 32-bit bitfields as a [`Vector`], with each element containing the value from 32
425/// invocations.
426/// Note that vector size will always be set to 4 even for `PLANE_DIM <= 64`, because we can't
427/// retrieve the actual plane size at expand time. Use the runtime`PLANE_DIM` to index appropriately.
428pub fn plane_ballot(_elem: bool) -> Vector<u32, Const<4>> {
429    unexpanded!()
430}
431
432/// Module containing the expand function for [`plane_ballot()`].
433pub mod plane_ballot {
434    use super::*;
435
436    /// Expand method of [`plane_ballot()`].
437    pub fn expand(scope: &Scope, elem: NativeExpand<bool>) -> NativeExpand<Vector<u32, Const<4>>> {
438        let value = elem.read_value(scope);
439        let op = BallotOp::new(scope.ctx_mut(), value);
440        scope.register_with_result(&op).into()
441    }
442}
cubecl_core/frontend/plane.rs

cubecl_core/frontend/
plane.rs
