cubecl_cpp/metal/
ty.rs

1use cubecl_core::ir::{
2    AddressSpace,
3    prelude::*,
4    types::{
5        AtomicType, PointerType,
6        scalar::{BFloat16Type, Float8E4M3Type, Float8E5M2Type, Float16Type},
7    },
8};
9
10use crate::{
11    shared::ty::{TypeExtCPP, TypeToCPP, UniformPointerType, ptr_constness},
12    target::Metal,
13};
14
15macro_rules! metal_ty {
16    ($ty: ty, $impl: expr) => {
17        #[type_interface_impl]
18        impl TypeToCPP<Metal> for $ty {
19            fn to_cpp(&self, ctx: &Context) -> String {
20                $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl)
21            }
22        }
23    };
24}
25pub(super) use metal_ty;
26
27metal_ty!(Float16Type, |_, _| "half".into());
28metal_ty!(BFloat16Type, |_, _| "bfloat".into());
29metal_ty!(Float8E4M3Type, |_, _| "uint8_t".into());
30metal_ty!(Float8E5M2Type, |_, _| "uint8_t".into());
31
32metal_ty!(PointerType, |ty, ctx| format!(
33    "{} {} {}*",
34    ptr_space(ty.address_space),
35    ty.inner.to_cpp(ctx),
36    ptr_constness(ctx, ty.address_space),
37));
38metal_ty!(UniformPointerType, |ty, ctx| format!(
39    "constant {} const*",
40    ty.inner.to_cpp(ctx)
41));
42
43pub fn ptr_space(addr_space: AddressSpace) -> &'static str {
44    match addr_space {
45        AddressSpace::Global(_) => "device",
46        AddressSpace::Shared => "threadgroup",
47        AddressSpace::Local => "thread",
48    }
49}
50
51metal_ty!(AtomicType, |ty, ctx| format!(
52    "atomic<{}>",
53    ty.inner.to_cpp(ctx)
54));
cubecl_cpp/metal/ty.rs

cubecl_cpp/metal/
ty.rs
