cubecl_cpp/shared/
ty.rs

1use core::cell::Ref;
2
3use cubecl_core::ir::{
4    AddressSpace, ContextExt, GlobalState,
5    attributes::{
6        ATTR_BUFFER_BINDING, ATTR_BUFFER_IO, BufferBindingAttr, BufferIOAttr, FuncInterface,
7    },
8    interfaces::TypedExt,
9    match_ty,
10    prelude::*,
11    types::{ArrayType, AtomicType, PointerType, RuntimeArrayType, VectorType, scalar::*},
12};
13use pliron::{
14    builtin::{
15        attributes::TypeAttr,
16        types::{IntegerType, Signedness, UnitType},
17    },
18    r#type::TypedHandle,
19};
20
21use crate::{
22    cuda::ty::*,
23    error::CompileError,
24    target::{Shared, dispatch_target},
25};
26
27macro_rules! shared_ty {
28    ($ty: ty, $impl: expr) => {
29        #[type_interface_impl]
30        impl TypeToCPP<Shared> for $ty {
31            fn to_cpp(&self, ctx: &Context) -> String {
32                $crate::shared::closure_inference_hack::<$ty, String>(self, ctx, $impl)
33            }
34        }
35    };
36}
37
38pub trait TypeExtCPP {
39    fn to_cpp(&self, ctx: &Context) -> String;
40}
41impl TypeExtCPP for Ref<'_, dyn Type> {
42    fn to_cpp(&self, ctx: &Context) -> String {
43        let target_cpp = dispatch_target!(ctx, {
44            let inner_cpp = type_cast::<dyn TypeToCPP<Target>>(&**self);
45            inner_cpp.map(|it| it.to_cpp(ctx))
46        });
47        if let Some(cpp) = target_cpp {
48            return cpp;
49        }
50        let shared = type_cast::<dyn TypeToCPP<Shared>>(&**self)
51            .ok_or_else(|| {
52                CompileError::UnsupportedType(format!("{}{}", self.get_type_id(), self.disp(ctx)))
53            })
54            .unwrap(); // Unwrap for now because handling errors everywhere is annoying
55        shared.to_cpp(ctx)
56    }
57}
58impl TypeExtCPP for TypeHandle {
59    fn to_cpp(&self, ctx: &Context) -> String {
60        self.deref(ctx).to_cpp(ctx)
61    }
62}
63impl<T: Type> TypeExtCPP for TypedHandle<T> {
64    fn to_cpp(&self, ctx: &Context) -> String {
65        self.to_handle().to_cpp(ctx)
66    }
67}
68impl TypeExtCPP for TypeAttr {
69    fn to_cpp(&self, ctx: &Context) -> String {
70        self.get_type(ctx).to_cpp(ctx)
71    }
72}
73
74impl<T: Type + ?Sized> TypeExt for T {}
75pub trait TypeExt: Type {
76    fn display(&self, ctx: &Context) -> String {
77        format!("{}{}", self.get_type_id(), self.disp(ctx))
78    }
79}
80
81macro_rules! is_one_of {
82    ($ty: expr; $($types: ty),*) => {
83        false $(|| $ty.is::<$types>())*
84    };
85}
86
87pub trait TypedExtCPP: Typed {
88    fn is_uniform_ptr(&self, ctx: &Context) -> bool {
89        let ty = self.get_type(ctx).deref(ctx);
90        ty.is::<UniformPointerType>()
91    }
92
93    fn is_complex(&self, ctx: &Context) -> bool {
94        let ty = self.scalar_ty(ctx).deref(ctx);
95        is_one_of!(ty; Complex32Type, Complex64Type)
96    }
97
98    fn is_half(&self, ctx: &Context) -> bool {
99        let ty = self.scalar_ty(ctx).deref(ctx);
100        is_one_of!(ty; Float16Type, BFloat16Type)
101    }
102
103    fn is_half2(&self, ctx: &Context) -> bool {
104        let ty = self.scalar_ty(ctx).deref(ctx);
105        is_one_of!(ty; Float16x2Type, BFloat16x2Type)
106    }
107
108    fn is_float8(&self, ctx: &Context) -> bool {
109        let ty = self.scalar_ty(ctx).deref(ctx);
110        is_one_of!(ty; Float8E8M0Type, Float8E5M2Type, Float8E4M3Type)
111    }
112
113    fn is_float8x2(&self, ctx: &Context) -> bool {
114        let ty = self.scalar_ty(ctx).deref(ctx);
115        is_one_of!(ty; Float8E8M0x2Type, Float8E5M2x2Type, Float8E4M3x2Type)
116    }
117
118    fn is_float6(&self, ctx: &Context) -> bool {
119        let ty = self.scalar_ty(ctx).deref(ctx);
120        is_one_of!(ty; Float6E3M2Type, Float6E2M3Type)
121    }
122
123    fn is_float6x2(&self, ctx: &Context) -> bool {
124        let ty = self.scalar_ty(ctx).deref(ctx);
125        is_one_of!(ty; Float6E3M2x2Type, Float6E2M3x2Type)
126    }
127
128    fn is_float4(&self, ctx: &Context) -> bool {
129        let ty = self.scalar_ty(ctx).deref(ctx);
130        ty.is::<Float4E2M1Type>()
131    }
132
133    fn is_float4x2(&self, ctx: &Context) -> bool {
134        let ty = self.scalar_ty(ctx).deref(ctx);
135        ty.is::<Float4E2M1x2Type>()
136    }
137
138    fn is_fp8_fp6_fp4(&self, ctx: &Context) -> bool {
139        self.is_float8(ctx) || self.is_float6(ctx) || self.is_float4(ctx)
140    }
141
142    fn is_packed_fp6_fp8_fp4(&self, ctx: &Context) -> bool {
143        self.is_float8x2(ctx) || self.is_float6x2(ctx) || self.is_float4x2(ctx)
144    }
145
146    fn can_pack(&self, ctx: &Context) -> bool {
147        if !self.is_vector(ctx) && !self.is_float4x2(ctx) {
148            return false;
149        }
150        let scalar = self.scalar_ty(ctx);
151        scalar.is_float16(ctx)
152            || scalar.is_bfloat16(ctx)
153            || scalar.is_fp8_fp6_fp4(ctx)
154            || scalar.is_float4x2(ctx)
155    }
156
157    fn packed_type(&self, ctx: &Context) -> TypeHandle {
158        let ty = self.get_type(ctx).deref(ctx);
159        if let Some(ptr) = ty.downcast_ref::<PointerType>() {
160            return PointerType::get(ctx, ptr.inner.packed_type(ctx), ptr.address_space)
161                .to_handle();
162        } else if let Some(atomic) = ty.downcast_ref::<AtomicType>() {
163            return AtomicType::get(ctx, atomic.inner.packed_type(ctx)).to_handle();
164        }
165
166        assert!(self.can_pack(ctx), "Should be packable");
167        // Already natively packed
168        if self.is_float4x2(ctx) {
169            return self.get_type(ctx);
170        }
171        let vec = ty.downcast_ref::<VectorType>().unwrap();
172        let scalar = vec.inner.deref(ctx);
173        let scalar = match_ty!((scalar) {
174            Float16Type => Float16x2Type::get(ctx).into(),
175            BFloat16Type => BFloat16x2Type::get(ctx).into(),
176            Float8E8M0Type => Float8E8M0x2Type::get(ctx).into(),
177            Float8E5M2Type => Float8E5M2x2Type::get(ctx).into(),
178            Float8E4M3Type => Float8E4M3x2Type::get(ctx).into(),
179            Float6E2M3Type => Float6E2M3x2Type::get(ctx).into(),
180            Float6E3M2Type => Float6E3M2x2Type::get(ctx).into(),
181            Float4E2M1Type => Float4E2M1x2Type::get(ctx).into(),;
182            _ => panic!("Unexpected type {}", scalar.display(ctx))
183        });
184        if vec.vectorization > 2 {
185            VectorType::get(ctx, scalar, vec.vectorization / 2).to_handle()
186        } else {
187            scalar
188        }
189    }
190
191    /// Whether the type is an integer that may be auto-promoted by C++
192    /// They need special handling
193    fn is_small_int(&self, ctx: &Context) -> bool {
194        self.is_int(ctx) && self.scalar_ty(ctx).size(ctx) < 4
195    }
196
197    fn is_small_signed_int(&self, ctx: &Context) -> bool {
198        self.is_signed_int(ctx) && self.scalar_ty(ctx).size(ctx) < 4
199    }
200
201    fn is_small_unsigned_int(&self, ctx: &Context) -> bool {
202        self.is_unsigned_int(ctx) && self.scalar_ty(ctx).size(ctx) < 4
203    }
204}
205impl<T: Typed> TypedExtCPP for T {}
206
207#[type_interface]
208pub trait TypeToCPP<T> {
209    verify_ty_succ!();
210    fn to_cpp(&self, ctx: &Context) -> String;
211}
212
213shared_ty!(VectorType, |ty, ctx| {
214    format!("{}_{}", ty.inner.to_cpp(ctx), ty.vectorization)
215});
216
217shared_ty!(AtomicType, |ty, ctx| ty.inner.to_cpp(ctx));
218shared_ty!(RuntimeArrayType, |ty, ctx| ty.inner.to_cpp(ctx));
219
220shared_ty!(ArrayType, |ty, ctx| {
221    format!("array<{}, {}>", ty.inner.to_cpp(ctx), ty.length)
222});
223
224pub fn ptr_constness(ctx: &Context, addr_space: AddressSpace) -> &'static str {
225    match addr_space {
226        AddressSpace::Global(idx) => match find_global_constness(ctx, idx) {
227            true => "const",
228            false => "",
229        },
230        AddressSpace::Shared | AddressSpace::Local => "",
231    }
232}
233
234fn find_global_constness(ctx: &Context, idx: usize) -> bool {
235    let func = ctx.aux_ty::<GlobalState>().entry_func;
236    let num_args = func.get_entry_block(ctx).deref(ctx).get_num_arguments();
237    let arg_pos = (0..num_args)
238        .filter_map(|i| Some((i, func.get_arg_attr(ctx, i, &ATTR_BUFFER_BINDING)?)))
239        .find(|(_, binding): &(_, Ref<'_, BufferBindingAttr>)| binding.buffer_pos == idx)
240        .expect("Should exist");
241    let io = func.get_arg_attr::<BufferIOAttr>(ctx, arg_pos.0, &ATTR_BUFFER_IO);
242    !io.expect("Should have IO attribute").is_writable()
243}
244
245#[pliron_type(
246    name = "cpp.info_ptr",
247    format = "`<` $inner `>`",
248    generate_get = true,
249    verifier = "succ"
250)]
251#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
252pub struct UniformPointerType {
253    pub inner: TypeHandle,
254}
255
256shared_ty!(UnitType, |_, _| "void".into());
257shared_ty!(IntegerType, |ty, _| match ty.signedness() {
258    Signedness::Signed => format!("int{}_t", ty.width()),
259    Signedness::Unsigned | Signedness::Signless => format!("uint{}_t", ty.width()),
260});
261shared_ty!(BoolType, |_, _| "bool".into());
262
263shared_ty!(FloatFlex32Type, |_, _| "float".into());
264shared_ty!(Float32Type, |_, _| "float".into());
265shared_ty!(Float64Type, |_, _| "double".into());
266
267shared_ty!(IndexType, |_, ctx| {
268    ctx.address_type().unsigned_type().to_type(ctx).to_cpp(ctx)
269});
270
271/// Vector of three unsigned integers. This is the only native vector type that's actually relevant
272/// for codegen, so we can just special case it and only use it where necessary (builtin types).
273#[pliron_type(
274    name = "cpp.uvec3",
275    format = "",
276    generate_get = true,
277    verifier = "succ"
278)]
279#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
280pub struct Uvec3Type;
281
282shared_ty!(Uvec3Type, |_, _| "uint3".into());
283
284#[pliron_type(
285    name = "cpp.info_st",
286    format = "",
287    generate_get = true,
288    verifier = "succ"
289)]
290#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
291pub struct InfoStructType;
292
293shared_ty!(InfoStructType, |_, _| "info_st".into());
cubecl_cpp/shared/ty.rs

cubecl_cpp/shared/
ty.rs
