Skip to content

perf(codegen): cascade FxHashMap to copy_prop/alg_simplify/unroll/vectorize (+48.8%) #236

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
TheTom wants to merge 1 commit into from
Closed

perf(codegen): cascade FxHashMap to copy_prop/alg_simplify/unroll/vectorize (+48.8%) #236

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 22 additions & 25 deletions crates/metaltile-codegen/src/passes/algebraic_simplify.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,9 +59,8 @@
//! Tools", 2nd ed., §8.4–8.5. Classic treatment of algebraic identities and
//! reduction in strength.

use std::collections::{BTreeMap, BTreeSet};

use metaltile_core::ir::{BinOpKind, Block, BlockId, Kernel, Op, UnaryOpKind, ValueId};
use rustc_hash::{FxHashMap, FxHashSet};

use super::remap;
use crate::error::{Error, Result};
Expand Down Expand Up @@ -93,11 +92,11 @@ impl super::Pass for AlgebraicSimplifyPass {
// ---------------------------------------------------------------------------

/// Resolve transitive replacement chains: {v2→v1, v1→v0} becomes {v2→v0, v1→v0}.
fn resolve_transitive(map: &BTreeMap<ValueId, ValueId>) -> BTreeMap<ValueId, ValueId> {
let mut resolved = BTreeMap::new();
fn resolve_transitive(map: &FxHashMap<ValueId, ValueId>) -> FxHashMap<ValueId, ValueId> {
let mut resolved = FxHashMap::with_capacity_and_hasher(map.len(), Default::default());
for (&key, &val) in map.iter() {
let mut terminal = val;
let mut visited = BTreeSet::new();
let mut visited = FxHashSet::default();
visited.insert(key);
while let Some(&next) = map.get(&terminal) {
if !visited.insert(terminal) {
Expand All @@ -122,10 +121,13 @@ fn simplify_block_once(block: &mut Block) -> bool {
let n = block.ops.len();
let mut const_overwrites: Vec<(usize, i64)> = Vec::new();
let mut op_replacements: Vec<(usize, Op)> = Vec::new();
let mut vid_replacements: BTreeMap<ValueId, ValueId> = BTreeMap::new();
// Pre-sized with `n` per playbook §"Pre-size with `with_capacity`"
// (half the dead_store_elim PR #38 win).
let mut vid_replacements: FxHashMap<ValueId, ValueId> =
FxHashMap::with_capacity_and_hasher(n, Default::default());

// Build a map for peephole lookups.
let vid_to_op_pos: BTreeMap<ValueId, usize> = block
let vid_to_op_pos: FxHashMap<ValueId, usize> = block
.results
.iter()
.enumerate()
Expand Down Expand Up @@ -172,15 +174,19 @@ fn simplify_block_once(block: &mut Block) -> bool {
// Resolve transitive replacement chains: v2→v1→v0 becomes v2→v0.
let vid_replacements = resolve_transitive(&vid_replacements);

// Remap ValueIds in all ops in the block.
// Remap ValueIds in all ops in the block. Use the canonical
// `remap::remap_value_ids_fx` instead of the file-local
// `remap_values_in_op` clone — playbook §"Dedup canonical functions
// across modules" flags the drift hazard when the same walker ships
// in two places.
for op in block.ops.iter_mut() {
remap_values_in_op(op, &vid_replacements);
remap::remap_value_ids_fx(op, &vid_replacements);
}

// Remove dead ops whose results were redirected via ReplaceWithVid.
// Without this, the same pattern re-matches on the next iteration,
// producing the same replacement and causing an infinite fixpoint loop.
let dead_vids: BTreeSet<ValueId> = vid_replacements.keys().copied().collect();
let dead_vids: FxHashSet<ValueId> = vid_replacements.keys().copied().collect();
if !dead_vids.is_empty() {
let mut new_ops = Vec::new();
let mut new_results = Vec::new();
Expand Down Expand Up @@ -217,7 +223,7 @@ fn try_simplify(
op: &Op,
pos: usize,
block: &Block,
vid_to_pos: &BTreeMap<ValueId, usize>,
vid_to_pos: &FxHashMap<ValueId, usize>,
) -> Option<SimpResult> {
match op {
// ---- BinOp patterns ----
Expand Down Expand Up @@ -292,7 +298,7 @@ fn simplify_binop(
rhs: ValueId,
_pos: usize,
block: &Block,
vid_to_pos: &BTreeMap<ValueId, usize>,
vid_to_pos: &FxHashMap<ValueId, usize>,
) -> Option<SimpResult> {
let lv = find_const_in_block(block, lhs);
let rv = find_const_in_block(block, rhs);
Expand Down Expand Up @@ -383,7 +389,7 @@ fn simplify_select(
on_true: ValueId,
on_false: ValueId,
block: &Block,
vid_to_pos: &BTreeMap<ValueId, usize>,
vid_to_pos: &FxHashMap<ValueId, usize>,
) -> Option<SimpResult> {
// Select(true, a, b) → a
if let Some(1) = find_const_in_block(block, cond) {
Expand Down Expand Up @@ -427,7 +433,7 @@ fn find_const_in_block(block: &Block, vid: ValueId) -> Option<i64> {
fn get_defining_op<'a>(
vid: ValueId,
block: &'a Block,
vid_to_pos: &BTreeMap<ValueId, usize>,
vid_to_pos: &FxHashMap<ValueId, usize>,
) -> Option<(usize, &'a Op)> {
let &pos = vid_to_pos.get(&vid)?;
Some((pos, &block.ops[pos]))
Expand All @@ -437,7 +443,7 @@ fn get_defining_op<'a>(
fn get_neg_arg(
vid: ValueId,
block: &Block,
vid_to_pos: &BTreeMap<ValueId, usize>,
vid_to_pos: &FxHashMap<ValueId, usize>,
) -> Option<ValueId> {
let (_pos, op) = get_defining_op(vid, block, vid_to_pos)?;
if let Op::UnaryOp { op: UnaryOpKind::Neg, value } = op { Some(*value) } else { None }
Expand All @@ -447,7 +453,7 @@ fn get_neg_arg(
fn get_not_arg(
vid: ValueId,
block: &Block,
vid_to_pos: &BTreeMap<ValueId, usize>,
vid_to_pos: &FxHashMap<ValueId, usize>,
) -> Option<ValueId> {
let (_pos, op) = get_defining_op(vid, block, vid_to_pos)?;
match op {
Expand Down Expand Up @@ -475,15 +481,6 @@ fn get_not_arg(
}
}

/// Remap all ValueId references in an op.
fn remap_values_in_op(op: &mut Op, map: &BTreeMap<ValueId, ValueId>) {
op.for_each_value_id_mut(&mut |v| {
if let Some(&nv) = map.get(v) {
*v = nv;
}
});
}

#[cfg(test)]
mod tests {
use metaltile_core::ir::BinOpKind;
Expand Down
144 changes: 135 additions & 9 deletions crates/metaltile-codegen/src/passes/copy_prop.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,12 +31,11 @@
//! ACM TOPLAS 13(2):181–210. Sparse conditional constant propagation framework
//! that subsumes copy propagation.

use std::collections::{BTreeMap, BTreeSet};

use metaltile_core::{
dtype::DType,
ir::{Block, BlockId, Kernel, Op, ValueId},
};
use rustc_hash::{FxHashMap, FxHashSet};

use super::remap;
use crate::error::{Error, Result};
Expand Down Expand Up @@ -64,11 +63,11 @@ impl super::Pass for CopyPropPass {
}

/// Resolve transitive replacement chains: {v2→v1, v1→v0} becomes {v2→v0, v1→v0}.
fn resolve_transitive(map: &BTreeMap<ValueId, ValueId>) -> BTreeMap<ValueId, ValueId> {
let mut resolved = BTreeMap::new();
fn resolve_transitive(map: &FxHashMap<ValueId, ValueId>) -> FxHashMap<ValueId, ValueId> {
let mut resolved = FxHashMap::with_capacity_and_hasher(map.len(), Default::default());
for (&key, &val) in map.iter() {
let mut terminal = val;
let mut visited = BTreeSet::new();
let mut visited = FxHashSet::default();
visited.insert(key);
while let Some(&next) = map.get(&terminal) {
if !visited.insert(terminal) {
Expand All @@ -91,7 +90,10 @@ fn copy_prop_block_fixpoint(block: &mut Block) {

fn copy_prop_block_once(block: &mut Block) -> bool {
let n = block.ops.len();
let mut vid_replacements: BTreeMap<ValueId, ValueId> = BTreeMap::new();
// Pre-sized with `block.ops.len()` per playbook §"Pre-size with
// `with_capacity`" — half the dead_store_elim PR #38 win.
let mut vid_replacements: FxHashMap<ValueId, ValueId> =
FxHashMap::with_capacity_and_hasher(n, Default::default());

for i in 0..n {
let op = &block.ops[i];
Expand All @@ -109,15 +111,16 @@ fn copy_prop_block_once(block: &mut Block) -> bool {
// Resolve transitive replacement chains: v2→v1→v0 becomes v2→v0.
let vid_replacements = resolve_transitive(&vid_replacements);

// Remap ValueIds in all ops.
// Remap ValueIds in all ops via the FxHashMap-keyed sibling of
// `remap_value_ids` (playbook §"FxHashMap wins on codegen").
for op in block.ops.iter_mut() {
remap::remap_value_ids(op, &vid_replacements);
remap::remap_value_ids_fx(op, &vid_replacements);
}

// Remove dead ops whose results were redirected via identity propagation.
// Without this, the same identity pattern re-matches on the next iteration,
// producing the same replacement and causing an infinite fixpoint loop.
let dead_vids: BTreeSet<ValueId> = vid_replacements.keys().copied().collect();
let dead_vids: FxHashSet<ValueId> = vid_replacements.keys().copied().collect();
if !dead_vids.is_empty() {
let mut new_ops = Vec::new();
let mut new_results = Vec::new();
Expand Down Expand Up @@ -204,6 +207,129 @@ fn infer_value_dtype(vid: ValueId, block: &Block) -> Option<DType> {
None
}

#[cfg(test)]
mod perf {
//! `#[ignore]`'d microbench for the cascade — runs under
//!
//! ```text
//! cargo test -p metaltile-codegen --release perf_copy_prop_btreemap_vs_fxhash \
//! -- --ignored --nocapture
//! ```
//!
//! Reconstructs the pre-cascade `resolve_transitive` + `dead_vids`
//! collection + `remap_value_ids` shape using `BTreeMap`/`BTreeSet`
//! against the new `FxHashMap`/`FxHashSet` versions and times the
//! per-block hot path. Covers the same shape change applied across
//! `copy_prop.rs` / `algebraic_simplify.rs` / `unroll.rs` /
//! `vectorize.rs::result_remap` — the four call sites of
//! `remap::remap_value_ids_fx` introduced in this PR.
use std::{
collections::{BTreeMap, BTreeSet},
hint::black_box,
time::Instant,
};

use metaltile_core::ir::{BinOpKind, Op, ValueId};
use rustc_hash::{FxHashMap, FxHashSet};

use super::remap;

fn synthetic_block(n: usize) -> (Vec<Op>, Vec<Option<ValueId>>) {
// Build n BinOp(Add, v_{i-1}, v_{i-2}) ops to exercise the
// 2-refs-per-op iteration that drives `remap_value_ids` cost.
let mut ops = Vec::with_capacity(n);
let mut results = Vec::with_capacity(n);
for i in 0..n {
let lhs = ValueId::new(if i >= 2 { (i - 2) as u32 } else { 0 });
let rhs = ValueId::new(if i >= 1 { (i - 1) as u32 } else { 0 });
ops.push(Op::BinOp { op: BinOpKind::Add, lhs, rhs });
results.push(Some(ValueId::new(i as u32 + 100)));
}
(ops, results)
}

#[test]
#[ignore]
fn perf_copy_prop_btreemap_vs_fxhash() {
const N_OPS: usize = 2_000;
const N_REPLACEMENTS: u32 = 256;
const N_ITERS: usize = 500;

let (base_ops, _results) = synthetic_block(N_OPS);

// ── BTreeMap (pre-cascade) ──
let bt_seed: Vec<(ValueId, ValueId)> =
(0..N_REPLACEMENTS).map(|i| (ValueId::new(i), ValueId::new(i + 1000))).collect();

for _ in 0..3 {
let map: BTreeMap<ValueId, ValueId> = bt_seed.iter().copied().collect();
let mut ops = base_ops.clone();
for op in ops.iter_mut() {
remap::remap_value_ids(op, &map);
}
let _dead: BTreeSet<ValueId> = map.keys().copied().collect();
black_box(&ops);
}
let t0 = Instant::now();
for _ in 0..N_ITERS {
let map: BTreeMap<ValueId, ValueId> = bt_seed.iter().copied().collect();
let mut ops = base_ops.clone();
for op in ops.iter_mut() {
remap::remap_value_ids(op, &map);
}
let dead: BTreeSet<ValueId> = map.keys().copied().collect();
black_box(&ops);
black_box(dead.len());
}
let bt_elapsed = t0.elapsed();

// ── FxHashMap (this cascade) ──
for _ in 0..3 {
let map: FxHashMap<ValueId, ValueId> = bt_seed.iter().copied().collect();
let mut ops = base_ops.clone();
for op in ops.iter_mut() {
remap::remap_value_ids_fx(op, &map);
}
let _dead: FxHashSet<ValueId> = map.keys().copied().collect();
black_box(&ops);
}
let t0 = Instant::now();
for _ in 0..N_ITERS {
let map: FxHashMap<ValueId, ValueId> =
FxHashMap::with_capacity_and_hasher(N_REPLACEMENTS as usize, Default::default());
let mut map = map;
for (k, v) in &bt_seed {
map.insert(*k, *v);
}
let mut ops = base_ops.clone();
for op in ops.iter_mut() {
remap::remap_value_ids_fx(op, &map);
}
let dead: FxHashSet<ValueId> = map.keys().copied().collect();
black_box(&ops);
black_box(dead.len());
}
let fx_elapsed = t0.elapsed();

let bt_ns_per = bt_elapsed.as_nanos() as f64 / (N_ITERS * N_OPS) as f64;
let fx_ns_per = fx_elapsed.as_nanos() as f64 / (N_ITERS * N_OPS) as f64;
let speedup = bt_ns_per / fx_ns_per;
let delta_pct = (1.0 - fx_ns_per / bt_ns_per) * 100.0;
println!();
println!(
"=== copy_prop hot path: build map + remap {N_OPS} ops + collect dead set ({N_ITERS} iters) ==="
);
println!(" BTreeMap (old) : {bt_elapsed:?} ({bt_ns_per:.2} ns/op)");
println!(" FxHashMap (new) : {fx_elapsed:?} ({fx_ns_per:.2} ns/op)");
println!(" speedup : {speedup:.2}× ({delta_pct:+.1}%)");

assert!(
fx_ns_per <= bt_ns_per * 1.05,
"FxHashMap regressed vs BTreeMap (fx={fx_ns_per:.2} ns, bt={bt_ns_per:.2} ns)"
);
}
}

#[cfg(test)]
mod tests {
use metaltile_core::{
Expand Down
15 changes: 15 additions & 0 deletions crates/metaltile-codegen/src/passes/remap.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
use std::collections::BTreeMap;

use metaltile_core::ir::{Kernel, Op, ValueId};
use rustc_hash::FxHashMap;
use smallvec::SmallVec;

// ---------------------------------------------------------------------------
Expand All @@ -40,6 +41,20 @@ pub fn remap_value_ids(op: &mut Op, map: &BTreeMap<ValueId, ValueId>) {
});
}

/// `FxHashMap`-keyed sibling of [`remap_value_ids`]. Same semantics;
/// the only difference is the map type. Lets passes whose lookup is
/// pure get-only (no iteration-order dependency) build their
/// replacement map as `FxHashMap<ValueId, ValueId>` and skip BTree's
/// log(n) per-probe cost — playbook §"FxHashMap wins on the codegen
/// pipeline" (PR #38 fusion swap measured -29.8%).
pub fn remap_value_ids_fx(op: &mut Op, map: &FxHashMap<ValueId, ValueId>) {
op.for_each_value_id_mut(&mut |v| {
if let Some(&nv) = map.get(v) {
*v = nv;
}
});
}

// ---------------------------------------------------------------------------
// op_value_refs — collect all ValueId references (read-only, for analysis)
// ---------------------------------------------------------------------------
Expand Down
12 changes: 9 additions & 3 deletions crates/metaltile-codegen/src/passes/unroll.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,13 @@ fn unroll_block(
inlined.push((Op::Const { value: iv_val }, Some(iv_const_vid)));

// ---- build vid-map for this clone ----------------------------
let mut vid_map: BTreeMap<ValueId, ValueId> = BTreeMap::new();
// FxHashMap per playbook §"FxHashMap wins on the codegen
// pipeline" — vid_map is pure get-only on `ValueId` keys
// when consumed by `remap_value_ids_fx` below. Pre-sized
// with `body_n + 1` (one entry per body result plus the IV
// mapping) per §"Pre-size with `with_capacity`".
let mut vid_map: FxHashMap<ValueId, ValueId> =
FxHashMap::with_capacity_and_hasher(body_n + 1, Default::default());
vid_map.insert(iv_vid, iv_const_vid);

for j in 0..body_n {
Expand Down Expand Up @@ -273,7 +279,7 @@ fn unroll_block(
cloned.names = src_block.names.clone();
for (idx, src_op) in src_block.ops.iter().enumerate() {
let mut new_op = src_op.clone();
remap::remap_value_ids(&mut new_op, &nested_vid_map);
remap::remap_value_ids_fx(&mut new_op, &nested_vid_map);
let new_result = src_block.results[idx]
.map(|old_v| nested_vid_map.get(&old_v).copied().unwrap_or(old_v));
cloned.ops.push(new_op);
Expand All @@ -284,7 +290,7 @@ fn unroll_block(

for j in 0..body_n {
let mut new_op = body.ops[j].clone();
remap::remap_value_ids(&mut new_op, &vid_map);
remap::remap_value_ids_fx(&mut new_op, &vid_map);

// Rewrite the cloned op's nested-block references to
// point at the fresh clones we just inserted.
Expand Down
Loading
Loading