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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use cubecl_core::{
    self as cubecl,
    ir::{dialect::vector::*, interfaces::TypedExt, prelude::*},
    prelude::*,
};
use itertools::Itertools;

use crate::{
    shared::{
        binary::lower_binop, scoped_block, shared_op_with_out, ty::TypeExtCPP, unary::lower_unop,
    },
    target::{CtxTarget, Target},
};

shared_op_with_out!(CompositeConstructOp, |op, ctx| {
    let values = op.values(ctx).iter().map(|it| it.name(ctx)).join(", ");
    let ty = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
    format!("{ty}{{{values}}}")
});

shared_op_with_out!(VectorBroadcastOp, |op, ctx| {
    let vec = op.get_result(ctx).vector_size(ctx);
    let values = (0..vec).map(|_| op.input(ctx).name(ctx)).join(", ");
    let ty = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
    format!("{ty}{{{values}}}")
});

shared_op_with_out!(CompositeInsertOp, |op, ctx| {
    assert!(op.composite(ctx).is_vector(ctx));
    let vector_size = op.composite(ctx).vector_size(ctx);
    let vector = op.composite(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    let index = op.index(ctx).0;
    let new_values = (0..vector_size)
        .map(|i| {
            if i == index {
                value.to_string()
            } else {
                format!("{vector}.i_{i}")
            }
        })
        .join(", ");
    format!("{{{new_values}}}")
});

shared_op_with_out!(CompositeExtractOp, |op, ctx| {
    assert!(op.composite(ctx).is_vector(ctx));
    let vector = op.composite(ctx).name(ctx);
    let index = op.index(ctx).0;
    format!("{vector}.i_{index}")
});

shared_op_with_out!(VectorInsertDynamicOp, |op, ctx| {
    let vector = op.vector(ctx).name(ctx);
    let value = op.value(ctx).name(ctx);
    let index = op.index(ctx).name(ctx);
    let elem_ty = op.value(ctx).get_type(ctx).to_cpp(ctx);
    let vector_ty = op.vector(ctx).get_type(ctx).to_cpp(ctx);
    scoped_block!(
        format!("{vector_ty} tmp = {vector};")
        format!("*(reinterpret_cast<{elem_ty}*>(&tmp) + {index}) = {value};")
        "return tmp;"
    )
});

shared_op_with_out!(VectorExtractDynamicOp, |op, ctx| {
    let vector = op.vector(ctx).name(ctx);
    let index = op.index(ctx).name(ctx);
    let elem_ty = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
    format!("reinterpret_cast<const {elem_ty}*>(&{vector})[{index}]")
});

#[cube]
fn magnitude<T: Float, N: Size>(input: Vector<T, N>) -> T {
    let mut out = T::from_int(0);
    #[unroll]
    for i in 0..input.vector_size() {
        let val = input.extract(i);
        out += val * val;
    }
    out.sqrt()
}

#[cube]
fn normalize<T: Float, N: Size>(input: Vector<T, N>) -> Vector<T, N> {
    let mut norm = T::from_int(0);
    #[unroll]
    for i in 0..input.vector_size() {
        let val = input.extract(i);
        norm += val * val;
    }
    input * Vector::cast_from(norm.inverse_sqrt())
}

#[cube]
fn sum<T: Float, N: Size>(input: Vector<T, N>) -> T {
    let mut sum = T::from_int(0);
    #[unroll]
    for i in 0..input.vector_size() {
        sum += input.extract(i);
    }
    sum
}

#[cube]
fn dot<T: Float, N: Size>(lhs: Vector<T, N>, rhs: Vector<T, N>) -> T {
    sum(lhs * rhs)
}

lower_unop!(MagnitudeOp, magnitude, |_, ctx| {
    ctx.target() != Target::Metal
});

lower_unop!(NormalizeOp, normalize, |_, ctx| {
    ctx.target() != Target::Metal
});

lower_unop!(ISumOp, sum);
lower_unop!(FSumOp, sum);
lower_binop!(SDotOp, dot);
lower_binop!(UDotOp, dot);
lower_binop!(FDotOp, dot, |_, ctx| ctx.target() != Target::Metal);