diff --git a/crates/metaltile-codegen/src/passes/const_fold.rs b/crates/metaltile-codegen/src/passes/const_fold.rs index a04ecd27..792a84a2 100644 --- a/crates/metaltile-codegen/src/passes/const_fold.rs +++ b/crates/metaltile-codegen/src/passes/const_fold.rs @@ -31,9 +31,8 @@ //! - Aho, Lam, Sethi & Ullman (2006), "Compilers: Principles, Techniques, and //! Tools", 2nd ed., §8.4 (constant folding), §9.1.2 (dead-code elimination). -use std::collections::BTreeSet; - use metaltile_core::ir::{BinOpKind, Block, BlockId, Kernel, Op, UnaryOpKind, ValueId}; +use rustc_hash::FxHashSet; use crate::error::Result; @@ -76,7 +75,7 @@ impl super::Pass for ConstFoldPass { // Collect all ValueIds referenced by any block in the kernel (cross-block uses). // These must be treated as "live" in each individual block's DCE pass to prevent // eliminating values that are consumed by sibling or child blocks. - let mut cross_block_refs: BTreeSet = BTreeSet::new(); + let mut cross_block_refs: FxHashSet = FxHashSet::default(); for block in kernel.blocks.values() { for op in &block.ops { collect_uses(op, &mut cross_block_refs); @@ -267,8 +266,9 @@ fn replace_value_in_op(op: &mut Op, old: ValueId, new: ValueId) { // Dead Code Elimination // --------------------------------------------------------------------------- -fn dce_block(block: &mut Block, cross_block_refs: &BTreeSet) { - let mut used: BTreeSet = BTreeSet::new(); +fn dce_block(block: &mut Block, cross_block_refs: &FxHashSet) { + let mut used: FxHashSet = + FxHashSet::with_capacity_and_hasher(block.ops.len(), Default::default()); for op in &block.ops { collect_uses(op, &mut used); } @@ -296,7 +296,7 @@ fn dce_block(block: &mut Block, cross_block_refs: &BTreeSet) { block.results = new_results; } -fn collect_uses(op: &Op, used: &mut BTreeSet) { +fn collect_uses(op: &Op, used: &mut FxHashSet) { for &v in op.value_refs().iter() { used.insert(*v); } diff --git a/crates/metaltile-codegen/src/passes/fusion.rs b/crates/metaltile-codegen/src/passes/fusion.rs index 5d1f7411..51dba629 100644 --- a/crates/metaltile-codegen/src/passes/fusion.rs +++ b/crates/metaltile-codegen/src/passes/fusion.rs @@ -35,8 +35,6 @@ //! TensorFlow blog. Production operator fusion for ML workloads. //! https://developers.googleblog.com/xla-tensorflow-compiled/ -use std::collections::BTreeSet; - use metaltile_core::ir::{Block, BlockId, Kernel, Op, ValueId}; use rustc_hash::{FxHashMap, FxHashSet}; @@ -99,7 +97,7 @@ impl super::Pass for FusionPass { // defined in B but used in at least one other block (i.e. a child block). // Pinned values must not be fused away — they need a standalone declaration // so that child blocks can reference the variable by name. - let mut pinned_per_block: FxHashMap> = FxHashMap::default(); + let mut pinned_per_block: FxHashMap> = FxHashMap::default(); for (vid, def_bid) in &def_block { if let Some(use_bids) = used_in.get(vid) { for &use_bid in use_bids { @@ -142,7 +140,7 @@ impl super::Pass for FusionPass { // Fusion on a single block // --------------------------------------------------------------------------- -fn fuse_block(block: &mut Block, pinned: &BTreeSet) -> Result<()> { +fn fuse_block(block: &mut Block, pinned: &FxHashSet) -> Result<()> { // Phase 1: build def-use graph. // uses[vid] = set of op indices that reference vid. let mut uses: FxHashMap> = FxHashMap::default(); diff --git a/crates/metaltile-codegen/src/passes/register_estimate.rs b/crates/metaltile-codegen/src/passes/register_estimate.rs index e2cbe84b..349d9447 100644 --- a/crates/metaltile-codegen/src/passes/register_estimate.rs +++ b/crates/metaltile-codegen/src/passes/register_estimate.rs @@ -64,7 +64,8 @@ pub fn estimate_registers(kernel: &Kernel) -> RegisterEstimate { /// Compute the maximum live ValueId count in a single block. fn block_max_live(block: &metaltile_core::ir::Block) -> usize { - let mut live: std::collections::BTreeSet = std::collections::BTreeSet::new(); + let mut live: rustc_hash::FxHashSet = + rustc_hash::FxHashSet::with_capacity_and_hasher(block.ops.len(), Default::default()); let mut max = 0usize; for (i, op) in block.ops.iter().enumerate() { @@ -112,6 +113,71 @@ mod tests { assert_eq!(est.max_live, 3); // v0, v1, v2 all live at the Add } + /// Side-by-side: BTreeSet live set (pre-PR impl) vs + /// FxHashSet live set (new impl). Same workload run in + /// one binary so the comparison is noise-bounded. + fn block_max_live_btree(block: &metaltile_core::ir::Block) -> usize { + let mut live: std::collections::BTreeSet = std::collections::BTreeSet::new(); + let mut max = 0usize; + for (i, op) in block.ops.iter().enumerate() { + for vid in crate::passes::remap::op_value_refs(op) { + live.insert(vid); + } + if let Some(Some(vid)) = block.results.get(i) { + live.insert(*vid); + } + max = max.max(live.len()); + } + max + } + + #[test] + #[ignore = "perf microbench"] + fn perf_block_max_live_btree_vs_fxhash() { + // Build a kernel with ~2000 ops and dense ValueIds — the live + // set grows continuously, exercising the hash/compare cost. + let mut k = Kernel::new("perf_live"); + for i in 0..2000u32 { + k.body.push_op(Op::Const { value: i as i64 }, ValueId::new(i)); + } + + const ITERS: usize = 10_000; + + for _ in 0..1_000 { + std::hint::black_box(block_max_live_btree(std::hint::black_box(&k.body))); + std::hint::black_box(block_max_live(std::hint::black_box(&k.body))); + } + + let t0 = std::time::Instant::now(); + for _ in 0..ITERS { + std::hint::black_box(block_max_live_btree(std::hint::black_box(&k.body))); + } + let bt_elapsed = t0.elapsed(); + let bt_ns_per = bt_elapsed.as_nanos() as f64 / ITERS as f64; + + let t1 = std::time::Instant::now(); + for _ in 0..ITERS { + std::hint::black_box(block_max_live(std::hint::black_box(&k.body))); + } + let fx_elapsed = t1.elapsed(); + let fx_ns_per = fx_elapsed.as_nanos() as f64 / ITERS as f64; + + println!(); + println!("=== block_max_live: 2000-op live-set walk × {ITERS} iters ==="); + println!(" BTreeSet (old): {bt_elapsed:>10.2?} ({bt_ns_per:>8.1} ns/call)"); + println!(" FxHashSet (new): {fx_elapsed:>10.2?} ({fx_ns_per:>8.1} ns/call)"); + let speedup = bt_ns_per / fx_ns_per; + println!( + " → speedup : {speedup:.2}× ({:+.1}%)", + (1.0 - fx_ns_per / bt_ns_per) * 100.0 + ); + + assert!( + fx_ns_per * 1.05 <= bt_ns_per, + "FxHashSet block_max_live ({fx_ns_per:.1} ns) should beat BTreeSet ({bt_ns_per:.1} ns)" + ); + } + #[test] fn estimates_regs_per_thread() { let mut k = Kernel::new("regs"); diff --git a/crates/metaltile-codegen/src/passes/value_sink.rs b/crates/metaltile-codegen/src/passes/value_sink.rs index 2647be65..840d40b2 100644 --- a/crates/metaltile-codegen/src/passes/value_sink.rs +++ b/crates/metaltile-codegen/src/passes/value_sink.rs @@ -39,9 +39,8 @@ //! - Aho, Lam, Sethi & Ullman (2006), "Compilers: Principles, Techniques, and //! Tools", 2nd ed., §9.5. Partial-redundancy elimination and code sinking. -use std::collections::{BTreeMap, BTreeSet}; - use metaltile_core::ir::{Block, BlockId, Kernel, Op, ValueId}; +use rustc_hash::{FxHashMap, FxHashSet}; use super::remap; use crate::error::{Error, Result}; @@ -76,8 +75,11 @@ impl super::Pass for ValueSinkPass { /// Count every use of every `ValueId` across the kernel's entry body and all /// stored blocks. Sinking is intra-block (it only reorders ops within one /// block) so this count never goes stale as sinking proceeds. -fn compute_global_use_count(kernel: &Kernel) -> BTreeMap { - let mut use_count: BTreeMap = BTreeMap::new(); +fn compute_global_use_count(kernel: &Kernel) -> FxHashMap { + // Pre-size from a rough op-count: each op contributes ≤ 3 refs. + let cap = kernel.body.ops.len() + kernel.blocks.values().map(|b| b.ops.len()).sum::(); + let mut use_count: FxHashMap = + FxHashMap::with_capacity_and_hasher(cap, Default::default()); let mut tally = |ops: &[Op]| { for op in ops { for vid in remap::op_value_refs(op) { @@ -116,7 +118,7 @@ struct SinkPlan { // main logic // --------------------------------------------------------------------------- -fn sink_in_block(block: &mut Block, global_use_count: &BTreeMap) { +fn sink_in_block(block: &mut Block, global_use_count: &FxHashMap) { let n = block.ops.len(); if n == 0 { return; @@ -197,8 +199,12 @@ fn rebuild_with_sinking(block: &mut Block, plans: &[SinkPlan]) { let n = old_ops.len(); // Build maps: from_pos → remove; to_pos → insert (Vec of ops). - let mut remove_at: BTreeSet = BTreeSet::new(); - let mut insert_before: BTreeMap)>> = BTreeMap::new(); + // Position-keyed sets/maps — pure get/contains, no ordered iteration + // (the rebuild loop walks positions 0..n itself). + let mut remove_at: FxHashSet = + FxHashSet::with_capacity_and_hasher(plans.len(), Default::default()); + let mut insert_before: FxHashMap)>> = + FxHashMap::with_capacity_and_hasher(plans.len(), Default::default()); for plan in plans { remove_at.insert(plan.from);