cubecl_cpp/cuda/
signature.rs

1use cubecl_core::ir::{
2    attributes::{ATTR_TENSOR_MAP_BINDING, EntrypointInterface, FuncInterface},
3    interfaces::TypedExt,
4    settings::Dim3,
5};
6use itertools::Itertools;
7use pliron::{
8    builtin::{
9        op_interfaces::{SingleBlockRegionInterface, SymbolOpInterface},
10        ops::{FuncOp, ModuleOp},
11        type_interfaces::FunctionTypeInterface,
12        types::FunctionType,
13    },
14    context::Context,
15    dict_key,
16    r#type::Typed,
17    value::Value,
18};
19
20use crate::{
21    cuda::cuda_op,
22    shared::{
23        CppValue,
24        branch::block_to_cpp,
25        define_array_polyfill, define_tensormap_opaque,
26        ty::{TypeExtCPP, TypedExtCPP},
27        type_definitions,
28    },
29};
30
31dict_key!(ATTR_GRID_CONSTANT, "grid_constant");
32
33cuda_op!(ModuleOp, |op, ctx| {
34    let mut out = String::new();
35    type_definitions(&mut out, "long long").unwrap();
36    define_array_polyfill(&mut out).unwrap();
37
38    // This is fine to generate even on old cards, it's just an opaque block of memory
39    // The headers are dumb and only work in NVCC so we just need to define it ourselves
40    define_tensormap_opaque(&mut out).unwrap();
41
42    out.push_str(&block_to_cpp(ctx, op.get_body(ctx, 0)));
43    out
44});
45
46cuda_op!(FuncOp, |op, ctx| {
47    let func_name = op.get_symbol_name(ctx);
48    let ty = op.get_type(ctx).deref(ctx);
49    let func_ty = ty.downcast_ref::<FunctionType>().unwrap();
50    let return_ty = func_ty.res_types()[0].to_cpp(ctx);
51    let attributes = if let Some(abi) = op.get_entrypoint_abi(ctx) {
52        let cluster_dim = match abi.cluster_dim {
53            Some(Dim3 { x, y, z }) => format!("__cluster_dims__({x}, {y}, {z})"),
54            None => "".into(),
55        };
56        format!(
57            r#"extern "C" __global__ {return_ty} __launch_bounds__({}) {cluster_dim}"#,
58            abi.cube_dim.num_elems(),
59        )
60    } else {
61        format!("__device__ {return_ty}")
62    };
63
64    let entry_block = op.get_entry_block(ctx);
65
66    let block = entry_block.deref(ctx);
67    let params = block.arguments().enumerate();
68    let params = params.map(|(i, arg)| gen_param(ctx, op, i, arg)).join(", ");
69
70    let body = block_to_cpp(ctx, entry_block);
71
72    format!("{attributes} {func_name}({params}) {{\n{body}\n}}\n")
73});
74
75fn gen_param(ctx: &Context, func: &FuncOp, i: usize, arg: Value) -> String {
76    let mut segments = vec![];
77
78    if func.has_arg_attr(ctx, i, &ATTR_GRID_CONSTANT) {
79        segments.push("__grid_constant__".into());
80    }
81    if func.has_arg_attr(ctx, i, &ATTR_TENSOR_MAP_BINDING) {
82        segments.extend(["__grid_constant__".into()]);
83    }
84    segments.push(arg.get_type(ctx).to_cpp(ctx));
85    segments.push("const".into());
86    if arg.is_ptr(ctx) || arg.is_uniform_ptr(ctx) {
87        segments.push("__restrict__".into());
88    }
89    segments.push(arg.name(ctx).to_string());
90    segments.join(" ")
91}
cubecl_cpp/cuda/signature.rs

cubecl_cpp/cuda/
signature.rs
