cubecl_cpp/hip/
plane.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3use cubecl_core::ir::{cube_op, dialect::plane::*};
4use pliron::{
5    builtin::types::{IntegerType, Signedness},
6    derive::op_interface_impl,
7    value::Value,
8};
9
10use crate::{
11    hip::hip_op_with_out,
12    shared::{elect, lowering::LowerOp},
13    target::Hip,
14};
15
16hip_op_with_out!(BroadcastOp, |op, ctx| {
17    let val = op.input(ctx).name(ctx);
18    let lane = op.lane(ctx).0;
19    format!("__shfl({val}, {lane});")
20});
21
22hip_op_with_out!(ShuffleOp, |op, ctx| {
23    let val = op.input(ctx).name(ctx);
24    let lane = op.lane(ctx).name(ctx);
25    format!("__shfl({val}, {lane});")
26});
27
28hip_op_with_out!(ShuffleXorOp, |op, ctx| {
29    let val = op.input(ctx).name(ctx);
30    let mask = op.mask(ctx).name(ctx);
31    format!("__shfl_xor({val}, {mask});")
32});
33
34hip_op_with_out!(ShuffleUpOp, |op, ctx| {
35    let val = op.input(ctx).name(ctx);
36    let delta = op.delta(ctx).name(ctx);
37    format!("__shfl_up({val}, {delta});")
38});
39
40hip_op_with_out!(ShuffleDownOp, |op, ctx| {
41    let val = op.input(ctx).name(ctx);
42    let delta = op.delta(ctx).name(ctx);
43    format!("__shfl_down({val}, {delta});")
44});
45
46hip_op_with_out!(AllOp, |op, ctx| {
47    let val = op.input(ctx).name(ctx);
48    format!("static_cast<bool>(__all({val}));")
49});
50
51hip_op_with_out!(AnyOp, |op, ctx| {
52    let val = op.input(ctx).name(ctx);
53    format!("static_cast<bool>(__any({val}));")
54});
55
56#[cube_op(name = "hip.ballot")]
57#[result_ty(fixed = IntegerType::get(ctx, 64, Signedness::Unsigned).to_handle())]
58pub struct HipBallotOp {
59    input: Value,
60}
61
62hip_op_with_out!(HipBallotOp, |op, ctx| {
63    let val = op.input(ctx).name(ctx);
64    format!("__ballot({val});")
65});
66
67#[cube]
68fn hip_ballot(value: bool) -> u64 {
69    intrinsic!(|scope| {
70        let value = value.read_value(scope);
71        let ballot = HipBallotOp::new(scope.ctx_mut(), value);
72        scope.register_with_result(&ballot).into()
73    })
74}
75
76/// Unlike CUDA's 32 bit `__ballot_sync`, HIP's `__ballot` returns a 64 bit mask so it can cover
77/// wave64. It has to be split across two of the result's 32 bit lanes instead of narrowed into one.
78#[cube]
79fn ballot(value: bool) -> Vector<u32, Const<4>> {
80    let mut out = Vector::<u64, Const<2>>::zero();
81    out.insert(0usize, hip_ballot(value));
82    Vector::reinterpret(out)
83}
84
85#[op_interface_impl]
86impl LowerOp<Hip> for BallotOp {
87    fn lower(&self, scope: &Scope) -> Vec<Value> {
88        let value = self.input(scope.ctx()).into();
89        vec![ballot::expand(scope, value).read_value(scope)]
90    }
91}
92
93#[op_interface_impl]
94impl LowerOp<Hip> for ElectOp {
95    fn lower(&self, scope: &Scope) -> Vec<Value> {
96        vec![elect::expand::<u64>(scope).read_value(scope)]
97    }
98}
cubecl_cpp/hip/plane.rs

cubecl_cpp/hip/
plane.rs
