cubecl_ir/interfaces/
mod.rs

1use crate::{
2    AddressSpace, CanMaterialize, ConstantValue, ElemType, NoMemoryEffect,
3    dialect::synchronization::SyncScope,
4    prelude::*,
5    types::{AtomicType, PointerType, VectorType, scalar::*},
6};
7use pliron::{
8    alloc::vec::Vec,
9    attribute::{AttrObj, AttributeDict},
10    builtin::{attr_interfaces::TypedAttrInterface, ops::ConstantOp, types::IntegerType},
11    context::Context,
12    derive::{op_interface, type_interface},
13    opts::dce::SideEffects,
14    printable::Printable,
15    r#type::{TypeHandle, type_cast},
16    utils::apint::APInt,
17    value::Use,
18};
19
20pub mod aliasing;
21pub mod control_flow;
22pub mod memory_slot;
23pub mod traits;
24pub mod uniformity;
25
26#[macro_export]
27macro_rules! verify_op_succ {
28    () => {
29        fn verify(
30            _op: &dyn pliron::op::Op,
31            _ctx: &pliron::context::Context,
32        ) -> pliron::result::Result<()>
33        where
34            Self: Sized,
35        {
36            Ok(())
37        }
38    };
39}
40
41#[macro_export]
42macro_rules! verify_ty_succ {
43    () => {
44        fn verify(
45            _op: &dyn pliron::r#type::Type,
46            _ctx: &pliron::context::Context,
47        ) -> pliron::result::Result<()>
48        where
49            Self: Sized,
50        {
51            Ok(())
52        }
53    };
54}
55
56#[macro_export]
57macro_rules! verify_attr_succ {
58    () => {
59        fn verify(
60            _op: &dyn pliron::attribute::Attribute,
61            _ctx: &pliron::context::Context,
62        ) -> pliron::result::Result<()>
63        where
64            Self: Sized,
65        {
66            Ok(())
67        }
68    };
69}
70
71// Pure marker
72#[op_interface]
73pub trait ReturnLike: IsTerminatorInterface + NResultsInterface<0> {
74    verify_op_succ!();
75}
76
77#[op_interface]
78pub trait TriviallyUnrollable: MaterializableOp {
79    verify_op_succ!();
80}
81
82/// Op that can be rematerialized from a set of operands.
83/// Should be implemented for anything that doesn't have regions or successors.
84#[op_interface]
85pub trait MaterializableOp {
86    verify_op_succ!();
87    fn materialize(
88        &self,
89        ctx: &mut Context,
90        result_ty: Vec<TypeHandle>,
91        operands: Vec<Value>,
92        attributes: AttributeDict,
93    ) -> Ptr<Operation>;
94}
95
96CanMaterialize!(ConstantOp);
97
98#[op_interface]
99pub trait Synchronizes: SideEffects {
100    verify_op_succ!();
101
102    /// Synchronizes at least at this scope. Should be used for optimizations where smaller scopes
103    /// are more conservative (i.e. marking shared memory as unused).
104    fn minimum_scope(&self, ctx: &Context) -> SyncScope;
105    /// Synchronizes at most at this scope. Should be used for optimizations where larger scopes
106    /// are more conservative (i.e. making memory writes visible).
107    fn maximum_scope(&self, ctx: &Context) -> SyncScope;
108}
109
110macro_rules! synchronizes {
111    ($ty: ty, $scope: expr) => {
112        #[::pliron::derive::op_interface_impl]
113        impl crate::interfaces::Synchronizes for $ty {
114            #[allow(unused_variables)]
115            fn minimum_scope(&self, ctx: &::pliron::context::Context) -> SyncScope {
116                $scope
117            }
118            #[allow(unused_variables)]
119            fn maximum_scope(&self, ctx: &::pliron::context::Context) -> SyncScope {
120                $scope
121            }
122        }
123        #[pliron::derive::op_interface_impl]
124        impl pliron::opts::dce::SideEffects for $ty {
125            fn has_side_effects(&self, _ctx: &Context) -> bool {
126                true
127            }
128        }
129    };
130}
131pub(crate) use synchronizes;
132
133#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
134pub enum MemoryEffect {
135    Read(Value),
136    Write(Value),
137    ReadAllInSpace(AddressSpace),
138    WriteAllInSpace(AddressSpace),
139    ReadAll,
140    WriteAll,
141    // Not analyzable, clobber the entire state
142    Opaque,
143}
144
145impl MemoryEffect {
146    pub fn value(&self) -> Option<Value> {
147        match self {
148            MemoryEffect::Read(value) | MemoryEffect::Write(value) => Some(*value),
149            MemoryEffect::ReadAllInSpace(_)
150            | MemoryEffect::WriteAllInSpace(_)
151            | MemoryEffect::ReadAll
152            | MemoryEffect::WriteAll
153            | MemoryEffect::Opaque => None,
154        }
155    }
156}
157
158impl Printable for MemoryEffect {
159    fn fmt(
160        &self,
161        ctx: &Context,
162        _state: &pliron::printable::State,
163        f: &mut core::fmt::Formatter<'_>,
164    ) -> core::fmt::Result {
165        match self {
166            MemoryEffect::Read(value) => write!(f, "Read({})", value.disp(ctx)),
167            MemoryEffect::Write(value) => write!(f, "Write({})", value.disp(ctx)),
168            MemoryEffect::ReadAllInSpace(address_space) => {
169                write!(f, "ReadAllInSpace({})", address_space.disp(ctx))
170            }
171            MemoryEffect::WriteAllInSpace(address_space) => {
172                write!(f, "WriteAllInSpace({})", address_space.disp(ctx))
173            }
174            MemoryEffect::ReadAll => write!(f, "ReadAll"),
175            MemoryEffect::WriteAll => write!(f, "WriteAll"),
176            MemoryEffect::Opaque => write!(f, "Opaque"),
177        }
178    }
179}
180
181#[op_interface]
182pub trait MemoryEffects {
183    verify_op_succ!();
184    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect>;
185    fn has_effects(&self, ctx: &Context) -> bool {
186        !self.memory_effects(ctx).is_empty()
187    }
188}
189
190NoMemoryEffect!(ConstantOp);
191
192#[type_interface]
193pub trait AlignedType {
194    verify_ty_succ!();
195
196    fn align(&self, ctx: &Context) -> usize;
197}
198
199#[macro_export]
200macro_rules! aligned {
201    ($ty: ty, $align: expr) => {
202        #[::pliron::derive::type_interface_impl]
203        impl $crate::interfaces::AlignedType for $ty {
204            #[allow(unused_variables)]
205            fn align(&self, ctx: &::pliron::context::Context) -> usize {
206                $align
207            }
208        }
209    };
210}
211
212#[type_interface]
213pub trait SizedType: AlignedType {
214    verify_ty_succ!();
215    fn size(&self, ctx: &Context) -> usize;
216    fn size_bits(&self, ctx: &Context) -> usize {
217        self.size(ctx) * 8
218    }
219}
220
221#[macro_export]
222macro_rules! sized {
223    ($ty: ty, $size: expr) => {
224        #[::pliron::derive::type_interface_impl]
225        impl $crate::interfaces::SizedType for $ty {
226            #[allow(unused_variables)]
227            fn size(&self, ctx: &::pliron::context::Context) -> usize {
228                $size
229            }
230        }
231    };
232}
233
234#[type_interface]
235pub trait MaybeVectorizedType {
236    verify_ty_succ!();
237
238    fn vector_size(&self, ctx: &Context) -> usize;
239    fn try_vector_size(&self, ctx: &Context) -> Option<usize> {
240        Some(self.vector_size(ctx))
241    }
242}
243
244#[macro_export]
245macro_rules! scalar {
246    ($ty: ty) => {
247        #[::pliron::derive::type_interface_impl]
248        impl $crate::interfaces::MaybeVectorizedType for $ty {
249            fn vector_size(&self, _ctx: &::pliron::context::Context) -> usize {
250                1
251            }
252        }
253
254        #[::pliron::derive::type_interface_impl]
255        impl $crate::interfaces::ScalarizableType for $ty {
256            fn scalar_type(&self, ctx: &Context) -> ::pliron::r#type::TypeHandle {
257                use ::pliron::r#type::Type;
258                self.get_self_handle(ctx)
259            }
260        }
261
262        #[::pliron::derive::type_interface_impl]
263        impl $crate::interfaces::HasElementType for $ty {
264            fn element_type(&self, ctx: &Context) -> Option<::pliron::r#type::TypeHandle> {
265                use ::pliron::r#type::Type;
266                Some(self.get_self_handle(ctx))
267            }
268        }
269    };
270}
271
272#[type_interface]
273pub trait MaybePackedType {
274    verify_ty_succ!();
275
276    fn packing_factor(&self, ctx: &Context) -> usize;
277}
278
279macro_rules! not_packed {
280    ($ty: ty) => {
281        #[::pliron::derive::type_interface_impl]
282        impl crate::interfaces::MaybePackedType for $ty {
283            fn packing_factor(&self, _ctx: &::pliron::context::Context) -> usize {
284                1
285            }
286        }
287    };
288}
289pub(crate) use not_packed;
290
291#[type_interface]
292pub trait ScalarizableType {
293    verify_ty_succ!();
294    fn scalar_type(&self, ctx: &Context) -> TypeHandle;
295}
296
297#[type_interface]
298pub trait ScalarType {
299    verify_ty_succ!();
300    fn elem_type(&self, ctx: &Context) -> ElemType;
301}
302
303#[type_interface]
304pub trait IndexableType {
305    verify_ty_succ!();
306
307    fn indexed_type(&self, ctx: &Context) -> TypeHandle;
308}
309
310#[type_interface]
311pub trait HasElementType {
312    verify_ty_succ!();
313    fn element_type(&self, ctx: &Context) -> Option<TypeHandle>;
314}
315
316#[op_interface]
317pub trait SimplifyInterface {
318    verify_op_succ!();
319    fn check_fold(&self, ctx: &Context, operand_attrs: &[Option<AttrObj>]) -> Option<Value>;
320}
321
322#[op_interface]
323pub trait CanonicalizeInterface {
324    verify_op_succ!();
325    fn canonicalize(&self, ctx: &mut Context, rewriter: &mut MatchRewriter) -> Result<()>;
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
