cubecl-opt 0.11.0-pre.1

Compiler optimizations for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

use alloc::{collections::VecDeque, vec::Vec};
use core::ops::Deref;

use crate::{Function, GlobalState, NodeIndex};
use hashbrown::{HashMap, HashSet};
use petgraph::algo::dominators;

use super::Analysis;

/// Dominator tree for the program graph
pub struct Dominators(dominators::Dominators<NodeIndex>);
/// Post dominator tree for the program graph
pub struct PostDominators(dominators::Dominators<NodeIndex>);

impl Deref for Dominators {
    type Target = dominators::Dominators<NodeIndex>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Deref for PostDominators {
    type Target = dominators::Dominators<NodeIndex>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Analysis for Dominators {
    fn init(func: &mut crate::Function, _: &GlobalState) -> Self {
        Dominators(dominators::simple_fast(&func.graph, func.root))
    }
}

impl Analysis for PostDominators {
    fn init(func: &mut crate::Function, _: &GlobalState) -> Self {
        let mut reversed = func.graph.clone();
        reversed.reverse();
        PostDominators(dominators::simple_fast(&reversed, func.ret))
    }
}

impl Dominators {
    pub fn breadth_first_nodes(&self) -> Vec<NodeIndex> {
        let mut out = Vec::new();
        let mut worklist = VecDeque::new();
        worklist.push_back(self.root());
        while let Some(node) = worklist.pop_front() {
            out.push(node);
            worklist.extend(self.immediately_dominated_by(node));
        }
        out
    }
}

/// Dominance frontiers for each block
pub struct DomFrontiers {
    /// The dominance frontiers of each block (where phi nodes must be inserted).
    dom_frontiers: HashMap<NodeIndex, HashSet<NodeIndex>>,
}

impl Deref for DomFrontiers {
    type Target = HashMap<NodeIndex, HashSet<NodeIndex>>;

    fn deref(&self) -> &Self::Target {
        &self.dom_frontiers
    }
}

impl DomFrontiers {
    /// Find dominance frontiers for each block
    pub fn new(func: &mut Function, state: &GlobalState) -> Self {
        let doms = func.analysis::<Dominators>(state);
        let nodes = func.node_ids().into_iter().map(|it| (it, HashSet::new()));
        let mut dom_frontiers: HashMap<NodeIndex, HashSet<NodeIndex>> = nodes.collect();

        for node in func.node_ids() {
            let predecessors = func.predecessors(node);
            if predecessors.len() >= 2 {
                for predecessor in predecessors {
                    let mut runner = predecessor;
                    while runner != doms.immediate_dominator(node).unwrap() {
                        dom_frontiers.get_mut(&runner).unwrap().insert(node);
                        runner = doms.immediate_dominator(runner).unwrap();
                    }
                }
            }
        }
        Self { dom_frontiers }
    }
}

impl Analysis for DomFrontiers {
    fn init(func: &mut Function, state: &GlobalState) -> Self {
        DomFrontiers::new(func, state)
    }
}