cubecl_ir/attributes/
mod.rs

1use core::{cell::Ref, fmt};
2
3use alloc::boxed::Box;
4
5use derive_more::From;
6use derive_new::new;
7use num_traits::{AsPrimitive, NumCast};
8use pliron::{
9    builtin::{
10        attr_interfaces::{MaterializableAttr, TypedAttrInterface},
11        attributes::IntegerAttr,
12        ops::ConstantOp,
13        types::IntegerType,
14    },
15    combine::{Parser, parser::char},
16    context::{Context, Ptr},
17    derive::{attr_interface_impl, pliron_attr},
18    irfmt::parsers::{spaced, type_parse},
19    op::Op,
20    operation::Operation,
21    parsable::{IntoParseResult, Parsable, ParseResult, StateStream},
22    printable::{self, Printable},
23    r#type::{TypeHandle, type_impls},
24    utils::apint::{APInt, bw},
25};
26
27use crate::{
28    ConstantValue,
29    apfloat::{APFloat, APFloatType},
30    interfaces::{ConstantAttr, TypedExt},
31    settings::Dim3,
32    try_cast_ty,
33    types::scalar::*,
34};
35
36mod entrypoint;
37
38pub use entrypoint::*;
39
40macro_rules! materialize_const {
41    ($ty: ty) => {
42        #[attr_interface_impl]
43        impl MaterializableAttr for $ty {
44            fn materialize(&self, ctx: &mut Context) -> Ptr<Operation> {
45                let const_op = ConstantOp::new(ctx, Box::new(self.clone()));
46                const_op.get_operation()
47            }
48        }
49    };
50}
51
52#[macro_export]
53macro_rules! ext_attribute {
54    ($name: ident: $ty: ty, $($implementors: ty),*) => {
55        paste::paste! {
56            dict_key!([<ATTR_KEY_ $name:upper>], stringify!($name));
57
58            #[op_interface]
59            pub trait [<$name:upper:camel> Interface] {
60                fn [<get_ $name>]<'a>(&self, ctx: &'a pliron::context::Context) -> Option<core::cell::Ref<'a, $ty>> {
61                    let self_op = self.get_operation().deref(ctx);
62                    Ref::filter_map(self_op, |self_op| {
63                        self_op
64                        .attributes
65                        .get::<$ty>(&[<ATTR_KEY_ $name:upper>])
66                    }).ok()
67                }
68
69                fn [<set_ $name>](&self, ctx: &mut Context, value: $ty) {
70                    let mut self_op = self.get_operation().deref_mut(ctx);
71                    self_op.attributes.set([<ATTR_KEY_ $name:upper>].clone(), value);
72                }
73
74                fn verify(_op: &dyn pliron::op::Op, _ctx: &pliron::context::Context) -> pliron::result::Result<()>
75                where
76                    Self: Sized,
77                {
78                    Ok(())
79                }
80            }
81        }
82    };
83}
84
85/// A zero-value attribute, used for zero-initializing arbitrary types with whatever "zero" means
86/// for it. Arrays get all fields zero-initialized, floats and ints initialize to zero, booleans
87/// to false, etc.
88#[pliron_attr(name = "cube.zero", format = "`[zero: ` $ty `]`", verifier = "succ")]
89#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)]
90pub struct ZeroAttr {
91    pub ty: TypeHandle,
92}
93materialize_const!(ZeroAttr);
94
95impl ZeroAttr {
96    pub fn new(ty: impl Into<TypeHandle>) -> Self {
97        Self { ty: ty.into() }
98    }
99}
100
101#[attr_interface_impl]
102impl TypedAttrInterface for ZeroAttr {
103    fn get_type(&self, _ctx: &Context) -> TypeHandle {
104        self.ty
105    }
106}
107
108#[attr_interface_impl]
109impl ConstantAttr for ZeroAttr {
110    fn as_const_val(&self, ctx: &Context) -> ConstantValue {
111        let ty = self.ty.deref(ctx);
112        if type_impls::<dyn APFloatType>(&*ty) {
113            ConstantValue::Float(0.0)
114        } else if self.ty.is_unsigned_int(ctx) || self.ty.is_index(ctx) {
115            ConstantValue::UInt(0)
116        } else if self.ty.is_signed_int(ctx) {
117            ConstantValue::Int(0)
118        } else if self.ty.is_bool(ctx) {
119            ConstantValue::Bool(false)
120        } else {
121            panic!("Invalid value type for `as_const_val`")
122        }
123    }
124    fn float_as_f64(&self, ctx: &Context) -> Option<f64> {
125        let ty = self.ty.deref(ctx);
126        if type_impls::<dyn APFloatType>(&*ty) {
127            Some(0.0)
128        } else {
129            None
130        }
131    }
132}
133
134#[pliron_attr(name = "cube.index", format = "$0", verifier = "succ")]
135#[derive(new, From, PartialEq, Eq, Clone, Copy, Debug, Hash, PartialOrd, Ord)]
136pub struct IndexAttr(pub usize);
137materialize_const!(IndexAttr);
138
139impl IndexAttr {
140    pub fn as_value(&self, _ctx: &Context) -> Option<usize> {
141        Some(self.0)
142    }
143
144    pub fn with_value(&self, _ctx: &Context, new_val: usize) -> Self {
145        Self::new(new_val)
146    }
147}
148
149#[attr_interface_impl]
150impl ConstantAttr for IndexAttr {
151    fn as_const_val(&self, _ctx: &Context) -> ConstantValue {
152        ConstantValue::UInt(self.0 as u64)
153    }
154}
155
156impl From<IndexAttr> for usize {
157    fn from(value: IndexAttr) -> Self {
158        value.0
159    }
160}
161
162#[attr_interface_impl]
163impl TypedAttrInterface for IndexAttr {
164    fn get_type(&self, ctx: &Context) -> TypeHandle {
165        IndexType::get(ctx).into()
166    }
167}
168
169/// A boolean attribute
170#[pliron_attr(name = "cube.bool", format = "$0", verifier = "succ")]
171#[derive(new, PartialEq, Eq, Clone, Copy, Debug, Hash)]
172pub struct BoolAttr(pub bool);
173materialize_const!(BoolAttr);
174
175impl BoolAttr {
176    pub fn as_value(&self, _ctx: &Context) -> Option<bool> {
177        Some(self.0)
178    }
179
180    pub fn with_value(&self, _ctx: &Context, new_val: bool) -> Self {
181        Self::new(new_val)
182    }
183}
184
185impl From<BoolAttr> for bool {
186    fn from(value: BoolAttr) -> Self {
187        value.0
188    }
189}
190
191impl From<bool> for BoolAttr {
192    fn from(value: bool) -> Self {
193        BoolAttr::new(value)
194    }
195}
196
197impl BoolAttr {
198    /// The answer for one lane of `result`, or [`None`] where `result` has more than one.
199    ///
200    /// A comparison answers once per lane, and [`BoolAttr`] carries no vectorization: it types
201    /// itself as a bare [`BoolType`]. Folding a vector comparison to a single `true` would put
202    /// one bool where a vector of them belongs, and the backends then emit a scalar into a slot
203    /// typed for a vector. Folds that answer per lane build their attribute here rather than with
204    /// [`BoolAttr::new`], and simply decline to fold a vector.
205    pub fn per_lane(
206        ctx: &Context,
207        result: impl pliron::r#type::Typed,
208        value: bool,
209    ) -> Option<Self> {
210        use crate::interfaces::TypedExt;
211        (result.vector_size(ctx) == 1).then(|| Self::new(value))
212    }
213}
214
215#[attr_interface_impl]
216impl TypedAttrInterface for BoolAttr {
217    fn get_type(&self, ctx: &Context) -> TypeHandle {
218        BoolType::get(ctx).into()
219    }
220}
221
222#[attr_interface_impl]
223impl ConstantAttr for BoolAttr {
224    fn as_const_val(&self, _ctx: &Context) -> ConstantValue {
225        ConstantValue::Bool(self.0)
226    }
227}
228
229pub trait IntAttrExt {
230    fn as_value<T>(&self, ctx: &Context) -> Option<T>
231    where
232        T: TypedLiteral + Copy + 'static,
233        i128: AsPrimitive<T>;
234
235    fn with_value<T: NumCast>(&self, ctx: &Context, new_val: T) -> Self;
236}
237
238impl IntAttrExt for IntegerAttr {
239    fn as_value<T>(&self, ctx: &Context) -> Option<T>
240    where
241        T: TypedLiteral + Copy + 'static,
242        i128: AsPrimitive<T>,
243    {
244        if T::is_same_type(ctx, self.get_type().into()) {
245            Some(self.value().to_i128().as_())
246        } else {
247            None
248        }
249    }
250
251    fn with_value<T: NumCast>(&self, ctx: &Context, new_val: T) -> Self {
252        let width = bw(self.get_type().deref(ctx).width() as usize);
253        let val = new_val.to_i128().expect("Should succeed");
254        Self::new(self.get_type(), APInt::from_i128(val, width))
255    }
256}
257
258#[attr_interface_impl]
259impl ConstantAttr for IntegerAttr {
260    fn as_const_val(&self, ctx: &Context) -> ConstantValue {
261        if self.get_type().deref(ctx).is_signed() {
262            ConstantValue::Int(self.value().to_i64())
263        } else {
264            ConstantValue::UInt(self.value().to_u64())
265        }
266    }
267}
268
269#[pliron_attr(name = "cube.float", verifier = "succ")]
270#[derive(new, PartialEq, Clone, Debug, Hash)]
271pub struct FloatAttr {
272    pub ty: TypeHandle,
273    pub val: APFloat,
274}
275materialize_const!(FloatAttr);
276
277impl Printable for FloatAttr {
278    fn fmt(
279        &self,
280        ctx: &Context,
281        state: &printable::State,
282        f: &mut fmt::Formatter<'_>,
283    ) -> fmt::Result {
284        write!(f, "{}: ", self.ty.disp(ctx))?;
285        self.float_type(ctx).disp_value(self.val, ctx, state, f)
286    }
287}
288
289impl Parsable for FloatAttr {
290    type Arg = ();
291    type Parsed = Self;
292
293    fn parse<'a>(input: &mut StateStream<'a>, _: Self::Arg) -> ParseResult<'a, Self::Parsed> {
294        let ty = type_parse(input)?.0;
295        spaced(char::char(':')).parse_stream(input).into_result()?;
296        // Safety: We know this context is not mutably borrowed for value parsing
297        let ctx = dupe_ref(input.state.ctx);
298        let val = try_cast_ty!(ty.deref(ctx), ctx, dyn APFloatType).parse_value(input)?;
299        Ok(FloatAttr::new(ty, val.0)).into_parse_result()
300    }
301}
302
303fn dupe_ref<'b>(ref_: &Context) -> &'b Context {
304    let ctx: *const Context = ref_;
305    unsafe { &*ctx }
306}
307
308impl FloatAttr {
309    pub fn as_value<T: NumCast + TypedLiteral>(&self, ctx: &Context) -> Option<T> {
310        if T::is_same_type(ctx, self.ty) {
311            Some(T::from(self.float_type(ctx).value_to_f64(self.val)).expect("Should succeed"))
312        } else {
313            None
314        }
315    }
316
317    pub fn with_value<T: NumCast>(&self, ctx: &Context, new_val: T) -> Self {
318        Self::from_f64(ctx, self.ty, new_val.to_f64().expect("Should convert"))
319    }
320
321    pub fn from_f64(ctx: &Context, ty: TypeHandle, val: f64) -> Self {
322        let val = try_cast_ty!(ty.deref(ctx), ctx, dyn APFloatType).value_from_f64(val);
323        Self::new(ty, val)
324    }
325
326    pub fn float_type<'a>(&self, ctx: &'a Context) -> Ref<'a, dyn APFloatType> {
327        Ref::map(self.ty.deref(ctx), |ty| {
328            try_cast_ty!(ty, ctx, dyn APFloatType)
329        })
330    }
331}
332
333#[pliron_attr(name = "cube.dim3", format, verifier = "succ")]
334#[derive(new, From, PartialEq, Clone, Debug, Hash)]
335pub struct Dim3Attr(pub Dim3);
336
337#[attr_interface_impl]
338impl TypedAttrInterface for FloatAttr {
339    fn get_type(&self, _ctx: &Context) -> TypeHandle {
340        self.ty
341    }
342}
343
344#[attr_interface_impl]
345impl ConstantAttr for FloatAttr {
346    fn as_const_val(&self, ctx: &Context) -> ConstantValue {
347        let value = self.float_type(ctx).value_to_f64(self.val);
348        ConstantValue::Float(value)
349    }
350    fn float_as_f64(&self, ctx: &Context) -> Option<f64> {
351        let val = self.float_type(ctx).value_to_f64(self.val);
352        Some(val)
353    }
354}
355
356pub trait TypedLiteral {
357    fn is_same_type(ctx: &Context, ty: TypeHandle) -> bool;
358}
359
360macro_rules! literal {
361    ($ty: ty, $ir_ty: ty, $pred: expr) => {
362        impl TypedLiteral for $ty {
363            fn is_same_type(ctx: &Context, ty: TypeHandle) -> bool {
364                ty.deref(ctx).downcast_ref::<$ir_ty>().is_some_and($pred)
365            }
366        }
367    };
368    ($ty: ty, $ir_ty: ty) => {
369        literal!($ty, $ir_ty, |_| true);
370    };
371}
372
373literal!(usize, IndexType);
374
375literal!(i8, IntegerType, |it| it.width() == 8);
376literal!(i16, IntegerType, |it| it.width() == 16);
377literal!(i32, IntegerType, |it| it.width() == 32);
378literal!(i64, IntegerType, |it| it.width() == 64);
379
380literal!(u8, IntegerType, |it| it.width() == 8);
381literal!(u16, IntegerType, |it| it.width() == 16);
382literal!(u32, IntegerType, |it| it.width() == 32);
383literal!(u64, IntegerType, |it| it.width() == 64);
384
385literal!(half::f16, Float16Type);
386literal!(half::bf16, BFloat16Type);
387literal!(f32, Float32Type);
388literal!(f64, Float64Type);
cubecl_ir/attributes/mod.rs

cubecl_ir/attributes/
mod.rs
