use std::hash::{Hash, Hasher};
pub mod add_sub_cancel;
pub mod aos_to_soa_promote;
mod arithmetic_combine;
pub mod bank_conflict_pad;
pub mod bitwise_combine;
mod body_index;
pub mod bitwise_idemp {
use crate::rewrites::self_binop::rewrite_self_binops;
use crate::KernelDescriptor;
use vyre_foundation::ir::BinOp;
#[must_use]
pub fn bitwise_idemp(desc: &KernelDescriptor) -> KernelDescriptor {
rewrite_self_binops(desc, |bin| matches!(bin, BinOp::BitAnd | BinOp::BitOr))
}
}
pub mod boolean_simplify;
pub mod branch_collapse;
pub mod canonicalize;
pub mod cmp_normalize;
pub mod cmp_self_false;
mod commutative_lit_chain;
pub mod const_buffer_promote;
mod dataflow_facts;
pub mod dead_store;
pub mod descriptor_const_fold;
pub mod descriptor_cse;
pub mod descriptor_dce;
pub mod div_combine;
pub mod drop_unused_bindings;
pub mod drop_unused_child_bodies;
pub mod drop_unused_literals;
pub mod egraph_saturation;
pub mod emit_order;
pub mod identity_elim;
pub mod licm;
mod literal;
pub mod load_forwarding;
pub mod loop_fission;
pub mod loop_fusion;
pub mod loop_unroll;
pub mod loop_zero_iter;
pub mod matmul_promote;
mod memory_address;
mod rhs_lit_chain;
mod self_binop;
pub mod min_max_idemp {
use crate::rewrites::self_binop::rewrite_self_binops;
use crate::KernelDescriptor;
use vyre_foundation::ir::BinOp;
#[must_use]
pub fn min_max_idemp(desc: &KernelDescriptor) -> KernelDescriptor {
rewrite_self_binops(desc, |bin| matches!(bin, BinOp::Min | BinOp::Max))
}
}
pub mod mod_idemp;
pub mod mul_add_to_fma;
pub mod negate_cancel;
pub mod select_fold;
pub mod shared_mem_promote;
pub mod shift_combine;
pub mod strength_reduce;
pub mod sub_combine;
pub mod tail_mask;
pub mod unary_idemp;
pub mod xor_self_zero;
pub use add_sub_cancel::add_sub_cancel;
pub use aos_to_soa_promote::{promote as aos_to_soa_promote, LayoutHint as AosSoaLayoutHint};
pub use arithmetic_combine::add_combine::add_combine;
pub use arithmetic_combine::mul_combine::mul_combine;
pub use arithmetic_combine::{add_combine, mul_combine};
pub use bank_conflict_pad::bank_conflict_pad;
pub use bitwise_combine::bitwise_combine;
pub use bitwise_idemp::bitwise_idemp;
pub use boolean_simplify::boolean_simplify;
pub use branch_collapse::branch_collapse;
pub use canonicalize::canonicalize;
pub use cmp_normalize::cmp_normalize;
pub use cmp_self_false::cmp_self_false;
pub use const_buffer_promote::const_buffer_promote;
pub use dead_store::{
dead_store, dead_store_with_alias_facts, dead_store_with_dataflow_analysis_facts,
dead_store_with_dataflow_facts, dead_store_with_weir_alias_facts,
};
pub use descriptor_const_fold::descriptor_const_fold;
pub use descriptor_cse::descriptor_cse;
pub use descriptor_dce::descriptor_dce;
pub use div_combine::div_combine;
pub use drop_unused_bindings::drop_unused_bindings;
pub use drop_unused_child_bodies::drop_unused_child_bodies;
pub use drop_unused_literals::drop_unused_literals;
pub use emit_order::emit_order;
pub use identity_elim::identity_elim;
pub use licm::{
licm, licm_with_alias_facts, licm_with_dataflow_analysis_facts, licm_with_dataflow_facts,
licm_with_weir_alias_facts,
};
pub use load_forwarding::{
load_forwarding, load_forwarding_with_alias_facts,
load_forwarding_with_dataflow_analysis_facts, load_forwarding_with_dataflow_facts,
load_forwarding_with_weir_alias_facts,
};
pub use loop_fission::{
loop_fission, loop_fission_with_alias_facts, loop_fission_with_dataflow_analysis_facts,
loop_fission_with_dataflow_facts, loop_fission_with_weir_alias_facts,
};
pub use loop_fusion::{
loop_fusion, loop_fusion_with_alias_facts, loop_fusion_with_dataflow_analysis_facts,
loop_fusion_with_dataflow_facts, loop_fusion_with_weir_alias_facts,
};
pub use loop_unroll::loop_unroll;
pub use loop_zero_iter::loop_zero_iter;
pub use matmul_promote::{infer_matmul_tile_loops, matmul_promote, MatmulTileLoopPlan};
pub use min_max_idemp::min_max_idemp;
pub use mod_idemp::mod_idemp;
pub use mul_add_to_fma::mul_add_to_fma;
pub use negate_cancel::negate_cancel;
pub use select_fold::select_fold;
pub use shared_mem_promote::shared_mem_promote;
pub use shift_combine::shift_combine;
pub use strength_reduce::strength_reduce;
pub use sub_combine::sub_combine;
pub use tail_mask::apply_tail_mask;
pub use unary_idemp::unary_idemp;
pub use xor_self_zero::xor_self_zero;
#[derive(Debug, Clone, Copy)]
pub struct DescriptorRewritePass {
pub name: &'static str,
pub rewrite: fn(&crate::KernelDescriptor) -> crate::KernelDescriptor,
}
impl DescriptorRewritePass {
#[must_use]
fn run(self, desc: &crate::KernelDescriptor) -> crate::KernelDescriptor {
(self.rewrite)(desc)
}
}
fn egraph_saturation_pass(desc: &crate::KernelDescriptor) -> crate::KernelDescriptor {
egraph_saturation::saturate_algebraic_descriptor(desc).0
}
const CANONICAL_REWRITE_PASSES: &[DescriptorRewritePass] = &[
DescriptorRewritePass {
name: "strength_reduce",
rewrite: strength_reduce,
},
DescriptorRewritePass {
name: "shift_combine",
rewrite: shift_combine,
},
DescriptorRewritePass {
name: "shared_mem_promote",
rewrite: shared_mem_promote,
},
DescriptorRewritePass {
name: "bank_conflict_pad",
rewrite: bank_conflict_pad,
},
DescriptorRewritePass {
name: "const_buffer_promote",
rewrite: const_buffer_promote,
},
DescriptorRewritePass {
name: "descriptor_const_fold",
rewrite: descriptor_const_fold,
},
DescriptorRewritePass {
name: "add_combine",
rewrite: add_combine,
},
DescriptorRewritePass {
name: "sub_combine",
rewrite: sub_combine,
},
DescriptorRewritePass {
name: "mul_combine",
rewrite: mul_combine,
},
DescriptorRewritePass {
name: "div_combine",
rewrite: div_combine,
},
DescriptorRewritePass {
name: "mod_idemp",
rewrite: mod_idemp,
},
DescriptorRewritePass {
name: "add_sub_cancel",
rewrite: add_sub_cancel,
},
DescriptorRewritePass {
name: "bitwise_combine",
rewrite: bitwise_combine,
},
DescriptorRewritePass {
name: "identity_elim",
rewrite: identity_elim,
},
DescriptorRewritePass {
name: "boolean_simplify",
rewrite: boolean_simplify,
},
DescriptorRewritePass {
name: "negate_cancel",
rewrite: negate_cancel,
},
DescriptorRewritePass {
name: "unary_idemp",
rewrite: unary_idemp,
},
DescriptorRewritePass {
name: "select_fold",
rewrite: select_fold,
},
DescriptorRewritePass {
name: "min_max_idemp",
rewrite: min_max_idemp,
},
DescriptorRewritePass {
name: "bitwise_idemp",
rewrite: bitwise_idemp,
},
DescriptorRewritePass {
name: "branch_collapse",
rewrite: branch_collapse,
},
DescriptorRewritePass {
name: "loop_fusion",
rewrite: loop_fusion,
},
DescriptorRewritePass {
name: "loop_unroll",
rewrite: loop_unroll,
},
DescriptorRewritePass {
name: "loop_zero_iter",
rewrite: loop_zero_iter,
},
DescriptorRewritePass {
name: "licm",
rewrite: licm,
},
DescriptorRewritePass {
name: "load_forwarding",
rewrite: load_forwarding,
},
DescriptorRewritePass {
name: "mul_add_to_fma",
rewrite: mul_add_to_fma,
},
DescriptorRewritePass {
name: "matmul_promote",
rewrite: matmul_promote,
},
DescriptorRewritePass {
name: "descriptor_dce_after_forwarding",
rewrite: descriptor_dce,
},
DescriptorRewritePass {
name: "dead_store",
rewrite: dead_store,
},
DescriptorRewritePass {
name: "descriptor_dce",
rewrite: descriptor_dce,
},
DescriptorRewritePass {
name: "cmp_normalize",
rewrite: cmp_normalize,
},
DescriptorRewritePass {
name: "cmp_self_false",
rewrite: cmp_self_false,
},
DescriptorRewritePass {
name: "xor_self_zero",
rewrite: xor_self_zero,
},
DescriptorRewritePass {
name: "canonicalize",
rewrite: canonicalize,
},
DescriptorRewritePass {
name: "descriptor_cse",
rewrite: descriptor_cse,
},
DescriptorRewritePass {
name: "egraph_saturation",
rewrite: egraph_saturation_pass,
},
DescriptorRewritePass {
name: "descriptor_const_fold_post_saturation",
rewrite: descriptor_const_fold,
},
DescriptorRewritePass {
name: "identity_elim_post_saturation",
rewrite: identity_elim,
},
DescriptorRewritePass {
name: "descriptor_dce_post_saturation",
rewrite: descriptor_dce,
},
DescriptorRewritePass {
name: "cmp_normalize_post_saturation",
rewrite: cmp_normalize,
},
DescriptorRewritePass {
name: "canonicalize_post_saturation",
rewrite: canonicalize,
},
DescriptorRewritePass {
name: "descriptor_cse_post_saturation",
rewrite: descriptor_cse,
},
DescriptorRewritePass {
name: "drop_unused_bindings",
rewrite: drop_unused_bindings,
},
DescriptorRewritePass {
name: "drop_unused_literals",
rewrite: drop_unused_literals,
},
DescriptorRewritePass {
name: "drop_unused_child_bodies",
rewrite: drop_unused_child_bodies,
},
DescriptorRewritePass {
name: "emit_order",
rewrite: emit_order,
},
];
#[must_use]
pub const fn canonical_rewrite_passes() -> &'static [DescriptorRewritePass] {
CANONICAL_REWRITE_PASSES
}
fn descriptor_hash(desc: &crate::KernelDescriptor) -> u64 {
let mut h = rustc_hash::FxHasher::default();
desc.hash(&mut h);
h.finish()
}
fn run_descriptor_passes(
desc: &crate::KernelDescriptor,
passes: &[DescriptorRewritePass],
) -> crate::KernelDescriptor {
let mut current = desc.clone();
for pass in passes {
let pre_hash = descriptor_hash(¤t);
let next = pass.run(¤t);
if descriptor_hash(&next) != pre_hash || next != current {
current = next;
}
#[cfg(debug_assertions)]
debug_verify_after_rewrite(¤t, pass.name);
}
current
}
#[must_use]
pub fn run_all_once(desc: &crate::KernelDescriptor) -> crate::KernelDescriptor {
run_descriptor_passes(desc, canonical_rewrite_passes())
}
#[cfg(debug_assertions)]
fn debug_verify_after_rewrite(desc: &crate::KernelDescriptor, pass: &str) {
if let Err(errors) = crate::verify::verify(desc) {
panic!(
"rewrite pass `{pass}` produced an invalid KernelDescriptor - {} violation(s):\n{errors:#?}",
errors.len()
);
}
}
#[must_use]
pub fn run_all_once_with_dataflow_facts(
desc: &crate::KernelDescriptor,
alias_facts: &crate::analyses::alias_facts::AliasFactSet,
reaching_defs: &crate::analyses::reaching_def_facts::ReachingDefFactSet,
) -> crate::KernelDescriptor {
let reduced = strength_reduce(desc);
let shared_promoted = shared_mem_promote(&reduced);
let padded = bank_conflict_pad(&shared_promoted);
let const_promoted = const_buffer_promote(&padded);
let folded = descriptor_const_fold(&const_promoted);
let identified = identity_elim(&folded);
let collapsed = branch_collapse(&identified);
let fused = loop_fusion_with_dataflow_facts(&collapsed, alias_facts, reaching_defs);
let unrolled = loop_unroll(&fused);
let hoisted = licm_with_dataflow_facts(&unrolled, alias_facts, reaching_defs);
let forwarded = load_forwarding_with_dataflow_facts(&hoisted, alias_facts, reaching_defs);
let cleaned = descriptor_dce(&forwarded);
let dse_done = dead_store_with_dataflow_facts(&cleaned, alias_facts, reaching_defs);
let dced = descriptor_dce(&dse_done);
let canon = canonicalize(&dced);
let merged = descriptor_cse(&canon);
let (saturated, _) = egraph_saturation::saturate_algebraic_descriptor(&merged);
let saturated_folded = descriptor_const_fold(&saturated);
let saturated_identified = identity_elim(&saturated_folded);
let saturated_dced = descriptor_dce(&saturated_identified);
let saturated_canon = canonicalize(&saturated_dced);
let saturated_merged = descriptor_cse(&saturated_canon);
let pruned_bindings = drop_unused_bindings(&saturated_merged);
let pruned_literals = drop_unused_literals(&pruned_bindings);
let pruned_children = drop_unused_child_bodies(&pruned_literals);
emit_order(&pruned_children)
}
pub const RUN_ALL_MAX_ITERS: usize = 4;
#[must_use]
pub fn run_all(desc: &crate::KernelDescriptor) -> crate::KernelDescriptor {
run_all_with_stats(desc).0
}
#[must_use]
pub fn run_all_with_dataflow_facts(
desc: &crate::KernelDescriptor,
alias_facts: &crate::analyses::alias_facts::AliasFactSet,
reaching_defs: &crate::analyses::reaching_def_facts::ReachingDefFactSet,
) -> crate::KernelDescriptor {
run_all_with_dataflow_stats(desc, alias_facts, reaching_defs).0
}
#[must_use]
pub fn run_all_with_dataflow_analysis_facts(
desc: &crate::KernelDescriptor,
alias_facts: &crate::analyses::weir_alias::AliasFactSet,
reaching_defs: &crate::analyses::weir_reaching_def::ReachingDefFactSet,
) -> crate::KernelDescriptor {
run_all_with_dataflow_facts(desc, alias_facts, reaching_defs)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct OptimizationStats {
pub ops_before: usize,
pub ops_after: usize,
pub bindings_before: usize,
pub bindings_after: usize,
pub literals_before: usize,
pub literals_after: usize,
pub iterations: usize,
pub converged: bool,
}
impl OptimizationStats {
pub fn ops_eliminated(&self) -> usize {
self.ops_before.saturating_sub(self.ops_after)
}
pub fn bindings_dropped(&self) -> usize {
self.bindings_before.saturating_sub(self.bindings_after)
}
pub fn is_no_op(&self) -> bool {
self.ops_before == self.ops_after
&& self.bindings_before == self.bindings_after
&& self.literals_before == self.literals_after
}
pub fn off_graph_dropped(&self) -> usize {
self.bindings_dropped() + self.literals_before.saturating_sub(self.literals_after)
}
pub fn merge(&mut self, other: OptimizationStats) {
self.ops_before = self.ops_before.saturating_add(other.ops_before);
self.ops_after = self.ops_after.saturating_add(other.ops_after);
self.bindings_before = self.bindings_before.saturating_add(other.bindings_before);
self.bindings_after = self.bindings_after.saturating_add(other.bindings_after);
self.literals_before = self.literals_before.saturating_add(other.literals_before);
self.literals_after = self.literals_after.saturating_add(other.literals_after);
self.iterations = self.iterations.saturating_add(other.iterations);
self.converged = self.converged && other.converged;
}
pub fn zero() -> Self {
OptimizationStats {
ops_before: 0,
ops_after: 0,
bindings_before: 0,
bindings_after: 0,
literals_before: 0,
literals_after: 0,
iterations: 0,
converged: true,
}
}
pub fn format_short(&self) -> String {
format!(
"ops {}→{} (-{}), bindings {}→{} (-{}), iters {} ({})",
self.ops_before,
self.ops_after,
self.ops_eliminated(),
self.bindings_before,
self.bindings_after,
self.bindings_dropped(),
self.iterations,
if self.converged {
"converged"
} else {
"stopped"
},
)
}
}
impl std::fmt::Display for OptimizationStats {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.format_short())
}
}
#[must_use]
pub fn run_all_with_stats(
desc: &crate::KernelDescriptor,
) -> (crate::KernelDescriptor, OptimizationStats) {
let ops_before = desc.body.ops.len();
let bindings_before = desc.bindings.slots.len();
let literals_before = desc.body.literals.len();
let mut current = run_all_once(desc);
#[cfg(debug_assertions)]
debug_verify_after_rewrite(¤t, "run_all_once (iter 1)");
let mut iterations = 1usize;
let mut current_hash = descriptor_hash(¤t);
let mut converged = current_hash == descriptor_hash(desc) && current == *desc;
while !converged && iterations < RUN_ALL_MAX_ITERS {
let next = run_all_once(¤t);
iterations += 1;
#[cfg(debug_assertions)]
debug_verify_after_rewrite(&next, &format!("run_all_once (iter {iterations})"));
let next_hash = descriptor_hash(&next);
converged = next_hash == current_hash && next == current;
current = next;
current_hash = next_hash;
}
let stats = OptimizationStats {
ops_before,
ops_after: current.body.ops.len(),
bindings_before,
bindings_after: current.bindings.slots.len(),
literals_before,
literals_after: current.body.literals.len(),
iterations,
converged,
};
(current, stats)
}
#[must_use]
pub fn run_all_with_dataflow_stats(
desc: &crate::KernelDescriptor,
alias_facts: &crate::analyses::alias_facts::AliasFactSet,
reaching_defs: &crate::analyses::reaching_def_facts::ReachingDefFactSet,
) -> (crate::KernelDescriptor, OptimizationStats) {
let ops_before = desc.body.ops.len();
let bindings_before = desc.bindings.slots.len();
let literals_before = desc.body.literals.len();
let mut current = run_all_once_with_dataflow_facts(desc, alias_facts, reaching_defs);
let mut iterations = 1usize;
let mut current_hash = descriptor_hash(¤t);
let mut converged = current_hash == descriptor_hash(desc) && current == *desc;
while !converged && iterations < RUN_ALL_MAX_ITERS {
let next = run_all_once_with_dataflow_facts(¤t, alias_facts, reaching_defs);
iterations += 1;
let next_hash = descriptor_hash(&next);
converged = next_hash == current_hash && next == current;
current = next;
current_hash = next_hash;
}
let stats = OptimizationStats {
ops_before,
ops_after: current.body.ops.len(),
bindings_before,
bindings_after: current.bindings.slots.len(),
literals_before,
literals_after: current.body.literals.len(),
iterations,
converged,
};
(current, stats)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
BindingLayout, Dispatch, KernelBody, KernelDescriptor, KernelOp, KernelOpKind, LiteralValue,
};
#[test]
fn run_all_on_empty_kernel_returns_empty() {
let desc = KernelDescriptor {
id: "k".into(),
bindings: BindingLayout { slots: vec![] },
dispatch: Dispatch::new(1, 1, 1),
body: KernelBody {
ops: vec![],
child_bodies: vec![],
literals: vec![],
},
};
let out = run_all(&desc);
assert!(out.body.ops.is_empty());
}
#[test]
fn run_all_is_idempotent() {
let desc = KernelDescriptor {
id: "k".into(),
bindings: BindingLayout { slots: vec![] },
dispatch: Dispatch::new(1, 1, 1),
body: KernelBody {
ops: vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(1),
}, KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(2),
}, ],
child_bodies: vec![],
literals: vec![LiteralValue::U32(5), LiteralValue::U32(99)],
},
};
let once = run_all(&desc);
let twice = run_all(&once);
assert_eq!(once.body.ops.len(), twice.body.ops.len());
assert_eq!(once.body.literals, twice.body.literals);
}
#[test]
fn run_all_collapses_kitchen_sink_kernel() {
let desc = KernelDescriptor {
id: "kitchen_sink".into(),
bindings: BindingLayout {
slots: vec![crate::BindingSlot {
slot: 0,
element_type: vyre_foundation::ir::DataType::U32,
element_count: None,
memory_class: crate::MemoryClass::Global,
visibility: crate::BindingVisibility::ReadWrite,
name: "buf".into(),
}],
},
dispatch: Dispatch::new(1, 1, 1),
body: KernelBody {
ops: vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![2],
result: Some(2),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![3],
result: Some(3),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(4),
},
KernelOp {
kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Add),
operands: vec![3, 0],
result: Some(5),
},
KernelOp {
kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul),
operands: vec![3, 2],
result: Some(6),
},
KernelOp {
kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul),
operands: vec![3, 1],
result: Some(7),
},
KernelOp {
kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul),
operands: vec![3, 0],
result: Some(8),
},
KernelOp {
kind: KernelOpKind::StoreGlobal,
operands: vec![0, 0, 5],
result: None,
},
KernelOp {
kind: KernelOpKind::StoreGlobal,
operands: vec![0, 0, 6],
result: None,
},
],
child_bodies: vec![],
literals: vec![
LiteralValue::U32(0),
LiteralValue::U32(1),
LiteralValue::U32(8),
LiteralValue::U32(7),
],
},
};
let before_op_count = desc.body.ops.len();
let out = run_all(&desc);
let store_count = out
.body
.ops
.iter()
.filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
.count();
assert_eq!(store_count, 1, "dead_store must drop the overwritten store");
let mul_count = out
.body
.ops
.iter()
.filter(|o| {
matches!(
o.kind,
KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul)
)
})
.count();
assert_eq!(mul_count, 0, "all 3 Mul ops should be eliminated");
let add_count = out
.body
.ops
.iter()
.filter(|o| {
matches!(
o.kind,
KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Add)
)
})
.count();
assert_eq!(add_count, 0, "Add(r3, 0) → r3 should be eliminated");
assert!(
out.body.ops.len() <= before_op_count - 5,
"expected op count to drop by ≥5 (was {before_op_count}, now {})",
out.body.ops.len()
);
let twice = run_all(&out);
assert_eq!(out.body.ops.len(), twice.body.ops.len());
}
#[test]
fn run_all_forwards_then_drops_redundant_load() {
let desc = KernelDescriptor {
id: "stl".into(),
bindings: BindingLayout {
slots: vec![crate::BindingSlot {
slot: 0,
element_type: vyre_foundation::ir::DataType::U32,
element_count: None,
memory_class: crate::MemoryClass::Global,
visibility: crate::BindingVisibility::ReadWrite,
name: "buf".into(),
}],
},
dispatch: Dispatch::new(1, 1, 1),
body: KernelBody {
ops: vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
}, KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
}, KernelOp {
kind: KernelOpKind::StoreGlobal,
operands: vec![0, 0, 1],
result: None,
},
KernelOp {
kind: KernelOpKind::LoadGlobal,
operands: vec![0, 0],
result: Some(2),
},
KernelOp {
kind: KernelOpKind::StoreGlobal,
operands: vec![0, 0, 2],
result: None,
},
],
child_bodies: vec![],
literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
},
};
let out = run_all(&desc);
let store_count = out
.body
.ops
.iter()
.filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
.count();
assert_eq!(
store_count, 1,
"dead_store must drop the redundant first store"
);
let load_count = out
.body
.ops
.iter()
.filter(|o| matches!(o.kind, KernelOpKind::LoadGlobal))
.count();
assert_eq!(load_count, 0, "descriptor_dce must drop the redundant load");
}
#[test]
fn run_all_unrolls_then_simplifies() {
let desc = KernelDescriptor {
id: "loop_then_dse".into(),
bindings: BindingLayout {
slots: vec![crate::BindingSlot {
slot: 0,
element_type: vyre_foundation::ir::DataType::U32,
element_count: None,
memory_class: crate::MemoryClass::Global,
visibility: crate::BindingVisibility::ReadWrite,
name: "buf".into(),
}],
},
dispatch: Dispatch::new(1, 1, 1),
body: KernelBody {
ops: vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(2),
},
KernelOp {
kind: KernelOpKind::StructuredForLoop {
loop_var: std::sync::Arc::from("i"),
},
operands: vec![0, 1, 0],
result: None,
},
],
child_bodies: vec![KernelBody {
ops: vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::StoreGlobal,
operands: vec![0, 0, 1],
result: None,
},
],
child_bodies: vec![],
literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
}],
literals: vec![LiteralValue::U32(0), LiteralValue::U32(2)],
},
};
let out = run_all(&desc);
let loop_count = out
.body
.ops
.iter()
.filter(|o| matches!(o.kind, KernelOpKind::StructuredForLoop { .. }))
.count();
assert_eq!(loop_count, 0, "loop should be unrolled");
let store_count = out
.body
.ops
.iter()
.filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
.count();
assert!(
store_count >= 1,
"at least one store from the unrolled body should survive"
);
}
#[test]
fn run_all_with_stats_reports_op_reduction() {
let desc = KernelDescriptor {
id: "stats".into(),
bindings: BindingLayout {
slots: vec![crate::BindingSlot {
slot: 0,
element_type: vyre_foundation::ir::DataType::U32,
element_count: None,
memory_class: crate::MemoryClass::Global,
visibility: crate::BindingVisibility::ReadWrite,
name: "buf".into(),
}],
},
dispatch: Dispatch::new(1, 1, 1),
body: KernelBody {
ops: vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Add),
operands: vec![1, 0],
result: Some(2),
}, KernelOp {
kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul),
operands: vec![1, 0],
result: Some(3),
}, KernelOp {
kind: KernelOpKind::StoreGlobal,
operands: vec![0, 0, 1],
result: None,
},
],
child_bodies: vec![],
literals: vec![LiteralValue::U32(0), LiteralValue::U32(99)],
},
};
let (_out, stats) = run_all_with_stats(&desc);
assert!(
stats.ops_eliminated() >= 2,
"expected ≥2 ops eliminated, got {} ({} → {})",
stats.ops_eliminated(),
stats.ops_before,
stats.ops_after
);
assert!(stats.iterations >= 1);
assert!(stats.iterations <= RUN_ALL_MAX_ITERS);
assert!(stats.converged, "pipeline must converge");
}
#[test]
fn optimization_stats_format_short_includes_all_fields() {
let s = OptimizationStats {
ops_before: 11,
ops_after: 3,
bindings_before: 3,
bindings_after: 1,
literals_before: 4,
literals_after: 2,
iterations: 2,
converged: true,
};
let f = s.format_short();
assert!(f.contains("ops 11→3 (-8)"));
assert!(f.contains("bindings 3→1 (-2)"));
assert!(f.contains("iters 2"));
assert!(f.contains("converged"));
}
#[test]
fn optimization_stats_merge_is_associative() {
let a = OptimizationStats {
ops_before: 10,
ops_after: 3,
bindings_before: 2,
bindings_after: 1,
literals_before: 5,
literals_after: 2,
iterations: 2,
converged: true,
};
let b = OptimizationStats {
ops_before: 7,
ops_after: 4,
bindings_before: 1,
bindings_after: 1,
literals_before: 3,
literals_after: 3,
iterations: 1,
converged: true,
};
let c = OptimizationStats {
ops_before: 5,
ops_after: 2,
bindings_before: 1,
bindings_after: 0,
literals_before: 2,
literals_after: 1,
iterations: 1,
converged: true,
};
let mut left = a;
left.merge(b);
left.merge(c);
let mut bc = b;
bc.merge(c);
let mut right = a;
right.merge(bc);
assert_eq!(left, right);
}
#[test]
fn optimization_stats_merge_aggregates() {
let mut acc = OptimizationStats::zero();
acc.merge(OptimizationStats {
ops_before: 10,
ops_after: 3,
bindings_before: 2,
bindings_after: 1,
literals_before: 5,
literals_after: 2,
iterations: 2,
converged: true,
});
acc.merge(OptimizationStats {
ops_before: 7,
ops_after: 4,
bindings_before: 1,
bindings_after: 1,
literals_before: 3,
literals_after: 3,
iterations: 1,
converged: false,
});
assert_eq!(acc.ops_before, 17);
assert_eq!(acc.ops_after, 7);
assert_eq!(acc.iterations, 3);
assert!(!acc.converged); }
#[test]
fn optimization_stats_zero_is_identity() {
let s = OptimizationStats {
ops_before: 5,
ops_after: 3,
bindings_before: 1,
bindings_after: 1,
literals_before: 2,
literals_after: 1,
iterations: 2,
converged: true,
};
let mut acc = OptimizationStats::zero();
acc.merge(s);
assert_eq!(acc, s);
}
#[test]
fn optimization_stats_format_short_marks_stopped() {
let s = OptimizationStats {
ops_before: 5,
ops_after: 5,
bindings_before: 1,
bindings_after: 1,
literals_before: 1,
literals_after: 1,
iterations: 4,
converged: false,
};
assert!(s.format_short().contains("stopped"));
}
#[test]
fn run_all_with_stats_reports_no_change_when_already_optimal() {
let desc = KernelDescriptor {
id: "minimal".into(),
bindings: BindingLayout {
slots: vec![crate::BindingSlot {
slot: 0,
element_type: vyre_foundation::ir::DataType::U32,
element_count: None,
memory_class: crate::MemoryClass::Global,
visibility: crate::BindingVisibility::ReadWrite,
name: "buf".into(),
}],
},
dispatch: Dispatch::new(1, 1, 1),
body: KernelBody {
ops: vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::StoreGlobal,
operands: vec![0, 0, 1],
result: None,
},
],
child_bodies: vec![],
literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
},
};
let (_out, stats) = run_all_with_stats(&desc);
assert_eq!(stats.ops_before, stats.ops_after);
assert_eq!(stats.bindings_before, stats.bindings_after);
assert_eq!(stats.iterations, 1);
assert!(stats.converged);
}
#[test]
fn run_all_with_stats_reports_dropped_bindings() {
let desc = KernelDescriptor {
id: "with_unused".into(),
bindings: BindingLayout {
slots: vec![
crate::BindingSlot {
slot: 0,
element_type: vyre_foundation::ir::DataType::U32,
element_count: None,
memory_class: crate::MemoryClass::Global,
visibility: crate::BindingVisibility::ReadWrite,
name: "used".into(),
},
crate::BindingSlot {
slot: 9,
element_type: vyre_foundation::ir::DataType::U32,
element_count: None,
memory_class: crate::MemoryClass::Global,
visibility: crate::BindingVisibility::ReadOnly,
name: "unused".into(),
},
],
},
dispatch: Dispatch::new(1, 1, 1),
body: KernelBody {
ops: vec![
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![0],
result: Some(0),
},
KernelOp {
kind: KernelOpKind::Literal,
operands: vec![1],
result: Some(1),
},
KernelOp {
kind: KernelOpKind::StoreGlobal,
operands: vec![0, 0, 1],
result: None,
},
],
child_bodies: vec![],
literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
},
};
let (_out, stats) = run_all_with_stats(&desc);
assert_eq!(stats.bindings_before, 2);
assert_eq!(stats.bindings_after, 2);
assert_eq!(stats.bindings_dropped(), 0);
}
}