cubecl_cpp/shared/
plane.rs

1use cubecl_core::{
2    self as cubecl,
3    ir::{dialect::plane, prelude::*},
4    prelude::*,
5};
6
7use cubecl_core::prelude::polyfills::plane::{
8    OpAdd, OpMax, OpMin, OpMul, plane_reduce, plane_reduce_exclusive, plane_reduce_inclusive,
9};
10
11use crate::{
12    cuda::packed_ops::packable,
13    shared::{lowering::LowerOp, shared_op_with_out, unroll::unrolling},
14    target::{CtxTarget, Target},
15};
16
17define_scalar!(T);
18define_size!(S);
19
20macro_rules! lower_unop {
21    ($ty: ty, $reduce: ident, $op: ty $(,$args: expr)*) => {
22        #[op_interface_impl]
23        impl LowerOp for $ty {
24            fn should_lower(&self, ctx: &Context) -> bool {
25                ctx.target() != Target::Metal
26            }
27            fn lower(&self, scope: &Scope) -> Vec<Value> {
28                let input = self.input(scope.ctx());
29                scope.register_value_type::<T, S>(input);
30                vec![$reduce::expand::<T, S, $op>(scope, input.into(), $($args),*).read_value(scope)]
31            }
32        }
33    };
34}
35
36lower_unop!(plane::ISumOp, plane_reduce, OpAdd);
37lower_unop!(plane::FSumOp, plane_reduce, OpAdd);
38lower_unop!(plane::IProdOp, plane_reduce, OpMul);
39lower_unop!(plane::FProdOp, plane_reduce, OpMul);
40lower_unop!(plane::SMinOp, plane_reduce, OpMin);
41lower_unop!(plane::UMinOp, plane_reduce, OpMin);
42lower_unop!(plane::FMinOp, plane_reduce, OpMin);
43lower_unop!(plane::SMaxOp, plane_reduce, OpMax);
44lower_unop!(plane::UMaxOp, plane_reduce, OpMax);
45lower_unop!(plane::FMaxOp, plane_reduce, OpMax);
46
47lower_unop!(plane::InclusiveISumOp, plane_reduce_inclusive, OpAdd);
48lower_unop!(plane::InclusiveFSumOp, plane_reduce_inclusive, OpAdd);
49lower_unop!(plane::InclusiveIProdOp, plane_reduce_inclusive, OpMul);
50lower_unop!(plane::InclusiveFProdOp, plane_reduce_inclusive, OpMul);
51
52lower_unop!(plane::ExclusiveISumOp, plane_reduce_exclusive, OpAdd, 0);
53lower_unop!(plane::ExclusiveFSumOp, plane_reduce_exclusive, OpAdd, 0);
54lower_unop!(plane::ExclusiveIProdOp, plane_reduce_exclusive, OpMul, 1);
55lower_unop!(plane::ExclusiveFProdOp, plane_reduce_exclusive, OpMul, 1);
56
57unrolling!(plane::BroadcastOp);
58packable!(plane::BroadcastOp);
59
60unrolling!(plane::ShuffleOp);
61packable!(plane::ShuffleOp);
62
63unrolling!(plane::ShuffleXorOp);
64packable!(plane::ShuffleXorOp);
65
66unrolling!(plane::ShuffleUpOp);
67packable!(plane::ShuffleUpOp);
68
69unrolling!(plane::ShuffleDownOp);
70packable!(plane::ShuffleDownOp);
71
72unrolling!(plane::AllOp);
73unrolling!(plane::AnyOp);
74
75#[cube_op(name = "cpp.activemask")]
76#[result_ty(argument)]
77struct ActiveMask {}
78shared_op_with_out!(ActiveMask, |_, _| "__activemask()".into());
79
80#[cube]
81fn activemask<T: Int>() -> T {
82    intrinsic!(|scope| {
83        let mask = ActiveMask::new(scope.ctx_mut(), T::__expand_as_type(scope));
84        scope.register_with_result(&mask).into()
85    })
86}
87
88// Lowest active lane, requires a generic because HIP uses u64 and CUDA uses u32 for `__activemask()`
89#[cube]
90pub fn elect<T: Int>() -> bool {
91    u32::cast_from(activemask::<T>().trailing_zeros()) == UNIT_POS_PLANE
92}
cubecl_cpp/shared/plane.rs

cubecl_cpp/shared/
plane.rs
