cubecl_cpp/hip/
signature.rs

1use cubecl_core::ir::{attributes::EntrypointInterface, interfaces::TypedExt, prelude::*};
2use itertools::Itertools;
3use pliron::builtin::{
4    ops::{FuncOp, ModuleOp},
5    types::FunctionType,
6};
7
8use crate::{
9    hip::{extension::compile_wmma_extensions, hip_op},
10    shared::{
11        CppValue,
12        branch::block_to_cpp,
13        define_array_polyfill,
14        ty::{TypeExtCPP, TypedExtCPP},
15        type_definitions,
16    },
17};
18
19hip_op!(ModuleOp, |op, ctx| {
20    let mut out = String::new();
21    type_definitions(&mut out, "long long").unwrap();
22    define_array_polyfill(&mut out).unwrap();
23    out.push_str(&compile_wmma_extensions(ctx, op.get_operation()));
24    out.push_str(&block_to_cpp(ctx, op.get_body(ctx, 0)));
25    out
26});
27
28hip_op!(FuncOp, |op, ctx| {
29    let func_name = op.get_symbol_name(ctx);
30    let ty = op.get_type(ctx).deref(ctx);
31    let func_ty = ty.downcast_ref::<FunctionType>().unwrap();
32    let return_ty = func_ty.res_types()[0].to_cpp(ctx);
33    let attributes = if let Some(abi) = op.get_entrypoint_abi(ctx) {
34        format!(
35            r#"extern "C" __global__ {return_ty} __launch_bounds__({})"#,
36            abi.cube_dim.num_elems(),
37        )
38    } else {
39        format!("__device__ {return_ty}")
40    };
41
42    let entry_block = op.get_entry_block(ctx);
43
44    let block = entry_block.deref(ctx);
45    let params = block.arguments();
46    let params = params.map(|arg| gen_param(ctx, arg)).join(", ");
47
48    let body = block_to_cpp(ctx, entry_block);
49
50    format!("{attributes} {func_name}({params}) {{\n{body}\n}}\n")
51});
52
53fn gen_param(ctx: &Context, arg: Value) -> String {
54    let mut segments = vec![];
55    segments.push(arg.get_type(ctx).to_cpp(ctx));
56    segments.push("const".into());
57    if arg.is_ptr(ctx) || arg.is_uniform_ptr(ctx) {
58        segments.push("__restrict__".into());
59    }
60    segments.push(arg.name(ctx).to_string());
61    segments.join(" ")
62}
cubecl_cpp/hip/signature.rs

cubecl_cpp/hip/
signature.rs
