cubecl_cpp/metal/
dialect.rs

1use super::arch::MetalArchitecture;
2use crate::shared::{CppValue, SupportedMmaCombinations};
3use cubecl_core::ir::{
4    ElemType, FloatKind,
5    dialect::{
6        general::PrintfOp,
7        synchronization::{SyncOp, SyncScope},
8    },
9    features::MmaConfig,
10};
11use itertools::Itertools;
12
13macro_rules! metal_op {
14    ($ty: ty, $impl: expr) => {
15        #[pliron::derive::op_interface_impl]
16        impl $crate::shared::operation::OpToCPP<$crate::target::Metal> for $ty {
17            fn to_cpp(&self, ctx: &pliron::context::Context) -> String {
18                $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl)
19            }
20        }
21    };
22}
23pub(super) use metal_op;
24
25macro_rules! metal_op_with_out {
26    ($ty: ty, $impl: expr) => {
27        #[pliron::derive::op_interface_impl]
28        impl $crate::shared::operation::OpToCPP<$crate::target::Metal> for $ty {
29            fn to_cpp(&self, ctx: &pliron::context::Context) -> String {
30                use cubecl_core::ir::prelude::*;
31                use $crate::shared::CppValue;
32                let op = $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl);
33                let out = self.get_result(ctx).fmt_left(ctx);
34                format!("{out} = {op};\n")
35            }
36        }
37    };
38}
39pub(super) use metal_op_with_out;
40
41metal_op!(PrintfOp, |op, ctx| {
42    let format_string = String::from(op.format_string(ctx).clone());
43    let args = op.args(ctx);
44    let args = args.iter().map(|it| format!(", {}", it.name(ctx))).join("");
45    format!("os_log_default.log({format_string:?}{});\n", args)
46});
47
48metal_op!(SyncOp, |op, ctx| {
49    match op.scope(ctx).0 {
50        SyncScope::Plane => "simdgroup_barrier(mem_flags::mem_threadgroup);\n",
51        SyncScope::Cube => "threadgroup_barrier(mem_flags::mem_threadgroup);\n",
52        // Three instructions, each measured on an M2 to be doing something. The barrier is the
53        // cube's half, and on its own it leaves one threadgroup's writes unseen by the next. The
54        // fence at `thread_scope_device` is the device half, and *which side of the barrier it
55        // sits on* decides what it covers: with the fence only after, a store made by one unit
56        // before the barrier still reached no other threadgroup, and the same kernel with one
57        // more `threadgroup_barrier` in front of it did. So the release is stated before the
58        // cube meets and the acquire after, and the barrier between them is what makes both the
59        // whole cube's. `atomic_thread_fence` is MSL 3.2, the version this backend compiles at.
60        SyncScope::Device => {
61            "atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);\n\
62             threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup);\n\
63             atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);\n"
64        }
65        SyncScope::Unit => "",
66    }
67    .into()
68});
69
70// Coop Matrices dialect
71
72pub fn supported_cmma_combinations_metal(_arch: &MetalArchitecture) -> SupportedMmaCombinations {
73    let types = vec![
74        (
75            ElemType::Float(FloatKind::F16),
76            ElemType::Float(FloatKind::F16),
77            ElemType::Float(FloatKind::F16),
78        ),
79        (
80            ElemType::Float(FloatKind::F16),
81            ElemType::Float(FloatKind::F16),
82            ElemType::Float(FloatKind::F32),
83        ),
84        (
85            ElemType::Float(FloatKind::BF16),
86            ElemType::Float(FloatKind::BF16),
87            ElemType::Float(FloatKind::BF16),
88        ),
89        (
90            ElemType::Float(FloatKind::F32),
91            ElemType::Float(FloatKind::F32),
92            ElemType::Float(FloatKind::F32),
93        ),
94    ];
95    types
96        .into_iter()
97        .map(|(a_type, b_type, cd_type)| MmaConfig {
98            a_type,
99            b_type,
100            cd_type,
101            m: 8,
102            n: 8,
103            k: 8,
104        })
105        .collect()
106}
cubecl_cpp/metal/dialect.rs

cubecl_cpp/metal/
dialect.rs
