Skip to main content

lift_opt/
common_subexpr.rs

1use lift_core::context::Context;
2use lift_core::pass::{AnalysisCache, Pass, PassResult};
3use std::collections::HashMap;
4
5/// Common Subexpression Elimination (CSE): detects identical operations
6/// with the same inputs and attributes, replacing duplicates with
7/// references to the first occurrence.
8#[derive(Debug)]
9pub struct CommonSubexprElimination;
10
11/// A fingerprint for an operation used to detect duplicates.
12#[derive(Hash, PartialEq, Eq, Clone, Debug)]
13struct OpFingerprint {
14    name: String,
15    inputs: Vec<lift_core::values::ValueKey>,
16    attrs_hash: String,
17}
18
19impl Pass for CommonSubexprElimination {
20    fn name(&self) -> &str {
21        "common-subexpr-elimination"
22    }
23
24    fn run(&self, ctx: &mut Context, _cache: &mut AnalysisCache) -> PassResult {
25        let mut eliminated = 0usize;
26        let mut ops_to_remove = Vec::new();
27
28        let block_keys: Vec<_> = ctx.blocks.keys().collect();
29
30        for block_key in block_keys {
31            let op_list = match ctx.blocks.get(block_key) {
32                Some(b) => b.ops.clone(),
33                None => continue,
34            };
35
36            // Map from fingerprint -> first op's result values
37            let mut seen: HashMap<OpFingerprint, Vec<lift_core::values::ValueKey>> = HashMap::new();
38
39            for &op_key in &op_list {
40                let fingerprint = match ctx.ops.get(op_key) {
41                    Some(op) => {
42                        // Skip ops with side effects (measurement, store, etc.)
43                        let name_str = ctx.strings.resolve(op.name).to_string();
44                        if name_str.contains("measure")
45                            || name_str.contains("store")
46                            || name_str.contains("send")
47                            || name_str.contains("receive")
48                            || name_str.contains("barrier")
49                            || name_str.contains("reset")
50                        {
51                            continue;
52                        }
53
54                        // Skip ops with no results (nothing to deduplicate)
55                        if op.results.is_empty() {
56                            continue;
57                        }
58
59                        let attrs_str = format!("{:?}", op.attrs);
60                        OpFingerprint {
61                            name: name_str,
62                            inputs: op.inputs.clone(),
63                            attrs_hash: attrs_str,
64                        }
65                    }
66                    None => continue,
67                };
68
69                if let Some(existing_results) = seen.get(&fingerprint) {
70                    // Found a duplicate - rewire users
71                    let dup_results = match ctx.ops.get(op_key) {
72                        Some(op) => op.results.clone(),
73                        None => continue,
74                    };
75
76                    if dup_results.len() != existing_results.len() {
77                        continue;
78                    }
79
80                    // Rewire all users of dup_results to use existing_results
81                    let all_ops: Vec<_> = ctx.ops.keys().collect();
82                    for &ok in &all_ops {
83                        if ok == op_key {
84                            continue;
85                        }
86                        if let Some(other) = ctx.ops.get_mut(ok) {
87                            for inp in &mut other.inputs {
88                                for (idx, dup_r) in dup_results.iter().enumerate() {
89                                    if *inp == *dup_r {
90                                        *inp = existing_results[idx];
91                                    }
92                                }
93                            }
94                        }
95                    }
96
97                    ops_to_remove.push(op_key);
98                    eliminated += 1;
99                } else {
100                    let results = match ctx.ops.get(op_key) {
101                        Some(op) => op.results.clone(),
102                        None => continue,
103                    };
104                    seen.insert(fingerprint, results);
105                }
106            }
107
108            // Remove from block
109            if !ops_to_remove.is_empty() {
110                if let Some(block) = ctx.blocks.get_mut(block_key) {
111                    block.ops.retain(|op| !ops_to_remove.contains(op));
112                }
113            }
114        }
115
116        // Remove from slotmap
117        for op_key in &ops_to_remove {
118            if let Some(op) = ctx.ops.remove(*op_key) {
119                for result in &op.results {
120                    ctx.values.remove(*result);
121                }
122            }
123        }
124
125        if eliminated > 0 {
126            tracing::info!(
127                pass = "cse",
128                eliminated = eliminated,
129                "Common subexpression elimination applied"
130            );
131            PassResult::Changed
132        } else {
133            PassResult::Unchanged
134        }
135    }
136
137    fn invalidates(&self) -> Vec<&str> {
138        vec!["analysis"]
139    }
140}