cubecl_opt/analyses/
uniformity.rs

1use cubecl_ir::{
2    Builtin, Operation, OperationReflect, Plane, Synchronization, Variable, VariableKind,
3};
4use petgraph::{graph::EdgeIndex, visit::EdgeRef};
5use std::collections::{HashMap, HashSet};
6
7use crate::{ControlFlow, NodeIndex, Optimizer};
8
9use super::Analysis;
10
11#[derive(Default, Clone)]
12pub struct Uniformity {
13    block_uniformity: HashMap<NodeIndex, bool>,
14    variable_uniformity: HashMap<Variable, bool>,
15    visited: HashSet<EdgeIndex>,
16}
17
18impl Analysis for Uniformity {
19    fn init(opt: &mut Optimizer) -> Self {
20        let mut this = Self::default();
21        this.run(opt);
22        this
23    }
24}
25
26impl Uniformity {
27    fn run(&mut self, opt: &Optimizer) {
28        let root = opt.entry();
29        self.block_uniformity.insert(root, true);
30        while self.analyze_block(opt, root).is_none() {}
31    }
32
33    fn analyze_block(&mut self, opt: &Optimizer, block_id: NodeIndex) -> Option<()> {
34        let block = opt.block(block_id);
35        let mut block_uniform = self.block_uniformity[&block_id];
36
37        for phi in block.phi_nodes.borrow().iter() {
38            let uniform = phi.entries.iter().all(|entry| {
39                let block_uniform = self.is_block_uniform(entry.block);
40                let value_uniform = self.is_var_uniform(entry.value);
41                block_uniform && value_uniform
42            }) && block_uniform;
43            self.mark_uniformity(phi.out, uniform && block_uniform)?;
44        }
45
46        for inst in block.ops.borrow().values() {
47            if inst.out.is_none() {
48                continue;
49            }
50            let out = inst.out.unwrap();
51            match &inst.operation {
52                Operation::Plane(plane) => match plane {
53                    // Elect returns true on only one unit, so it's always non-uniform
54                    // Inclusive/exclusive scans are non-uniform by definition
55                    Plane::Elect
56                    | Plane::ExclusiveSum(_)
57                    | Plane::InclusiveSum(_)
58                    | Plane::ExclusiveProd(_)
59                    | Plane::InclusiveProd(_) => self.mark_uniformity(out, false)?,
60                    // Reductions are always uniform if executed in uniform control flow
61                    Plane::Sum(_)
62                    | Plane::Prod(_)
63                    | Plane::Min(_)
64                    | Plane::Max(_)
65                    | Plane::All(_)
66                    | Plane::Any(_)
67                    | Plane::Ballot(_) => self.mark_uniformity(out, block_uniform)?,
68                    // Broadcast maps to shuffle or broadcast, if id or value is uniform, so will
69                    // the output, otherwise not.
70                    Plane::Broadcast(op) => {
71                        let input_uniform =
72                            self.is_var_uniform(op.lhs) || self.is_var_uniform(op.rhs);
73                        self.mark_uniformity(out, input_uniform && block_uniform)?;
74                    }
75                },
76                Operation::Synchronization(sync) => match sync {
77                    Synchronization::SyncUnits | Synchronization::SyncStorage => {
78                        block_uniform = true;
79                    }
80                    Synchronization::SyncProxyShared => {}
81                },
82                op => {
83                    let is_uniform =
84                        op.is_pure() && self.is_all_uniform(op.args()) && block_uniform;
85                    self.mark_uniformity(out, is_uniform)?;
86                }
87            }
88        }
89
90        match &*block.control_flow.borrow() {
91            ControlFlow::IfElse {
92                cond,
93                then,
94                or_else,
95                merge,
96            } => {
97                let is_uniform = self.is_var_uniform(*cond);
98                self.block_uniformity
99                    .insert(*then, is_uniform && block_uniform);
100                self.block_uniformity
101                    .insert(*or_else, is_uniform && block_uniform);
102                if let Some(merge) = merge {
103                    self.block_uniformity.insert(*merge, block_uniform);
104                }
105            }
106            ControlFlow::Switch {
107                value,
108                default,
109                branches,
110                merge,
111            } => {
112                let is_uniform = self.is_var_uniform(*value);
113                self.block_uniformity
114                    .insert(*default, is_uniform && block_uniform);
115                for branch in branches {
116                    self.block_uniformity
117                        .insert(branch.1, is_uniform && block_uniform);
118                }
119                if let Some(merge) = merge {
120                    self.block_uniformity.insert(*merge, block_uniform);
121                }
122            }
123            ControlFlow::Loop {
124                body,
125                continue_target,
126                merge,
127            } => {
128                // If we don't know the break condition, we can't detect whether it's uniform
129                self.block_uniformity.insert(block_id, false);
130                self.block_uniformity.insert(*body, false);
131                self.block_uniformity.insert(*continue_target, false);
132                self.block_uniformity.insert(*merge, false);
133            }
134            ControlFlow::LoopBreak {
135                break_cond,
136                body,
137                continue_target,
138                merge,
139            } => {
140                let is_uniform = self.is_var_uniform(*break_cond);
141                self.block_uniformity
142                    .insert(block_id, is_uniform && block_uniform);
143                self.block_uniformity
144                    .insert(*body, is_uniform && block_uniform);
145                self.block_uniformity
146                    .insert(*continue_target, is_uniform && block_uniform);
147                self.block_uniformity
148                    .insert(*merge, is_uniform && block_uniform);
149            }
150            ControlFlow::Return => {}
151            ControlFlow::None => {
152                let successor = opt.successors(block_id)[0];
153                self.block_uniformity
154                    .entry(successor)
155                    .and_modify(|it| {
156                        *it |= block_uniform;
157                    })
158                    .or_insert(block_uniform);
159            }
160        }
161
162        for edge in opt.program.edges(block_id) {
163            if !self.visited.contains(&edge.id()) {
164                self.visited.insert(edge.id());
165                self.analyze_block(opt, edge.target())?;
166            }
167        }
168
169        Some(())
170    }
171
172    fn mark_uniformity(&mut self, var: Variable, value: bool) -> Option<()> {
173        if let Some(val) = self.variable_uniformity.get_mut(&var) {
174            // If the value was already set before and has been invalidated, we need to revisit
175            // all edges. This only happens for loopback edges, where an uninitialized variable
176            // was assumed to be uniform but actually isn't
177            let invalidate = !value && *val;
178            *val = *val && value;
179            if invalidate {
180                self.visited.clear();
181                return None;
182            }
183        } else {
184            self.variable_uniformity.insert(var, value);
185        }
186        Some(())
187    }
188
189    fn is_all_uniform(&self, args: Option<Vec<Variable>>) -> bool {
190        args.map(|it| it.iter().all(|it| self.is_var_uniform(*it)))
191            .unwrap_or(false)
192    }
193
194    /// Whether a variable is plane uniform
195    pub fn is_var_uniform(&self, var: Variable) -> bool {
196        match var.kind {
197            VariableKind::ConstantArray { .. }
198            | VariableKind::SharedMemory { .. }
199            | VariableKind::GlobalInputArray(_)
200            | VariableKind::GlobalOutputArray(_)
201            | VariableKind::GlobalScalar(_)
202            | VariableKind::ConstantScalar(_) => true,
203            VariableKind::Builtin(builtin) => match builtin {
204                Builtin::UnitPosPlane
205                | Builtin::AbsolutePos
206                | Builtin::AbsolutePosX
207                | Builtin::AbsolutePosY
208                | Builtin::AbsolutePosZ
209                | Builtin::UnitPos
210                | Builtin::UnitPosX
211                | Builtin::UnitPosY
212                | Builtin::UnitPosZ => false,
213                Builtin::CubePos
214                | Builtin::CubePosX
215                | Builtin::CubePosY
216                | Builtin::CubePosZ
217                | Builtin::CubePosCluster
218                | Builtin::CubePosClusterX
219                | Builtin::CubePosClusterY
220                | Builtin::CubePosClusterZ
221                | Builtin::CubeDim
222                | Builtin::CubeDimX
223                | Builtin::CubeDimY
224                | Builtin::CubeDimZ
225                | Builtin::CubeClusterDim
226                | Builtin::CubeClusterDimX
227                | Builtin::CubeClusterDimY
228                | Builtin::CubeClusterDimZ
229                | Builtin::CubeCount
230                | Builtin::CubeCountX
231                | Builtin::CubeCountY
232                | Builtin::CubeCountZ
233                | Builtin::PlaneDim => true,
234            },
235            VariableKind::LocalMut { .. } => false,
236            VariableKind::LocalArray { .. }
237            | VariableKind::LocalConst { .. }
238            | VariableKind::Versioned { .. }
239            | VariableKind::Matrix { .. }
240            | VariableKind::Slice { .. }
241            | VariableKind::Barrier { .. }
242            | VariableKind::Pipeline { .. } => {
243                self.variable_uniformity.get(&var).copied().unwrap_or(true)
244            }
245            VariableKind::TensorMap(_) => true,
246        }
247    }
248
249    pub fn is_block_uniform(&self, block: NodeIndex) -> bool {
250        self.block_uniformity.get(&block).copied().unwrap_or(true)
251    }
252}
cubecl_opt/analyses/uniformity.rs

cubecl_opt/analyses/
uniformity.rs
