cubecl_ir/types/
mod.rs

1use pliron::{
2    attribute::AttrObj,
3    derive::{pliron_type, type_interface_impl},
4    utils::table::HMap,
5};
6
7use crate::{
8    AddressSpace, aligned,
9    interfaces::{
10        AlignedType, HasElementType, IndexableType, MaybePackedType, MaybeVectorizedType,
11        ScalarizableType, SizedType, TypedExt, memory_slot::DestructurableTypeInterface,
12    },
13    prelude::*,
14    sized,
15    types::aggregate::index_attr,
16};
17
18pub mod aggregate;
19pub mod barrier;
20pub mod cuda;
21pub mod fp8;
22pub mod matrix;
23pub mod scalar;
24pub mod spirv;
25
26pub use fp8::Fp8Format;
27pub use matrix::{MatrixIdent, MatrixLayout, MatrixScope, MatrixShape};
28
29#[pliron_type(
30    name = "vector.vector",
31    format = "`<` $vectorization ` x ` $inner `>`",
32    generate_get = true,
33    verifier = "succ"
34)]
35#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
36pub struct VectorType {
37    pub inner: TypeHandle,
38    pub vectorization: usize,
39}
40
41#[type_interface_impl]
42impl MaybeVectorizedType for VectorType {
43    fn vector_size(&self, _ctx: &Context) -> usize {
44        self.vectorization
45    }
46}
47
48#[type_interface_impl]
49impl MaybePackedType for VectorType {
50    fn packing_factor(&self, ctx: &Context) -> usize {
51        self.inner.packing_factor(ctx)
52    }
53}
54
55#[type_interface_impl]
56impl AlignedType for VectorType {
57    fn align(&self, ctx: &Context) -> usize {
58        self.inner.align(ctx) * self.vectorization
59    }
60}
61
62#[type_interface_impl]
63impl SizedType for VectorType {
64    fn size(&self, ctx: &Context) -> usize {
65        self.inner.size(ctx) * self.vectorization
66    }
67}
68
69#[type_interface_impl]
70impl ScalarizableType for VectorType {
71    fn scalar_type(&self, _ctx: &Context) -> TypeHandle {
72        self.inner
73    }
74}
75
76#[type_interface_impl]
77impl HasElementType for VectorType {
78    fn element_type(&self, ctx: &Context) -> Option<TypeHandle> {
79        Some(self.get_self_handle(ctx))
80    }
81}
82
83#[type_interface_impl]
84impl DestructurableTypeInterface for VectorType {
85    fn subelement_index_map(&self, _ctx: &Context) -> Option<HMap<AttrObj, TypeHandle>> {
86        let mut out = HMap::new();
87        for i in 0..self.vectorization {
88            out.insert(index_attr(i), self.inner);
89        }
90        Some(out)
91    }
92
93    fn type_at_index(&self, _ctx: &Context, _index: &AttrObj) -> TypeHandle {
94        self.inner
95    }
96}
97
98#[pliron_type(
99    name = "atomic.atomic",
100    format = "`<` $inner `>`",
101    generate_get = true,
102    verifier = "succ"
103)]
104#[derive(Hash, PartialEq, Eq, Debug, Clone)]
105pub struct AtomicType {
106    pub inner: TypeHandle,
107}
108
109#[type_interface_impl]
110impl MaybeVectorizedType for AtomicType {
111    fn vector_size(&self, ctx: &Context) -> usize {
112        self.inner.vector_size(ctx)
113    }
114    fn try_vector_size(&self, ctx: &Context) -> Option<usize> {
115        self.inner.try_get_vector_size(ctx)
116    }
117}
118
119#[type_interface_impl]
120impl AlignedType for AtomicType {
121    fn align(&self, ctx: &Context) -> usize {
122        self.inner.align(ctx)
123    }
124}
125
126#[type_interface_impl]
127impl SizedType for AtomicType {
128    fn size(&self, ctx: &Context) -> usize {
129        self.inner.size(ctx)
130    }
131}
132
133#[type_interface_impl]
134impl HasElementType for AtomicType {
135    fn element_type(&self, ctx: &Context) -> Option<TypeHandle> {
136        type_cast::<dyn HasElementType>(&*self.inner.deref(ctx))?.element_type(ctx)
137    }
138}
139
140#[pliron_type(
141    name = "cube.ptr",
142    format = "`<` $inner `, ` $address_space `>`",
143    generate_get = true,
144    verifier = "succ"
145)]
146#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
147pub struct PointerType {
148    pub inner: TypeHandle,
149    pub address_space: AddressSpace,
150}
151aligned!(PointerType, align_of::<u64>());
152sized!(PointerType, size_of::<u64>());
153
154#[type_interface_impl]
155impl MaybeVectorizedType for PointerType {
156    fn vector_size(&self, ctx: &Context) -> usize {
157        self.inner.vector_size(ctx)
158    }
159    fn try_vector_size(&self, ctx: &Context) -> Option<usize> {
160        self.inner.try_get_vector_size(ctx)
161    }
162}
163
164#[type_interface_impl]
165impl MaybePackedType for PointerType {
166    fn packing_factor(&self, ctx: &Context) -> usize {
167        self.inner.packing_factor(ctx)
168    }
169}
170
171#[type_interface_impl]
172impl HasElementType for PointerType {
173    fn element_type(&self, ctx: &Context) -> Option<TypeHandle> {
174        type_cast::<dyn HasElementType>(&*self.inner.deref(ctx))?.element_type(ctx)
175    }
176}
177
178#[pliron_type(
179    name = "cube.array",
180    format = "`[` $inner `; ` $length `]`",
181    generate_get = true,
182    verifier = "succ"
183)]
184#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
185pub struct ArrayType {
186    pub inner: TypeHandle,
187    pub length: usize,
188}
189
190#[type_interface_impl]
191impl MaybeVectorizedType for ArrayType {
192    fn vector_size(&self, ctx: &Context) -> usize {
193        self.inner.vector_size(ctx)
194    }
195    fn try_vector_size(&self, ctx: &Context) -> Option<usize> {
196        self.inner.try_get_vector_size(ctx)
197    }
198}
199
200#[type_interface_impl]
201impl AlignedType for ArrayType {
202    fn align(&self, ctx: &Context) -> usize {
203        self.inner.align(ctx)
204    }
205}
206
207#[type_interface_impl]
208impl SizedType for ArrayType {
209    fn size(&self, ctx: &Context) -> usize {
210        self.inner.size(ctx) * self.length
211    }
212}
213
214#[type_interface_impl]
215impl IndexableType for ArrayType {
216    fn indexed_type(&self, _ctx: &Context) -> TypeHandle {
217        self.inner
218    }
219}
220
221#[type_interface_impl]
222impl MaybePackedType for ArrayType {
223    fn packing_factor(&self, ctx: &Context) -> usize {
224        self.inner.packing_factor(ctx)
225    }
226}
227
228#[type_interface_impl]
229impl HasElementType for ArrayType {
230    fn element_type(&self, ctx: &Context) -> Option<TypeHandle> {
231        type_cast::<dyn HasElementType>(&*self.inner.deref(ctx))?.element_type(ctx)
232    }
233}
234
235#[type_interface_impl]
236impl DestructurableTypeInterface for ArrayType {
237    fn subelement_index_map(&self, _ctx: &Context) -> Option<HMap<AttrObj, TypeHandle>> {
238        let mut out = HMap::new();
239        for i in 0..self.length {
240            out.insert(index_attr(i), self.inner);
241        }
242        Some(out)
243    }
244
245    fn type_at_index(&self, _ctx: &Context, _index: &AttrObj) -> TypeHandle {
246        self.inner
247    }
248}
249
250/// Raw byte array.
251/// Separate to mark it as semantically opaque and not valid as an input to ops that take a normal
252/// array.
253#[pliron_type(
254    name = "cube.bytes",
255    format = "",
256    generate_get = true,
257    verifier = "succ"
258)]
259#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
260pub struct BytesType;
261aligned!(BytesType, align_of::<u8>());
262
263#[pliron_type(
264    name = "cube.runtime_array",
265    format = "`[` $inner `]`",
266    generate_get = true,
267    verifier = "succ"
268)]
269#[derive(Hash, PartialEq, Eq, Debug, Clone, Copy)]
270pub struct RuntimeArrayType {
271    pub inner: TypeHandle,
272}
273
274#[type_interface_impl]
275impl MaybeVectorizedType for RuntimeArrayType {
276    fn vector_size(&self, ctx: &Context) -> usize {
277        self.inner.vector_size(ctx)
278    }
279    fn try_vector_size(&self, ctx: &Context) -> Option<usize> {
280        self.inner.try_get_vector_size(ctx)
281    }
282}
283
284#[type_interface_impl]
285impl AlignedType for RuntimeArrayType {
286    fn align(&self, ctx: &Context) -> usize {
287        self.inner.align(ctx)
288    }
289}
290
291#[type_interface_impl]
292impl IndexableType for RuntimeArrayType {
293    fn indexed_type(&self, _ctx: &Context) -> TypeHandle {
294        self.inner
295    }
296}
297
298#[type_interface_impl]
299impl MaybePackedType for RuntimeArrayType {
300    fn packing_factor(&self, ctx: &Context) -> usize {
301        self.inner.packing_factor(ctx)
302    }
303}
304
305#[type_interface_impl]
306impl HasElementType for RuntimeArrayType {
307    fn element_type(&self, ctx: &Context) -> Option<TypeHandle> {
308        type_cast::<dyn HasElementType>(&*self.inner.deref(ctx))?.element_type(ctx)
309    }
310}
cubecl_ir/types/mod.rs

cubecl_ir/types/
mod.rs
