cubecl_ir/attributes/
mod.rs

1use core::{cell::Ref, fmt};
2
3use alloc::{boxed::Box, vec::Vec};
4
5use derive_more::{From, Into};
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, control_flow::SymbolVisibility},
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#[macro_export]
86macro_rules! typed_vec_attr {
87    ($ty: ty, $name: literal, $vec_ty: ident) => {
88        /// A vector of other attributes.
89        #[pliron::derive::def_attribute($name)]
90        #[pliron::derive::format_attribute("`[` vec($0, CharSpace(`,`)) `]`")]
91        #[pliron::derive::verify_succ]
92        #[derive(PartialEq, Eq, Clone, Debug, Hash, Default, derive_more::From)]
93        pub struct $vec_ty(pub Vec<$ty>);
94
95        impl $vec_ty {
96            pub fn new(value: alloc::vec::Vec<$ty>) -> Self {
97                $vec_ty(value)
98            }
99        }
100    };
101}
102/// A zero-value attribute, used for zero-initializing arbitrary types with whatever "zero" means
103/// for it. Arrays get all fields zero-initialized, floats and ints initialize to zero, booleans
104/// to false, etc.
105#[pliron_attr(name = "cube.zero", format = "`[zero: ` $ty `]`", verifier = "succ")]
106#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)]
107pub struct ZeroAttr {
108    pub ty: TypeHandle,
109}
110materialize_const!(ZeroAttr);
111
112impl ZeroAttr {
113    pub fn new(ty: impl Into<TypeHandle>) -> Self {
114        Self { ty: ty.into() }
115    }
116}
117
118#[attr_interface_impl]
119impl TypedAttrInterface for ZeroAttr {
120    fn get_type(&self, _ctx: &Context) -> TypeHandle {
121        self.ty
122    }
123}
124
125#[attr_interface_impl]
126impl ConstantAttr for ZeroAttr {
127    fn as_const_val(&self, ctx: &Context) -> ConstantValue {
128        let ty = self.ty.deref(ctx);
129        if type_impls::<dyn APFloatType>(&*ty) {
130            ConstantValue::Float(0.0)
131        } else if self.ty.is_unsigned_int(ctx) || self.ty.is_index(ctx) {
132            ConstantValue::UInt(0)
133        } else if self.ty.is_signed_int(ctx) {
134            ConstantValue::Int(0)
135        } else if self.ty.is_bool(ctx) {
136            ConstantValue::Bool(false)
137        } else {
138            panic!("Invalid value type for `as_const_val`")
139        }
140    }
141    fn as_int(&self, ctx: &Context) -> Option<APInt> {
142        if self.ty.is_int(ctx) || self.ty.is_index(ctx) {
143            Some(APInt::zero(bw(self.ty.size_bits(ctx))))
144        } else {
145            None
146        }
147    }
148    fn float_as_f64(&self, ctx: &Context) -> Option<f64> {
149        let ty = self.ty.deref(ctx);
150        if type_impls::<dyn APFloatType>(&*ty) {
151            Some(0.0)
152        } else {
153            None
154        }
155    }
156}
157
158#[pliron_attr(name = "cube.index", format = "$0", verifier = "succ")]
159#[derive(new, From, PartialEq, Eq, Clone, Copy, Debug, Hash, PartialOrd, Ord)]
160pub struct IndexAttr(pub usize);
161materialize_const!(IndexAttr);
162
163impl IndexAttr {
164    pub fn as_value(&self, _ctx: &Context) -> Option<usize> {
165        Some(self.0)
166    }
167
168    pub fn with_value(&self, _ctx: &Context, new_val: usize) -> Self {
169        Self::new(new_val)
170    }
171}
172
173#[attr_interface_impl]
174impl ConstantAttr for IndexAttr {
175    fn as_const_val(&self, _ctx: &Context) -> ConstantValue {
176        ConstantValue::UInt(self.0 as u64)
177    }
178    fn as_int(&self, ctx: &Context) -> Option<APInt> {
179        let size = IndexType::get(ctx).to_handle().size_bits(ctx);
180        Some(APInt::from_usize(self.0, bw(size)))
181    }
182}
183
184impl From<IndexAttr> for usize {
185    fn from(value: IndexAttr) -> Self {
186        value.0
187    }
188}
189
190#[attr_interface_impl]
191impl TypedAttrInterface for IndexAttr {
192    fn get_type(&self, ctx: &Context) -> TypeHandle {
193        IndexType::get(ctx).into()
194    }
195}
196
197/// A boolean attribute
198#[pliron_attr(name = "cube.bool", format = "$0", verifier = "succ")]
199#[derive(new, PartialEq, Eq, Clone, Copy, Debug, Hash)]
200pub struct BoolAttr(pub bool);
201materialize_const!(BoolAttr);
202
203impl BoolAttr {
204    pub fn as_value(&self, _ctx: &Context) -> Option<bool> {
205        Some(self.0)
206    }
207
208    pub fn with_value(&self, _ctx: &Context, new_val: bool) -> Self {
209        Self::new(new_val)
210    }
211}
212
213impl From<BoolAttr> for bool {
214    fn from(value: BoolAttr) -> Self {
215        value.0
216    }
217}
218
219impl From<bool> for BoolAttr {
220    fn from(value: bool) -> Self {
221        BoolAttr::new(value)
222    }
223}
224
225impl BoolAttr {
226    /// The answer for one lane of `result`, or [`None`] where `result` has more than one.
227    ///
228    /// A comparison answers once per lane, and [`BoolAttr`] carries no vectorization: it types
229    /// itself as a bare [`BoolType`]. Folding a vector comparison to a single `true` would put
230    /// one bool where a vector of them belongs, and the backends then emit a scalar into a slot
231    /// typed for a vector. Folds that answer per lane build their attribute here rather than with
232    /// [`BoolAttr::new`], and simply decline to fold a vector.
233    pub fn per_lane(
234        ctx: &Context,
235        result: impl pliron::r#type::Typed,
236        value: bool,
237    ) -> Option<Self> {
238        use crate::interfaces::TypedExt;
239        (result.vector_size(ctx) == 1).then(|| Self::new(value))
240    }
241}
242
243#[attr_interface_impl]
244impl TypedAttrInterface for BoolAttr {
245    fn get_type(&self, ctx: &Context) -> TypeHandle {
246        BoolType::get(ctx).into()
247    }
248}
249
250#[attr_interface_impl]
251impl ConstantAttr for BoolAttr {
252    fn as_const_val(&self, _ctx: &Context) -> ConstantValue {
253        ConstantValue::Bool(self.0)
254    }
255}
256
257pub trait IntAttrExt {
258    fn as_value<T>(&self, ctx: &Context) -> Option<T>
259    where
260        T: TypedLiteral + Copy + 'static,
261        i128: AsPrimitive<T>;
262
263    fn with_value<T: NumCast>(&self, ctx: &Context, new_val: T) -> Self;
264}
265
266impl IntAttrExt for IntegerAttr {
267    fn as_value<T>(&self, ctx: &Context) -> Option<T>
268    where
269        T: TypedLiteral + Copy + 'static,
270        i128: AsPrimitive<T>,
271    {
272        if T::is_same_type(ctx, self.get_type().into()) {
273            Some(self.value().to_i128().as_())
274        } else {
275            None
276        }
277    }
278
279    fn with_value<T: NumCast>(&self, ctx: &Context, new_val: T) -> Self {
280        let width = bw(self.get_type().deref(ctx).width() as usize);
281        let val = new_val.to_i128().expect("Should succeed");
282        Self::new(self.get_type(), APInt::from_i128(val, width))
283    }
284}
285
286#[attr_interface_impl]
287impl ConstantAttr for IntegerAttr {
288    fn as_const_val(&self, ctx: &Context) -> ConstantValue {
289        if self.get_type().deref(ctx).is_signed() {
290            ConstantValue::Int(self.value().to_i64())
291        } else {
292            ConstantValue::UInt(self.value().to_u64())
293        }
294    }
295    fn as_int(&self, _ctx: &Context) -> Option<APInt> {
296        Some(self.value())
297    }
298}
299
300typed_vec_attr!(IntegerAttr, "cube.integer_vec", IntegerVecAttr);
301
302#[pliron_attr(name = "cube.float", verifier = "succ")]
303#[derive(new, PartialEq, Clone, Debug, Hash)]
304pub struct FloatAttr {
305    pub ty: TypeHandle,
306    pub val: APFloat,
307}
308materialize_const!(FloatAttr);
309
310impl Printable for FloatAttr {
311    fn fmt(
312        &self,
313        ctx: &Context,
314        state: &printable::State,
315        f: &mut fmt::Formatter<'_>,
316    ) -> fmt::Result {
317        write!(f, "{}: ", self.ty.disp(ctx))?;
318        self.float_type(ctx).disp_value(self.val, ctx, state, f)
319    }
320}
321
322impl Parsable for FloatAttr {
323    type Arg = ();
324    type Parsed = Self;
325
326    fn parse<'a>(input: &mut StateStream<'a>, _: Self::Arg) -> ParseResult<'a, Self::Parsed> {
327        let ty = type_parse(input)?.0;
328        spaced(char::char(':')).parse_stream(input).into_result()?;
329        // Safety: We know this context is not mutably borrowed for value parsing
330        let ctx = dupe_ref(input.state.ctx);
331        let val = try_cast_ty!(ty.deref(ctx), ctx, dyn APFloatType).parse_value(input)?;
332        Ok(FloatAttr::new(ty, val.0)).into_parse_result()
333    }
334}
335
336fn dupe_ref<'b>(ref_: &Context) -> &'b Context {
337    let ctx: *const Context = ref_;
338    unsafe { &*ctx }
339}
340
341impl FloatAttr {
342    pub fn as_value<T: NumCast + TypedLiteral>(&self, ctx: &Context) -> Option<T> {
343        if T::is_same_type(ctx, self.ty) {
344            Some(T::from(self.float_type(ctx).value_to_f64(self.val)).expect("Should succeed"))
345        } else {
346            None
347        }
348    }
349
350    pub fn with_value<T: NumCast>(&self, ctx: &Context, new_val: T) -> Self {
351        Self::from_f64(ctx, self.ty, new_val.to_f64().expect("Should convert"))
352    }
353
354    pub fn from_f64(ctx: &Context, ty: TypeHandle, val: f64) -> Self {
355        let val = try_cast_ty!(ty.deref(ctx), ctx, dyn APFloatType).value_from_f64(val);
356        Self::new(ty, val)
357    }
358
359    pub fn float_type<'a>(&self, ctx: &'a Context) -> Ref<'a, dyn APFloatType> {
360        Ref::map(self.ty.deref(ctx), |ty| {
361            try_cast_ty!(ty, ctx, dyn APFloatType)
362        })
363    }
364}
365
366#[pliron_attr(name = "cube.complex", verifier = "succ")]
367#[derive(new, PartialEq, Eq, Clone, Copy, Debug, Hash)]
368pub struct ComplexAttr {
369    pub ty: TypeHandle,
370    pub re: APFloat,
371    pub im: APFloat,
372}
373materialize_const!(ComplexAttr);
374
375impl Printable for ComplexAttr {
376    fn fmt(
377        &self,
378        ctx: &Context,
379        state: &printable::State,
380        f: &mut fmt::Formatter<'_>,
381    ) -> fmt::Result {
382        write!(f, "{}, ", self.ty.disp(ctx))?;
383        let float_ty = self.float_type(ctx);
384        float_ty.disp_value(self.re, ctx, state, f)?;
385        write!(f, ", ")?;
386        float_ty.disp_value(self.im, ctx, state, f)
387    }
388}
389
390impl Parsable for ComplexAttr {
391    type Arg = ();
392    type Parsed = Self;
393
394    fn parse<'a>(input: &mut StateStream<'a>, _: Self::Arg) -> ParseResult<'a, Self::Parsed> {
395        let ty = type_parse(input)?.0;
396        spaced(char::char(',')).parse_stream(input).into_result()?;
397        let ctx = dupe_ref(input.state.ctx);
398        let float_ty = Self::float_type_for(ctx, ty);
399        let re = float_ty.parse_value(input)?.0;
400        spaced(char::char(',')).parse_stream(input).into_result()?;
401        let im = float_ty.parse_value(input)?.0;
402        Ok(Self::new(ty, re, im)).into_parse_result()
403    }
404}
405
406impl ComplexAttr {
407    pub fn from_f64(ctx: &Context, ty: TypeHandle, re: f64, im: f64) -> Self {
408        let float_ty = Self::float_type_for(ctx, ty);
409        Self::new(ty, float_ty.value_from_f64(re), float_ty.value_from_f64(im))
410    }
411
412    fn float_type_for(ctx: &Context, ty: TypeHandle) -> Ref<'_, dyn APFloatType> {
413        let ty: TypeHandle = if ty.deref(ctx).is::<Complex32Type>() {
414            Float32Type::get(ctx).into()
415        } else if ty.deref(ctx).is::<Complex64Type>() {
416            Float64Type::get(ctx).into()
417        } else {
418            panic!("expected complex type")
419        };
420        Ref::map(ty.deref(ctx), |ty| try_cast_ty!(ty, ctx, dyn APFloatType))
421    }
422
423    pub fn float_type<'a>(&self, ctx: &'a Context) -> Ref<'a, dyn APFloatType> {
424        Self::float_type_for(ctx, self.ty)
425    }
426}
427
428#[attr_interface_impl]
429impl TypedAttrInterface for ComplexAttr {
430    fn get_type(&self, _ctx: &Context) -> TypeHandle {
431        self.ty
432    }
433}
434
435#[attr_interface_impl]
436impl ConstantAttr for ComplexAttr {
437    fn as_const_val(&self, ctx: &Context) -> ConstantValue {
438        let float_ty = self.float_type(ctx);
439        ConstantValue::Complex(
440            float_ty.value_to_f64(self.re),
441            float_ty.value_to_f64(self.im),
442        )
443    }
444}
445
446#[pliron_attr(name = "cube.dim3", format, verifier = "succ")]
447#[derive(new, From, PartialEq, Clone, Debug, Hash)]
448pub struct Dim3Attr(pub Dim3);
449
450#[attr_interface_impl]
451impl TypedAttrInterface for FloatAttr {
452    fn get_type(&self, _ctx: &Context) -> TypeHandle {
453        self.ty
454    }
455}
456
457#[attr_interface_impl]
458impl ConstantAttr for FloatAttr {
459    fn as_const_val(&self, ctx: &Context) -> ConstantValue {
460        let value = self.float_type(ctx).value_to_f64(self.val);
461        ConstantValue::Float(value)
462    }
463    fn float_as_f64(&self, ctx: &Context) -> Option<f64> {
464        let val = self.float_type(ctx).value_to_f64(self.val);
465        Some(val)
466    }
467}
468
469pub trait TypedLiteral {
470    fn is_same_type(ctx: &Context, ty: TypeHandle) -> bool;
471}
472
473macro_rules! literal {
474    ($ty: ty, $ir_ty: ty, $pred: expr) => {
475        impl TypedLiteral for $ty {
476            fn is_same_type(ctx: &Context, ty: TypeHandle) -> bool {
477                ty.deref(ctx).downcast_ref::<$ir_ty>().is_some_and($pred)
478            }
479        }
480    };
481    ($ty: ty, $ir_ty: ty) => {
482        literal!($ty, $ir_ty, |_| true);
483    };
484}
485
486literal!(usize, IndexType);
487
488literal!(i8, IntegerType, |it| it.width() == 8);
489literal!(i16, IntegerType, |it| it.width() == 16);
490literal!(i32, IntegerType, |it| it.width() == 32);
491literal!(i64, IntegerType, |it| it.width() == 64);
492
493literal!(u8, IntegerType, |it| it.width() == 8);
494literal!(u16, IntegerType, |it| it.width() == 16);
495literal!(u32, IntegerType, |it| it.width() == 32);
496literal!(u64, IntegerType, |it| it.width() == 64);
497
498literal!(half::f16, Float16Type);
499literal!(half::bf16, BFloat16Type);
500literal!(f32, Float32Type);
501literal!(f64, Float64Type);
502
503/// Symbol visibility
504#[pliron_attr(name = "cube.sym_visibility", format = "$0", verifier = "succ")]
505#[derive(new, PartialEq, Eq, Clone, Copy, Debug, Hash, From, Into)]
506pub struct SymbolVisibilityAttr(pub SymbolVisibility);
cubecl_ir/attributes/mod.rs

cubecl_ir/attributes/
mod.rs
