Skip to main content

cubecl_cpp/hip/
ty.rs

1use cubecl_core::{
2    cmma::MatrixType,
3    ir::{
4        prelude::*,
5        types::{
6            PointerType,
7            scalar::{BFloat16Type, Float8E4M3Type, Float8E5M2Type, Float16Type},
8        },
9    },
10};
11
12use crate::{
13    shared::{
14        signature::ty_includes,
15        ty::{TypeExtCPP, TypeToCPP, UniformPointerType, ptr_constness},
16    },
17    target::Hip,
18};
19
20macro_rules! hip_ty {
21    ($ty: ty, $impl: expr) => {
22        #[type_interface_impl]
23        impl TypeToCPP<Hip> for $ty {
24            fn to_cpp(&self, ctx: &Context) -> String {
25                $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl)
26            }
27        }
28    };
29}
30pub(super) use hip_ty;
31
32hip_ty!(Float16Type, |_, _| "__half".into());
33hip_ty!(BFloat16Type, |_, _| "__hip_bfloat16".into());
34hip_ty!(Float8E4M3Type, |_, _| "uint8_t".into());
35hip_ty!(Float8E5M2Type, |_, _| "uint8_t".into());
36
37hip_ty!(PointerType, |ty, ctx| format!(
38    "{} {}*",
39    ty.inner.to_cpp(ctx),
40    ptr_constness(ctx, ty.address_space),
41));
42hip_ty!(UniformPointerType, |ty, ctx| format!(
43    "{} const*",
44    ty.inner.to_cpp(ctx)
45));
46
47ty_includes!(Hip, [MatrixType] => "rocwmma/rocwmma.hpp");
48// hiprtc's builtin header only declares the legacy `hip_bfloat16` struct;
49// `__hip_bfloat16` (the CUDA-compatible type emitted above) needs the real
50// header, exactly as the pre-pliron dialect included it whenever bf16 appeared.
51// (`__half` needs no include: the builtin header covers it.)
52ty_includes!(Hip, [BFloat16Type, crate::cuda::ty::BFloat16x2Type] => "hip/hip_bf16.h");