cubecl_opt/analyses/
dataflow_solver.rs

1use core::{
2    any::{TypeId, type_name},
3    cell::RefCell,
4    hash::{BuildHasher, Hash, Hasher},
5    marker::PhantomData,
6    ops::BitOrAssign,
7};
8
9use alloc::{boxed::Box, collections::VecDeque, rc::Rc, vec::Vec};
10use cubecl_environment::collections::HashMap;
11use cubecl_ir::prelude::*;
12use downcast_rs::{Downcast, impl_downcast};
13use foldhash::fast::FixedState;
14use pliron::{
15    basic_block::BasicBlock,
16    graph::HasLabel,
17    linked_list::{ContainsLinkedList, LinkedList},
18    printable::Printable,
19    verify_err_noloc,
20};
21
22use smallvec::SmallVec;
23pub use solver::*;
24
25pub mod control_flow_uniformity;
26pub mod dead_code;
27pub mod sccp;
28pub mod sparse;
29pub mod value_uniformity;
30
31pub type SmallPtrVec<T> = SmallVec<[T; 8]>;
32
33#[derive(Clone, Copy, PartialEq, Eq, Hash)]
34pub enum ChangeResult {
35    Changed,
36    Unchanged,
37}
38
39impl BitOrAssign for ChangeResult {
40    fn bitor_assign(&mut self, rhs: Self) {
41        if matches!(rhs, ChangeResult::Changed) {
42            *self = ChangeResult::Changed;
43        }
44    }
45}
46
47/// Nested module to ensure none of the unsafe abstractions leave this scope.
48mod solver {
49    use core::cell::{Ref, RefMut};
50
51    use super::*;
52
53    pub type SolverWorkItem = (ProgramPoint, TypeId);
54
55    /// Read-only ref, can read but not update state.
56    pub struct ReadRef<'a, T: PrintableState> {
57        value: Rc<RefCell<dyn PrintableState>>,
58        _ty: PhantomData<&'a T>,
59    }
60
61    impl<T: PrintableState> ReadRef<'_, T> {
62        #[track_caller]
63        pub fn deref(&self) -> Ref<'_, T> {
64            Ref::map(self.value.borrow(), |value| {
65                value.downcast_ref::<T>().unwrap()
66            })
67        }
68    }
69
70    /// Write-only ref, required for soundness in cases where multiple aliasing lattice elements are
71    /// referenced at the same time. Mutation is only allowed through `update_state`.
72    pub struct WriteRef<'a, T: PrintableState> {
73        value: Rc<RefCell<dyn PrintableState>>,
74        _ty: PhantomData<&'a T>,
75    }
76
77    impl<T: PrintableState> WriteRef<'_, T> {
78        #[track_caller]
79        fn deref(&self) -> RefMut<'_, T> {
80            RefMut::map(self.value.borrow_mut(), |value| {
81                value.downcast_mut::<T>().unwrap()
82            })
83        }
84    }
85
86    impl<T: PrintableState> PartialEq<WriteRef<'_, T>> for ReadRef<'_, T> {
87        fn eq(&self, other: &WriteRef<'_, T>) -> bool {
88            Rc::ptr_eq(&self.value, &other.value)
89        }
90    }
91
92    impl<T: PrintableState> PartialEq<ReadRef<'_, T>> for WriteRef<'_, T> {
93        fn eq(&self, other: &ReadRef<'_, T>) -> bool {
94            Rc::ptr_eq(&self.value, &other.value)
95        }
96    }
97
98    pub struct SolverConfig {
99        pub is_interprocedural: bool,
100    }
101
102    impl Default for SolverConfig {
103        fn default() -> Self {
104            Self {
105                is_interprocedural: true,
106            }
107        }
108    }
109
110    type AnalysisStates = HashMap<u64, States>;
111    type States = HashMap<TypeId, StateEntry>;
112    type StateEntry = Rc<RefCell<dyn PrintableState>>;
113
114    pub struct DataflowSolver {
115        child_analyses: HashMap<TypeId, Box<dyn DataflowAnalysis>>,
116        worklist: RefCell<VecDeque<SolverWorkItem>>,
117        anchor_hash: FixedState,
118        analysis_states: RefCell<AnalysisStates>,
119        config: SolverConfig,
120    }
121
122    impl DataflowSolver {
123        pub fn new(config: SolverConfig) -> Self {
124            Self {
125                child_analyses: Default::default(),
126                worklist: Default::default(),
127                anchor_hash: Default::default(),
128                analysis_states: Default::default(),
129                config,
130            }
131        }
132
133        #[track_caller]
134        pub fn update_state<A: AnalysisState>(
135            &self,
136            ctx: &Context,
137            state: &WriteRef<A>,
138            update: impl FnOnce(&mut A) -> ChangeResult,
139        ) {
140            let mut state = state.deref();
141            let changed = update(&mut state);
142            if changed == ChangeResult::Changed {
143                state.on_update(ctx, self);
144            }
145        }
146
147        pub fn get_or_create<T: AnalysisState>(&self, anchor: T::Anchor) -> ReadRef<'_, T> {
148            let anchor_hash = self.hash_anchor(&anchor);
149            let mut states = self.analysis_states.borrow_mut();
150            let state = states.entry(anchor_hash).or_default();
151            let id = TypeId::of::<T>();
152            let state = state
153                .entry(id)
154                .or_insert_with(|| Rc::new(RefCell::new(T::create(anchor))))
155                .clone();
156            ReadRef {
157                value: state,
158                _ty: PhantomData,
159            }
160        }
161
162        pub fn get_or_create_mut<T: AnalysisState>(&self, anchor: T::Anchor) -> WriteRef<'_, T> {
163            let anchor_hash = self.hash_anchor(&anchor);
164            let mut states = self.analysis_states.borrow_mut();
165            let state = states.entry(anchor_hash).or_default();
166            let id = TypeId::of::<T>();
167            let state = state
168                .entry(id)
169                .or_insert_with(|| Rc::new(RefCell::new(T::create(anchor))))
170                .clone();
171            WriteRef {
172                value: state,
173                _ty: PhantomData,
174            }
175        }
176
177        pub fn get_or_create_for<A: 'static, T: AnalysisState>(
178            &self,
179            dependent: ProgramPoint,
180            anchor: T::Anchor,
181        ) -> ReadRef<'_, T> {
182            let is_equivalent = self.is_equivalent::<T>(&anchor, dependent);
183            let state = self.get_or_create::<T>(anchor);
184            if !is_equivalent {
185                state.deref().add_dependency::<A>(dependent);
186            }
187            state
188        }
189
190        pub fn lookup_state<T: AnalysisState>(&self, anchor: T::Anchor) -> Option<ReadRef<'_, T>> {
191            let anchor_hash = self.hash_anchor(&anchor);
192            let states = self.analysis_states.borrow();
193            let state = states.get(&anchor_hash)?.get(&TypeId::of::<T>())?.clone();
194            Some(ReadRef {
195                value: state,
196                _ty: PhantomData,
197            })
198        }
199    }
200
201    impl DataflowSolver {
202        pub fn load<A: DataflowAnalysis>(&mut self, analysis: A) {
203            let key = TypeId::of::<A>();
204            let existing = self.child_analyses.insert(key, Box::new(analysis));
205            assert!(
206                existing.is_none(),
207                "Tried loading {} twice",
208                type_name::<A>()
209            )
210        }
211
212        pub fn require_loaded<A: DataflowAnalysis>(&self) -> Result<()> {
213            if !self.child_analyses.contains_key(&TypeId::of::<A>()) {
214                return verify_err_noloc!(
215                    "Missing required dataflow analysis {}",
216                    type_name::<A>()
217                );
218            }
219            Ok(())
220        }
221
222        pub fn initialize_and_run(&mut self, ctx: &Context, root: Ptr<Operation>) -> Result<()> {
223            let is_interprocedural = self.config.is_interprocedural;
224            if is_interprocedural && !root.impls::<dyn SymbolTableInterface>(ctx) {
225                self.config.is_interprocedural = false;
226            }
227
228            // Take it temporarily so we get mutable access without borrowing self
229            let mut child_analyses = core::mem::take(&mut self.child_analyses);
230
231            // Initialize equivalent lattice anchors.
232            for analysis in child_analyses.values() {
233                analysis.initialize_equivalent_lattice_anchor(self, ctx, root);
234            }
235
236            // Initialize the analyses.
237            for analysis in child_analyses.values_mut() {
238                if let Err(err) = analysis.initialize(self, ctx, root) {
239                    self.child_analyses = child_analyses;
240                    self.config.is_interprocedural = is_interprocedural;
241                    return Err(err);
242                }
243            }
244
245            self.child_analyses = child_analyses;
246
247            while let Some((point, analysis)) = {
248                let mut worklist = self.worklist.borrow_mut();
249                worklist.pop_front()
250            } {
251                let analysis = self.child_analyses.get(&analysis).expect("Should exist");
252                if let Err(err) = analysis.visit(self, ctx, point) {
253                    self.config.is_interprocedural = is_interprocedural;
254                    return Err(err);
255                }
256            }
257
258            self.config.is_interprocedural = is_interprocedural;
259            Ok(())
260        }
261
262        pub fn enqueue(&self, work_item: SolverWorkItem) {
263            self.worklist.borrow_mut().push_back(work_item);
264        }
265
266        pub fn config(&self) -> &SolverConfig {
267            &self.config
268        }
269
270        pub fn is_equivalent<T: AnalysisState>(
271            &self,
272            _lhs: &T::Anchor,
273            _rhs: ProgramPoint,
274        ) -> bool {
275            // TODO: Support equivalence classes
276            false
277        }
278
279        // Includes `TypeID` so different anchor types hash to different values even if the inner data
280        // is the same.
281        fn hash_anchor<T: Clone + Printable + Hash + 'static>(&self, anchor: &T) -> u64 {
282            let mut hasher = self.anchor_hash.build_hasher();
283            TypeId::of::<T>().hash(&mut hasher);
284            anchor.hash(&mut hasher);
285            hasher.finish()
286        }
287    }
288
289    impl Printable for DataflowSolver {
290        fn fmt(
291            &self,
292            ctx: &Context,
293            _state: &pliron::printable::State,
294            f: &mut core::fmt::Formatter<'_>,
295        ) -> core::fmt::Result {
296            writeln!(f, "DataflowSolver {{")?;
297            let states = self.analysis_states.borrow();
298            let mut entries = states
299                .values()
300                .flat_map(|states| states.iter())
301                .collect::<Vec<_>>();
302            entries.sort_by_key(|it| it.0);
303
304            for (_, state) in entries {
305                writeln!(f, "    {},", state.borrow().disp(ctx))?;
306            }
307            writeln!(f, "}}")
308        }
309    }
310}
311
312#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
313pub enum ProgramPoint {
314    Operation(Ptr<Operation>),
315    BeforeOpInBlock(Ptr<BasicBlock>, Ptr<Operation>),
316    EndOfBlock(Ptr<BasicBlock>),
317}
318
319impl Printable for ProgramPoint {
320    fn fmt(
321        &self,
322        ctx: &Context,
323        _state: &pliron::printable::State,
324        f: &mut core::fmt::Formatter<'_>,
325    ) -> core::fmt::Result {
326        match self {
327            ProgramPoint::Operation(op) => write!(f, "ProgramPoint::Operation({})", op.disp(ctx)),
328            ProgramPoint::BeforeOpInBlock(block, _) if self.is_block_start(ctx) => {
329                write!(f, "ProgramPoint::StartOfBlock({})", block.label(ctx))
330            }
331            ProgramPoint::BeforeOpInBlock(block, op) => {
332                write!(
333                    f,
334                    "ProgramPoint::BeforeOpInBlock({}, {})",
335                    block.label(ctx),
336                    op.disp(ctx)
337                )
338            }
339            ProgramPoint::EndOfBlock(block) => {
340                write!(f, "ProgramPoint::EndOfBlock({})", block.label(ctx))
341            }
342        }
343    }
344}
345
346impl ProgramPoint {
347    pub fn before_op(ctx: &Context, op: Ptr<Operation>) -> ProgramPoint {
348        if let Some(block) = op.deref(ctx).get_parent_block() {
349            ProgramPoint::BeforeOpInBlock(block, op)
350        } else {
351            ProgramPoint::Operation(op)
352        }
353    }
354
355    pub fn at_block_start(ctx: &Context, block: Ptr<BasicBlock>) -> ProgramPoint {
356        if let Some(first) = block.deref(ctx).iter(ctx).next() {
357            ProgramPoint::BeforeOpInBlock(block, first)
358        } else {
359            ProgramPoint::EndOfBlock(block)
360        }
361    }
362
363    pub fn after_op(ctx: &Context, op: Ptr<Operation>) -> ProgramPoint {
364        if let Some(block) = op.deref(ctx).get_parent_block() {
365            if let Some(next) = op.deref(ctx).get_next() {
366                ProgramPoint::BeforeOpInBlock(block, next)
367            } else {
368                ProgramPoint::EndOfBlock(block)
369            }
370        } else {
371            ProgramPoint::Operation(op)
372        }
373    }
374
375    pub fn at_block_end(_ctx: &Context, block: Ptr<BasicBlock>) -> ProgramPoint {
376        ProgramPoint::EndOfBlock(block)
377    }
378
379    pub fn next_op(&self, _ctx: &Context) -> Option<Ptr<Operation>> {
380        match self {
381            ProgramPoint::Operation(op) => Some(*op),
382            ProgramPoint::BeforeOpInBlock(_, op) => Some(*op),
383            ProgramPoint::EndOfBlock(_) => None,
384        }
385    }
386
387    pub fn prev_op(&self, ctx: &Context) -> Option<Ptr<Operation>> {
388        match self {
389            ProgramPoint::Operation(op) => Some(*op),
390            ProgramPoint::BeforeOpInBlock(_, op) => op.deref(ctx).get_prev(),
391            ProgramPoint::EndOfBlock(block) => block.deref(ctx).get_tail(),
392        }
393    }
394
395    pub fn is_block_start(&self, ctx: &Context) -> bool {
396        match self {
397            ProgramPoint::Operation(_) => false,
398            // true if no preceding op
399            ProgramPoint::BeforeOpInBlock(_, op) => op.deref(ctx).get_prev().is_none(),
400            // true if Empty block
401            ProgramPoint::EndOfBlock(ptr) => ptr.deref(ctx).get_head().is_none(),
402        }
403    }
404
405    pub fn is_block_end(&self, _ctx: &Context) -> bool {
406        match self {
407            ProgramPoint::Operation(_) | ProgramPoint::BeforeOpInBlock(..) => false,
408            ProgramPoint::EndOfBlock(_) => true,
409        }
410    }
411
412    pub fn block(&self) -> Option<Ptr<BasicBlock>> {
413        match self {
414            ProgramPoint::Operation(_) => None,
415            ProgramPoint::BeforeOpInBlock(block, _) => Some(*block),
416            ProgramPoint::EndOfBlock(block) => Some(*block),
417        }
418    }
419}
420
421pub trait AnalysisState: PrintableState + Sized + 'static {
422    type Anchor: Printable + Clone + Hash + 'static;
423
424    fn create(anchor: Self::Anchor) -> Self;
425    fn add_dependency<A: 'static>(&self, point: ProgramPoint);
426    fn on_update(&self, ctx: &Context, solver: &DataflowSolver);
427}
428
429pub trait DataflowAnalysis: Downcast {
430    /// Verify analysis can be run on the solver. Should be used to verify required analyses are
431    /// loaded.
432    fn verify(&self, solver: &DataflowSolver, ctx: &Context, root: Ptr<Operation>) -> Result<()> {
433        let _ = (solver, ctx, root);
434        Ok(())
435    }
436
437    #[allow(clippy::result_unit_err)]
438    fn initialize(
439        &mut self,
440        solver: &mut DataflowSolver,
441        ctx: &Context,
442        root: Ptr<Operation>,
443    ) -> Result<()>;
444
445    #[allow(clippy::result_unit_err)]
446    fn visit(&self, solver: &DataflowSolver, ctx: &Context, point: ProgramPoint) -> Result<()>;
447
448    fn initialize_equivalent_lattice_anchor(
449        &self,
450        solver: &mut DataflowSolver,
451        ctx: &Context,
452        root: Ptr<Operation>,
453    ) {
454        let _ = (solver, ctx, root);
455    }
456}
457impl_downcast!(DataflowAnalysis);
458
459pub trait PrintableState: Printable + Downcast {}
460impl<T: Printable + 'static> PrintableState for T {}
461impl_downcast!(PrintableState);
cubecl_opt/analyses/dataflow_solver.rs

cubecl_opt/analyses/
dataflow_solver.rs
