Skip to content

Commit c272297

Browse files
authored
ZJIT: Use worklist algorithm for type inference (ruby#16122)
## Summary - Replace the RPO fixed-point iteration in `infer_types` with a worklist-based approach - Instead of re-scanning every reachable block on each iteration, only re-process blocks whose incoming types actually changed - Add `BitSet::remove` helper for worklist membership tracking ## Benchmark On lobsters (release+stats build, 10 iterations): | Stat | Branch | Master | Delta | |------|--------|--------|-------| | compile_time | 1,604ms | 1,808ms | -204ms (-11.3%) | | Avg iter time | 766ms | 807ms | -41ms (-5.1%) |
1 parent 09ae953 commit c272297

File tree

2 files changed

+68
-45
lines changed

2 files changed

+68
-45
lines changed

zjit/src/bitset.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,16 @@ impl<T: Into<usize> + Copy> BitSet<T> {
3737
}
3838
}
3939

40+
/// Clear a bit. Returns whether the bit was previously set.
41+
pub fn remove(&mut self, idx: T) -> bool {
42+
debug_assert!(idx.into() < self.num_bits);
43+
let entry_idx = idx.into() / ENTRY_NUM_BITS;
44+
let bit_idx = idx.into() % ENTRY_NUM_BITS;
45+
let was_set = (self.entries[entry_idx] & (1 << bit_idx)) != 0;
46+
self.entries[entry_idx] &= !(1 << bit_idx);
47+
was_set
48+
}
49+
4050
pub fn get(&self, idx: T) -> bool {
4151
debug_assert!(idx.into() < self.num_bits);
4252
let entry_idx = idx.into() / ENTRY_NUM_BITS;

zjit/src/hir.rs

Lines changed: 58 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -2610,60 +2610,73 @@ impl Function {
26102610
self.copy_param_types();
26112611

26122612
let mut reachable = BlockSet::with_capacity(self.blocks.len());
2613+
2614+
// Maintain both a worklist and a fast membership check to avoid linear search
2615+
let mut worklist: VecDeque<BlockId> = VecDeque::with_capacity(self.blocks.len());
2616+
let mut in_worklist = BlockSet::with_capacity(self.blocks.len());
2617+
macro_rules! worklist_add {
2618+
($block:expr) => {
2619+
if in_worklist.insert($block) {
2620+
worklist.push_back($block);
2621+
}
2622+
};
2623+
}
2624+
26132625
for entry_block in self.entry_blocks() {
26142626
reachable.insert(entry_block);
2627+
worklist_add!(entry_block);
2628+
}
2629+
2630+
// Helper to propagate types along a branch edge and enqueue the target if anything changed
2631+
macro_rules! enqueue {
2632+
($self:ident, $target:expr) => {
2633+
let newly_reachable = reachable.insert($target.target);
2634+
let mut target_changed = newly_reachable;
2635+
for (idx, arg) in $target.args.iter().enumerate() {
2636+
let param = $self.blocks[$target.target.0].params[idx];
2637+
let new = self.insn_types[param.0].union(self.insn_types[arg.0]);
2638+
if !self.insn_types[param.0].bit_equal(new) {
2639+
self.insn_types[param.0] = new;
2640+
target_changed = true;
2641+
}
2642+
}
2643+
if target_changed {
2644+
worklist_add!($target.target);
2645+
}
2646+
};
26152647
}
26162648

2617-
// Walk the graph, computing types until fixpoint
2618-
let rpo = self.rpo();
2619-
loop {
2620-
let mut changed = false;
2621-
for &block in &rpo {
2622-
if !reachable.get(block) { continue; }
2623-
for insn_id in &self.blocks[block.0].insns {
2624-
let insn_type = match self.find(*insn_id) {
2625-
Insn::IfTrue { val, target: BranchEdge { target, args } } => {
2626-
assert!(!self.type_of(val).bit_equal(types::Empty));
2627-
if self.type_of(val).could_be(Type::from_cbool(true)) {
2628-
reachable.insert(target);
2629-
for (idx, arg) in args.iter().enumerate() {
2630-
let param = self.blocks[target.0].params[idx];
2631-
self.insn_types[param.0] = self.type_of(param).union(self.type_of(*arg));
2632-
}
2633-
}
2634-
continue;
2635-
}
2636-
Insn::IfFalse { val, target: BranchEdge { target, args } } => {
2637-
assert!(!self.type_of(val).bit_equal(types::Empty));
2638-
if self.type_of(val).could_be(Type::from_cbool(false)) {
2639-
reachable.insert(target);
2640-
for (idx, arg) in args.iter().enumerate() {
2641-
let param = self.blocks[target.0].params[idx];
2642-
self.insn_types[param.0] = self.type_of(param).union(self.type_of(*arg));
2643-
}
2644-
}
2645-
continue;
2649+
// Walk the graph, computing types until worklist is empty
2650+
while let Some(block) = worklist.pop_front() {
2651+
in_worklist.remove(block);
2652+
if !reachable.get(block) { continue; }
2653+
for insn_id in &self.blocks[block.0].insns {
2654+
let insn_type = match self.find(*insn_id) {
2655+
Insn::IfTrue { val, target } => {
2656+
assert!(!self.type_of(val).bit_equal(types::Empty));
2657+
if self.type_of(val).could_be(Type::from_cbool(true)) {
2658+
enqueue!(self, target);
26462659
}
2647-
Insn::Jump(BranchEdge { target, args }) => {
2648-
reachable.insert(target);
2649-
for (idx, arg) in args.iter().enumerate() {
2650-
let param = self.blocks[target.0].params[idx];
2651-
self.insn_types[param.0] = self.type_of(param).union(self.type_of(*arg));
2652-
}
2653-
continue;
2660+
continue;
2661+
}
2662+
Insn::IfFalse { val, target } => {
2663+
assert!(!self.type_of(val).bit_equal(types::Empty));
2664+
if self.type_of(val).could_be(Type::from_cbool(false)) {
2665+
enqueue!(self, target);
26542666
}
2655-
insn if insn.has_output() => self.infer_type(*insn_id),
2656-
_ => continue,
2657-
};
2658-
if !self.type_of(*insn_id).bit_equal(insn_type) {
2659-
self.insn_types[insn_id.0] = insn_type;
2660-
changed = true;
2667+
continue;
26612668
}
2669+
Insn::Jump(target) => {
2670+
enqueue!(self, target);
2671+
continue;
2672+
}
2673+
insn if insn.has_output() => self.infer_type(*insn_id),
2674+
_ => continue,
2675+
};
2676+
if !self.type_of(*insn_id).bit_equal(insn_type) {
2677+
self.insn_types[insn_id.0] = insn_type;
26622678
}
26632679
}
2664-
if !changed {
2665-
break;
2666-
}
26672680
}
26682681
}
26692682

0 commit comments

Comments
 (0)