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::GetStorageKey(_) => None,
108        InstOp::GetElemPtr {
109            base,
110            elem_ptr_ty,
111            indices,
112        } => Some(Expr::GetElemPtr {
113            base: vntable.value_map.get(base).cloned().unwrap(),
114            elem_ptr_ty: *elem_ptr_ty,
115            indices: indices
116                .iter()
117                .map(|idx| vntable.value_map.get(idx).cloned().unwrap())
118                .collect(),
119        }),
120        InstOp::IntToPtr(val, ty) => Some(Expr::IntToPtr(
121            vntable.value_map.get(val).cloned().unwrap(),
122            *ty,
123        )),
124        InstOp::Load(_) => None,
125        InstOp::MemCopyBytes { .. } => None,
126        InstOp::MemCopyVal { .. } => None,
127        InstOp::MemClearVal { .. } => None,
128        InstOp::Nop => None,
129        InstOp::PtrToInt(val, ty) => Some(Expr::PtrToInt(
130            vntable.value_map.get(val).cloned().unwrap(),
131            *ty,
132        )),
133        InstOp::Ret(_, _) => None,
134        InstOp::Store { .. } => None,
135    }
136}
137
138/// Convert a PHI argument to Expr
139fn phi_to_expr(context: &Context, vntable: &VNTable, phi_arg: Value) -> Expr {
140    let phi_arg = phi_arg.get_argument(context).unwrap();
141    let phi_args = phi_arg
142        .block
143        .pred_iter(context)
144        .map(|pred| {
145            let incoming_val = phi_arg
146                .get_val_coming_from(context, pred)
147                .expect("No parameter from predecessor");
148            vntable.value_map.get(&incoming_val).cloned().unwrap()
149        })
150        .collect();
151    Expr::Phi(phi_args)
152}
153
154#[derive(Default)]
155struct VNTable {
156    value_map: FxHashMap<Value, ValueNumber>,
157    expr_map: FxHashMap<Expr, ValueNumber>,
158}
159
160impl Debug for VNTable {
161    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162        writeln!(f, "value_map:")?;
163        self.value_map.iter().for_each(|(key, value)| {
164            if format!("v{:?}", key.0.data()) == "v620v3" {
165                writeln!(f, "\tv{:?} -> {:?}", key.0.data(), value).expect("writeln! failed");
166            }
167        });
168        Ok(())
169    }
170}
171
172/// Wrapper around [DomTree::dominates] to work at instruction level.
173/// Does `inst1` dominate `inst2` ?
174fn dominates(context: &Context, dom_tree: &DomTree, inst1: Value, inst2: Value) -> bool {
175    let block1 = match &context.values[inst1.0].value {
176        crate::ValueDatum::Argument(arg) => arg.block,
177        crate::ValueDatum::Constant(_) => {
178            panic!("Shouldn't be querying dominance info for constants")
179        }
180        crate::ValueDatum::Instruction(i) => i.parent,
181    };
182    let block2 = match &context.values[inst2.0].value {
183        crate::ValueDatum::Argument(arg) => arg.block,
184        crate::ValueDatum::Constant(_) => {
185            panic!("Shouldn't be querying dominance info for constants")
186        }
187        crate::ValueDatum::Instruction(i) => i.parent,
188    };
189
190    if block1 == block2 {
191        let inst1_idx = block1
192            .instruction_iter(context)
193            .position(|inst| inst == inst1)
194            // Not found indicates a block argument
195            .unwrap_or(0);
196        let inst2_idx = block1
197            .instruction_iter(context)
198            .position(|inst| inst == inst2)
199            // Not found indicates a block argument
200            .unwrap_or(0);
201        inst1_idx < inst2_idx
202    } else {
203        dom_tree.dominates(block1, block2)
204    }
205}
206
207pub fn cse(
208    context: &mut Context,
209    analyses: &AnalysisResults,
210    function: Function,
211) -> Result<bool, IrError> {
212    let mut vntable = VNTable::default();
213
214    // Function arg values map to themselves.
215    for arg in function.args_iter(context) {
216        vntable.value_map.insert(arg.1, ValueNumber::Number(arg.1));
217    }
218
219    // Map all other arg values map to Top.
220    for block in function.block_iter(context).skip(1) {
221        for arg in block.arg_iter(context) {
222            vntable.value_map.insert(*arg, ValueNumber::Top);
223        }
224    }
225
226    // Initialize all instructions and constants. Constants need special treatment.
227    // They don't have PartialEq implemented. So we need to value number them manually.
228    // This map maps the hash of a constant value to all possible collisions of it.
229    let mut const_map = FxHashMap::<u64, Vec<Value>>::default();
230    for (_, inst) in function.instruction_iter(context) {
231        vntable.value_map.insert(inst, ValueNumber::Top);
232        for (const_opd_val, const_opd_const) in inst
233            .get_instruction(context)
234            .unwrap()
235            .op
236            .get_operands()
237            .iter()
238            .filter_map(|opd| opd.get_constant(context).map(|copd| (opd, copd)))
239        {
240            let mut state = FxHasher::default();
241            const_opd_const.hash(&mut state);
242            let hash = state.finish();
243            if let Some(existing_const) = const_map.get(&hash).and_then(|consts| {
244                consts.iter().find(|val| {
245                    let c = val
246                        .get_constant(context)
247                        .expect("const_map can only contain consts");
248                    const_opd_const == c
249                })
250            }) {
251                vntable
252                    .value_map
253                    .insert(*const_opd_val, ValueNumber::Number(*existing_const));
254            } else {
255                const_map
256                    .entry(hash)
257                    .and_modify(|consts| consts.push(*const_opd_val))
258                    .or_insert_with(|| vec![*const_opd_val]);
259                vntable
260                    .value_map
261                    .insert(*const_opd_val, ValueNumber::Number(*const_opd_val));
262            }
263        }
264    }
265
266    // We need to iterate over the blocks in RPO.
267    let post_order: &PostOrder = analyses.get_analysis_result(function);
268
269    // RPO based value number (Sec 4.2).
270    let mut changed = true;
271    while changed {
272        changed = false;
273        // For each block in RPO:
274        for (block_idx, block) in post_order.po_to_block.iter().rev().enumerate() {
275            // Process PHIs and then the other instructions.
276            if block_idx != 0 {
277                // Entry block arguments are not PHIs.
278                for (phi, expr_opt) in block
279                    .arg_iter(context)
280                    .map(|arg| (*arg, Some(phi_to_expr(context, &vntable, *arg))))
281                    .collect_vec()
282                {
283                    let expr = expr_opt.expect("PHIs must always translate to a valid Expr");
284                    // We first try to see if PHIs can be simplified into a single value.
285                    let vn = {
286                        let Expr::Phi(ref phi_args) = expr else {
287                            panic!("Expr must be a PHI")
288                        };
289                        phi_args
290                            .iter()
291                            .map(|vn| Some(*vn))
292                            .reduce(|vn1, vn2| {
293                                // Here `None` indicates Bottom of the lattice.
294                                if let (Some(vn1), Some(vn2)) = (vn1, vn2) {
295                                    match (vn1, vn2) {
296                                        (ValueNumber::Top, ValueNumber::Top) => {
297                                            Some(ValueNumber::Top)
298                                        }
299                                        (ValueNumber::Top, ValueNumber::Number(vn))
300                                        | (ValueNumber::Number(vn), ValueNumber::Top) => {
301                                            Some(ValueNumber::Number(vn))
302                                        }
303                                        (ValueNumber::Number(vn1), ValueNumber::Number(vn2)) => {
304                                            (vn1 == vn2).then_some(ValueNumber::Number(vn1))
305                                        }
306                                    }
307                                } else {
308                                    None
309                                }
310                            })
311                            .flatten()
312                            // The PHI couldn't be simplified to a single ValueNumber.
313                            .unwrap_or(ValueNumber::Number(phi))
314                    };
315
316                    match vntable.value_map.entry(phi) {
317                        hash_map::Entry::Occupied(occ) if *occ.get() == vn => {}
318                        _ => {
319                            changed = true;
320                            vntable.value_map.insert(phi, vn);
321                        }
322                    }
323                }
324            }
325
326            for (inst, expr_opt) in block
327                .instruction_iter(context)
328                .map(|instr| (instr, instr_to_expr(context, &vntable, instr)))
329                .collect_vec()
330            {
331                // lookup(expr, x)
332                let vn = if let Some(expr) = expr_opt {
333                    match vntable.expr_map.entry(expr) {
334                        hash_map::Entry::Occupied(occ) => *occ.get(),
335                        hash_map::Entry::Vacant(vac) => *(vac.insert(ValueNumber::Number(inst))),
336                    }
337                } else {
338                    // Instructions that always map to their own value number
339                    // (i.e., they can never be equal to some other instruction).
340                    ValueNumber::Number(inst)
341                };
342                match vntable.value_map.entry(inst) {
343                    hash_map::Entry::Occupied(occ) if *occ.get() == vn => {}
344                    _ => {
345                        changed = true;
346                        vntable.value_map.insert(inst, vn);
347                    }
348                }
349            }
350        }
351        vntable.expr_map.clear();
352    }
353
354    // create a partition of congruent (equal) values.
355    let mut partition = FxHashMap::<ValueNumber, FxHashSet<Value>>::default();
356    vntable.value_map.iter().for_each(|(v, vn)| {
357        // If v is a constant or its value number is itself, don't add to the partition.
358        // The latter condition is so that we have only > 1 sized partitions.
359        if v.is_constant(context)
360            || matches!(vn, ValueNumber::Top)
361            || matches!(vn, ValueNumber::Number(v2) if (v == v2 || v2.is_constant(context)))
362        {
363            return;
364        }
365        partition
366            .entry(*vn)
367            .and_modify(|part| {
368                part.insert(*v);
369            })
370            .or_insert(vec![*v].into_iter().collect());
371    });
372
373    // For convenience, now add back back `v` into `partition[VN[v]]` if it isn't already there.
374    partition.iter_mut().for_each(|(vn, v_part)| {
375        let ValueNumber::Number(v) = vn else {
376            panic!("We cannot have Top at this point");
377        };
378        v_part.insert(*v);
379        assert!(
380            v_part.len() > 1,
381            "We've only created partitions with size greater than 1"
382        );
383    });
384
385    // There are two ways to replace congruent values (see the paper cited, Sec 5).
386    // 1. Dominator based. If v1 and v2 are equal, v1 dominates v2, we just remove v2
387    // and replace its uses with v1. Simple, and what we're going to do.
388    // 2. AVAIL based. More powerful, but also requires a data-flow-analysis for AVAIL
389    // and later on, mem2reg again since replacements will need breaking SSA.
390    let dom_tree: &DomTree = analyses.get_analysis_result(function);
391    let mut replace_map = FxHashMap::<Value, Value>::default();
392    let mut modified = false;
393    // Check every set in the partition.
394    partition.iter().for_each(|(_leader, vals)| {
395        // Iterate over every pair of values, checking if one can replace the other.
396        for v_pair in vals.iter().combinations(2) {
397            let (v1, v2) = (*v_pair[0], *v_pair[1]);
398            if dominates(context, dom_tree, v1, v2) {
399                modified = true;
400                replace_map.insert(v2, v1);
401            } else if dominates(context, dom_tree, v2, v1) {
402                modified = true;
403                replace_map.insert(v1, v2);
404            }
405        }
406    });
407
408    function.replace_values(context, &replace_map, None);
409
410    Ok(modified)
411}
sway_ir/optimize/cse.rs

sway_ir/optimize/
cse.rs
