cubecl-cpp 0.11.0-pre.4

CPP transpiler for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

use cubecl_core::{
    frontend::cast_value,
    ir::{ContextExt, dialect::plane::*, types::scalar::BoolType},
    prelude::*,
};
use pliron::{
    builtin::types::{IntegerType, Signedness},
    derive::op_interface_impl,
    value::Value,
};

use crate::{
    cuda::{cuda_op_with_out, ptx::InlinePtxOp},
    ptx_block,
    shared::{CompilationOptions, elect, lowering::LowerOp},
    target::Cuda,
};

cuda_op_with_out!(BroadcastOp, |op, ctx| {
    let val = op.input(ctx).name(ctx);
    let lane = op.lane(ctx).0;
    format!("__shfl_sync(__activemask(), {val}, {lane})")
});

cuda_op_with_out!(ShuffleOp, |op, ctx| {
    let val = op.input(ctx).name(ctx);
    let lane = op.lane(ctx).name(ctx);
    format!("__shfl_sync(__activemask(), {val}, {lane})")
});

cuda_op_with_out!(ShuffleXorOp, |op, ctx| {
    let val = op.input(ctx).name(ctx);
    let mask = op.mask(ctx).name(ctx);
    format!("__shfl_xor_sync(__activemask(), {val}, {mask})")
});

cuda_op_with_out!(ShuffleUpOp, |op, ctx| {
    let val = op.input(ctx).name(ctx);
    let delta = op.delta(ctx).name(ctx);
    format!("__shfl_up_sync(__activemask(), {val}, {delta})")
});

cuda_op_with_out!(ShuffleDownOp, |op, ctx| {
    let val = op.input(ctx).name(ctx);
    let delta = op.delta(ctx).name(ctx);
    format!("__shfl_down_sync(__activemask(), {val}, {delta})")
});

cuda_op_with_out!(AllOp, |op, ctx| {
    let val = op.input(ctx).name(ctx);
    format!("__all_sync(__activemask(), {val})")
});

cuda_op_with_out!(AnyOp, |op, ctx| {
    let val = op.input(ctx).name(ctx);
    format!("__any_sync(__activemask(), {val})")
});

cuda_op_with_out!(BallotOp, |op, ctx| {
    let val = op.input(ctx).name(ctx);
    format!("{{__ballot_sync(__activemask(), {val}), 0, 0, 0}}")
});

#[op_interface_impl]
impl LowerOp<Cuda> for ElectOp {
    fn lower(&self, scope: &cubecl_core::ir::Scope) -> Vec<Value> {
        let ctx = scope.ctx_mut();
        let opts = ctx.aux_ty::<CompilationOptions>();
        let native_elect = opts.supports_features.elect_sync;
        if native_elect {
            let u32 = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle();
            let ptx = ptx_block! {
                ".reg .pred %%px;"
                ".reg .b32 %mask;"
                "activemask.b32 %mask;"
                "elect.sync _|%%px, %mask;"
                "selp.b32 $0, 1, 0, %%px;"
            };
            let op = InlinePtxOp::new_volatile(ctx, Some(u32), ptx, vec![]);
            scope.register(&op);
            let cast = cast_value(scope, op.result(ctx).unwrap(), BoolType::get(ctx).into());
            vec![cast]
        } else {
            vec![elect::expand::<u32>(scope).read_value(scope)]
        }
    }
}