cubecl_ir/
rewrite.rs

1use core::{
2    fmt::Debug,
3    hash::Hash,
4    marker::PhantomData,
5    ops::{Deref, Range},
6};
7
8use cubecl_macros_internal::NamedRewrite;
9use derive_more::{Deref, DerefMut, From};
10use derive_new::new;
11use pliron::{
12    attribute::{AttrObj, Attribute},
13    builtin::{
14        given_names::{get_operation_result_name, set_operation_result_name},
15        ops::ConstantOp,
16    },
17    graph::walkers::{
18        WalkConfig,
19        uninterruptible::{
20            immutable::{self},
21            mutable,
22        },
23    },
24    irbuild::{
25        dialect_conversion::apply_dialect_conversion,
26        match_rewrite::{RewriterOrder, apply_match_rewrite},
27    },
28    op::{OpInterfaceMarker, OpObj},
29    value::Use,
30    verify_err_noloc,
31};
32
33use crate::{
34    dialect::BlockPtrExt,
35    interfaces::{CanonicalizeInterface, SimplifyInterface},
36    prelude::*,
37};
38
39/// A preset config when order doesn't matter
40pub const WALKCONFIG_ANY: WalkConfig = WALKCONFIG_PREORDER_FORWARD;
41
42pub trait NamedRewrite {
43    fn name(&self) -> &str;
44}
45
46#[derive(new, From, Clone, Debug, Default)]
47pub struct DialectConversionPass<T: DialectConversion>(T);
48
49impl<T: DialectConversion + NamedRewrite> Pass for DialectConversionPass<T> {
50    fn name(&self) -> &str {
51        self.0.name()
52    }
53
54    fn run(
55        &mut self,
56        op: Ptr<Operation>,
57        ctx: &mut Context,
58        _analyses: &mut AnalysisManager,
59    ) -> Result<PassResult> {
60        let mut res = PassResult::default();
61        res.ir_changed = apply_dialect_conversion(ctx, &mut self.0, op)?;
62        Ok(res)
63    }
64}
65
66#[derive(new, From, Clone, Debug, Default)]
67pub struct MatchRewritePass<T: MatchRewrite>(pub T);
68
69impl<T: MatchRewrite + NamedRewrite> Pass for MatchRewritePass<T> {
70    fn name(&self) -> &str {
71        self.0.name()
72    }
73
74    fn run(
75        &mut self,
76        op: Ptr<Operation>,
77        ctx: &mut Context,
78        _analyses: &mut AnalysisManager,
79    ) -> Result<PassResult> {
80        let mut res = PassResult::default();
81        res.ir_changed = apply_match_rewrite(ctx, &mut self.0, RewriterOrder::default(), op)?;
82        Ok(res)
83    }
84}
85
86#[derive(new, Clone, Copy, Default, Debug)]
87pub struct CombinedPass<P1: Pass, P2: Pass> {
88    pass_1: P1,
89    pass_2: P2,
90}
91
92#[pass_name]
93impl<P1: Pass, P2: Pass> Pass for CombinedPass<P1, P2> {
94    fn run(
95        &mut self,
96        op: Ptr<Operation>,
97        ctx: &mut Context,
98        analyses: &mut AnalysisManager,
99    ) -> Result<PassResult> {
100        let mut res = PassResult::default();
101        res.ir_changed |= self.pass_1.run(op, ctx, analyses)?.ir_changed;
102        res.ir_changed |= self.pass_2.run(op, ctx, analyses)?.ir_changed;
103        Ok(res)
104    }
105}
106
107pub type SimplifyOpsPass = MatchRewritePass<SimplifyOps>;
108
109#[derive(Default, Clone, Copy, NamedRewrite)]
110pub struct SimplifyOps;
111
112impl MatchRewrite for SimplifyOps {
113    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
114        op.impls::<dyn SimplifyInterface>(ctx)
115    }
116
117    fn rewrite(
118        &mut self,
119        ctx: &mut Context,
120        rewriter: &mut MatchRewriter,
121        op: Ptr<Operation>,
122    ) -> Result<()> {
123        let dyn_op = op.dyn_op(ctx);
124        let operand_attrs = const_operands(ctx, op);
125        let simplify = op_cast::<dyn SimplifyInterface>(&*dyn_op).unwrap();
126        if let Some(value) = simplify.check_fold(ctx, &operand_attrs) {
127            rewriter.replace_operation_with_values(ctx, op, vec![value]);
128        }
129        Ok(())
130    }
131}
132
133pub fn const_operand<T: Attribute>(ctx: &Context, op: Ptr<Operation>, idx: usize) -> Option<T> {
134    Some(*const_operands(ctx, op).remove(idx)?.downcast().ok()?)
135}
136
137pub fn const_operands(ctx: &Context, op: Ptr<Operation>) -> Vec<Option<AttrObj>> {
138    op.deref(ctx)
139        .operands()
140        .map(|opd| {
141            Some(
142                opd.defining_op()?
143                    .as_op::<ConstantOp>(ctx)?
144                    .get_attr_builtin_constant_value(ctx)?
145                    .clone(),
146            )
147        })
148        .collect()
149}
150
151pub type CanonicalizePass = MatchRewritePass<Canonicalize>;
152
153#[derive(Default, Clone, Copy, NamedRewrite)]
154pub struct Canonicalize;
155
156impl MatchRewrite for Canonicalize {
157    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
158        op.impls::<dyn CanonicalizeInterface>(ctx)
159    }
160
161    fn rewrite(
162        &mut self,
163        ctx: &mut Context,
164        rewriter: &mut MatchRewriter,
165        op: Ptr<Operation>,
166    ) -> Result<()> {
167        let dyn_op = op.dyn_op(ctx);
168        let canonicalize = op_cast::<dyn CanonicalizeInterface>(&*dyn_op).unwrap();
169        canonicalize.canonicalize(ctx, rewriter)?;
170        Ok(())
171    }
172}
173
174pub type VisitOpsCallback<T, State> = fn(&Context, &mut State, T);
175pub type VisitOpsMutCallback<T, State> = fn(&mut Context, &mut State, T);
176
177pub fn visit_all_ops_of_type<T: Op, State>(
178    ctx: &Context,
179    state: &mut State,
180    root: Ptr<Operation>,
181    callback: VisitOpsCallback<T, State>,
182) {
183    immutable::walk_op(
184        ctx,
185        &mut (state, callback),
186        &WALKCONFIG_PREORDER_FORWARD,
187        root,
188        |ctx, (state, callback), node| {
189            if let IRNode::Operation(op) = node
190                && let Some(op) = op.as_op::<T>(ctx)
191            {
192                callback(ctx, state, op)
193            }
194        },
195    );
196}
197
198pub fn visit_all_ops_of_type_mut<T: Op, State>(
199    ctx: &mut Context,
200    state: &mut State,
201    root: Ptr<Operation>,
202    callback: VisitOpsMutCallback<T, State>,
203) {
204    mutable::walk_op(
205        ctx,
206        &mut (state, callback),
207        &WALKCONFIG_PREORDER_FORWARD,
208        root,
209        |ctx, (state, callback), node| {
210            if let IRNode::Operation(op) = node
211                && let Some(op) = op.as_op::<T>(ctx)
212            {
213                callback(ctx, state, op)
214            }
215        },
216    );
217}
218
219pub fn visit_all_ops_with_interface<T: ?Sized + OpInterfaceMarker + 'static, State>(
220    ctx: &Context,
221    state: &mut State,
222    root: Ptr<Operation>,
223    callback: for<'a> fn(&Context, &mut State, &'a T),
224) {
225    immutable::walk_op(
226        ctx,
227        &mut (state, callback),
228        &WALKCONFIG_ANY,
229        root,
230        |ctx, (state, callback), node| {
231            if let IRNode::Operation(op) = node {
232                let dyn_op = op.dyn_op(ctx);
233                if let Some(op) = op_cast::<T>(&*dyn_op) {
234                    callback(ctx, state, op)
235                }
236            }
237        },
238    );
239}
240
241pub fn visit_all_values<State>(
242    ctx: &Context,
243    state: &mut State,
244    root: Ptr<Operation>,
245    callback: for<'a> fn(&Context, &mut State, Value),
246) {
247    immutable::walk_op(
248        ctx,
249        &mut (state, callback),
250        &WALKCONFIG_PREORDER_FORWARD,
251        root,
252        |ctx, (state, callback), node| match node {
253            IRNode::Operation(ptr) => {
254                for res in ptr.results(ctx) {
255                    callback(ctx, state, res)
256                }
257            }
258            IRNode::BasicBlock(ptr) => {
259                for res in ptr.arguments(ctx) {
260                    callback(ctx, state, res)
261                }
262            }
263            IRNode::Region(_) => {}
264        },
265    );
266}
267
268pub trait RewriteOp<T: Op> {
269    fn should_rewrite(&self, _ctx: &Context, _op: T) -> bool {
270        true
271    }
272
273    fn rewrite(&mut self, ctx: &mut Context, rewriter: &mut MatchRewriter, op: T);
274}
275
276#[derive(new, Deref, DerefMut)]
277pub struct MatchRewriteOp<T: Op, R: RewriteOp<T>> {
278    #[deref]
279    #[deref_mut]
280    pub inner: R,
281    _t: PhantomData<T>,
282}
283
284impl<T: Op, R: RewriteOp<T>> From<R> for MatchRewriteOp<T, R> {
285    fn from(value: R) -> Self {
286        Self::new(value)
287    }
288}
289
290impl<T: Op, R: RewriteOp<T>> MatchRewrite for MatchRewriteOp<T, R> {
291    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
292        op.as_op::<T>(ctx)
293            .is_some_and(|op| self.should_rewrite(ctx, op))
294    }
295
296    fn rewrite(
297        &mut self,
298        ctx: &mut Context,
299        rewriter: &mut MatchRewriter,
300        op: Ptr<Operation>,
301    ) -> Result<()> {
302        let op = op.as_op::<T>(ctx).unwrap();
303        RewriteOp::rewrite(&mut self.inner, ctx, rewriter, op);
304        Ok(())
305    }
306}
307
308pub trait RewriterExt: Rewriter {
309    fn replace_op_with(&mut self, ctx: &mut Context, op: Ptr<Operation>, new_op: Ptr<Operation>) {
310        new_op.insert_before(ctx, op);
311        transfer_result_names(ctx, op, &new_op.results(ctx));
312        self.replace_operation(ctx, op, new_op);
313    }
314    fn append_op_with_result(&mut self, ctx: &mut Context, op: &impl OneResultInterface) -> Value {
315        self.append_op(ctx, op);
316        op.get_result(ctx)
317    }
318}
319impl<R: Rewriter> RewriterExt for R {}
320
321pub fn transfer_result_names(ctx: &Context, old_op: Ptr<Operation>, values: &[Value]) {
322    for (idx, value) in values.iter().enumerate() {
323        transfer_result_name(ctx, old_op, *value, idx);
324    }
325}
326
327pub fn transfer_result_name(ctx: &Context, old_op: Ptr<Operation>, value: Value, idx: usize) {
328    if let Some(new_op) = value.defining_op() {
329        set_operation_result_name(
330            ctx,
331            new_op,
332            idx,
333            get_operation_result_name(ctx, old_op, idx),
334        );
335    }
336}
337
338pub struct TraitOp<T: OpInterfaceMarker + ?Sized> {
339    obj: OpObj,
340    _marker: PhantomData<T>,
341}
342
343impl<T: OpInterfaceMarker + ?Sized> TraitOp<T> {
344    pub fn dyn_op(&self) -> &dyn Op {
345        &*self.obj
346    }
347}
348
349impl<T: OpInterfaceMarker + 'static + ?Sized> TraitOp<T> {
350    pub fn try_from_op(op: Ptr<Operation>, ctx: &Context) -> Option<Self> {
351        let op = op.dyn_op(ctx);
352        if !op_impls::<T>(&*op) {
353            None
354        } else {
355            Some(TraitOp {
356                obj: op,
357                _marker: PhantomData,
358            })
359        }
360    }
361}
362
363impl<T: OpInterfaceMarker + 'static + ?Sized> TryFrom<OpObj> for TraitOp<T> {
364    type Error = pliron::result::Error;
365
366    fn try_from(value: OpObj) -> Result<Self> {
367        if !op_impls::<T>(&*value) {
368            verify_err_noloc!("Op doesn't implement trait")
369        } else {
370            Ok(TraitOp {
371                obj: value,
372                _marker: PhantomData,
373            })
374        }
375    }
376}
377
378impl<T: OpInterfaceMarker + 'static + ?Sized> Deref for TraitOp<T> {
379    type Target = T;
380
381    fn deref(&self) -> &Self::Target {
382        op_cast(self.obj.as_ref()).unwrap()
383    }
384}
385
386impl<T: OpInterfaceMarker + 'static + ?Sized> Eq for TraitOp<T> {}
387impl<T: OpInterfaceMarker + 'static + ?Sized> PartialEq for TraitOp<T> {
388    fn eq(&self, other: &Self) -> bool {
389        self.obj == other.obj && self._marker == other._marker
390    }
391}
392
393impl<T: OpInterfaceMarker + 'static + ?Sized> Hash for TraitOp<T> {
394    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
395        self.obj.hash(state);
396    }
397}
398
399impl<T: OpInterfaceMarker + 'static + ?Sized> Copy for TraitOp<T> {}
400impl<T: OpInterfaceMarker + 'static + ?Sized> Clone for TraitOp<T> {
401    fn clone(&self) -> Self {
402        *self
403    }
404}
405
406impl<T: OpInterfaceMarker + 'static + ?Sized> Debug for TraitOp<T> {
407    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
408        let op = self.obj.get_operation();
409        Debug::fmt(&op, f)
410    }
411}
412
413pub(crate) fn operand_range_to_uses(
414    ctx: &Context,
415    op: Ptr<Operation>,
416    range: Range<usize>,
417) -> Vec<Use<Value>> {
418    let op = op.deref(ctx);
419    range.map(|idx| op.get_operand_as_use(idx)).collect()
420}
421
422#[macro_export]
423macro_rules! small_map {
424    {$($k: expr => $v: expr),* $(,)?} => {{
425        let mut out = $crate::pliron::utils::table::SmallMap::new();
426        $(out.insert($k, $v);)*
427        out
428    }};
429}
430
431#[macro_export]
432macro_rules! small_set {
433    {$($v: expr),* $(,)?} => {{
434        let mut out = $crate::pliron::utils::table::SmallSet::new();
435        $(out.insert($v);)*
436        out
437    }};
438}
cubecl_ir/rewrite.rs

cubecl_ir/
rewrite.rs
