cubecl_ir/dialect/
base.rs

1use core::cell::Ref;
2
3use crate::{
4    prelude::*,
5    types::{AtomicType, PointerType},
6};
7
8#[allow(clippy::wrong_self_convention)]
9pub trait OperationPtrExt: Sized {
10    fn is_op<T: Op>(self, ctx: &Context) -> bool {
11        self.as_op::<T>(ctx).is_some()
12    }
13    fn is_terminator(self, ctx: &Context) -> bool {
14        self.impls::<dyn IsTerminatorInterface>(ctx)
15    }
16    fn impls<T: ?Sized + OpInterfaceMarker + 'static>(self, ctx: &Context) -> bool {
17        op_impls::<T>(&*self.dyn_op(ctx))
18    }
19    fn cast<T: ?Sized + OpInterfaceMarker + 'static>(self, ctx: &Context) -> Option<TraitOp<T>>;
20    fn as_op<T: Op>(self, ctx: &Context) -> Option<T>;
21    fn dyn_op(self, ctx: &Context) -> OpObj;
22    fn operand(self, ctx: &Context, idx: usize) -> Value;
23    fn operand_as_use(self, ctx: &Context, idx: usize) -> Use<Value>;
24    fn operands(self, ctx: &Context) -> Vec<Value>;
25    fn operands_as_uses(self, ctx: &Context) -> Vec<Use<Value>>;
26    fn result(self, ctx: &Context) -> Value;
27    fn results(self, ctx: &Context) -> Vec<Value>;
28    fn result_types(self, ctx: &Context) -> Vec<TypeHandle>;
29    fn regions(self, ctx: &Context) -> Vec<Ptr<Region>>;
30    fn result_names(self, ctx: &Context) -> Vec<Option<Identifier>>;
31    fn opt_result(self, ctx: &Context) -> Option<Value>;
32    fn get_attr<'a, T: Attribute>(self, ctx: &'a Context, key: &Identifier) -> Option<Ref<'a, T>>;
33    fn set_attr<T: Attribute>(self, ctx: &Context, key: &Identifier, value: T);
34    fn parent_module(self, ctx: &Context) -> ModuleOp;
35
36    fn is_ancestor_of(self, ctx: &Context, other: Ptr<Operation>) -> bool;
37}
38
39pub trait BlockPtrExt: Sized {
40    fn arguments(self, ctx: &Context) -> Vec<Value>;
41    fn is_empty(&self, ctx: &Context) -> bool;
42    fn is_entry_block(&self, ctx: &Context) -> bool;
43    fn ops_with_interface<T: OpInterfaceMarker + ?Sized + 'static>(
44        &self,
45        ctx: &Context,
46    ) -> impl Iterator<Item = TraitOp<T>>;
47}
48
49pub trait RegionPtrExt: Sized {
50    fn arguments(self, ctx: &Context) -> Vec<Value>;
51    fn is_empty(&self, ctx: &Context) -> bool;
52}
53
54pub trait OperationExt {
55    fn immediately_nested_ops(&self, ctx: &Context) -> impl Iterator<Item = Ptr<Operation>>;
56}
57
58impl OperationPtrExt for Ptr<Operation> {
59    fn as_op<T: Op>(self, ctx: &Context) -> Option<T> {
60        Operation::get_op(self, ctx)
61    }
62    fn dyn_op(self, ctx: &Context) -> OpObj {
63        Operation::get_op_dyn(self, ctx)
64    }
65    fn cast<T: ?Sized + OpInterfaceMarker + 'static>(self, ctx: &Context) -> Option<TraitOp<T>> {
66        TraitOp::try_from_op(self, ctx)
67    }
68    fn operand(self, ctx: &Context, idx: usize) -> Value {
69        Operation::get_operand(&self.deref(ctx), idx)
70    }
71    fn operand_as_use(self, ctx: &Context, idx: usize) -> Use<Value> {
72        Operation::get_operand_as_use(&self.deref(ctx), idx)
73    }
74    fn operands(self, ctx: &Context) -> Vec<Value> {
75        self.deref(ctx).operands().collect()
76    }
77    fn operands_as_uses(self, ctx: &Context) -> Vec<Use<Value>> {
78        self.deref(ctx).operands_as_uses().collect()
79    }
80    fn result(self, ctx: &Context) -> Value {
81        Operation::get_result(&self.deref(ctx), 0)
82    }
83    fn results(self, ctx: &Context) -> Vec<Value> {
84        self.deref(ctx).results().collect()
85    }
86    fn result_types(self, ctx: &Context) -> Vec<TypeHandle> {
87        self.deref(ctx).result_types().collect()
88    }
89    fn regions(self, ctx: &Context) -> Vec<Ptr<Region>> {
90        self.deref(ctx).regions().collect()
91    }
92    fn result_names(self, ctx: &Context) -> Vec<Option<Identifier>> {
93        self.deref(ctx)
94            .results()
95            .map(|it| it.given_name(ctx))
96            .collect()
97    }
98    fn opt_result(self, ctx: &Context) -> Option<Value> {
99        self.deref(ctx).results().next()
100    }
101    fn get_attr<'a, T: Attribute>(self, ctx: &'a Context, key: &Identifier) -> Option<Ref<'a, T>> {
102        Ref::filter_map(self.deref(ctx), |op| op.attributes.get(key)).ok()
103    }
104    fn set_attr<T: Attribute>(self, ctx: &Context, key: &Identifier, value: T) {
105        self.deref_mut(ctx).attributes.set(key.clone(), value);
106    }
107    fn parent_module(self, ctx: &Context) -> ModuleOp {
108        if self.is_op::<ModuleOp>(ctx) {
109            return self.as_op(ctx).unwrap();
110        }
111        let mut op = self;
112        while let Some(parent) = op.deref(ctx).get_parent_op(ctx) {
113            if parent.is_op::<ModuleOp>(ctx) {
114                return parent.as_op(ctx).unwrap();
115            }
116            op = parent;
117        }
118        panic!("Op is not contained in any module")
119    }
120
121    fn is_ancestor_of(self, ctx: &Context, mut child: Ptr<Operation>) -> bool {
122        while let Some(parent) = child.deref(ctx).get_parent_op(ctx) {
123            if parent == self {
124                return true;
125            }
126            child = parent;
127        }
128        false
129    }
130}
131
132impl OperationExt for Operation {
133    fn immediately_nested_ops(&self, ctx: &Context) -> impl Iterator<Item = Ptr<Operation>> {
134        self.regions()
135            .flat_map(|region| region.deref(ctx).iter(ctx))
136            .flat_map(|block| block.deref(ctx).iter(ctx))
137    }
138}
139
140impl BlockPtrExt for Ptr<BasicBlock> {
141    fn arguments(self, ctx: &Context) -> Vec<Value> {
142        self.deref(ctx).arguments().collect()
143    }
144
145    fn is_empty(&self, ctx: &Context) -> bool {
146        self.deref(ctx).iter(ctx).count() == 0
147    }
148
149    fn is_entry_block(&self, ctx: &Context) -> bool {
150        self.deref(ctx)
151            .get_parent_region()
152            .is_some_and(|region| region.deref(ctx).get_entry_block() == Some(*self))
153    }
154
155    fn ops_with_interface<T: OpInterfaceMarker + ?Sized + 'static>(
156        &self,
157        ctx: &Context,
158    ) -> impl Iterator<Item = TraitOp<T>> {
159        self.deref(ctx)
160            .iter(ctx)
161            .filter_map(|op| TraitOp::try_from_op(op, ctx))
162    }
163}
164
165impl RegionPtrExt for Ptr<Region> {
166    fn arguments(self, ctx: &Context) -> Vec<Value> {
167        match self.entry_node(ctx) {
168            Some(entry) => entry.arguments(ctx),
169            None => vec![],
170        }
171    }
172
173    fn is_empty(&self, ctx: &Context) -> bool {
174        self.entry_node(ctx).is_none()
175    }
176}
177
178macro_rules! pure_unop {
179    ($name: literal, $ty: ident) => {
180        #[cubecl_macros_internal::cube_op(name = $name)]
181        #[result_ty(same_as = input)]
182        #[$crate::prelude::op_interfaces(
183            SameOperandsType,
184            SameOperandsAndResultType,
185            $crate::interfaces::TriviallyUnrollable
186        )]
187        #[$crate::prelude::op_traits(
188            $crate::CanMaterialize,
189            $crate::Pure,
190            $crate::PropagatesUniformity
191        )]
192        pub struct $ty {
193            pub input: Value,
194        }
195    };
196}
197use pliron::{
198    attribute::Attribute,
199    basic_block::BasicBlock,
200    builtin::ops::ModuleOp,
201    common_traits::Named,
202    graph::ControlFlowGraph,
203    identifier::Identifier,
204    linked_list::ContainsLinkedList,
205    op::{OpInterfaceMarker, OpObj, op_impls},
206    region::Region,
207    r#type::TypeHandle,
208    value::Use,
209};
210pub(crate) use pure_unop;
211
212macro_rules! pure_binop {
213    ($name: literal, $ty: ident) => {
214        #[cubecl_macros_internal::cube_op(name = $name)]
215        #[result_ty(same_as = lhs)]
216        #[$crate::prelude::op_interfaces(
217            SameOperandsType,
218            SameOperandsAndResultType,
219            $crate::interfaces::TriviallyUnrollable
220        )]
221        #[$crate::prelude::op_traits(
222            $crate::CanMaterialize,
223            $crate::Pure,
224            $crate::PropagatesUniformity
225        )]
226        pub struct $ty {
227            pub lhs: Value,
228            pub rhs: Value,
229        }
230    };
231}
232pub(crate) use pure_binop;
233
234pub fn ptr_value_ty(ctx: &Context, input: &Value) -> TypeHandle {
235    let in_ty = input.get_type(ctx).deref(ctx);
236    let ptr_ty = in_ty.downcast_ref::<PointerType>();
237    let mut inner = ptr_ty.expect("Should be a pointer").inner;
238    // If type is atomic, it's the value inside it after load
239    {
240        let inner_ref = inner.deref(ctx);
241        if let Some(AtomicType {
242            inner: inner_val, ..
243        }) = inner_ref.downcast_ref()
244        {
245            inner = *inner_val;
246        }
247    }
248    inner
249}
cubecl_ir/dialect/base.rs

cubecl_ir/dialect/
base.rs
