lift_opt/
dce.rs

1use lift_core::context::Context;
2use lift_core::pass::{AnalysisCache, Pass, PassResult};
3use std::collections::HashSet;
4
5#[derive(Debug)]
6pub struct DeadCodeElimination;
7
8impl Pass for DeadCodeElimination {
9    fn name(&self) -> &str {
10        "dce"
11    }
12
13    fn run(&self, ctx: &mut Context, _cache: &mut AnalysisCache) -> PassResult {
14        let mut used_values: HashSet<lift_core::values::ValueKey> = HashSet::new();
15
16        // Mark phase: find all values used as inputs
17        for (_op_key, op) in &ctx.ops {
18            for &input in &op.inputs {
19                used_values.insert(input);
20            }
21        }
22
23        // Sweep phase: find ops whose results are never used
24        let mut dead_ops: Vec<lift_core::operations::OpKey> = Vec::new();
25        for (op_key, op) in &ctx.ops {
26            let op_name = ctx.strings.resolve(op.name);
27            // Never remove terminators or side-effecting ops
28            if op_name == "core.return" || op_name == "core.br" || op_name == "core.cond_br" {
29                continue;
30            }
31            // Don't remove quantum ops (they have side effects on qubits)
32            if op_name.starts_with("quantum.") {
33                continue;
34            }
35
36            if !op.results.is_empty() && op.results.iter().all(|r| !used_values.contains(r)) {
37                dead_ops.push(op_key);
38            }
39        }
40
41        if dead_ops.is_empty() {
42            return PassResult::Unchanged;
43        }
44
45        // Remove dead ops from their parent blocks
46        let dead_set: HashSet<_> = dead_ops.iter().copied().collect();
47        for (_block_key, block) in &mut ctx.blocks {
48            block.ops.retain(|op| !dead_set.contains(op));
49        }
50
51        // Remove the ops and their result values
52        for op_key in &dead_ops {
53            if let Some(op) = ctx.ops.remove(*op_key) {
54                for result in &op.results {
55                    ctx.values.remove(*result);
56                }
57            }
58        }
59
60        tracing::info!("DCE removed {} dead operations", dead_ops.len());
61        PassResult::Changed
62    }
63
64    fn invalidates(&self) -> Vec<&str> {
65        vec!["analysis"]
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72    use lift_core::attributes::Attributes;
73    use lift_core::location::Location;
74
75    #[test]
76    fn test_dce_removes_unused_ops() {
77        let mut ctx = Context::new();
78        let f32_ty = ctx.make_float_type(32);
79
80        let block = ctx.create_block();
81        let arg = ctx.create_block_arg(block, f32_ty);
82
83        // This op's result is never used -> dead
84        let (dead_op, _dead_results) = ctx.create_op(
85            "tensor.relu",
86            "tensor",
87            vec![arg],
88            vec![f32_ty],
89            Attributes::new(),
90            Location::unknown(),
91        );
92        ctx.add_op_to_block(block, dead_op);
93
94        // Return the original arg
95        let (ret_op, _) = ctx.create_op(
96            "core.return",
97            "core",
98            vec![arg],
99            vec![],
100            Attributes::new(),
101            Location::unknown(),
102        );
103        ctx.add_op_to_block(block, ret_op);
104
105        let mut cache = AnalysisCache::new();
106        let result = DeadCodeElimination.run(&mut ctx, &mut cache);
107        assert_eq!(result, PassResult::Changed);
108
109        // Only the return should remain
110        let block_data = ctx.get_block(block).unwrap();
111        assert_eq!(block_data.ops.len(), 1);
112    }
113}
lift_opt/dce.rs

lift_opt/
dce.rs
