perf(codegen): FxHashMap cascade pt2 — value_sink/const_fold/fusion/register_estimate (+57.6%)

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<ValueId> = BTreeSet::new();
+        let mut cross_block_refs: FxHashSet<ValueId> = 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<ValueId>) {
-    let mut used: BTreeSet<ValueId> = BTreeSet::new();
+fn dce_block(block: &mut Block, cross_block_refs: &FxHashSet<ValueId>) {
+    let mut used: FxHashSet<ValueId> =
+        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<ValueId>) {
     block.results = new_results;
 }
 
-fn collect_uses(op: &Op, used: &mut BTreeSet<ValueId>) {
+fn collect_uses(op: &Op, used: &mut FxHashSet<ValueId>) {
     for &v in op.value_refs().iter() {
         used.insert(*v);
     }

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<BlockId, BTreeSet<ValueId>> = FxHashMap::default();
+        let mut pinned_per_block: FxHashMap<BlockId, FxHashSet<ValueId>> = 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<ValueId>) -> Result<()> {
+fn fuse_block(block: &mut Block, pinned: &FxHashSet<ValueId>) -> Result<()> {
     // Phase 1: build def-use graph.
     // uses[vid] = set of op indices that reference vid.
     let mut uses: FxHashMap<ValueId, Vec<usize>> = FxHashMap::default();

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<ValueId> = std::collections::BTreeSet::new();
+    let mut live: rustc_hash::FxHashSet<ValueId> =
+        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<ValueId> live set (pre-PR impl) vs
+    /// FxHashSet<ValueId> 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<ValueId> = 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");

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<ValueId, usize> {
-    let mut use_count: BTreeMap<ValueId, usize> = BTreeMap::new();
+fn compute_global_use_count(kernel: &Kernel) -> FxHashMap<ValueId, usize> {
+    // 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::<usize>();
+    let mut use_count: FxHashMap<ValueId, usize> =
+        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<ValueId, usize>) {
+fn sink_in_block(block: &mut Block, global_use_count: &FxHashMap<ValueId, usize>) {
     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<usize> = BTreeSet::new();
-    let mut insert_before: BTreeMap<usize, Vec<(Op, Option<ValueId>)>> = 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<usize> =
+        FxHashSet::with_capacity_and_hasher(plans.len(), Default::default());
+    let mut insert_before: FxHashMap<usize, Vec<(Op, Option<ValueId>)>> =
+        FxHashMap::with_capacity_and_hasher(plans.len(), Default::default());
 
     for plan in plans {
         remove_at.insert(plan.from);