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

use cubecl_core::ir::{attributes::EntrypointInterface, interfaces::TypedExt, prelude::*};
use itertools::Itertools;
use pliron::builtin::{
    ops::{FuncOp, ModuleOp},
    types::FunctionType,
};

use crate::{
    hip::{extension::compile_wmma_extensions, hip_op},
    shared::{
        CppValue,
        branch::block_to_cpp,
        define_array_polyfill,
        ty::{TypeExtCPP, TypedExtCPP},
        type_definitions,
    },
};

hip_op!(ModuleOp, |op, ctx| {
    let mut out = String::new();
    type_definitions(&mut out, "long long").unwrap();
    define_array_polyfill(&mut out).unwrap();
    out.push_str(&compile_wmma_extensions(ctx, op.get_operation()));
    out.push_str(&block_to_cpp(ctx, op.get_body(ctx, 0)));
    out
});

hip_op!(FuncOp, |op, ctx| {
    let func_name = op.get_symbol_name(ctx);
    let ty = op.get_type(ctx).deref(ctx);
    let func_ty = ty.downcast_ref::<FunctionType>().unwrap();
    let return_ty = func_ty.res_types()[0].to_cpp(ctx);
    let attributes = if let Some(abi) = op.get_entrypoint_abi(ctx) {
        format!(
            r#"extern "C" __global__ {return_ty} __launch_bounds__({})"#,
            abi.cube_dim.num_elems(),
        )
    } else {
        format!("__device__ {return_ty}")
    };

    let entry_block = op.get_entry_block(ctx);

    let block = entry_block.deref(ctx);
    let params = block.arguments();
    let params = params.map(|arg| gen_param(ctx, arg)).join(", ");

    let body = block_to_cpp(ctx, entry_block);

    format!("{attributes} {func_name}({params}) {{\n{body}\n}}\n")
});

fn gen_param(ctx: &Context, arg: Value) -> String {
    let mut segments = vec![];
    segments.push(arg.get_type(ctx).to_cpp(ctx));
    segments.push("const".into());
    if arg.is_ptr(ctx) || arg.is_uniform_ptr(ctx) {
        segments.push("__restrict__".into());
    }
    segments.push(arg.name(ctx).to_string());
    segments.join(" ")
}