From 8adb74c9d756b418a8e04fcdfdcd9ff75eafcc56 Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Fri, 24 Jul 2026 10:54:38 -0700 Subject: [PATCH 1/2] Optimize localify (regalloc onto locals): - Use a bitset rather than a HashSet for the live-set fixpoint computation; - Use a dense index space consisting *only* of block-crossing values for those bitsets, to keep them smaller; - Use a local dense liveset datastructure that uses epoch-numbers-per-slot to allow O(1) clearing between blocks. Overall, speeds up compilation on a very large input from ~173 seconds to ~7 seconds. --- src/backend/bitset.rs | 38 ++++++ src/backend/cross_block_ids.rs | 204 +++++++++++++++++++++++++++++++++ src/backend/dense_live_set.rs | 104 +++++++++++++++++ src/backend/localify.rs | 167 ++++++++++++++++++++------- src/backend/mod.rs | 3 + 5 files changed, 476 insertions(+), 40 deletions(-) create mode 100644 src/backend/bitset.rs create mode 100644 src/backend/cross_block_ids.rs create mode 100644 src/backend/dense_live_set.rs diff --git a/src/backend/bitset.rs b/src/backend/bitset.rs new file mode 100644 index 0000000..230f426 --- /dev/null +++ b/src/backend/bitset.rs @@ -0,0 +1,38 @@ +use std::convert::TryFrom; + +/// A lazily-grown bitset. +#[derive(Clone, Debug, Default)] +pub struct Bitset { + words: Vec, +} + +impl Bitset { + /// Set bit `i`; returns whether it was newly set. + pub fn insert(&mut self, i: u32) -> bool { + let word = usize::try_from(i / 64).unwrap(); + let bit = 1u64 << (i % 64); + if word >= self.words.len() { + self.words.resize(word + 1, 0); + } + let w = &mut self.words[word]; + let new = *w & bit == 0; + *w |= bit; + new + } + + /// Iterate over all set bit-indices. + pub fn iter(&self) -> impl Iterator + '_ { + self.words.iter().enumerate().flat_map(|(wi, &w)| { + let mut w = w; + std::iter::from_fn(move || { + if w == 0 { + return None; + } + let bit = w.trailing_zeros(); + w &= w - 1; + let limb_start = u32::try_from(wi).unwrap() * 64; + Some(limb_start + bit) + }) + }) + } +} diff --git a/src/backend/cross_block_ids.rs b/src/backend/cross_block_ids.rs new file mode 100644 index 0000000..f574019 --- /dev/null +++ b/src/backend/cross_block_ids.rs @@ -0,0 +1,204 @@ +//! Cross-block value id management for localification. +//! +//! Only SSA values that have a use in a different block from where they +//! are defined can appear in live-in/live-out sets. Those "cross-block" +//! values are assigned compact dense indices (`CrossBlockId`) so that +//! inter-block liveness bitsets and dense sets operate over a much +//! smaller space than the full `Value` space. + +use std::num::NonZeroU32; +use std::ops::Index; + +use crate::backend::treeify::Trees; +use crate::entity::EntityRef; +use crate::ir::{FunctionBody, Value, ValueDef}; + +/// A dense index into the compact set of cross-block SSA values. +/// +/// Only values with a use outside their defining block get a +/// `CrossBlockId`; purely block-local values do not participate in +/// inter-block liveness and are represented by the absence of a value +/// in the two-way map. +/// +/// The inner `NonZeroU32` stores the external 0-based ID plus one, +/// so that `Option` benefits from Rust's niche +/// optimization and packs into a single `u32`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct CrossBlockId(NonZeroU32); + +impl CrossBlockId { + /// Construct a `CrossBlockId` from a raw u32. Used when the id + /// comes from iterating over a bitset or dense array. + pub fn new(n: u32) -> Self { + // SAFETY: n+1 is always >= 1, so never zero. + Self(NonZeroU32::new(n + 1).unwrap()) + } + + /// The inner numeric value, used for indexing into bitsets and + /// dense arrays. + pub fn as_u32(self) -> u32 { + self.0.get() - 1 + } +} + +/// Two-way map between `Value` (SSA value) and `CrossBlockId` (compact index). +/// +/// Tracks which values are "cross-block" (used in a different block from +/// where they are defined). Only cross-block values can appear in +/// live-in/live-out sets. +#[derive(Clone, Debug, Default)] +pub struct CrossBlockValues { + /// Compact id per value (`None` = block-local). Indexed by `Value` index. + ids: Vec>, + /// Compact id -> Value. Indexed by `CrossBlockId.0`. + values: Vec, +} + +impl CrossBlockValues { + /// The number of cross-block values registered. + pub fn len(&self) -> usize { + self.values.len() + } + + /// Whether no values have been registered yet. + pub fn is_empty(&self) -> bool { + self.values.is_empty() + } + + /// Pre-allocate the `ids` array for the given number of total SSA values. + pub fn init(&mut self, total_values: usize) { + self.ids = vec![None; total_values]; + } + + /// Look up the cross-block id for a value, if it is cross-block. + pub fn get(&self, value: Value) -> Option { + self.ids[value.index()] + } + + /// Look up the value for a cross-block id. + pub fn value(&self, id: CrossBlockId) -> Value { + self.values[id.0.get() as usize - 1] + } + + /// Register a value as cross-block, assigning it a new compact id. + /// Returns the assigned id. If already registered, returns the existing id. + pub fn insert(&mut self, value: Value) -> CrossBlockId { + if let Some(id) = self.ids[value.index()] { + return id; + } + let id = CrossBlockId::new(self.values.len() as u32); + self.ids[value.index()] = Some(id); + self.values.push(value); + id + } + + /// Check whether a value has been registered as cross-block. + pub fn contains(&self, value: Value) -> bool { + self.ids[value.index()].is_some() + } + + pub fn ids_mut(&mut self) -> &mut Vec> { + &mut self.ids + } + + pub fn ids(&self) -> &Vec> { + &self.ids + } + + pub fn iter_values(&self) -> impl Iterator + '_ { + self.values.iter().copied() + } + + /// Build the cross-block value map by scanning all uses in the + /// function body. + /// + /// Only values that have a (proper) use outside the block where + /// their def is visited are recorded. A def the block walk skips + /// (treeified or remat insts) never cancels its uses, so such + /// values stay tracked everywhere, matching the walk's behavior. + pub fn build(body: &FunctionBody, trees: &Trees) -> Self { + // 1. Compute def_block for every value. + let mut def_block = vec![u32::MAX; body.values.len()]; + for (block, block_def) in body.blocks.entries() { + for &(_, param) in &block_def.params { + def_block[param.index()] = block.index() as u32; + } + for &inst in &block_def.insts { + if is_tree_or_remat(inst, trees) { + continue; + } + def_block[inst.index()] = block.index() as u32; + } + } + + // 2. Walk every block and record uses that cross block boundaries. + let mut cross_blocks = CrossBlockValues::default(); + cross_blocks.init(body.values.len()); + + for (block, block_def) in body.blocks.entries() { + let cur = block.index() as u32; + + // Visit terminator uses. + block_def.terminator.visit_uses(|u| { + visit_use(u, body, trees, cur, &def_block, &mut cross_blocks); + }); + + // Visit instruction args (in reverse order per the body walk). + for &inst in block_def.insts.iter().rev() { + if is_tree_or_remat(inst, trees) { + continue; + } + if let ValueDef::Operator(_, args, _) = &body.values[inst] { + for &arg in &body.arg_pool[*args] { + visit_use(arg, body, trees, cur, &def_block, &mut cross_blocks); + } + } + } + } + + cross_blocks + } +} + +/// Check whether a value is treeified (owned by another) or rematerialized. +fn is_tree_or_remat(value: Value, trees: &Trees) -> bool { + trees.owner.contains_key(&value) || trees.remat.contains(&value) +} + +/// Recursively resolve a use through aliases and treeified/remat values, +/// recording cross-block values into the map. +fn visit_use( + value: Value, + body: &FunctionBody, + trees: &Trees, + cur_block: u32, + def_block: &[u32], + cross_blocks: &mut CrossBlockValues, +) { + let value = body.resolve_alias(value); + if let ValueDef::PickOutput(value, _, _) = body.values[value] { + visit_use(value, body, trees, cur_block, def_block, cross_blocks); + return; + } + if is_tree_or_remat(value, trees) { + // Treeified or remat: visit its args without recording the value itself. + if let ValueDef::Operator(_, args, _) = body.values[value] { + for &arg in &body.arg_pool[args] { + visit_use(arg, body, trees, cur_block, def_block, cross_blocks); + } + } + return; + } + // Proper use: record if it crosses block boundaries. + if def_block[value.index()] != cur_block { + cross_blocks.insert(value); + } +} + +impl Index for CrossBlockValues { + type Output = Value; + + fn index(&self, id: CrossBlockId) -> &Self::Output { + &self.values[id.0.get() as usize - 1] + } +} diff --git a/src/backend/dense_live_set.rs b/src/backend/dense_live_set.rs new file mode 100644 index 0000000..07c94e8 --- /dev/null +++ b/src/backend/dense_live_set.rs @@ -0,0 +1,104 @@ +//! A dense, epoch-based live-set used during block walks. +//! +//! # Epoch-based O(1) clearing +//! +//! The `DenseLiveSet` is reused across multiple block visits without +//! reallocating or clearing its backing storage. +//! +//! We want a dense set because accesses are much faster than a HashSet. +//! We would not want to allocate a dense set separately for every block +//! we process, however (as we get further down in the function, there +//! would be a larger and larger "prefix" of unused IDs). Instead we +//! want to use one allocation for a whole scan through a function body. +//! +//! But we cannot clear a dense set in O(1) if its per-index entries are +//! just booleans; that entails scanning through the whole thing with a +//! bulk memset (or at least, a walk through the dirty-index list and a +//! store at each index). So instead, this is achieved through an epoch +//! counter: +//! +//! * Each slot in the `mark` array stores `epoch * 2 + state`, where +//! `state` is `0` (live) or `1` (dead). +//! * When a new block walk begins, `next_epoch()` increments the epoch +//! counter and clears the `touched` list — both O(1) operations. +//! * Because the old epoch values in `mark` no longer match the current +//! epoch, stale entries from previous block visits are automatically +//! treated as absent without needing to zero out the array. + +/// A dense live-set with O(1) clearing via epoch tagging. +#[derive(Debug)] +pub struct DenseLiveSet { + /// Per-slot mark: `epoch * 2 + state` where state is 0 (live) or 1 + /// (dead). Slots with a mark from a previous epoch are treated as + /// absent. + mark: Vec, + /// The current epoch counter. Incremented on `next_epoch()` to + /// logically clear the set. + epoch: u64, + /// Indices touched (set_live or set_dead) during the current epoch. + /// Cleared on each `next_epoch()` call. + touched: Vec, +} + +impl DenseLiveSet { + /// Create a new set with space for `n` indices. + pub fn new(n: usize) -> Self { + Self { + mark: vec![0; n], + epoch: 0, + touched: Vec::new(), + } + } + + /// Start a new epoch, logically clearing the set in O(1). + /// + /// Increments the epoch counter and clears the touched-index list. + pub fn next_epoch(&mut self) { + self.epoch += 1; + self.touched.clear(); + } + + /// Mark index `i` as live. + /// + /// Returns `true` if it was already live in this epoch (i.e., this + /// is not a first-use). + pub fn set_live(&mut self, i: u32) -> bool { + let m = &mut self.mark[i as usize]; + let newly_touched = *m / 2 != self.epoch; + if newly_touched { + self.touched.push(i); + } + let was_live = !newly_touched && *m % 2 == 0; + *m = self.epoch * 2; + !was_live + } + + /// Mark index `i` as dead (retired by a def). + /// + /// Returns `true` if it was live in this epoch. + pub fn set_dead(&mut self, i: u32) -> bool { + let m = &mut self.mark[i as usize]; + if *m / 2 != self.epoch { + self.touched.push(i); + *m = self.epoch * 2 + 1; + return false; + } + let was_live = *m % 2 == 0; + *m = self.epoch * 2 + 1; + was_live + } + + /// Check if index `i` is currently live in this epoch. + pub fn is_live(&self, i: u32) -> bool { + let m = self.mark[i as usize]; + m / 2 == self.epoch && m % 2 == 0 + } + + /// Returns the list of indices touched during the current epoch. + /// + /// Useful for iterating only over entries that were modified in the + /// current epoch, avoiding a full scan of all possible indices. + pub fn touched(&self) -> &[u32] { + &self.touched + } +} diff --git a/src/backend/localify.rs b/src/backend/localify.rs index 62d8790..a06bfbb 100644 --- a/src/backend/localify.rs +++ b/src/backend/localify.rs @@ -1,11 +1,21 @@ //! Localification: a simple form of register allocation that picks //! locations for SSA values in Wasm locals. +//! +//! Performance notes: liveness sets and live-range maps are keyed by +//! `Value`, a dense index space, so all per-value state lives in flat +//! arrays. The liveness fixpoint additionally runs over a compacted +//! index space of only the cross-block values (values with a use +//! outside their defining block), and we track cross-block liveness +//! with bitsets indexed by this index space. +use crate::backend::bitset::Bitset; +use crate::backend::cross_block_ids::{CrossBlockId, CrossBlockValues}; +use crate::backend::dense_live_set::DenseLiveSet; use crate::backend::treeify::Trees; use crate::cfg::CFGInfo; -use crate::entity::{EntityVec, PerEntity}; +use crate::entity::{EntityRef, EntityVec, PerEntity}; use crate::ir::{Block, FunctionBody, Local, Type, Value, ValueDef}; -use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; +use fxhash::FxHashMap as HashMap; use smallvec::{smallvec, SmallVec}; use std::ops::Range; @@ -27,13 +37,19 @@ struct Context<'a> { trees: &'a Trees, results: Localifier, - /// Precise liveness for each block: live Values at the end. - block_end_live: PerEntity>, + /// Two-way map between `Value` and `CrossBlockId` for the compacted + /// inter-block value index space. + cross_blocks: CrossBlockValues, + + /// Precise liveness for each block: cross-block values live at the + /// end, as a bitset over compact ids. + block_end_live: PerEntity, /// Liveranges for each Value, in an arbitrary index space /// (concretely, the span of first to last instruction visit step - /// index in an RPO walk over the function body). - ranges: HashMap>, + /// index in an RPO walk over the function body). `usize::MAX` + /// start means "no range recorded". + ranges: Vec>, /// Number of points. points: usize, } @@ -49,6 +65,10 @@ trait Visitor { fn pre_params(&mut self) {} } +fn is_tree_or_remat(value: Value, trees: &Trees) -> bool { + trees.owner.contains_key(&value) || trees.remat.contains(&value) +} + struct BlockVisitor<'a, V: Visitor> { body: &'a FunctionBody, trees: &'a Trees, @@ -74,7 +94,7 @@ impl<'a, V: Visitor> BlockVisitor<'a, V> { self.visitor.pre_term(); for &inst in self.body.blocks[block].insts.iter().rev() { - if self.trees.owner.contains_key(&inst) || self.trees.remat.contains(&inst) { + if is_tree_or_remat(inst, self.trees) { continue; } self.visitor.post_inst(inst); @@ -107,9 +127,14 @@ impl<'a, V: Visitor> BlockVisitor<'a, V> { self.visit_use(value); return; } - if self.trees.owner.contains_key(&value) { - // If this is a treeified value, then don't process the use, - // but process the instruction directly here. + if is_tree_or_remat(value, self.trees) { + // If this is a treeified or rematerialized value, then + // don't process the use, but process the instruction + // directly here. (Lowering rematerializes remat values at + // each use and never lowers their defs, so a remat use + // must not keep a local alive: its def is never visited + // and it would otherwise be tracked as live everywhere, + // holding a dead local for the whole function.) self.visit_inst(value, /* root = */ false); } else { // Otherwise, this is a proper use. @@ -133,43 +158,71 @@ impl<'a> Context<'a> { cfg, trees, results, + cross_blocks: CrossBlockValues::default(), block_end_live: PerEntity::default(), - ranges: HashMap::default(), + ranges: vec![usize::MAX..usize::MAX; body.values.len()], points: 0, } } + /// Assign compact ids to the values with a (proper) use outside + /// the block where their def is VISITED; only those can appear in + /// a liveness set. A def the block walk skips (treeified or remat + /// insts) never cancels its uses, so such values stay tracked + /// everywhere, matching the walk's behavior. + fn find_cross_block_values(&mut self) { + self.cross_blocks = CrossBlockValues::build(self.body, self.trees); + } + fn compute_liveness(&mut self) { - struct LivenessVisitor { - live: HashSet, + struct LivenessVisitor<'b> { + cross_blocks: &'b CrossBlockValues, + live: DenseLiveSet, } - impl Visitor for LivenessVisitor { + impl<'b> Visitor for LivenessVisitor<'b> { fn visit_use(&mut self, value: Value) { - self.live.insert(value); + if let Some(id) = self.cross_blocks.get(value) { + self.live.set_live(id.as_u32()); + } } fn visit_def(&mut self, value: Value) { - self.live.remove(&value); + if let Some(id) = self.cross_blocks.get(value) { + self.live.set_dead(id.as_u32()); + } } } + let mut visitor = BlockVisitor::new( + self.body, + self.trees, + LivenessVisitor { + cross_blocks: &self.cross_blocks, + live: DenseLiveSet::new(self.cross_blocks.len()), + }, + ); let mut workqueue: Vec = self.cfg.rpo.values().cloned().collect(); - let mut workqueue_set: HashSet = workqueue.iter().cloned().collect(); + let mut workqueue_set = vec![true; self.body.blocks.len()]; while let Some(block) = workqueue.pop() { - workqueue_set.remove(&block); - let live = self.block_end_live[block].clone(); - let mut visitor = BlockVisitor::new(self.body, self.trees, LivenessVisitor { live }); + workqueue_set[block.index()] = false; + visitor.visitor.live.next_epoch(); + for id in self.block_end_live[block].iter() { + visitor.visitor.live.set_live(id); + } visitor.visit_block(block); - let live = visitor.visitor.live; + // Insert the live-in ids straight into each pred's + // live-out bitset. + let live = &visitor.visitor.live; for &pred in &self.body.blocks[block].preds { let pred_live = &mut self.block_end_live[pred]; let mut changed = false; - for &value in &live { - if pred_live.insert(value) { + for &id in live.touched() { + if live.is_live(id) && pred_live.insert(id) { changed = true; } } - if changed && workqueue_set.insert(pred) { + if changed && !workqueue_set[pred.index()] { + workqueue_set[pred.index()] = true; workqueue.push(pred); } } @@ -181,8 +234,22 @@ impl<'a> Context<'a> { struct LiveRangeVisitor<'b> { point: &'b mut usize, - live: HashMap, - ranges: &'b mut HashMap>, + /// Live values in the current block walk; the range start + /// (visit point of the last use seen) rides in `start`. + live: DenseLiveSet, + start: &'b mut [usize], + ranges: &'b mut [Range], + } + impl<'b> LiveRangeVisitor<'b> { + fn record_def(&mut self, value: Value, range: Range) { + let existing = &mut self.ranges[value.index()]; + if existing.start == usize::MAX { + *existing = range; + } else { + existing.start = std::cmp::min(existing.start, range.start); + existing.end = std::cmp::max(existing.end, range.end); + } + } } impl<'b> Visitor for LiveRangeVisitor<'b> { fn pre_params(&mut self) { @@ -195,30 +262,35 @@ impl<'a> Context<'a> { *self.point += 1; } fn visit_use(&mut self, value: Value) { - self.live.entry(value).or_insert(*self.point); + if self.live.set_live(value.index() as u32) { + self.start[value.index()] = *self.point; + } } fn visit_def(&mut self, value: Value) { - let range = if let Some(start) = self.live.remove(&value) { - start..(*self.point + 1) + let range = if self.live.set_dead(value.index() as u32) { + self.start[value.index()]..(*self.point + 1) } else { *self.point..(*self.point + 1) }; - let existing_range = self.ranges.entry(value).or_insert(range.clone()); - existing_range.start = std::cmp::min(existing_range.start, range.start); - existing_range.end = std::cmp::max(existing_range.end, range.end); + self.record_def(value, range); } } + let mut start = vec![0usize; self.body.values.len()]; + let mut live = DenseLiveSet::new(self.body.values.len()); for &block in self.cfg.rpo.values().rev() { + live.next_epoch(); let visitor = LiveRangeVisitor { - live: HashMap::default(), + live, + start: &mut start, point: &mut point, ranges: &mut self.ranges, }; - let mut visitor = BlockVisitor::new(&self.body, &self.trees, visitor); + let mut visitor = BlockVisitor::new(&self.body, self.trees, visitor); // Live-outs to succ blocks: in this block-local // handling, model them as uses as the end of the block. - for &livein in &self.block_end_live[block] { + for id in self.block_end_live[block].iter() { + let livein = self.cross_blocks.value(CrossBlockId::new(id)); let livein = self.body.resolve_alias(livein); visitor.visitor.visit_use(livein); } @@ -226,10 +298,19 @@ impl<'a> Context<'a> { visitor.visit_block(block); // Live-ins from pred blocks: anything still live has a // virtual def at top of block. - let still_live = visitor.visitor.live.keys().cloned().collect::>(); - for live in still_live { - visitor.visitor.visit_def(live); + let still_live: Vec = visitor + .visitor + .live + .touched() + .iter() + .copied() + .filter(|&i| visitor.visitor.live.is_live(i)) + .map(|i| Value::new(i as usize)) + .collect(); + for v in still_live { + visitor.visitor.visit_def(v); } + live = visitor.visitor.live; } self.points = point + 1; @@ -237,8 +318,13 @@ impl<'a> Context<'a> { fn allocate(&mut self) { // Sort values by ranges' starting points, then value to break ties. - let mut ranges: Vec<(Value, std::ops::Range)> = - self.ranges.iter().map(|(k, v)| (*k, v.clone())).collect(); + let mut ranges: Vec<(Value, std::ops::Range)> = self + .ranges + .iter() + .enumerate() + .filter(|(_, r)| r.start != usize::MAX) + .map(|(i, r)| (Value::new(i), r.clone())) + .collect(); ranges.sort_unstable_by_key(|(val, range)| (range.start, *val)); // Keep a list of expiring Locals by expiry point. @@ -302,6 +388,7 @@ impl<'a> Context<'a> { } fn compute(mut self) -> Localifier { + self.find_cross_block_values(); self.compute_liveness(); self.find_ranges(); self.allocate(); diff --git a/src/backend/mod.rs b/src/backend/mod.rs index 1f933cd..d02c45d 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -10,6 +10,9 @@ use rayon::prelude::*; use std::borrow::Cow; use std::convert::TryFrom; +pub(crate) mod bitset; +pub(crate) mod cross_block_ids; +pub(crate) mod dense_live_set; pub mod reducify; use reducify::Reducifier; pub mod stackify; From d8809bd31a3d4821452d073726bf5ce471497b01 Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Fri, 24 Jul 2026 14:22:06 -0700 Subject: [PATCH 2/2] Release 0.3.1. Includes #20 (optimizations to localify), #19 (use FxHashMap/FxHashSet everywhere), and #18 (a treeifier soundness fix). --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index f0bdc8e..eb71815 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "waffle" -version = "0.3.0" +version = "0.3.1" description = "Wasm Analysis Framework For Lightweight Experiments" authors = ["Chris Fallin "] license = "Apache-2.0 WITH LLVM-exception"