Skip to content

Commit bf54196

Browse files
benchclaude
andcommitted
feat: add redundancy elimination and loop optimization passes
Add four post-lowering optimization passes: - local_cse: Local common subexpression elimination within blocks - gvn: Global value numbering across blocks using dominator tree - licm: Loop invariant code motion - branch_fold: Branch condition simplification Also includes shared optimizer utilities and refactored structural passes. WIP: Tests and full integration in progress. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
1 parent 136373b commit bf54196

10 files changed

Lines changed: 4636 additions & 28 deletions

File tree

Lines changed: 365 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,365 @@
1+
//! Branch condition folding.
2+
//!
3+
//! Simplifies `BranchIf` terminators by looking at the instruction that
4+
//! defines the condition variable:
5+
//!
6+
//! - `Eqz(x)` as condition → swap branch targets, use `x` directly
7+
//! - `Ne(x, 0)` as condition → use `x` directly
8+
//! - `Eq(x, 0)` as condition → swap branch targets, use `x` directly
9+
//!
10+
//! After substitution, the defining instruction becomes dead (single use was
11+
//! the branch) and is cleaned up by `dead_instrs`.
12+
13+
use super::utils::{build_global_use_count, instr_dest};
14+
use crate::ir::{BinOp, IrFunction, IrInstr, IrTerminator, IrValue, UnOp, VarId};
15+
use std::collections::HashMap;
16+
17+
pub fn eliminate(func: &mut IrFunction) {
18+
loop {
19+
let global_uses = build_global_use_count(func);
20+
if !fold_one(func, &global_uses) {
21+
break;
22+
}
23+
}
24+
}
25+
26+
/// Attempt a single branch fold across the function. Returns `true` if a
27+
/// change was made.
28+
fn fold_one(func: &mut IrFunction, global_uses: &HashMap<VarId, usize>) -> bool {
29+
// Build a map of VarId → defining instruction info.
30+
// We only care about single-use vars defined by Eqz, Ne(x,0), or Eq(x,0).
31+
let mut var_defs: HashMap<VarId, VarDef> = HashMap::new();
32+
33+
// Also build a global constant map for checking if an operand is zero.
34+
let global_consts = build_global_const_map(func);
35+
36+
for block in &func.blocks {
37+
let mut local_consts = global_consts.clone();
38+
for instr in &block.instructions {
39+
if let IrInstr::Const { dest, value } = instr {
40+
local_consts.insert(*dest, *value);
41+
}
42+
43+
if let Some(dest) = instr_dest(instr) {
44+
match instr {
45+
IrInstr::UnOp {
46+
op: UnOp::I32Eqz | UnOp::I64Eqz,
47+
operand,
48+
..
49+
} => {
50+
var_defs.insert(dest, VarDef::Eqz(*operand));
51+
}
52+
IrInstr::BinOp {
53+
op: BinOp::I32Ne | BinOp::I64Ne,
54+
lhs,
55+
rhs,
56+
..
57+
} => {
58+
if is_zero(rhs, &local_consts) {
59+
var_defs.insert(dest, VarDef::NeZero(*lhs));
60+
} else if is_zero(lhs, &local_consts) {
61+
var_defs.insert(dest, VarDef::NeZero(*rhs));
62+
}
63+
}
64+
IrInstr::BinOp {
65+
op: BinOp::I32Eq | BinOp::I64Eq,
66+
lhs,
67+
rhs,
68+
..
69+
} => {
70+
if is_zero(rhs, &local_consts) {
71+
var_defs.insert(dest, VarDef::EqZero(*lhs));
72+
} else if is_zero(lhs, &local_consts) {
73+
var_defs.insert(dest, VarDef::EqZero(*rhs));
74+
}
75+
}
76+
_ => {}
77+
}
78+
}
79+
}
80+
}
81+
82+
// Now scan terminators for BranchIf with a foldable condition.
83+
for block in &mut func.blocks {
84+
let condition = match &block.terminator {
85+
IrTerminator::BranchIf { condition, .. } => *condition,
86+
_ => continue,
87+
};
88+
89+
// Only fold if the condition has exactly one use (the BranchIf).
90+
if global_uses.get(&condition).copied().unwrap_or(0) != 1 {
91+
continue;
92+
}
93+
94+
let def = match var_defs.get(&condition) {
95+
Some(d) => d,
96+
None => continue,
97+
};
98+
99+
match def {
100+
VarDef::Eqz(inner) | VarDef::EqZero(inner) => {
101+
// eqz(x) != 0 ≡ x == 0, so swap targets and use x
102+
if let IrTerminator::BranchIf {
103+
condition: cond,
104+
if_true,
105+
if_false,
106+
} = &mut block.terminator
107+
{
108+
*cond = *inner;
109+
std::mem::swap(if_true, if_false);
110+
}
111+
return true;
112+
}
113+
VarDef::NeZero(inner) => {
114+
// ne(x, 0) != 0 ≡ x != 0, so just use x
115+
if let IrTerminator::BranchIf {
116+
condition: cond, ..
117+
} = &mut block.terminator
118+
{
119+
*cond = *inner;
120+
}
121+
return true;
122+
}
123+
}
124+
}
125+
126+
false
127+
}
128+
129+
#[derive(Clone, Copy)]
130+
enum VarDef {
131+
Eqz(VarId),
132+
NeZero(VarId),
133+
EqZero(VarId),
134+
}
135+
136+
fn is_zero(var: &VarId, consts: &HashMap<VarId, IrValue>) -> bool {
137+
matches!(
138+
consts.get(var),
139+
Some(IrValue::I32(0)) | Some(IrValue::I64(0))
140+
)
141+
}
142+
143+
fn build_global_const_map(func: &IrFunction) -> HashMap<VarId, IrValue> {
144+
let mut total_defs: HashMap<VarId, usize> = HashMap::new();
145+
let mut const_defs: HashMap<VarId, IrValue> = HashMap::new();
146+
147+
for block in &func.blocks {
148+
for instr in &block.instructions {
149+
if let Some(dest) = instr_dest(instr) {
150+
*total_defs.entry(dest).or_insert(0) += 1;
151+
if let IrInstr::Const { dest, value } = instr {
152+
const_defs.insert(*dest, *value);
153+
}
154+
}
155+
}
156+
}
157+
158+
const_defs
159+
.into_iter()
160+
.filter(|(v, _)| total_defs.get(v).copied().unwrap_or(0) == 1)
161+
.collect()
162+
}
163+
164+
// ── Tests ─────────────────────────────────────────────────────────────────────
165+
166+
#[cfg(test)]
167+
mod tests {
168+
use super::*;
169+
use crate::ir::{BlockId, IrBlock, TypeIdx};
170+
171+
fn make_func(blocks: Vec<IrBlock>) -> IrFunction {
172+
IrFunction {
173+
params: vec![],
174+
locals: vec![],
175+
blocks,
176+
entry_block: BlockId(0),
177+
return_type: None,
178+
type_idx: TypeIdx::new(0),
179+
}
180+
}
181+
182+
#[test]
183+
fn eqz_swaps_targets() {
184+
// v1 = Eqz(v0); BranchIf(v1, B1, B2) → BranchIf(v0, B2, B1)
185+
let mut func = make_func(vec![IrBlock {
186+
id: BlockId(0),
187+
instructions: vec![IrInstr::UnOp {
188+
dest: VarId(1),
189+
op: UnOp::I32Eqz,
190+
operand: VarId(0),
191+
}],
192+
terminator: IrTerminator::BranchIf {
193+
condition: VarId(1),
194+
if_true: BlockId(1),
195+
if_false: BlockId(2),
196+
},
197+
}]);
198+
eliminate(&mut func);
199+
match &func.blocks[0].terminator {
200+
IrTerminator::BranchIf {
201+
condition,
202+
if_true,
203+
if_false,
204+
} => {
205+
assert_eq!(*condition, VarId(0));
206+
assert_eq!(*if_true, BlockId(2), "targets should be swapped");
207+
assert_eq!(*if_false, BlockId(1));
208+
}
209+
other => panic!("expected BranchIf, got {other:?}"),
210+
}
211+
}
212+
213+
#[test]
214+
fn ne_zero_simplifies() {
215+
// v1 = 0; v2 = Ne(v0, v1); BranchIf(v2, B1, B2) → BranchIf(v0, B1, B2)
216+
let mut func = make_func(vec![IrBlock {
217+
id: BlockId(0),
218+
instructions: vec![
219+
IrInstr::Const {
220+
dest: VarId(1),
221+
value: IrValue::I32(0),
222+
},
223+
IrInstr::BinOp {
224+
dest: VarId(2),
225+
op: BinOp::I32Ne,
226+
lhs: VarId(0),
227+
rhs: VarId(1),
228+
},
229+
],
230+
terminator: IrTerminator::BranchIf {
231+
condition: VarId(2),
232+
if_true: BlockId(1),
233+
if_false: BlockId(2),
234+
},
235+
}]);
236+
eliminate(&mut func);
237+
match &func.blocks[0].terminator {
238+
IrTerminator::BranchIf {
239+
condition,
240+
if_true,
241+
if_false,
242+
} => {
243+
assert_eq!(*condition, VarId(0));
244+
assert_eq!(*if_true, BlockId(1), "targets should NOT be swapped");
245+
assert_eq!(*if_false, BlockId(2));
246+
}
247+
other => panic!("expected BranchIf, got {other:?}"),
248+
}
249+
}
250+
251+
#[test]
252+
fn eq_zero_swaps() {
253+
// v1 = 0; v2 = Eq(v0, v1); BranchIf(v2, B1, B2) → BranchIf(v0, B2, B1)
254+
let mut func = make_func(vec![IrBlock {
255+
id: BlockId(0),
256+
instructions: vec![
257+
IrInstr::Const {
258+
dest: VarId(1),
259+
value: IrValue::I32(0),
260+
},
261+
IrInstr::BinOp {
262+
dest: VarId(2),
263+
op: BinOp::I32Eq,
264+
lhs: VarId(0),
265+
rhs: VarId(1),
266+
},
267+
],
268+
terminator: IrTerminator::BranchIf {
269+
condition: VarId(2),
270+
if_true: BlockId(1),
271+
if_false: BlockId(2),
272+
},
273+
}]);
274+
eliminate(&mut func);
275+
match &func.blocks[0].terminator {
276+
IrTerminator::BranchIf {
277+
condition,
278+
if_true,
279+
if_false,
280+
} => {
281+
assert_eq!(*condition, VarId(0));
282+
assert_eq!(*if_true, BlockId(2), "targets should be swapped");
283+
assert_eq!(*if_false, BlockId(1));
284+
}
285+
other => panic!("expected BranchIf, got {other:?}"),
286+
}
287+
}
288+
289+
#[test]
290+
fn multi_use_not_folded() {
291+
// v1 = Eqz(v0); use(v1) elsewhere → don't fold
292+
let mut func = make_func(vec![
293+
IrBlock {
294+
id: BlockId(0),
295+
instructions: vec![IrInstr::UnOp {
296+
dest: VarId(1),
297+
op: UnOp::I32Eqz,
298+
operand: VarId(0),
299+
}],
300+
terminator: IrTerminator::BranchIf {
301+
condition: VarId(1),
302+
if_true: BlockId(1),
303+
if_false: BlockId(2),
304+
},
305+
},
306+
IrBlock {
307+
id: BlockId(1),
308+
instructions: vec![],
309+
terminator: IrTerminator::Return {
310+
value: Some(VarId(1)), // second use of v1
311+
},
312+
},
313+
IrBlock {
314+
id: BlockId(2),
315+
instructions: vec![],
316+
terminator: IrTerminator::Return { value: None },
317+
},
318+
]);
319+
eliminate(&mut func);
320+
// Should NOT fold — v1 has 2 uses
321+
match &func.blocks[0].terminator {
322+
IrTerminator::BranchIf { condition, .. } => {
323+
assert_eq!(*condition, VarId(1), "should not have been folded");
324+
}
325+
other => panic!("expected BranchIf, got {other:?}"),
326+
}
327+
}
328+
329+
#[test]
330+
fn cross_block_zero_const() {
331+
// B0: v1 = 0; Jump(B1)
332+
// B1: v2 = Ne(v0, v1); BranchIf(v2, B2, B3) → BranchIf(v0, B2, B3)
333+
let mut func = make_func(vec![
334+
IrBlock {
335+
id: BlockId(0),
336+
instructions: vec![IrInstr::Const {
337+
dest: VarId(1),
338+
value: IrValue::I32(0),
339+
}],
340+
terminator: IrTerminator::Jump { target: BlockId(1) },
341+
},
342+
IrBlock {
343+
id: BlockId(1),
344+
instructions: vec![IrInstr::BinOp {
345+
dest: VarId(2),
346+
op: BinOp::I32Ne,
347+
lhs: VarId(0),
348+
rhs: VarId(1),
349+
}],
350+
terminator: IrTerminator::BranchIf {
351+
condition: VarId(2),
352+
if_true: BlockId(2),
353+
if_false: BlockId(3),
354+
},
355+
},
356+
]);
357+
eliminate(&mut func);
358+
match &func.blocks[1].terminator {
359+
IrTerminator::BranchIf { condition, .. } => {
360+
assert_eq!(*condition, VarId(0));
361+
}
362+
other => panic!("expected BranchIf, got {other:?}"),
363+
}
364+
}
365+
}

0 commit comments

Comments
 (0)