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#[derive(new, Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
271pub struct ExpressionValue(u64);
272
273impl Printable for ExpressionValue {
274    fn fmt(&self, _: &Context, _: &printable::State, f: &mut fmt::Formatter<'_>) -> fmt::Result {
275        write!(f, "e{}", self.0)
276    }
277}
278
279#[derive(new, Clone, PartialEq, Eq, Hash)]
280pub struct Expression {
281    pub result_type: TypeHandle,
282    #[new(default)]
283    pub result_idx: usize,
284    pub op_id: OpId,
285    pub operands: Vec<ExpressionValue>,
286    pub attributes: AttributeDict,
287    #[new(default)]
288    pub mem_value: Option<MemoryValue>,
289}
290
291impl Printable for Expression {
292    fn fmt(&self, ctx: &Context, _: &printable::State, f: &mut fmt::Formatter<'_>) -> fmt::Result {
293        let mut opds = self.operands.iter().map(|it| it.disp(ctx).to_string());
294        write!(
295            f,
296            "result{} = {} ({}{}) {} : {}",
297            self.result_idx,
298            self.op_id,
299            opds.join(", "),
300            self.mem_value
301                .map(|it| alloc::format!(", {it}"))
302                .unwrap_or_default(),
303            self.attributes.disp(ctx),
304            self.result_type.disp(ctx)
305        )
306    }
307}
308
309impl Debug for Expression {
310    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
311        f.debug_struct("Expression")
312            .field("result_type", &self.result_type)
313            .field("result_idx", &self.result_idx)
314            .field("op_id", &self.op_id.to_string())
315            .field("operands", &self.operands)
316            .field("attributes", &self.attributes)
317            .field("mem_value", &self.mem_value)
318            .finish()
319    }
320}
321
322#[op_interface]
323pub trait ExpressionCanonicalize: OneResultInterface {
324    verify_op_succ!();
325    fn canonical_expression(&self, ctx: &Context, operands: Vec<ExpressionValue>) -> Expression;
326}
327
328#[attr_interface]
329pub trait ConstantAttr: TypedAttrInterface {
330    verify_attr_succ!();
331    fn as_const_val(&self, ctx: &Context) -> ConstantValue;
332    fn as_int(&self, ctx: &Context) -> Option<APInt> {
333        let _ = ctx;
334        None
335    }
336    fn float_as_f64(&self, _ctx: &Context) -> Option<f64> {
337        None
338    }
339}
340
341#[macro_export]
342macro_rules! try_cast_ty {
343    ($ty: expr, $ctx: expr, $interface: ty) => {
344        $crate::prelude::type_cast::<$interface>(&*$ty)
345            .ok_or_else(|| {
346                $crate::alloc::format!(
347                    "Expected type {} {} to implement {}",
348                    $ty.get_type_id(),
349                    $ty.disp($ctx),
350                    stringify!($interface)
351                )
352            })
353            .unwrap()
354    };
355}
356
357#[macro_export]
358macro_rules! try_cast_op {
359    ($op: expr, $ctx: expr, $interface: ty) => {
360        $crate::prelude::op_cast::<$interface>(&*$op)
361            .ok_or_else(|| {
362                $crate::alloc::format!(
363                    "Expected op {} {} to implement {}",
364                    $op.get_opid(),
365                    $op.disp($ctx),
366                    stringify!($interface)
367                )
368            })
369            .unwrap()
370    };
371}
372
373#[macro_export]
374macro_rules! match_ty {
375    (($handle: expr) { $($ty: ty => $body: expr,)*; _ => $default: expr }) => {
376        (|| {
377            $(if $handle.is::<$ty>() {
378                return $body;
379            })*
380            $default
381        })()
382    };
383    (($handle: expr) { $($ty: ty => $body: expr,)* }) => {
384        (|| {
385            $(if $handle.is::<$ty>() {
386                return $body;
387            })*
388            unreachable!()
389        })()
390    };
391}
392
393pub trait TypedExt: Typed {
394    fn size(&self, ctx: &Context) -> usize {
395        let ty = self.get_type(ctx).deref(ctx);
396        let sized = try_cast_ty!(ty, ctx, dyn SizedType);
397        sized.size(ctx)
398    }
399
400    fn size_bits(&self, ctx: &Context) -> usize {
401        let ty = self.get_type(ctx).deref(ctx);
402        let sized = try_cast_ty!(ty, ctx, dyn SizedType);
403        sized.size_bits(ctx)
404    }
405
406    fn unpacked_size_bits(&self, ctx: &Context) -> usize {
407        self.element_ty(ctx).scalar_ty(ctx).size_bits(ctx) / self.packing_factor(ctx)
408    }
409
410    fn align(&self, ctx: &Context) -> usize {
411        let ty = self.get_type(ctx).deref(ctx);
412        let aligned = try_cast_ty!(ty, ctx, dyn AlignedType);
413        aligned.align(ctx)
414    }
415
416    fn is_ptr(&self, ctx: &Context) -> bool {
417        let ty = self.get_type(ctx).deref(ctx);
418        ty.is::<PointerType>()
419    }
420
421    fn is_atomic(&self, ctx: &Context) -> bool {
422        let ty = self.get_type(ctx).deref(ctx);
423        ty.is::<AtomicType>()
424    }
425
426    fn is_vector(&self, ctx: &Context) -> bool {
427        let ty = self.get_type(ctx).deref(ctx);
428        ty.is::<VectorType>()
429    }
430
431    fn is_vector_of_size(&self, ctx: &Context, size: usize) -> bool {
432        let ty = self.get_type(ctx).deref(ctx);
433        ty.downcast_ref::<VectorType>()
434            .is_some_and(|it| it.vectorization == size)
435    }
436
437    fn is_immutable(&self, ctx: &Context) -> bool {
438        !self.is_ptr(ctx)
439    }
440
441    fn vector_size(&self, ctx: &Context) -> usize {
442        let ty = self.get_type(ctx).deref(ctx);
443        let maybe_vec = try_cast_ty!(ty, ctx, dyn MaybeVectorizedType);
444        maybe_vec.vector_size(ctx)
445    }
446
447    fn try_get_vector_size(&self, ctx: &Context) -> Option<usize> {
448        let ty = self.get_type(ctx).deref(ctx);
449        let maybe_vec = type_cast::<dyn MaybeVectorizedType>(&*ty)?;
450        maybe_vec.try_vector_size(ctx)
451    }
452
453    fn packing_factor(&self, ctx: &Context) -> usize {
454        let ty = self.get_type(ctx).deref(ctx);
455        let maybe_packed = try_cast_ty!(ty, ctx, dyn MaybePackedType);
456        maybe_packed.packing_factor(ctx)
457    }
458
459    fn scalar_ty(&self, ctx: &Context) -> TypeHandle {
460        let ty = self.element_ty(ctx).deref(ctx);
461        let scalarizable = try_cast_ty!(ty, ctx, dyn ScalarizableType);
462        scalarizable.scalar_type(ctx)
463    }
464
465    fn element_ty(&self, ctx: &Context) -> TypeHandle {
466        let ty = self.get_type(ctx).deref(ctx);
467        let has_element_type = try_cast_ty!(ty, ctx, dyn HasElementType);
468        has_element_type
469            .element_type(ctx)
470            .expect("Expected element type to be some")
471    }
472
473    fn unwrap_ptr(&self, ctx: &Context) -> TypeHandle {
474        if let Some(ptr) = self.get_type(ctx).deref(ctx).downcast_ref::<PointerType>() {
475            ptr.inner
476        } else {
477            self.get_type(ctx)
478        }
479    }
480
481    fn try_get_scalar_ty(&self, ctx: &Context) -> Option<TypeHandle> {
482        let ty = self.get_type(ctx).deref(ctx);
483        let scalarizable = type_cast::<dyn ScalarizableType>(&*ty)?;
484        Some(scalarizable.scalar_type(ctx))
485    }
486
487    fn try_get_scalar_elem_ty(&self, ctx: &Context) -> Option<TypeHandle> {
488        let ty = self.get_type(ctx).deref(ctx);
489        let has_elem = type_cast::<dyn HasElementType>(&*ty)?;
490        let ty = has_elem.element_type(ctx)?.deref(ctx);
491        let scalarizable = type_cast::<dyn ScalarizableType>(&*ty)?;
492        Some(scalarizable.scalar_type(ctx))
493    }
494
495    fn is_index(&self, ctx: &Context) -> bool {
496        let ty = self.get_type(ctx).deref(ctx);
497        ty.is::<IndexType>()
498    }
499
500    fn is_int(&self, ctx: &Context) -> bool {
501        let ty = self.get_type(ctx).deref(ctx);
502        ty.is::<IntegerType>()
503    }
504
505    fn is_signed_int(&self, ctx: &Context) -> bool {
506        let ty = self.get_type(ctx).deref(ctx);
507        ty.downcast_ref::<IntegerType>()
508            .is_some_and(|it| it.is_signed())
509    }
510
511    fn is_unsigned_int(&self, ctx: &Context) -> bool {
512        let ty = self.get_type(ctx).deref(ctx);
513        ty.downcast_ref::<IntegerType>()
514            .is_some_and(|it| !it.is_signed())
515    }
516
517    fn is_int_of_width(&self, ctx: &Context, width: usize) -> bool {
518        let ty = self.get_type(ctx).deref(ctx);
519        ty.downcast_ref::<IntegerType>()
520            .is_some_and(|it| it.width() as usize == width)
521    }
522
523    fn is_float64(&self, ctx: &Context) -> bool {
524        self.get_type(ctx).deref(ctx).is::<Float64Type>()
525    }
526
527    fn is_float32(&self, ctx: &Context) -> bool {
528        self.get_type(ctx).deref(ctx).is::<Float32Type>()
529    }
530
531    fn is_tfloat32(&self, ctx: &Context) -> bool {
532        self.get_type(ctx).deref(ctx).is::<TFloat32Type>()
533    }
534
535    fn is_float16(&self, ctx: &Context) -> bool {
536        self.get_type(ctx).deref(ctx).is::<Float16Type>()
537    }
538
539    fn is_bfloat16(&self, ctx: &Context) -> bool {
540        self.get_type(ctx).deref(ctx).is::<BFloat16Type>()
541    }
542
543    fn is_float(&self, ctx: &Context) -> bool {
544        self.is_float16(ctx)
545            | self.is_float32(ctx)
546            | self.is_float64(ctx)
547            | self.is_tfloat32(ctx)
548            | self.is_bfloat16(ctx)
549    }
550
551    fn is_bool(&self, ctx: &Context) -> bool {
552        self.get_type(ctx).deref(ctx).is::<BoolType>()
553    }
554}
555
556impl<T: Typed> TypedExt for T {}
557
558pub trait TypeExt {
559    fn as_ptr(&self, ctx: &Context) -> PointerType;
560}
561
562impl TypeExt for TypeHandle {
563    fn as_ptr(&self, ctx: &Context) -> PointerType {
564        *TypedHandle::from_handle(*self, ctx)
565            .expect("Should be pointer")
566            .deref(ctx)
567    }
568}
569
570pub trait ValueExt {
571    fn replace_all_uses_except_with(&self, ctx: &Context, except: Use<Value>, other: &Value);
572}
573
574impl ValueExt for Value {
575    fn replace_all_uses_except_with(&self, ctx: &Context, except: Use<Value>, other: &Value) {
576        self.replace_some_uses_with(ctx, |_, r#use| r#use != &except, other);
577    }
578}
cubecl_ir/interfaces/mod.rs

cubecl_ir/interfaces/
mod.rs
