cubecl_ir/types/
aggregate.rs

1use alloc::boxed::Box;
2use pliron::{
3    attribute::AttrObj,
4    derive::{pliron_type, type_interface_impl},
5    utils::table::HMap,
6};
7
8use crate::{
9    attributes::IndexAttr,
10    interfaces::memory_slot::DestructurableTypeInterface,
11    prelude::*,
12    types::scalar::{BoolType, IndexType},
13};
14
15#[pliron_type(
16    name = "cube.slice",
17    format = "`<` $base_ty `>`",
18    generate_get = true,
19    verifier = "succ"
20)]
21#[derive(Hash, PartialEq, Eq, Debug, Clone)]
22pub struct SliceType {
23    pub base_ty: TypeHandle,
24}
25
26#[type_interface_impl]
27impl DestructurableTypeInterface for SliceType {
28    fn subelement_index_map(&self, ctx: &Context) -> Option<HMap<AttrObj, TypeHandle>> {
29        let mut out = HMap::new();
30        out.insert(index_attr(0), self.base_ty);
31        out.insert(index_attr(1), IndexType::get(ctx).to_handle());
32        out.insert(index_attr(2), IndexType::get(ctx).to_handle());
33        Some(out)
34    }
35
36    fn type_at_index(&self, ctx: &Context, index: &AttrObj) -> TypeHandle {
37        match index.downcast_ref::<IndexAttr>().unwrap().0 {
38            0 => self.base_ty,
39            1 | 2 => IndexType::get(ctx).to_handle(),
40            _ => unreachable!(),
41        }
42    }
43}
44
45#[pliron_type(
46    name = "cube.checked_ptr",
47    format = "`<` $base_ty `>`",
48    generate_get = true,
49    verifier = "succ"
50)]
51#[derive(Hash, PartialEq, Eq, Debug, Clone)]
52pub struct CheckedPtrType {
53    pub base_ty: TypeHandle,
54}
55
56#[type_interface_impl]
57impl DestructurableTypeInterface for CheckedPtrType {
58    fn subelement_index_map(&self, ctx: &Context) -> Option<HMap<AttrObj, TypeHandle>> {
59        let mut out = HMap::new();
60        out.insert(index_attr(0), self.base_ty);
61        out.insert(index_attr(1), BoolType::get(ctx).to_handle());
62        Some(out)
63    }
64
65    fn type_at_index(&self, ctx: &Context, index: &AttrObj) -> TypeHandle {
66        match index.downcast_ref::<IndexAttr>().unwrap().0 {
67            0 => self.base_ty,
68            1 => BoolType::get(ctx).to_handle(),
69            _ => unreachable!(),
70        }
71    }
72}
73
74pub fn index_attr(idx: usize) -> AttrObj {
75    Box::new(IndexAttr(idx))
76}
77
78pub struct BoundsCheckMetadata;
79impl BoundsCheckMetadata {
80    pub const POINTER: usize = 0;
81    pub const IS_IN_BOUNDS: usize = 1;
82}
83
84pub struct SliceMetadata;
85impl SliceMetadata {
86    pub const LIST: usize = 0;
87    pub const OFFSET: usize = 1;
88    pub const LENGTH: usize = 2;
89}
cubecl_ir/types/aggregate.rs

cubecl_ir/types/
aggregate.rs
