cubecl_cpp/metal/
attribute.rs

1use derive_more::Display;
2use pliron::{
3    builtin::types::{IntegerType, Signedness},
4    context::Context,
5    derive::pliron_attr,
6    r#type::TypeHandle,
7};
8use std::fmt::Display;
9
10use crate::shared::ty::Uvec3Type;
11
12pub enum BufferAttribute {
13    Buffer,
14    ThreadGroup,
15    None,
16}
17
18impl BufferAttribute {
19    pub fn indexed_fmt(&self, index: usize, f: &mut std::fmt::Formatter) -> std::fmt::Result {
20        write!(f, " [[{self}({index})]]")
21    }
22}
23
24impl Display for BufferAttribute {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        match self {
27            Self::Buffer => f.write_str("buffer"),
28            Self::ThreadGroup => f.write_str("threadgroup"),
29            Self::None => Ok(()),
30        }
31    }
32}
33
34#[pliron_attr(name = "msl.builtin", format, verifier = "succ")]
35#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy, Display)]
36pub enum BuiltInAttr {
37    #[display("simdgroup_index_in_threadgroup")]
38    SIMDgroupIndexInThreadgroup,
39    #[display("thread_index_in_simdgroup")]
40    ThreadIndexInSIMDgroup,
41    #[display("thread_index_in_threadgroup")]
42    ThreadIndexInThreadgroup,
43    #[display("thread_position_in_grid")]
44    ThreadPositionInGrid,
45    #[display("thread_position_in_threadgroup")]
46    ThreadPositionInThreadgroup,
47    #[display("threadgroup_position_in_grid")]
48    ThreadgroupPositionInGrid,
49    #[display("threadgroups_per_grid")]
50    ThreadgroupsPerGrid,
51    #[display("threads_per_simdgroup")]
52    ThreadsPerSIMDgroup,
53}
54
55impl BuiltInAttr {
56    pub fn ty(&self, ctx: &Context) -> TypeHandle {
57        match self {
58            BuiltInAttr::SIMDgroupIndexInThreadgroup
59            | BuiltInAttr::ThreadIndexInSIMDgroup
60            | BuiltInAttr::ThreadIndexInThreadgroup
61            | BuiltInAttr::ThreadsPerSIMDgroup => {
62                IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle()
63            }
64
65            BuiltInAttr::ThreadPositionInGrid
66            | BuiltInAttr::ThreadPositionInThreadgroup
67            | BuiltInAttr::ThreadgroupPositionInGrid
68            | BuiltInAttr::ThreadgroupsPerGrid => Uvec3Type::get(ctx).to_handle(),
69        }
70    }
71}
cubecl_cpp/metal/attribute.rs

cubecl_cpp/metal/
attribute.rs
