sway_ir/optimize/
cse.rs

1//! Value numbering based common subexpression elimination.
2//! Reference: Value Driven Redundancy Elimination - Loren Taylor Simpson.
3
4use core::panic;
5use itertools::Itertools;
6use rustc_hash::{FxHashMap, FxHashSet, FxHasher};
7use slotmap::Key;
8use std::{
9    collections::hash_map,
10    fmt::Debug,
11    hash::{Hash, Hasher},
12};
13
14use crate::{
15    AnalysisResults, BinaryOpKind, Context, DebugWithContext, DomTree, Function, InstOp, IrError,
16    Pass, PassMutability, PostOrder, Predicate, ScopedPass, Type, UnaryOpKind, Value,
17    DOMINATORS_NAME, POSTORDER_NAME,
18};
19
20pub const CSE_NAME: &str = "cse";
21
22pub fn create_cse_pass() -> Pass {
23    Pass {
24        name: CSE_NAME,
25        descr: "Common subexpression elimination",
26        runner: ScopedPass::FunctionPass(PassMutability::Transform(cse)),
27        deps: vec![POSTORDER_NAME, DOMINATORS_NAME],
28    }
29}
30
31#[derive(Clone, Copy, Eq, PartialEq, Hash, DebugWithContext)]
32enum ValueNumber {
33    // Top of the lattice = Don't know = uninitialized
34    Top,
35    // Belongs to a congruence class represented by the inner value.
36    Number(Value),
37}
38
39impl Debug for ValueNumber {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        match self {
42            Self::Top => write!(f, "Top"),
43            Self::Number(arg0) => write!(f, "v{:?}", arg0.0.data()),
44        }
45    }
46}
47
48#[derive(Clone, Debug, Eq, PartialEq, Hash, DebugWithContext)]
49enum Expr {
50    Phi(Vec<ValueNumber>),
51    UnaryOp {
52        op: UnaryOpKind,
53        arg: ValueNumber,
54    },
55    BinaryOp {
56        op: BinaryOpKind,
57        arg1: ValueNumber,
58        arg2: ValueNumber,
59    },
60    BitCast(ValueNumber, Type),
61    CastPtr(ValueNumber, Type),
62    Cmp(Predicate, ValueNumber, ValueNumber),
63    GetElemPtr {
64        base: ValueNumber,
65        elem_ptr_ty: Type,
66        indices: Vec<ValueNumber>,
67    },
68    IntToPtr(ValueNumber, Type),
69    PtrToInt(ValueNumber, Type),
70}
71
72/// Convert an instruction to an expression for hashing
73/// Instructions that we don't handle will have their value numbers be equal to themselves.
74fn instr_to_expr(context: &Context, vntable: &VNTable, instr: Value) -> Option<Expr> {
75    match &instr.get_instruction(context).unwrap().op {
76        InstOp::AsmBlock(_, _) => None,
77        InstOp::UnaryOp { op, arg } => Some(Expr::UnaryOp {
78            op: *op,
79            arg: vntable.value_map.get(arg).cloned().unwrap(),
80        }),
81        InstOp::BinaryOp { op, arg1, arg2 } => Some(Expr::BinaryOp {
82            op: *op,
83            arg1: vntable.value_map.get(arg1).cloned().unwrap(),
84            arg2: vntable.value_map.get(arg2).cloned().unwrap(),
85        }),
86        InstOp::BitCast(val, ty) => Some(Expr::BitCast(
87            vntable.value_map.get(val).cloned().unwrap(),
88            *ty,
89        )),
90        InstOp::Branch(_) => None,
91        InstOp::Call(_, _) => None,
92        InstOp::CastPtr(val, ty) => Some(Expr::CastPtr(
93            vntable.value_map.get(val).cloned().unwrap(),
94            *ty,
95        )),
96        InstOp::Cmp(pred, val1, val2) => Some(Expr::Cmp(
97            *pred,
98            vntable.value_map.get(val1).cloned().unwrap(),
99            vntable.value_map.get(val2).cloned().unwrap(),
100        )),
101        InstOp::ConditionalBranch { .. } => None,
102        InstOp::ContractCall { .. } => None,
103        InstOp::FuelVm(_) => None,
104        InstOp::GetLocal(_) => None,
105        InstOp::GetGlobal(_) => None,
106        InstOp::GetConfig(_, _) => None,
107        InstOp::GetElemPtr {
108            base,
109            elem_ptr_ty,
110            indices,
111        } => Some(Expr::GetElemPtr {
112            base: vntable.value_map.get(base).cloned().unwrap(),
113            elem_ptr_ty: *elem_ptr_ty,
114            indices: indices
115                .iter()
116                .map(|idx| vntable.value_map.get(idx).cloned().unwrap())
117                .collect(),
118        }),
119        InstOp::IntToPtr(val, ty) => Some(Expr::IntToPtr(
120            vntable.value_map.get(val).cloned().unwrap(),
121            *ty,
122        )),
123        InstOp::Load(_) => None,
124        InstOp::MemCopyBytes { .. } => None,
125        InstOp::MemCopyVal { .. } => None,
126        InstOp::Nop => None,
127        InstOp::PtrToInt(val, ty) => Some(Expr::PtrToInt(
128            vntable.value_map.get(val).cloned().unwrap(),
129            *ty,
130        )),
131        InstOp::Ret(_, _) => None,
132        InstOp::Store { .. } => None,
133    }
134}
135
136/// Convert a PHI argument to Expr
137fn phi_to_expr(context: &Context, vntable: &VNTable, phi_arg: Value) -> Expr {
138    let phi_arg = phi_arg.get_argument(context).unwrap();
139    let phi_args = phi_arg
140        .block
141        .pred_iter(context)
142        .map(|pred| {
143            let incoming_val = phi_arg
144                .get_val_coming_from(context, pred)
145                .expect("No parameter from predecessor");
146            vntable.value_map.get(&incoming_val).cloned().unwrap()
147        })
148        .collect();
149    Expr::Phi(phi_args)
150}
151
152#[derive(Default)]
153struct VNTable {
154    value_map: FxHashMap<Value, ValueNumber>,
155    expr_map: FxHashMap<Expr, ValueNumber>,
156}
157
158impl Debug for VNTable {
159    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
160        writeln!(f, "value_map:")?;
161        self.value_map.iter().for_each(|(key, value)| {
162            if format!("v{:?}", key.0.data()) == "v620v3" {
163                writeln!(f, "\tv{:?} -> {:?}", key.0.data(), value).expect("writeln! failed");
164            }
165        });
166        Ok(())
167    }
168}
169
170/// Wrapper around [DomTree::dominates] to work at instruction level.
171/// Does `inst1` dominate `inst2` ?
172fn dominates(context: &Context, dom_tree: &DomTree, inst1: Value, inst2: Value) -> bool {
173    let block1 = match &context.values[inst1.0].value {
174        crate::ValueDatum::Argument(arg) => arg.block,
175        crate::ValueDatum::Constant(_) => {
176            panic!("Shouldn't be querying dominance info for constants")
177        }
178        crate::ValueDatum::Instruction(i) => i.parent,
179    };
180    let block2 = match &context.values[inst2.0].value {
181        crate::ValueDatum::Argument(arg) => arg.block,
182        crate::ValueDatum::Constant(_) => {
183            panic!("Shouldn't be querying dominance info for constants")
184        }
185        crate::ValueDatum::Instruction(i) => i.parent,
186    };
187
188    if block1 == block2 {
189        let inst1_idx = block1
190            .instruction_iter(context)
191            .position(|inst| inst == inst1)
192            // Not found indicates a block argument
193            .unwrap_or(0);
194        let inst2_idx = block1
195            .instruction_iter(context)
196            .position(|inst| inst == inst2)
197            // Not found indicates a block argument
198            .unwrap_or(0);
199        inst1_idx < inst2_idx
200    } else {
201        dom_tree.dominates(block1, block2)
202    }
203}
204
205pub fn cse(
206    context: &mut Context,
207    analyses: &AnalysisResults,
208    function: Function,
209) -> Result<bool, IrError> {
210    let mut vntable = VNTable::default();
211
212    // Function arg values map to themselves.
213    for arg in function.args_iter(context) {
214        vntable.value_map.insert(arg.1, ValueNumber::Number(arg.1));
215    }
216
217    // Map all other arg values map to Top.
218    for block in function.block_iter(context).skip(1) {
219        for arg in block.arg_iter(context) {
220            vntable.value_map.insert(*arg, ValueNumber::Top);
221        }
222    }
223
224    // Initialize all instructions and constants. Constants need special treatment.
225    // They don't have PartialEq implemented. So we need to value number them manually.
226    // This map maps the hash of a constant value to all possible collisions of it.
227    let mut const_map = FxHashMap::<u64, Vec<Value>>::default();
228    for (_, inst) in function.instruction_iter(context) {
229        vntable.value_map.insert(inst, ValueNumber::Top);
230        for (const_opd_val, const_opd_const) in inst
231            .get_instruction(context)
232            .unwrap()
233            .op
234            .get_operands()
235            .iter()
236            .filter_map(|opd| opd.get_constant(context).map(|copd| (opd, copd)))
237        {
238            let mut state = FxHasher::default();
239            const_opd_const.hash(&mut state);
240            let hash = state.finish();
241            if let Some(existing_const) = const_map.get(&hash).and_then(|consts| {
242                consts.iter().find(|val| {
243                    let c = val
244                        .get_constant(context)
245                        .expect("const_map can only contain consts");
246                    const_opd_const == c
247                })
248            }) {
249                vntable
250                    .value_map
251                    .insert(*const_opd_val, ValueNumber::Number(*existing_const));
252            } else {
253                const_map
254                    .entry(hash)
255                    .and_modify(|consts| consts.push(*const_opd_val))
256                    .or_insert_with(|| vec![*const_opd_val]);
257                vntable
258                    .value_map
259                    .insert(*const_opd_val, ValueNumber::Number(*const_opd_val));
260            }
261        }
262    }
263
264    // We need to iterate over the blocks in RPO.
265    let post_order: &PostOrder = analyses.get_analysis_result(function);
266
267    // RPO based value number (Sec 4.2).
268    let mut changed = true;
269    while changed {
270        changed = false;
271        // For each block in RPO:
272        for (block_idx, block) in post_order.po_to_block.iter().rev().enumerate() {
273            // Process PHIs and then the other instructions.
274            if block_idx != 0 {
275                // Entry block arguments are not PHIs.
276                for (phi, expr_opt) in block
277                    .arg_iter(context)
278                    .map(|arg| (*arg, Some(phi_to_expr(context, &vntable, *arg))))
279                    .collect_vec()
280                {
281                    let expr = expr_opt.expect("PHIs must always translate to a valid Expr");
282                    // We first try to see if PHIs can be simplified into a single value.
283                    let vn = {
284                        let Expr::Phi(ref phi_args) = expr else {
285                            panic!("Expr must be a PHI")
286                        };
287                        phi_args
288                            .iter()
289                            .map(|vn| Some(*vn))
290                            .reduce(|vn1, vn2| {
291                                // Here `None` indicates Bottom of the lattice.
292                                if let (Some(vn1), Some(vn2)) = (vn1, vn2) {
293                                    match (vn1, vn2) {
294                                        (ValueNumber::Top, ValueNumber::Top) => {
295                                            Some(ValueNumber::Top)
296                                        }
297                                        (ValueNumber::Top, ValueNumber::Number(vn))
298                                        | (ValueNumber::Number(vn), ValueNumber::Top) => {
299                                            Some(ValueNumber::Number(vn))
300                                        }
301                                        (ValueNumber::Number(vn1), ValueNumber::Number(vn2)) => {
302                                            (vn1 == vn2).then_some(ValueNumber::Number(vn1))
303                                        }
304                                    }
305                                } else {
306                                    None
307                                }
308                            })
309                            .flatten()
310                            // The PHI couldn't be simplified to a single ValueNumber.
311                            .unwrap_or(ValueNumber::Number(phi))
312                    };
313
314                    match vntable.value_map.entry(phi) {
315                        hash_map::Entry::Occupied(occ) if *occ.get() == vn => {}
316                        _ => {
317                            changed = true;
318                            vntable.value_map.insert(phi, vn);
319                        }
320                    }
321                }
322            }
323
324            for (inst, expr_opt) in block
325                .instruction_iter(context)
326                .map(|instr| (instr, instr_to_expr(context, &vntable, instr)))
327                .collect_vec()
328            {
329                // lookup(expr, x)
330                let vn = if let Some(expr) = expr_opt {
331                    match vntable.expr_map.entry(expr) {
332                        hash_map::Entry::Occupied(occ) => *occ.get(),
333                        hash_map::Entry::Vacant(vac) => *(vac.insert(ValueNumber::Number(inst))),
334                    }
335                } else {
336                    // Instructions that always map to their own value number
337                    // (i.e., they can never be equal to some other instruction).
338                    ValueNumber::Number(inst)
339                };
340                match vntable.value_map.entry(inst) {
341                    hash_map::Entry::Occupied(occ) if *occ.get() == vn => {}
342                    _ => {
343                        changed = true;
344                        vntable.value_map.insert(inst, vn);
345                    }
346                }
347            }
348        }
349        vntable.expr_map.clear();
350    }
351
352    // create a partition of congruent (equal) values.
353    let mut partition = FxHashMap::<ValueNumber, FxHashSet<Value>>::default();
354    vntable.value_map.iter().for_each(|(v, vn)| {
355        // If v is a constant or its value number is itself, don't add to the partition.
356        // The latter condition is so that we have only > 1 sized partitions.
357        if v.is_constant(context)
358            || matches!(vn, ValueNumber::Top)
359            || matches!(vn, ValueNumber::Number(v2) if (v == v2 || v2.is_constant(context)))
360        {
361            return;
362        }
363        partition
364            .entry(*vn)
365            .and_modify(|part| {
366                part.insert(*v);
367            })
368            .or_insert(vec![*v].into_iter().collect());
369    });
370
371    // For convenience, now add back back `v` into `partition[VN[v]]` if it isn't already there.
372    partition.iter_mut().for_each(|(vn, v_part)| {
373        let ValueNumber::Number(v) = vn else {
374            panic!("We cannot have Top at this point");
375        };
376        v_part.insert(*v);
377        assert!(
378            v_part.len() > 1,
379            "We've only created partitions with size greater than 1"
380        );
381    });
382
383    // There are two ways to replace congruent values (see the paper cited, Sec 5).
384    // 1. Dominator based. If v1 and v2 are equal, v1 dominates v2, we just remove v2
385    // and replace its uses with v1. Simple, and what we're going to do.
386    // 2. AVAIL based. More powerful, but also requires a data-flow-analysis for AVAIL
387    // and later on, mem2reg again since replacements will need breaking SSA.
388    let dom_tree: &DomTree = analyses.get_analysis_result(function);
389    let mut replace_map = FxHashMap::<Value, Value>::default();
390    let mut modified = false;
391    // Check every set in the partition.
392    partition.iter().for_each(|(_leader, vals)| {
393        // Iterate over every pair of values, checking if one can replace the other.
394        for v_pair in vals.iter().combinations(2) {
395            let (v1, v2) = (*v_pair[0], *v_pair[1]);
396            if dominates(context, dom_tree, v1, v2) {
397                modified = true;
398                replace_map.insert(v2, v1);
399            } else if dominates(context, dom_tree, v2, v1) {
400                modified = true;
401                replace_map.insert(v1, v2);
402            }
403        }
404    });
405
406    function.replace_values(context, &replace_map, None);
407
408    Ok(modified)
409}
sway_ir/optimize/cse.rs

sway_ir/optimize/
cse.rs
