Skip to main content

cubecl_opt/analyses/
base.rs

1use core::{any::Any, cell::RefCell};
2
3use alloc::rc::Rc;
4use type_map::TypeMap;
5
6use crate::{Function, GlobalState};
7
8use super::{
9    dominance::{Dominators, PostDominators},
10    liveness::Liveness,
11    post_order::PostOrder,
12    uniformity::Uniformity,
13};
14
15/// An analysis used by optimization passes. Unlike optimization passes, analyses can have state
16/// and persist until they're invalidated.
17pub trait Analysis {
18    /// Perform the analysis for the current optimizer state and return the persistent analysis state
19    fn init(opt: &mut Function, state: &GlobalState) -> Self;
20}
21
22#[derive(Default, Clone, Debug)]
23pub struct AnalysisCache {
24    cache: Rc<RefCell<TypeMap>>,
25}
26
27impl AnalysisCache {
28    pub fn get<A: Analysis + Any>(&self, func: &mut Function, state: &GlobalState) -> Rc<A> {
29        let analysis = self.cache.borrow().get::<Rc<A>>().cloned();
30        if let Some(analysis) = analysis {
31            analysis
32        } else {
33            let analysis = Rc::new(A::init(func, state));
34            self.cache.borrow_mut().insert(analysis.clone());
35            analysis
36        }
37    }
38
39    pub fn try_get<A: Any>(&self) -> Option<Rc<A>> {
40        self.cache.borrow().get().cloned()
41    }
42
43    pub fn invalidate<A: Analysis + Any>(&self) {
44        self.cache.borrow_mut().remove::<Rc<A>>();
45    }
46}
47
48impl Function {
49    /// Fetch an analysis if cached, or run it if not.
50    pub fn analysis<A: Analysis + Any>(&mut self, state: &GlobalState) -> Rc<A> {
51        let analyses = self.analysis_cache.clone();
52        analyses.get(self, state)
53    }
54
55    /// Invalidate an analysis by removing it from the cache. The analysis is rerun when requested
56    /// again.
57    pub fn invalidate_analysis<A: Analysis + Any>(&self) {
58        self.analysis_cache.invalidate::<A>();
59    }
60
61    /// Invalidate all analyses that rely on the structure of the control flow graph.
62    pub fn invalidate_structure(&self) {
63        self.invalidate_analysis::<PostOrder>();
64        self.invalidate_analysis::<Dominators>();
65        self.invalidate_analysis::<PostDominators>();
66        self.invalidate_analysis::<Liveness>();
67        self.invalidate_analysis::<Uniformity>();
68    }
69}