cubecl_ir/interfaces/
side_effects.rs

1use core::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign};
2
3use pliron::{builtin::ops::ConstantOp, opts::dce::SideEffects, printable::Printable};
4
5use crate::{AddressSpace, NoMemoryEffect, dialect::OperationExt, prelude::*};
6
7#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
8pub enum ModRefResult {
9    NoModRef,
10    Ref,
11    Mod,
12    ModRef,
13}
14
15impl ModRefResult {
16    pub fn join(&self, other: ModRefResult) -> ModRefResult {
17        match (*self, other) {
18            (ModRefResult::NoModRef, other) | (other, ModRefResult::NoModRef) => other,
19            (_, ModRefResult::ModRef) | (ModRefResult::ModRef, _) => ModRefResult::ModRef,
20            (ModRefResult::Ref, ModRefResult::Ref) => ModRefResult::Ref,
21            (ModRefResult::Mod, ModRefResult::Mod) => ModRefResult::Mod,
22            (ModRefResult::Ref, ModRefResult::Mod) | (ModRefResult::Mod, ModRefResult::Ref) => {
23                ModRefResult::ModRef
24            }
25        }
26    }
27
28    pub fn meet(&self, other: ModRefResult) -> ModRefResult {
29        match (*self, other) {
30            (ModRefResult::NoModRef, _) | (_, ModRefResult::NoModRef) => ModRefResult::NoModRef,
31            (ModRefResult::ModRef, other) | (other, ModRefResult::ModRef) => other,
32            (ModRefResult::Ref, ModRefResult::Ref) => ModRefResult::Ref,
33            (ModRefResult::Mod, ModRefResult::Mod) => ModRefResult::Mod,
34            (ModRefResult::Mod, ModRefResult::Ref) | (ModRefResult::Ref, ModRefResult::Mod) => {
35                ModRefResult::NoModRef
36            }
37        }
38    }
39
40    pub fn contains_mod(&self) -> bool {
41        match self {
42            ModRefResult::NoModRef | ModRefResult::Ref => false,
43            ModRefResult::Mod | ModRefResult::ModRef => true,
44        }
45    }
46}
47
48impl BitOr for ModRefResult {
49    type Output = Self;
50    fn bitor(self, rhs: Self) -> Self::Output {
51        self.join(rhs)
52    }
53}
54
55impl BitAnd for ModRefResult {
56    type Output = Self;
57    fn bitand(self, rhs: Self) -> Self::Output {
58        self.meet(rhs)
59    }
60}
61
62impl BitOrAssign for ModRefResult {
63    fn bitor_assign(&mut self, rhs: Self) {
64        *self = self.join(rhs);
65    }
66}
67
68impl BitAndAssign for ModRefResult {
69    fn bitand_assign(&mut self, rhs: Self) {
70        *self = self.meet(rhs);
71    }
72}
73
74#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
75pub enum MemoryEffect {
76    Read(Value),
77    Write(Value),
78    ReadAllInSpace(AddressSpace),
79    WriteAllInSpace(AddressSpace),
80    ReadAll,
81    WriteAll,
82    // Not analyzable, clobber the entire state
83    Opaque,
84}
85
86impl MemoryEffect {
87    pub fn value(&self) -> Option<Value> {
88        match self {
89            MemoryEffect::Read(value) | MemoryEffect::Write(value) => Some(*value),
90            MemoryEffect::ReadAllInSpace(_)
91            | MemoryEffect::WriteAllInSpace(_)
92            | MemoryEffect::ReadAll
93            | MemoryEffect::WriteAll
94            | MemoryEffect::Opaque => None,
95        }
96    }
97
98    pub fn mod_ref_mask(&self) -> ModRefResult {
99        match self {
100            MemoryEffect::Read(_) | MemoryEffect::ReadAllInSpace(_) | MemoryEffect::ReadAll => {
101                ModRefResult::Ref
102            }
103            MemoryEffect::Write(_) | MemoryEffect::WriteAllInSpace(_) | MemoryEffect::WriteAll => {
104                ModRefResult::Mod
105            }
106            MemoryEffect::Opaque => ModRefResult::ModRef,
107        }
108    }
109
110    pub fn modifies(&self) -> bool {
111        self.mod_ref_mask().contains_mod()
112    }
113}
114
115impl Printable for MemoryEffect {
116    fn fmt(
117        &self,
118        ctx: &Context,
119        _state: &pliron::printable::State,
120        f: &mut core::fmt::Formatter<'_>,
121    ) -> core::fmt::Result {
122        match self {
123            MemoryEffect::Read(value) => write!(f, "Read({})", value.disp(ctx)),
124            MemoryEffect::Write(value) => write!(f, "Write({})", value.disp(ctx)),
125            MemoryEffect::ReadAllInSpace(address_space) => {
126                write!(f, "ReadAllInSpace({})", address_space.disp(ctx))
127            }
128            MemoryEffect::WriteAllInSpace(address_space) => {
129                write!(f, "WriteAllInSpace({})", address_space.disp(ctx))
130            }
131            MemoryEffect::ReadAll => write!(f, "ReadAll"),
132            MemoryEffect::WriteAll => write!(f, "WriteAll"),
133            MemoryEffect::Opaque => write!(f, "Opaque"),
134        }
135    }
136}
137
138#[op_interface]
139pub trait MemoryEffectsOp {
140    verify_op_succ!();
141    fn memory_effects(&self, ctx: &Context) -> Vec<MemoryEffect>;
142    fn has_effects(&self, ctx: &Context) -> bool {
143        !self.memory_effects(ctx).is_empty()
144    }
145}
146
147NoMemoryEffect!(ConstantOp);
148
149pub fn has_memory_side_effects(effects: &[MemoryEffect]) -> bool {
150    effects.iter().any(|effect| effect.modifies())
151}
152
153pub fn get_nested_side_effects(ctx: &Context, op: Ptr<Operation>) -> bool {
154    op.deref(ctx).immediately_nested_ops(ctx).any(|op| {
155        op_cast::<dyn SideEffects>(&*op.dyn_op(ctx)).is_none_or(|op| op.has_side_effects(ctx))
156    })
157}
158
159pub fn get_nested_memory_effects(ctx: &Context, op: Ptr<Operation>) -> Vec<MemoryEffect> {
160    let mut out = vec![];
161    for op in op.deref(ctx).immediately_nested_ops(ctx) {
162        out.extend(match op_cast::<dyn MemoryEffectsOp>(&*op.dyn_op(ctx)) {
163            Some(mem_op) => mem_op.memory_effects(ctx),
164            None => vec![MemoryEffect::Opaque],
165        });
166    }
167    out
168}
169
170pub enum Speculatability {
171    /// The Operation in question cannot be speculatively executed. This could be
172    /// because it may invoke undefined behavior or have other side effects.
173    NotSpeculatable,
174
175    // The Operation in question can be speculatively executed. It does not have
176    // any side effects or undefined behavior.
177    Speculatable,
178
179    // The Operation in question can be speculatively executed if all the
180    // operations in all attached regions can also be speculatively executed.
181    RecursivelySpeculatable,
182}
183
184#[op_interface]
185pub trait ConditionallySpeculatable {
186    verify_op_succ!();
187    fn speculatability(&self, ctx: &Context) -> Speculatability;
188}
cubecl_ir/interfaces/side_effects.rs

cubecl_ir/interfaces/
side_effects.rs
