cubecl_ir/interfaces/
aliasing.rs

1use pliron::{basic_block::BasicBlock, value::DefiningEntity};
2
3use crate::prelude::*;
4
5#[op_interface]
6pub trait AliasingOp: OneResultInterface {
7    verify_op_succ!();
8    fn source_ptr(&self, ctx: &Context) -> Option<Value>;
9}
10
11pub trait PointerExt {
12    fn find_root_index(&self, ctx: &Context) -> usize;
13    fn get_root_defining_entity(&self, ctx: &Context) -> DefiningEntity;
14    fn get_root_defining_op(&self, ctx: &Context) -> Option<Ptr<Operation>> {
15        match self.get_root_defining_entity(ctx) {
16            DefiningEntity::Op(ptr) => Some(ptr),
17            DefiningEntity::Block(_) => None,
18        }
19    }
20    fn get_root_defining_block(&self, ctx: &Context) -> Option<Ptr<BasicBlock>> {
21        match self.get_root_defining_entity(ctx) {
22            DefiningEntity::Op(_) => None,
23            DefiningEntity::Block(ptr) => Some(ptr),
24        }
25    }
26}
27
28impl PointerExt for Value {
29    fn get_root_defining_entity(&self, ctx: &Context) -> DefiningEntity {
30        match self.defining_entity() {
31            DefiningEntity::Op(op) => {
32                if let Some(aliasing) = op_cast::<dyn AliasingOp>(&*op.dyn_op(ctx))
33                    && let Some(source) = aliasing.source_ptr(ctx)
34                {
35                    source.get_root_defining_entity(ctx)
36                } else {
37                    DefiningEntity::Op(op)
38                }
39            }
40            block @ DefiningEntity::Block(_) => block,
41        }
42    }
43
44    fn find_root_index(&self, ctx: &Context) -> usize {
45        match self.defining_entity() {
46            DefiningEntity::Op(op)
47                if let Some(aliasing) = op_cast::<dyn AliasingOp>(&*op.dyn_op(ctx))
48                    && let Some(source) = aliasing.source_ptr(ctx) =>
49            {
50                source.find_root_index(ctx)
51            }
52            _ => self.find_index(ctx),
53        }
54    }
55}
cubecl_ir/interfaces/aliasing.rs

cubecl_ir/interfaces/
aliasing.rs
