cubecl_ir/interfaces/
mod.rs

1use core::fmt::{self, Debug};
2
3use crate::{
4    CanMaterialize, ConstantValue, ElemType,
5    dialect::synchronization::SyncScope,
6    interfaces::memory_slot::MemoryValue,
7    prelude::*,
8    types::{AtomicType, PointerType, VectorType, scalar::*},
9};
10use alloc::string::ToString;
11use derive_new::new;
12use itertools::Itertools;
13use pliron::{
14    alloc::vec::Vec,
15    attribute::{AttrObj, AttributeDict},
16    builtin::{attr_interfaces::TypedAttrInterface, ops::ConstantOp, types::IntegerType},
17    context::Context,
18    derive::{op_interface, type_interface},
19    op::OpId,
20    printable::{self, Printable},
21    r#type::{TypeHandle, type_cast},
22    utils::apint::APInt,
23    value::Use,
24};
25
26pub mod aliasing;
27pub mod control_flow;
28pub mod memory_slot;
29pub mod side_effects;
30pub mod traits;
31pub mod uniformity;
32
33#[macro_export]
34macro_rules! verify_op_succ {
35    () => {
36        fn verify(
37            _op: &dyn pliron::op::Op,
38            _ctx: &pliron::context::Context,
39        ) -> pliron::result::Result<()>
40        where
41            Self: Sized,
42        {
43            Ok(())
44        }
45    };
46}
47
48#[macro_export]
49macro_rules! verify_ty_succ {
50    () => {
51        fn verify(
52            _op: &dyn pliron::r#type::Type,
53            _ctx: &pliron::context::Context,
54        ) -> pliron::result::Result<()>
55        where
56            Self: Sized,
57        {
58            Ok(())
59        }
60    };
61}
62
63#[macro_export]
64macro_rules! verify_attr_succ {
65    () => {
66        fn verify(
67            _op: &dyn pliron::attribute::Attribute,
68            _ctx: &pliron::context::Context,
69        ) -> pliron::result::Result<()>
70        where
71            Self: Sized,
72        {
73            Ok(())
74        }
75    };
76}
77
78// Pure marker
79#[op_interface]
80pub trait ReturnLike: IsTerminatorInterface + NResultsInterface<0> {
81    verify_op_succ!();
82}
83
84#[op_interface]
85pub trait TriviallyUnrollable: MaterializableOp {
86    verify_op_succ!();
87}
88
89/// Op that can be rematerialized from a set of operands.
90/// Should be implemented for anything that doesn't have regions or successors.
91#[op_interface]
92pub trait MaterializableOp {
93    verify_op_succ!();
94    fn materialize(
95        &self,
96        ctx: &mut Context,
97        result_ty: Vec<TypeHandle>,
98        operands: Vec<Value>,
99        attributes: AttributeDict,
100    ) -> Ptr<Operation>;
101}
102
103CanMaterialize!(ConstantOp);
104
105#[op_interface]
106pub trait Synchronizes {
107    verify_op_succ!();
108
109    /// Synchronizes at least at this scope. Should be used for optimizations where smaller scopes
110    /// are more conservative (i.e. marking shared memory as unused).
111    fn minimum_scope(&self, ctx: &Context) -> SyncScope;
112    /// Synchronizes at most at this scope. Should be used for optimizations where larger scopes
113    /// are more conservative (i.e. making memory writes visible).
114    fn maximum_scope(&self, ctx: &Context) -> SyncScope;
115}
116
117macro_rules! synchronizes {
118    ($ty: ty, $scope: expr) => {
119        #[::pliron::derive::op_interface_impl]
120        impl crate::interfaces::Synchronizes for $ty {
121            #[allow(unused_variables)]
122            fn minimum_scope(&self, ctx: &::pliron::context::Context) -> SyncScope {
123                $scope
124            }
125            #[allow(unused_variables)]
126            fn maximum_scope(&self, ctx: &::pliron::context::Context) -> SyncScope {
127                $scope
128            }
129        }
130    };
131}
132pub(crate) use synchronizes;
133
134#[type_interface]
135pub trait AlignedType {
136    verify_ty_succ!();
137
138    fn align(&self, ctx: &Context) -> usize;
139}
140
141#[macro_export]
142macro_rules! aligned {
143    ($ty: ty, $align: expr) => {
144        #[::pliron::derive::type_interface_impl]
145        impl $crate::interfaces::AlignedType for $ty {
146            #[allow(unused_variables)]
147            fn align(&self, ctx: &::pliron::context::Context) -> usize {
148                $align
149            }
150        }
151    };
152}
153
154#[type_interface]
155pub trait SizedType: AlignedType {
156    verify_ty_succ!();
157    fn size(&self, ctx: &Context) -> usize;
158    fn size_bits(&self, ctx: &Context) -> usize {
159        self.size(ctx) * 8
160    }
161}
162
163#[macro_export]
164macro_rules! sized {
165    ($ty: ty, $size: expr) => {
166        #[::pliron::derive::type_interface_impl]
167        impl $crate::interfaces::SizedType for $ty {
168            #[allow(unused_variables)]
169            fn size(&self, ctx: &::pliron::context::Context) -> usize {
170                $size
171            }
172        }
173    };
174}
175
176#[type_interface]
177pub trait MaybeVectorizedType {
178    verify_ty_succ!();
179
180    fn vector_size(&self, ctx: &Context) -> usize;
181    fn try_vector_size(&self, ctx: &Context) -> Option<usize> {
182        Some(self.vector_size(ctx))
183    }
184}
185
186#[macro_export]
187macro_rules! scalar {
188    ($ty: ty) => {
189        #[::pliron::derive::type_interface_impl]
190        impl $crate::interfaces::MaybeVectorizedType for $ty {
191            fn vector_size(&self, _ctx: &::pliron::context::Context) -> usize {
192                1
193            }
194        }
195
196        #[::pliron::derive::type_interface_impl]
197        impl $crate::interfaces::ScalarizableType for $ty {
198            fn scalar_type(&self, ctx: &Context) -> ::pliron::r#type::TypeHandle {
199                use ::pliron::r#type::Type;
200                self.get_self_handle(ctx)
201            }
202        }
203
204        #[::pliron::derive::type_interface_impl]
205        impl $crate::interfaces::HasElementType for $ty {
206            fn element_type(&self, ctx: &Context) -> Option<::pliron::r#type::TypeHandle> {
207                use ::pliron::r#type::Type;
208                Some(self.get_self_handle(ctx))
209            }
210        }
211    };
212}
213
214#[type_interface]
215pub trait MaybePackedType {
216    verify_ty_succ!();
217
218    fn packing_factor(&self, ctx: &Context) -> usize;
219}
220
221macro_rules! not_packed {
222    ($ty: ty) => {
223        #[::pliron::derive::type_interface_impl]
224        impl crate::interfaces::MaybePackedType for $ty {
225            fn packing_factor(&self, _ctx: &::pliron::context::Context) -> usize {
226                1
227            }
228        }
229    };
230}
231pub(crate) use not_packed;
232
233#[type_interface]
234pub trait ScalarizableType {
235    verify_ty_succ!();
236    fn scalar_type(&self, ctx: &Context) -> TypeHandle;
237}
238
239#[type_interface]
240pub trait ScalarType {
241    verify_ty_succ!();
242    fn elem_type(&self, ctx: &Context) -> ElemType;
243}
244
245#[type_interface]
246pub trait IndexableType {
247    verify_ty_succ!();
248
249    fn indexed_type(&self, ctx: &Context) -> TypeHandle;
250}
251
252#[type_interface]
253pub trait HasElementType {
254    verify_ty_succ!();
255    fn element_type(&self, ctx: &Context) -> Option<TypeHandle>;
256}
257
258#[op_interface]
259pub trait SimplifyInterface {
260    verify_op_succ!();
261    fn check_fold(&self, ctx: &Context, operand_attrs: &[Option<AttrObj>]) -> Option<Value>;
262}
263
264#[op_interface]
265pub trait CanonicalizeInterface {
266    verify_op_succ!();
267    fn canonicalize(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()>;
268}
269
270/// An op of a dialect the backends don't know, such as the tile DSL's. Every backend runs
271/// [`LowerOpsPass`](crate::rewrite::LowerOpsPass) ahead of its own passes, which lowers it to ops
272/// the backend knows.
273#[op_interface]
274pub trait LowerInterface {
275    verify_op_succ!();
276    fn lower(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()>;
277}
278
279#[derive(new, Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
280pub struct ExpressionValue(u64);
281
282impl Printable for ExpressionValue {
283    fn fmt(&self, _: &Context, _: &printable::State, f: &mut fmt::Formatter<'_>) -> fmt::Result {
284        write!(f, "e{}", self.0)
285    }
286}
287
288#[derive(new, Clone, PartialEq, Eq, Hash)]
289pub struct Expression {
290    pub result_type: TypeHandle,
291    #[new(default)]
292    pub result_idx: usize,
293    pub op_id: OpId,
294    pub operands: Vec<ExpressionValue>,
295    pub attributes: AttributeDict,
296    #[new(default)]
297    pub mem_value: Option<MemoryValue>,
298}
299
300impl Printable for Expression {
301    fn fmt(&self, ctx: &Context, _: &printable::State, f: &mut fmt::Formatter<'_>) -> fmt::Result {
302        let mut opds = self.operands.iter().map(|it| it.disp(ctx).to_string());
303        write!(
304            f,
305            "result{} = {} ({}{}) {} : {}",
306            self.result_idx,
307            self.op_id,
308            opds.join(", "),
309            self.mem_value
310                .map(|it| alloc::format!(", {it}"))
311                .unwrap_or_default(),
312            self.attributes.disp(ctx),
313            self.result_type.disp(ctx)
314        )
315    }
316}
317
318impl Debug for Expression {
319    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
320        f.debug_struct("Expression")
321            .field("result_type", &self.result_type)
322            .field("result_idx", &self.result_idx)
323            .field("op_id", &self.op_id.to_string())
324            .field("operands", &self.operands)
325            .field("attributes", &self.attributes)
326            .field("mem_value", &self.mem_value)
327            .finish()
328    }
329}
330
331#[op_interface]
332pub trait ExpressionCanonicalize: OneResultInterface {
333    verify_op_succ!();
334    fn canonical_expression(&self, ctx: &Context, operands: Vec<ExpressionValue>) -> Expression;
335}
336
337#[attr_interface]
338pub trait ConstantAttr: TypedAttrInterface {
339    verify_attr_succ!();
340    fn as_const_val(&self, ctx: &Context) -> ConstantValue;
341    fn as_int(&self, ctx: &Context) -> Option<APInt> {
342        let _ = ctx;
343        None
344    }
345    fn float_as_f64(&self, _ctx: &Context) -> Option<f64> {
346        None
347    }
348    /// The value of a scalar bool constant, without the panic `as_const_val` has for a zero of an
349    /// unsupported type such as a vector.
350    fn as_bool(&self, _ctx: &Context) -> Option<bool> {
351        None
352    }
353}
354
355#[macro_export]
356macro_rules! try_cast_ty {
357    ($ty: expr, $ctx: expr, $interface: ty) => {
358        $crate::prelude::type_cast::<$interface>(&*$ty)
359            .ok_or_else(|| {
360                $crate::alloc::format!(
361                    "Expected type {} {} to implement {}",
362                    $ty.get_type_id(),
363                    $ty.disp($ctx),
364                    stringify!($interface)
365                )
366            })
367            .unwrap()
368    };
369}
370
371#[macro_export]
372macro_rules! try_cast_op {
373    ($op: expr, $ctx: expr, $interface: ty) => {
374        $crate::prelude::op_cast::<$interface>(&*$op)
375            .ok_or_else(|| {
376                $crate::alloc::format!(
377                    "Expected op {} {} to implement {}",
378                    $op.get_opid(),
379                    $op.disp($ctx),
380                    stringify!($interface)
381                )
382            })
383            .unwrap()
384    };
385}
386
387#[macro_export]
388macro_rules! match_ty {
389    (($handle: expr) { $($ty: ty => $body: expr,)*; _ => $default: expr }) => {
390        (|| {
391            $(if $handle.is::<$ty>() {
392                return $body;
393            })*
394            $default
395        })()
396    };
397    (($handle: expr) { $($ty: ty => $body: expr,)* }) => {
398        (|| {
399            $(if $handle.is::<$ty>() {
400                return $body;
401            })*
402            unreachable!()
403        })()
404    };
405}
406
407pub trait TypedExt: Typed {
408    fn size(&self, ctx: &Context) -> usize {
409        let ty = self.get_type(ctx).deref(ctx);
410        let sized = try_cast_ty!(ty, ctx, dyn SizedType);
411        sized.size(ctx)
412    }
413
414    fn size_bits(&self, ctx: &Context) -> usize {
415        let ty = self.get_type(ctx).deref(ctx);
416        let sized = try_cast_ty!(ty, ctx, dyn SizedType);
417        sized.size_bits(ctx)
418    }
419
420    fn unpacked_size_bits(&self, ctx: &Context) -> usize {
421        self.element_ty(ctx).scalar_ty(ctx).size_bits(ctx) / self.packing_factor(ctx)
422    }
423
424    fn align(&self, ctx: &Context) -> usize {
425        let ty = self.get_type(ctx).deref(ctx);
426        let aligned = try_cast_ty!(ty, ctx, dyn AlignedType);
427        aligned.align(ctx)
428    }
429
430    fn is_ptr(&self, ctx: &Context) -> bool {
431        let ty = self.get_type(ctx).deref(ctx);
432        ty.is::<PointerType>()
433    }
434
435    fn is_atomic(&self, ctx: &Context) -> bool {
436        let ty = self.get_type(ctx).deref(ctx);
437        ty.is::<AtomicType>()
438    }
439
440    fn is_vector(&self, ctx: &Context) -> bool {
441        let ty = self.get_type(ctx).deref(ctx);
442        ty.is::<VectorType>()
443    }
444
445    fn is_vector_of_size(&self, ctx: &Context, size: usize) -> bool {
446        let ty = self.get_type(ctx).deref(ctx);
447        ty.downcast_ref::<VectorType>()
448            .is_some_and(|it| it.vectorization == size)
449    }
450
451    fn is_immutable(&self, ctx: &Context) -> bool {
452        !self.is_ptr(ctx)
453    }
454
455    fn vector_size(&self, ctx: &Context) -> usize {
456        let ty = self.get_type(ctx).deref(ctx);
457        let maybe_vec = try_cast_ty!(ty, ctx, dyn MaybeVectorizedType);
458        maybe_vec.vector_size(ctx)
459    }
460
461    fn try_get_vector_size(&self, ctx: &Context) -> Option<usize> {
462        let ty = self.get_type(ctx).deref(ctx);
463        let maybe_vec = type_cast::<dyn MaybeVectorizedType>(&*ty)?;
464        maybe_vec.try_vector_size(ctx)
465    }
466
467    fn packing_factor(&self, ctx: &Context) -> usize {
468        let ty = self.get_type(ctx).deref(ctx);
469        let maybe_packed = try_cast_ty!(ty, ctx, dyn MaybePackedType);
470        maybe_packed.packing_factor(ctx)
471    }
472
473    fn scalar_ty(&self, ctx: &Context) -> TypeHandle {
474        let ty = self.element_ty(ctx).deref(ctx);
475        let scalarizable = try_cast_ty!(ty, ctx, dyn ScalarizableType);
476        scalarizable.scalar_type(ctx)
477    }
478
479    fn element_ty(&self, ctx: &Context) -> TypeHandle {
480        let ty = self.get_type(ctx).deref(ctx);
481        let has_element_type = try_cast_ty!(ty, ctx, dyn HasElementType);
482        has_element_type
483            .element_type(ctx)
484            .expect("Expected element type to be some")
485    }
486
487    fn unwrap_ptr(&self, ctx: &Context) -> TypeHandle {
488        if let Some(ptr) = self.get_type(ctx).deref(ctx).downcast_ref::<PointerType>() {
489            ptr.inner
490        } else {
491            self.get_type(ctx)
492        }
493    }
494
495    fn try_get_scalar_ty(&self, ctx: &Context) -> Option<TypeHandle> {
496        let ty = self.get_type(ctx).deref(ctx);
497        let scalarizable = type_cast::<dyn ScalarizableType>(&*ty)?;
498        Some(scalarizable.scalar_type(ctx))
499    }
500
501    fn try_get_scalar_elem_ty(&self, ctx: &Context) -> Option<TypeHandle> {
502        let ty = self.get_type(ctx).deref(ctx);
503        let has_elem = type_cast::<dyn HasElementType>(&*ty)?;
504        let ty = has_elem.element_type(ctx)?.deref(ctx);
505        let scalarizable = type_cast::<dyn ScalarizableType>(&*ty)?;
506        Some(scalarizable.scalar_type(ctx))
507    }
508
509    fn is_index(&self, ctx: &Context) -> bool {
510        let ty = self.get_type(ctx).deref(ctx);
511        ty.is::<IndexType>()
512    }
513
514    fn is_int(&self, ctx: &Context) -> bool {
515        let ty = self.get_type(ctx).deref(ctx);
516        ty.is::<IntegerType>()
517    }
518
519    fn is_signed_int(&self, ctx: &Context) -> bool {
520        let ty = self.get_type(ctx).deref(ctx);
521        ty.downcast_ref::<IntegerType>()
522            .is_some_and(|it| it.is_signed())
523    }
524
525    fn is_unsigned_int(&self, ctx: &Context) -> bool {
526        let ty = self.get_type(ctx).deref(ctx);
527        ty.downcast_ref::<IntegerType>()
528            .is_some_and(|it| !it.is_signed())
529    }
530
531    fn is_int_of_width(&self, ctx: &Context, width: usize) -> bool {
532        let ty = self.get_type(ctx).deref(ctx);
533        ty.downcast_ref::<IntegerType>()
534            .is_some_and(|it| it.width() as usize == width)
535    }
536
537    fn is_float64(&self, ctx: &Context) -> bool {
538        self.get_type(ctx).deref(ctx).is::<Float64Type>()
539    }
540
541    fn is_float32(&self, ctx: &Context) -> bool {
542        self.get_type(ctx).deref(ctx).is::<Float32Type>()
543    }
544
545    fn is_tfloat32(&self, ctx: &Context) -> bool {
546        self.get_type(ctx).deref(ctx).is::<TFloat32Type>()
547    }
548
549    fn is_float16(&self, ctx: &Context) -> bool {
550        self.get_type(ctx).deref(ctx).is::<Float16Type>()
551    }
552
553    fn is_bfloat16(&self, ctx: &Context) -> bool {
554        self.get_type(ctx).deref(ctx).is::<BFloat16Type>()
555    }
556
557    fn is_float(&self, ctx: &Context) -> bool {
558        self.is_float16(ctx)
559            | self.is_float32(ctx)
560            | self.is_float64(ctx)
561            | self.is_tfloat32(ctx)
562            | self.is_bfloat16(ctx)
563    }
564
565    fn is_bool(&self, ctx: &Context) -> bool {
566        self.get_type(ctx).deref(ctx).is::<BoolType>()
567    }
568}
569
570impl<T: Typed> TypedExt for T {}
571
572pub trait TypeExt {
573    fn as_ptr(&self, ctx: &Context) -> PointerType;
574}
575
576impl TypeExt for TypeHandle {
577    fn as_ptr(&self, ctx: &Context) -> PointerType {
578        *TypedHandle::from_handle(*self, ctx)
579            .expect("Should be pointer")
580            .deref(ctx)
581    }
582}
583
584pub trait ValueExt {
585    fn replace_all_uses_except_with(&self, ctx: &Context, except: Use<Value>, other: &Value);
586}
587
588impl ValueExt for Value {
589    fn replace_all_uses_except_with(&self, ctx: &Context, except: Use<Value>, other: &Value) {
590        self.replace_some_uses_with(ctx, |_, r#use| r#use != &except, other);
591    }
592}
cubecl_ir/interfaces/mod.rs

cubecl_ir/interfaces/
mod.rs
