cubecl_cpp/metal/
signature.rs

1use cubecl_core::ir::{
2    attributes::{ATTR_BUFFER_BINDING, BufferBindingAttr, EntrypointInterface, FuncInterface},
3    prelude::*,
4};
5use cubecl_opt::passes::alloc_shared_memory::AllocSharedOp;
6use itertools::Itertools;
7use pliron::{
8    builtin::{
9        ops::{FuncOp, ModuleOp},
10        types::FunctionType,
11    },
12    dict_key,
13};
14
15use crate::{
16    metal::{BuiltInAttr, metal_op},
17    shared::{
18        CompilationOptions, CppValue, branch::block_to_cpp, signature::LoadInfoOp, ty::TypeExtCPP,
19        type_definitions,
20    },
21};
22
23dict_key!(ATTR_BUILTIN_ATTRIBUTE, "metal_builtin");
24
25const IMPORT: &str = "
26#include <metal_stdlib>
27using namespace metal;
28";
29
30metal_op!(ModuleOp, |op, ctx| {
31    let mut out = IMPORT.to_string();
32    type_definitions(&mut out, "long").unwrap();
33    out.push_str(&block_to_cpp(ctx, op.get_body(ctx, 0)));
34    out
35});
36
37metal_op!(FuncOp, |op, ctx| {
38    let func_name = op.get_symbol_name(ctx);
39    let ty = op.get_type(ctx).deref(ctx);
40    let func_ty = ty.downcast_ref::<FunctionType>().unwrap();
41    let return_ty = func_ty.res_types()[0].to_cpp(ctx);
42    let attributes = if let Some(abi) = op.get_entrypoint_abi(ctx) {
43        let threads_per_simdgroup = ctx.aux_ty::<CompilationOptions>().warp_size as u32;
44        format!(
45            r#"[[max_total_threads_per_threadgroup({})]] [[kernel]] {return_ty}"#,
46            max_total_threads_never_declaring_a_single_simdgroup(
47                abi.cube_dim.num_elems(),
48                threads_per_simdgroup,
49            ),
50        )
51    } else {
52        return_ty
53    };
54
55    let entry_block = op.get_entry_block(ctx);
56
57    let block = entry_block.deref(ctx);
58    let params = block.arguments().enumerate();
59    let params = params.map(|(i, arg)| gen_param(ctx, op, i, arg)).join(", ");
60
61    let body = block_to_cpp(ctx, entry_block);
62
63    format!("{attributes} {func_name}({params}) {{\n{body}\n}}\n")
64});
65
66fn max_total_threads_never_declaring_a_single_simdgroup(
67    cube_dim_total: u32,
68    threads_per_simdgroup: u32,
69) -> u32 {
70    let smallest_bound_compiled_correctly = 2 * threads_per_simdgroup;
71    cube_dim_total.max(smallest_bound_compiled_correctly)
72}
73
74fn gen_param(ctx: &Context, func: &FuncOp, i: usize, arg: Value) -> String {
75    let mut segments = vec![];
76    segments.push(arg.get_type(ctx).to_cpp(ctx));
77    segments.push("const".into());
78    segments.push(arg.name(ctx).to_string());
79    if let Some(binding) = func.get_arg_attr::<BufferBindingAttr>(ctx, i, &ATTR_BUFFER_BINDING) {
80        segments.push(format!("[[buffer({})]]", binding.buffer_pos));
81    }
82    if let Some(builtin) = func.get_arg_attr::<BuiltInAttr>(ctx, i, &ATTR_BUILTIN_ATTRIBUTE) {
83        segments.push(format!("[[{}]]", builtin));
84    }
85    segments.join(" ")
86}
87
88// Metal does support dynamically sized shared memory, but it can't be used from WGPU. Metal allows
89// allocating the full size statically, unlike CUDA, so it should be fine.
90metal_op!(AllocSharedOp, |op, ctx| {
91    let name = op.get_result(ctx).name(ctx);
92    let align = op.alignment(ctx).0;
93    let size = op.size(ctx).0;
94    format!("alignas({align}) threadgroup char {name}[{size}];\n")
95});
96
97metal_op!(LoadInfoOp, |op, ctx| {
98    let ptr = op.ptr(ctx).name(ctx);
99    let out = op.get_result(ctx);
100    let out_ty = out.get_type(ctx).to_cpp(ctx);
101    format!("constant {out_ty}& {} = *{ptr};\n", out.name(ctx))
102});
cubecl_cpp/metal/signature.rs

cubecl_cpp/metal/
signature.rs
