cubecl_cpp/cuda/
ty.rs

1use cubecl_core::{
2    cmma::MatrixType,
3    ir::{
4        aligned,
5        pliron::context::Context,
6        scalar, sized,
7        types::{
8            PointerType,
9            barrier::{BarrierLevel, BarrierTokenType, BarrierType},
10            cuda::TensorMapType,
11            scalar::*,
12        },
13    },
14};
15use pliron::derive::{format, pliron_type, type_interface_impl};
16
17use crate::{
18    shared::{
19        signature::{RequiresIncludesType, ty_includes},
20        ty::{TypeExtCPP, UniformPointerType, ptr_constness},
21    },
22    target::Cuda,
23};
24
25macro_rules! cuda_ty {
26    ($ty: ty, $impl: expr) => {
27        #[type_interface_impl]
28        impl crate::shared::ty::TypeToCPP<crate::target::Cuda> for $ty {
29            fn to_cpp(&self, ctx: &Context) -> String {
30                $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl)
31            }
32        }
33    };
34}
35pub(super) use cuda_ty;
36
37cuda_ty!(TensorMapType, |_, _| "CUtensorMap".into());
38cuda_ty!(BarrierType, |ty, _| match ty.0 {
39    BarrierLevel::Unit => "cuda::barrier<cuda::thread_scope_thread>".into(),
40    BarrierLevel::Cube => "cuda::barrier<cuda::thread_scope_block>".into(),
41});
42cuda_ty!(BarrierTokenType, |ty, ctx| {
43    format!("{}::arrival_token", ty.0.to_cpp(ctx))
44});
45
46#[type_interface_impl]
47impl RequiresIncludesType<Cuda> for BarrierType {
48    fn includes(&self, _ctx: &Context) -> Vec<String> {
49        vec![
50            "cuda/barrier".into(),
51            "cooperative_groups.h".into(),
52            "cooperative_groups/memcpy_async.h".into(),
53        ]
54    }
55}
56
57cuda_ty!(PointerType, |ty, ctx| format!(
58    "{} {}*",
59    ty.inner.to_cpp(ctx),
60    ptr_constness(ctx, ty.address_space),
61));
62cuda_ty!(UniformPointerType, |ty, ctx| format!(
63    "{} const*",
64    ty.inner.to_cpp(ctx)
65));
66
67#[pliron_type(
68    name = "cpp.f16x2",
69    format = "",
70    generate_get = true,
71    verifier = "succ"
72)]
73#[derive(new, Hash, PartialEq, Eq, Debug, Clone, Copy)]
74pub struct Float16x2Type;
75sized!(Float16x2Type, size_of::<u32>());
76aligned!(Float16x2Type, align_of::<u32>());
77scalar!(Float16x2Type);
78
79#[pliron_type(
80    name = "cpp.bf16x2",
81    format = "",
82    generate_get = true,
83    verifier = "succ"
84)]
85#[derive(new, Hash, PartialEq, Eq, Debug, Clone, Copy)]
86pub struct BFloat16x2Type;
87sized!(BFloat16x2Type, size_of::<u32>());
88aligned!(BFloat16x2Type, align_of::<u32>());
89scalar!(BFloat16x2Type);
90
91#[pliron_type(
92    name = "cuda.ue8m0x2",
93    format = "",
94    generate_get = true,
95    verifier = "succ"
96)]
97#[derive(Debug, Hash, PartialEq, Eq)]
98pub struct Float8E8M0x2Type;
99sized!(Float8E8M0x2Type, size_of::<u16>());
100aligned!(Float8E8M0x2Type, align_of::<u16>());
101scalar!(Float8E8M0x2Type);
102
103#[pliron_type(
104    name = "cuda.e4m3x2",
105    format = "",
106    generate_get = true,
107    verifier = "succ"
108)]
109#[derive(Debug, Hash, PartialEq, Eq)]
110pub struct Float8E4M3x2Type;
111sized!(Float8E4M3x2Type, size_of::<u16>());
112aligned!(Float8E4M3x2Type, align_of::<u16>());
113scalar!(Float8E4M3x2Type);
114
115#[pliron_type(
116    name = "cuda.e5m2x2",
117    format = "",
118    generate_get = true,
119    verifier = "succ"
120)]
121#[derive(Debug, Hash, PartialEq, Eq)]
122pub struct Float8E5M2x2Type;
123sized!(Float8E5M2x2Type, size_of::<u16>());
124aligned!(Float8E5M2x2Type, align_of::<u16>());
125scalar!(Float8E5M2x2Type);
126
127#[pliron_type(
128    name = "cuda.e3m2x2",
129    format = "",
130    generate_get = true,
131    verifier = "succ"
132)]
133#[derive(Debug, Hash, PartialEq, Eq)]
134pub struct Float6E3M2x2Type;
135sized!(Float6E3M2x2Type, size_of::<u16>());
136aligned!(Float6E3M2x2Type, align_of::<u16>());
137scalar!(Float6E3M2x2Type);
138
139#[pliron_type(
140    name = "cuda.e2m3x2",
141    format = "",
142    generate_get = true,
143    verifier = "succ"
144)]
145#[derive(Debug, Hash, PartialEq, Eq)]
146pub struct Float6E2M3x2Type;
147sized!(Float6E2M3x2Type, size_of::<u16>());
148aligned!(Float6E2M3x2Type, align_of::<u16>());
149scalar!(Float6E2M3x2Type);
150
151ty_includes!(Cuda, [Complex32Type, Complex64Type] => "cuComplex.h");
152ty_includes!(Cuda, [MatrixType, TFloat32Type] => "mma.h");
153ty_includes!(Cuda, [Float16Type, Float16x2Type] => "cuda_fp16.h");
154ty_includes!(Cuda, [BFloat16Type, BFloat16x2Type] => "cuda_bf16.h");
155ty_includes!(Cuda, [Float8E4M3Type, Float8E5M2Type, Float8E8M0Type] => "cuda_fp8.h");
156ty_includes!(Cuda, [Float8E4M3x2Type, Float8E5M2x2Type, Float8E8M0x2Type] => "cuda_fp8.h");
157ty_includes!(Cuda, [Float6E3M2Type, Float6E2M3Type, Float6E3M2x2Type, Float6E2M3x2Type] => "cuda_fp6.h");
158ty_includes!(Cuda, [Float4E2M1Type, Float4E2M1x2Type] => "cuda_fp4.h");
159
160cuda_ty!(Complex32Type, |_, _| "cuFloatComplex".into());
161cuda_ty!(Complex64Type, |_, _| "cuDoubleComplex".into());
162cuda_ty!(TFloat32Type, |_, _| "float".into());
163
164cuda_ty!(Float16x2Type, |_, _| "__half2".into());
165cuda_ty!(BFloat16x2Type, |_, _| "__nv_bfloat162".into());
166
167cuda_ty!(Float8E4M3x2Type, |_, _| "__nv_fp8x2_storage_t".into());
168cuda_ty!(Float8E5M2x2Type, |_, _| "__nv_fp8x2_storage_t".into());
169cuda_ty!(Float8E8M0x2Type, |_, _| "__nv_fp8x2_storage_t".into());
170
171cuda_ty!(Float6E3M2x2Type, |_, _| "__nv_fp6x2_storage_t".into());
172cuda_ty!(Float6E2M3x2Type, |_, _| "__nv_fp6x2_storage_t".into());
173
174cuda_ty!(Float4E2M1x2Type, |_, _| "__nv_fp4x2_storage_t".into());
175
176cuda_ty!(Float16Type, |_, _| "__half".into());
177cuda_ty!(BFloat16Type, |_, _| "__nv_bfloat16".into());
178
179cuda_ty!(Float8E4M3Type, |_, _| "__nv_fp8_storage_t".into());
180cuda_ty!(Float8E5M2Type, |_, _| "__nv_fp8_storage_t".into());
181cuda_ty!(Float8E8M0Type, |_, _| "__nv_fp8_storage_t".into());
182
183cuda_ty!(Float6E3M2Type, |_, _| "__nv_fp6_storage_t".into());
184cuda_ty!(Float6E2M3Type, |_, _| "__nv_fp6_storage_t".into());
185
186cuda_ty!(Float4E2M1Type, |_, _| "__nv_fp4_storage_t".into());
cubecl_cpp/cuda/ty.rs

cubecl_cpp/cuda/
ty.rs
