cubecl_cpp/cuda/
plane.rs

1use cubecl_core::{
2    frontend::cast_value,
3    ir::{ContextExt, dialect::plane::*, types::scalar::BoolType},
4    prelude::*,
5};
6use pliron::{
7    builtin::types::{IntegerType, Signedness},
8    derive::op_interface_impl,
9    value::Value,
10};
11
12use crate::{
13    cuda::{cuda_op_with_out, ptx::InlinePtxOp},
14    ptx_block,
15    shared::{CompilationOptions, elect, lowering::LowerOp},
16    target::Cuda,
17};
18
19cuda_op_with_out!(BroadcastOp, |op, ctx| {
20    let val = op.input(ctx).name(ctx);
21    let lane = op.lane(ctx).0;
22    format!("__shfl_sync(__activemask(), {val}, {lane})")
23});
24
25cuda_op_with_out!(ShuffleOp, |op, ctx| {
26    let val = op.input(ctx).name(ctx);
27    let lane = op.lane(ctx).name(ctx);
28    format!("__shfl_sync(__activemask(), {val}, {lane})")
29});
30
31cuda_op_with_out!(ShuffleXorOp, |op, ctx| {
32    let val = op.input(ctx).name(ctx);
33    let mask = op.mask(ctx).name(ctx);
34    format!("__shfl_xor_sync(__activemask(), {val}, {mask})")
35});
36
37cuda_op_with_out!(ShuffleUpOp, |op, ctx| {
38    let val = op.input(ctx).name(ctx);
39    let delta = op.delta(ctx).name(ctx);
40    format!("__shfl_up_sync(__activemask(), {val}, {delta})")
41});
42
43cuda_op_with_out!(ShuffleDownOp, |op, ctx| {
44    let val = op.input(ctx).name(ctx);
45    let delta = op.delta(ctx).name(ctx);
46    format!("__shfl_down_sync(__activemask(), {val}, {delta})")
47});
48
49cuda_op_with_out!(AllOp, |op, ctx| {
50    let val = op.input(ctx).name(ctx);
51    format!("__all_sync(__activemask(), {val})")
52});
53
54cuda_op_with_out!(AnyOp, |op, ctx| {
55    let val = op.input(ctx).name(ctx);
56    format!("__any_sync(__activemask(), {val})")
57});
58
59cuda_op_with_out!(BallotOp, |op, ctx| {
60    let val = op.input(ctx).name(ctx);
61    format!("{{__ballot_sync(__activemask(), {val}), 0, 0, 0}}")
62});
63
64#[op_interface_impl]
65impl LowerOp<Cuda> for ElectOp {
66    fn lower(&self, scope: &cubecl_core::ir::Scope) -> Vec<Value> {
67        let ctx = scope.ctx_mut();
68        let opts = ctx.aux_ty::<CompilationOptions>();
69        let native_elect = opts.supports_features.elect_sync;
70        if native_elect {
71            let u32 = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle();
72            let ptx = ptx_block! {
73                ".reg .pred %%px;"
74                ".reg .b32 %mask;"
75                "activemask.b32 %mask;"
76                "elect.sync _|%%px, %mask;"
77                "selp.b32 $0, 1, 0, %%px;"
78            };
79            let op = InlinePtxOp::new_volatile(ctx, Some(u32), ptx, vec![]);
80            scope.register(&op);
81            let cast = cast_value(scope, op.result(ctx).unwrap(), BoolType::get(ctx).into());
82            vec![cast]
83        } else {
84            vec![elect::expand::<u32>(scope).read_value(scope)]
85        }
86    }
87}
cubecl_cpp/cuda/plane.rs

cubecl_cpp/cuda/
plane.rs
