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::SyncCube | Synchronization::SyncStorage => {
78                        block_uniform = true;
79                    }
80                    Synchronization::SyncProxyShared => {}
81                    Synchronization::SyncPlane => {
82                        // TODO: not sure
83                    }
84                },
85                op => {
86                    let is_uniform =
87                        op.is_pure() && self.is_all_uniform(op.args()) && block_uniform;
88                    self.mark_uniformity(out, is_uniform)?;
89                }
90            }
91        }
92
93        match &*block.control_flow.borrow() {
94            ControlFlow::IfElse {
95                cond,
96                then,
97                or_else,
98                merge,
99            } => {
100                let is_uniform = self.is_var_uniform(*cond);
101                self.block_uniformity
102                    .insert(*then, is_uniform && block_uniform);
103                self.block_uniformity
104                    .insert(*or_else, is_uniform && block_uniform);
105                if let Some(merge) = merge {
106                    self.block_uniformity.insert(*merge, block_uniform);
107                }
108            }
109            ControlFlow::Switch {
110                value,
111                default,
112                branches,
113                merge,
114            } => {
115                let is_uniform = self.is_var_uniform(*value);
116                self.block_uniformity
117                    .insert(*default, is_uniform && block_uniform);
118                for branch in branches {
119                    self.block_uniformity
120                        .insert(branch.1, is_uniform && block_uniform);
121                }
122                if let Some(merge) = merge {
123                    self.block_uniformity.insert(*merge, block_uniform);
124                }
125            }
126            ControlFlow::Loop {
127                body,
128                continue_target,
129                merge,
130            } => {
131                // If we don't know the break condition, we can't detect whether it's uniform
132                self.block_uniformity.insert(block_id, false);
133                self.block_uniformity.insert(*body, false);
134                self.block_uniformity.insert(*continue_target, false);
135                self.block_uniformity.insert(*merge, false);
136            }
137            ControlFlow::LoopBreak {
138                break_cond,
139                body,
140                continue_target,
141                merge,
142            } => {
143                let is_uniform = self.is_var_uniform(*break_cond);
144                self.block_uniformity
145                    .insert(block_id, is_uniform && block_uniform);
146                self.block_uniformity
147                    .insert(*body, is_uniform && block_uniform);
148                self.block_uniformity
149                    .insert(*continue_target, is_uniform && block_uniform);
150                self.block_uniformity
151                    .insert(*merge, is_uniform && block_uniform);
152            }
153            ControlFlow::Return => {}
154            ControlFlow::None => {
155                let successor = opt.successors(block_id)[0];
156                self.block_uniformity
157                    .entry(successor)
158                    .and_modify(|it| {
159                        *it |= block_uniform;
160                    })
161                    .or_insert(block_uniform);
162            }
163        }
164
165        for edge in opt.program.edges(block_id) {
166            if !self.visited.contains(&edge.id()) {
167                self.visited.insert(edge.id());
168                self.analyze_block(opt, edge.target())?;
169            }
170        }
171
172        Some(())
173    }
174
175    fn mark_uniformity(&mut self, var: Variable, value: bool) -> Option<()> {
176        if let Some(val) = self.variable_uniformity.get_mut(&var) {
177            // If the value was already set before and has been invalidated, we need to revisit
178            // all edges. This only happens for loopback edges, where an uninitialized variable
179            // was assumed to be uniform but actually isn't
180            let invalidate = !value && *val;
181            *val = *val && value;
182            if invalidate {
183                self.visited.clear();
184                return None;
185            }
186        } else {
187            self.variable_uniformity.insert(var, value);
188        }
189        Some(())
190    }
191
192    fn is_all_uniform(&self, args: Option<Vec<Variable>>) -> bool {
193        args.map(|it| it.iter().all(|it| self.is_var_uniform(*it)))
194            .unwrap_or(false)
195    }
196
197    /// Whether a variable is plane uniform
198    pub fn is_var_uniform(&self, var: Variable) -> bool {
199        match var.kind {
200            VariableKind::ConstantArray { .. }
201            | VariableKind::SharedMemory { .. }
202            | VariableKind::GlobalInputArray(_)
203            | VariableKind::GlobalOutputArray(_)
204            | VariableKind::GlobalScalar(_)
205            | VariableKind::ConstantScalar(_) => true,
206            VariableKind::Builtin(builtin) => match builtin {
207                Builtin::UnitPosPlane
208                | Builtin::AbsolutePos
209                | Builtin::AbsolutePosX
210                | Builtin::AbsolutePosY
211                | Builtin::AbsolutePosZ
212                | Builtin::UnitPos
213                | Builtin::UnitPosX
214                | Builtin::UnitPosY
215                | Builtin::UnitPosZ => false,
216                Builtin::CubePos
217                | Builtin::CubePosX
218                | Builtin::CubePosY
219                | Builtin::CubePosZ
220                | Builtin::CubePosCluster
221                | Builtin::CubePosClusterX
222                | Builtin::CubePosClusterY
223                | Builtin::CubePosClusterZ
224                | Builtin::CubeDim
225                | Builtin::CubeDimX
226                | Builtin::CubeDimY
227                | Builtin::CubeDimZ
228                | Builtin::CubeClusterDim
229                | Builtin::CubeClusterDimX
230                | Builtin::CubeClusterDimY
231                | Builtin::CubeClusterDimZ
232                | Builtin::CubeCount
233                | Builtin::CubeCountX
234                | Builtin::CubeCountY
235                | Builtin::CubeCountZ
236                | Builtin::PlaneDim => true,
237            },
238            VariableKind::LocalMut { .. } => false,
239            VariableKind::LocalArray { .. }
240            | VariableKind::LocalConst { .. }
241            | VariableKind::Versioned { .. }
242            | VariableKind::Matrix { .. }
243            | VariableKind::Barrier { .. }
244            | VariableKind::Pipeline { .. } => {
245                self.variable_uniformity.get(&var).copied().unwrap_or(true)
246            }
247            VariableKind::TensorMap(_) => true,
248        }
249    }
250
251    pub fn is_block_uniform(&self, block: NodeIndex) -> bool {
252        self.block_uniformity.get(&block).copied().unwrap_or(true)
253    }
254}
cubecl_opt/analyses/uniformity.rs

cubecl_opt/analyses/
uniformity.rs
