Skip to main content

vyre_lower/rewrites/
mod.rs

1//! Real rewrite passes on `KernelDescriptor`.
2//!
3//! Until this module, every analysis in `vyre-lower` was read-only  -
4//! they detected patterns and returned reports. This module's passes
5//! are the OTHER side of that line: they take a `KernelDescriptor`,
6//! apply a transformation, and return an improved equivalent.
7//!
8//! ## Pass shape
9//!
10//! ```text
11//! pub fn rewrite(desc: &KernelDescriptor) -> KernelDescriptor;
12//! ```
13//!
14//! Every pass:
15//! - Is total (no `Result`)  -  if a pass can't apply, it returns the
16//!   input unchanged.
17//! - Preserves semantic equivalence  -  running both the input and the
18//!   output through `audit()` gives compatible reports for surviving
19//!   ops.
20//! - Renumbers operand ids when needed; the output's ids are dense
21//!   `0..N` for operand id space.
22//! - Is idempotent  -  applying twice gives the same result as once.
23//!
24//! Descriptor-cleanup passes use a `descriptor_*` prefix so they are not
25//! confused with foundation's Program-IR semantic optimizer passes.
26//! Phase 1 shipped three descriptor cleanups:
27//! - `descriptor_dce`  -  dead-op elimination
28//! - `descriptor_cse`  -  common-subexpression elimination
29//! - `descriptor_const_fold`  -  fold compile-time-constant arithmetic
30//!
31//! The canonical release pipeline includes memory-layout rewrites,
32//! external dataflow-aware rewrites, and bounded e-graph-family algebraic
33//! saturation. Saturation is followed by immediate cleanup so a single
34//! pass application does not leave reassociated inner chains alive until
35//! a later fixed-point iteration.
36
37use std::hash::{Hash, Hasher};
38
39pub mod add_sub_cancel;
40pub mod aos_to_soa_promote;
41mod arithmetic_combine;
42pub mod bank_conflict_pad;
43pub mod bitwise_combine;
44mod body_index;
45/// Bitwise idempotence: fold `BitAnd(x, x)` and `BitOr(x, x)` into `Copy(x)`.
46pub mod bitwise_idemp {
47    use crate::rewrites::self_binop::rewrite_self_binops;
48    use crate::KernelDescriptor;
49    use vyre_foundation::ir::BinOp;
50
51    #[must_use]
52    pub fn bitwise_idemp(desc: &KernelDescriptor) -> KernelDescriptor {
53        rewrite_self_binops(desc, |bin| matches!(bin, BinOp::BitAnd | BinOp::BitOr))
54    }
55}
56pub mod boolean_simplify;
57pub mod branch_collapse;
58pub mod canonicalize;
59pub mod cmp_normalize;
60pub mod cmp_self_false;
61mod commutative_lit_chain;
62pub mod const_buffer_promote;
63mod dataflow_facts;
64pub mod dead_store;
65pub mod descriptor_const_fold;
66pub mod descriptor_cse;
67pub mod descriptor_dce;
68pub mod div_combine;
69pub mod drop_unused_bindings;
70pub mod drop_unused_child_bodies;
71pub mod drop_unused_literals;
72pub mod egraph_saturation;
73pub mod emit_order;
74pub mod identity_elim;
75pub mod licm;
76mod literal;
77pub mod load_forwarding;
78pub mod loop_fission;
79pub mod loop_fusion;
80pub mod loop_unroll;
81pub mod loop_zero_iter;
82pub mod matmul_promote;
83mod memory_address;
84mod rhs_lit_chain;
85mod self_binop;
86/// Min/max idempotence: fold `Min(x, x)` and `Max(x, x)` into `Copy(x)`.
87pub mod min_max_idemp {
88    use crate::rewrites::self_binop::rewrite_self_binops;
89    use crate::KernelDescriptor;
90    use vyre_foundation::ir::BinOp;
91
92    #[must_use]
93    pub fn min_max_idemp(desc: &KernelDescriptor) -> KernelDescriptor {
94        rewrite_self_binops(desc, |bin| matches!(bin, BinOp::Min | BinOp::Max))
95    }
96}
97pub mod mod_idemp;
98pub mod mul_add_to_fma;
99pub mod negate_cancel;
100pub mod select_fold;
101pub mod shared_mem_promote;
102pub mod shift_combine;
103pub mod strength_reduce;
104pub mod sub_combine;
105pub mod tail_mask;
106pub mod unary_idemp;
107pub mod xor_self_zero;
108
109pub use add_sub_cancel::add_sub_cancel;
110pub use aos_to_soa_promote::{promote as aos_to_soa_promote, LayoutHint as AosSoaLayoutHint};
111pub use arithmetic_combine::add_combine::add_combine;
112pub use arithmetic_combine::mul_combine::mul_combine;
113pub use arithmetic_combine::{add_combine, mul_combine};
114pub use bank_conflict_pad::bank_conflict_pad;
115pub use bitwise_combine::bitwise_combine;
116pub use bitwise_idemp::bitwise_idemp;
117pub use boolean_simplify::boolean_simplify;
118pub use branch_collapse::branch_collapse;
119pub use canonicalize::canonicalize;
120pub use cmp_normalize::cmp_normalize;
121pub use cmp_self_false::cmp_self_false;
122pub use const_buffer_promote::const_buffer_promote;
123pub use dead_store::{
124    dead_store, dead_store_with_alias_facts, dead_store_with_dataflow_analysis_facts,
125    dead_store_with_dataflow_facts, dead_store_with_weir_alias_facts,
126};
127pub use descriptor_const_fold::descriptor_const_fold;
128pub use descriptor_cse::descriptor_cse;
129pub use descriptor_dce::descriptor_dce;
130pub use div_combine::div_combine;
131pub use drop_unused_bindings::drop_unused_bindings;
132pub use drop_unused_child_bodies::drop_unused_child_bodies;
133pub use drop_unused_literals::drop_unused_literals;
134pub use emit_order::emit_order;
135pub use identity_elim::identity_elim;
136pub use licm::{
137    licm, licm_with_alias_facts, licm_with_dataflow_analysis_facts, licm_with_dataflow_facts,
138    licm_with_weir_alias_facts,
139};
140pub use load_forwarding::{
141    load_forwarding, load_forwarding_with_alias_facts,
142    load_forwarding_with_dataflow_analysis_facts, load_forwarding_with_dataflow_facts,
143    load_forwarding_with_weir_alias_facts,
144};
145pub use loop_fission::{
146    loop_fission, loop_fission_with_alias_facts, loop_fission_with_dataflow_analysis_facts,
147    loop_fission_with_dataflow_facts, loop_fission_with_weir_alias_facts,
148};
149pub use loop_fusion::{
150    loop_fusion, loop_fusion_with_alias_facts, loop_fusion_with_dataflow_analysis_facts,
151    loop_fusion_with_dataflow_facts, loop_fusion_with_weir_alias_facts,
152};
153pub use loop_unroll::loop_unroll;
154pub use loop_zero_iter::loop_zero_iter;
155pub use matmul_promote::{infer_matmul_tile_loops, matmul_promote, MatmulTileLoopPlan};
156pub use min_max_idemp::min_max_idemp;
157pub use mod_idemp::mod_idemp;
158pub use mul_add_to_fma::mul_add_to_fma;
159pub use negate_cancel::negate_cancel;
160pub use select_fold::select_fold;
161pub use shared_mem_promote::shared_mem_promote;
162pub use shift_combine::shift_combine;
163pub use strength_reduce::strength_reduce;
164pub use sub_combine::sub_combine;
165pub use tail_mask::apply_tail_mask;
166pub use unary_idemp::unary_idemp;
167pub use xor_self_zero::xor_self_zero;
168
169/// One substrate-neutral lowered-IR rewrite in the canonical pipeline.
170#[derive(Debug, Clone, Copy)]
171pub struct DescriptorRewritePass {
172    /// Stable pass name used in stats, diagnostics, and benchmark attribution.
173    pub name: &'static str,
174    /// Total rewrite function. Returns the input unchanged when it cannot apply.
175    pub rewrite: fn(&crate::KernelDescriptor) -> crate::KernelDescriptor,
176}
177
178impl DescriptorRewritePass {
179    #[must_use]
180    fn run(self, desc: &crate::KernelDescriptor) -> crate::KernelDescriptor {
181        (self.rewrite)(desc)
182    }
183}
184
185fn egraph_saturation_pass(desc: &crate::KernelDescriptor) -> crate::KernelDescriptor {
186    egraph_saturation::saturate_algebraic_descriptor(desc).0
187}
188
189const CANONICAL_REWRITE_PASSES: &[DescriptorRewritePass] = &[
190    DescriptorRewritePass {
191        name: "strength_reduce",
192        rewrite: strength_reduce,
193    },
194    DescriptorRewritePass {
195        name: "shift_combine",
196        rewrite: shift_combine,
197    },
198    DescriptorRewritePass {
199        name: "shared_mem_promote",
200        rewrite: shared_mem_promote,
201    },
202    DescriptorRewritePass {
203        name: "bank_conflict_pad",
204        rewrite: bank_conflict_pad,
205    },
206    DescriptorRewritePass {
207        name: "const_buffer_promote",
208        rewrite: const_buffer_promote,
209    },
210    DescriptorRewritePass {
211        name: "descriptor_const_fold",
212        rewrite: descriptor_const_fold,
213    },
214    DescriptorRewritePass {
215        name: "add_combine",
216        rewrite: add_combine,
217    },
218    DescriptorRewritePass {
219        name: "sub_combine",
220        rewrite: sub_combine,
221    },
222    DescriptorRewritePass {
223        name: "mul_combine",
224        rewrite: mul_combine,
225    },
226    DescriptorRewritePass {
227        name: "div_combine",
228        rewrite: div_combine,
229    },
230    DescriptorRewritePass {
231        name: "mod_idemp",
232        rewrite: mod_idemp,
233    },
234    DescriptorRewritePass {
235        name: "add_sub_cancel",
236        rewrite: add_sub_cancel,
237    },
238    DescriptorRewritePass {
239        name: "bitwise_combine",
240        rewrite: bitwise_combine,
241    },
242    DescriptorRewritePass {
243        name: "identity_elim",
244        rewrite: identity_elim,
245    },
246    DescriptorRewritePass {
247        name: "boolean_simplify",
248        rewrite: boolean_simplify,
249    },
250    DescriptorRewritePass {
251        name: "negate_cancel",
252        rewrite: negate_cancel,
253    },
254    DescriptorRewritePass {
255        name: "unary_idemp",
256        rewrite: unary_idemp,
257    },
258    DescriptorRewritePass {
259        name: "select_fold",
260        rewrite: select_fold,
261    },
262    DescriptorRewritePass {
263        name: "min_max_idemp",
264        rewrite: min_max_idemp,
265    },
266    DescriptorRewritePass {
267        name: "bitwise_idemp",
268        rewrite: bitwise_idemp,
269    },
270    DescriptorRewritePass {
271        name: "branch_collapse",
272        rewrite: branch_collapse,
273    },
274    DescriptorRewritePass {
275        name: "loop_fusion",
276        rewrite: loop_fusion,
277    },
278    DescriptorRewritePass {
279        name: "loop_unroll",
280        rewrite: loop_unroll,
281    },
282    DescriptorRewritePass {
283        name: "loop_zero_iter",
284        rewrite: loop_zero_iter,
285    },
286    DescriptorRewritePass {
287        name: "licm",
288        rewrite: licm,
289    },
290    DescriptorRewritePass {
291        name: "load_forwarding",
292        rewrite: load_forwarding,
293    },
294    DescriptorRewritePass {
295        name: "mul_add_to_fma",
296        rewrite: mul_add_to_fma,
297    },
298    DescriptorRewritePass {
299        name: "matmul_promote",
300        rewrite: matmul_promote,
301    },
302    DescriptorRewritePass {
303        name: "descriptor_dce_after_forwarding",
304        rewrite: descriptor_dce,
305    },
306    DescriptorRewritePass {
307        name: "dead_store",
308        rewrite: dead_store,
309    },
310    DescriptorRewritePass {
311        name: "descriptor_dce",
312        rewrite: descriptor_dce,
313    },
314    DescriptorRewritePass {
315        name: "cmp_normalize",
316        rewrite: cmp_normalize,
317    },
318    DescriptorRewritePass {
319        name: "cmp_self_false",
320        rewrite: cmp_self_false,
321    },
322    DescriptorRewritePass {
323        name: "xor_self_zero",
324        rewrite: xor_self_zero,
325    },
326    DescriptorRewritePass {
327        name: "canonicalize",
328        rewrite: canonicalize,
329    },
330    DescriptorRewritePass {
331        name: "descriptor_cse",
332        rewrite: descriptor_cse,
333    },
334    DescriptorRewritePass {
335        name: "egraph_saturation",
336        rewrite: egraph_saturation_pass,
337    },
338    DescriptorRewritePass {
339        name: "descriptor_const_fold_post_saturation",
340        rewrite: descriptor_const_fold,
341    },
342    DescriptorRewritePass {
343        name: "identity_elim_post_saturation",
344        rewrite: identity_elim,
345    },
346    DescriptorRewritePass {
347        name: "descriptor_dce_post_saturation",
348        rewrite: descriptor_dce,
349    },
350    DescriptorRewritePass {
351        name: "cmp_normalize_post_saturation",
352        rewrite: cmp_normalize,
353    },
354    DescriptorRewritePass {
355        name: "canonicalize_post_saturation",
356        rewrite: canonicalize,
357    },
358    DescriptorRewritePass {
359        name: "descriptor_cse_post_saturation",
360        rewrite: descriptor_cse,
361    },
362    DescriptorRewritePass {
363        name: "drop_unused_bindings",
364        rewrite: drop_unused_bindings,
365    },
366    DescriptorRewritePass {
367        name: "drop_unused_literals",
368        rewrite: drop_unused_literals,
369    },
370    DescriptorRewritePass {
371        name: "drop_unused_child_bodies",
372        rewrite: drop_unused_child_bodies,
373    },
374    DescriptorRewritePass {
375        name: "emit_order",
376        rewrite: emit_order,
377    },
378];
379
380/// Canonical lowered-IR rewrite pipeline as data, not a second hand-coded compiler.
381#[must_use]
382pub const fn canonical_rewrite_passes() -> &'static [DescriptorRewritePass] {
383    CANONICAL_REWRITE_PASSES
384}
385
386/// Cheap descriptor fingerprint for convergence checks.
387///
388/// Uses FxHasher for speed  -  collision resistance isn't needed because
389/// a hash match is always confirmed by one final deep equality check
390/// before declaring convergence.
391fn descriptor_hash(desc: &crate::KernelDescriptor) -> u64 {
392    let mut h = rustc_hash::FxHasher::default();
393    desc.hash(&mut h);
394    h.finish()
395}
396
397fn run_descriptor_passes(
398    desc: &crate::KernelDescriptor,
399    passes: &[DescriptorRewritePass],
400) -> crate::KernelDescriptor {
401    let mut current = desc.clone();
402    for pass in passes {
403        let pre_hash = descriptor_hash(&current);
404        let next = pass.run(&current);
405        // Skip the allocation when the pass was a no-op.
406        if descriptor_hash(&next) != pre_hash || next != current {
407            current = next;
408        }
409        #[cfg(debug_assertions)]
410        debug_verify_after_rewrite(&current, pass.name);
411    }
412    current
413}
414
415/// Apply every shipped rewrite in canonical order. The ordering is
416/// chosen so that each pass exposes work for the passes that follow:
417///
418/// 1. `strength_reduce`  -  turns `mul/div/mod` by power-of-2 into
419///    shift/and. Synthesizes new Literal ops for the shift counts.
420/// 2. `shared_mem_promote`  -  stages proven repeated U32 global tile loads
421///    through workgroup memory.
422/// 3. `bank_conflict_pad`  -  pads simple shared-memory strided layouts whose
423///    pitch aliases shared-memory banks.
424/// 4. `const_buffer_promote`  -  promotes small fixed-size read-only global
425///    buffers with repeated loads to constant bindings.
426/// 5. `descriptor_const_fold`  -  folds `BinOp(Lit, Lit)` into a single Literal.
427///    Catches both the original literal-pair arithmetic and any new
428///    constant-shift produced by strength_reduce.
429/// 6. `identity_elim`  -  substitutes `Add(x, 0)`, `Mul(x, 1)`, etc. with
430///    `x` (and absorbing-zero patterns with `0`). Runs after descriptor_const_fold
431///    so its left/right-identity rules see the post-folded literals.
432/// 7. `branch_collapse`  -  replaces `If(Lit_bool, then, else)` with the
433///    selected arm inlined. Runs after descriptor_const_fold + identity_elim so
434///    conditions like `Add(x, 0) != 0` simplify down to literals first.
435/// 8. `loop_fusion`  -  canonical launch/loop overhead reducer for
436///    safe disjoint-write loops. `loop_fission` is exposed as an
437///    explicit cost-model transform, but is not run unconditionally
438///    because it intentionally opposes fusion.
439/// 9. `loop_unroll`  -  unrolls small constant-bound loops. Runs after
440///    branch_collapse (which often removes the guards that prevented
441///    unrolling) but before licm (no point hoisting from a loop that
442///    will get unrolled away).
443/// 10. `licm`  -  hoists loop-invariant ops out of remaining loops.
444/// 11. `load_forwarding`  -  store-to-load + load-to-load forwarding.
445///    Runs after licm because hoisted loads may reveal new forwardable
446///    pairs in straight-line code.
447/// 12. `dead_store`  -  drops stores whose value is overwritten before any
448///    observation. Runs after load_forwarding (forwarded loads may have
449///    removed the only observers that kept stores alive).
450/// 13. `descriptor_dce`  -  drops result-producing ops with no users. Cleans up
451///    everything orphaned by the substitutions in (3) and (7).
452/// 14. `descriptor_cse`  -  merges remaining structurally-equivalent ops before
453///     saturation so the equality surface is canonical.
454/// 15. `egraph_saturation`  -  reassociates algebraic constant chains under a
455///     bounded saturation contract.
456/// 16. post-saturation fold/DCE/CSE cleanup  -  immediately removes dead inner
457///     chain ops and merges the final reassociated shape.
458/// Single application of the canonical pass sequence  -  see [`run_all`]
459/// for the iterating wrapper. This exists so callers that want exactly
460/// one pass (e.g. for diagnostics) can have it.
461#[must_use]
462
463pub fn run_all_once(desc: &crate::KernelDescriptor) -> crate::KernelDescriptor {
464    run_descriptor_passes(desc, canonical_rewrite_passes())
465}
466
467/// Internal: assert the descriptor's structural invariants hold after a
468/// rewrite pass. Only active in debug builds. Catches rewrite bugs early
469/// instead of letting them propagate to the emitter as opaque naga errors.
470#[cfg(debug_assertions)]
471fn debug_verify_after_rewrite(desc: &crate::KernelDescriptor, pass: &str) {
472    if let Err(errors) = crate::verify::verify(desc) {
473        let violation_count = errors.len();
474        assert!(
475            violation_count == 0,
476            "rewrite pass `{pass}` produced an invalid KernelDescriptor  -  {} violation(s):\n{errors:#?}",
477            violation_count
478        );
479    }
480}
481
482/// Single canonical pass sequence using External alias and
483/// reaching-definition facts where they unlock stronger legality.
484#[must_use]
485pub fn run_all_once_with_dataflow_facts(
486    desc: &crate::KernelDescriptor,
487    alias_facts: &crate::analyses::alias_facts::AliasFactSet,
488    reaching_defs: &crate::analyses::reaching_def_facts::ReachingDefFactSet,
489) -> crate::KernelDescriptor {
490    let reduced = strength_reduce(desc);
491    let shared_promoted = shared_mem_promote(&reduced);
492    let padded = bank_conflict_pad(&shared_promoted);
493    let const_promoted = const_buffer_promote(&padded);
494    let folded = descriptor_const_fold(&const_promoted);
495    let identified = identity_elim(&folded);
496    let collapsed = branch_collapse(&identified);
497    let fused = loop_fusion_with_dataflow_facts(&collapsed, alias_facts, reaching_defs);
498    let unrolled = loop_unroll(&fused);
499    let hoisted = licm_with_dataflow_facts(&unrolled, alias_facts, reaching_defs);
500    let forwarded = load_forwarding_with_dataflow_facts(&hoisted, alias_facts, reaching_defs);
501    let cleaned = descriptor_dce(&forwarded);
502    let dse_done = dead_store_with_dataflow_facts(&cleaned, alias_facts, reaching_defs);
503    let dced = descriptor_dce(&dse_done);
504    let canon = canonicalize(&dced);
505    let merged = descriptor_cse(&canon);
506    let (saturated, _) = egraph_saturation::saturate_algebraic_descriptor(&merged);
507    let saturated_folded = descriptor_const_fold(&saturated);
508    let saturated_identified = identity_elim(&saturated_folded);
509    let saturated_dced = descriptor_dce(&saturated_identified);
510    let saturated_canon = canonicalize(&saturated_dced);
511    let saturated_merged = descriptor_cse(&saturated_canon);
512    let pruned_bindings = drop_unused_bindings(&saturated_merged);
513    let pruned_literals = drop_unused_literals(&pruned_bindings);
514    let pruned_children = drop_unused_child_bodies(&pruned_literals);
515    emit_order(&pruned_children)
516}
517
518/// Maximum number of `run_all_once` iterations before giving up and
519/// returning the latest output. In practice fixed point is reached
520/// in 1–2 iterations on every shape we've fuzzed.
521pub const RUN_ALL_MAX_ITERS: usize = 4;
522
523/// Apply [`run_all_once`] repeatedly until the descriptor reaches a full
524/// fixed point or [`RUN_ALL_MAX_ITERS`] is reached. Necessary because passes
525/// late in the pipeline (notably `descriptor_cse`) can expose opportunities
526/// for earlier passes (notably `dead_store`): two stores at distinct
527/// id-but-equal-value indices look distinct to dead_store in iteration 1, but
528/// after CSE merges the index ids in iteration 1, dead_store catches them in
529/// iteration 2.
530#[must_use]
531pub fn run_all(desc: &crate::KernelDescriptor) -> crate::KernelDescriptor {
532    run_all_with_stats(desc).0
533}
534
535/// Apply the canonical fixed-point rewrite pipeline with external dataflow facts.
536#[must_use]
537pub fn run_all_with_dataflow_facts(
538    desc: &crate::KernelDescriptor,
539    alias_facts: &crate::analyses::alias_facts::AliasFactSet,
540    reaching_defs: &crate::analyses::reaching_def_facts::ReachingDefFactSet,
541) -> crate::KernelDescriptor {
542    run_all_with_dataflow_stats(desc, alias_facts, reaching_defs).0
543}
544
545/// Apply the canonical fixed-point rewrite pipeline with Weir dataflow-analysis facts.
546#[must_use]
547pub fn run_all_with_dataflow_analysis_facts(
548    desc: &crate::KernelDescriptor,
549    alias_facts: &crate::analyses::weir_alias::AliasFactSet,
550    reaching_defs: &crate::analyses::weir_reaching_def::ReachingDefFactSet,
551) -> crate::KernelDescriptor {
552    run_all_with_dataflow_facts(desc, alias_facts, reaching_defs)
553}
554
555/// Per-pipeline-run statistics. Useful for benchmarks, regression
556/// diagnostics, and `--verbose` emit modes.
557#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
558pub struct OptimizationStats {
559    /// Top-level body op count before optimization.
560    pub ops_before: usize,
561    /// Top-level body op count after optimization.
562    pub ops_after: usize,
563    /// Number of bindings before. After equals `bindings_before` minus
564    /// what `drop_unused_bindings` stripped.
565    pub bindings_before: usize,
566    pub bindings_after: usize,
567    /// Number of literals in the top-level body's pool, before/after.
568    pub literals_before: usize,
569    pub literals_after: usize,
570    /// How many `run_all_once` iterations actually fired (1 to
571    /// `RUN_ALL_MAX_ITERS`). 1 means the pipeline converged on the
572    /// first pass; higher means later-pass changes exposed work for
573    /// earlier passes and the fixed-point loop kicked in.
574    pub iterations: usize,
575    /// True iff the fixed-point converged within the cap. False means
576    /// the pipeline stopped mid-way; output is still valid (every
577    /// individual pass is total) but may have residual optimization
578    /// opportunities.
579    pub converged: bool,
580}
581
582impl OptimizationStats {
583    /// Total ops eliminated (top-level body only). Saturating, so the
584    /// rare case where output exceeds input (e.g. loop_unroll inlining)
585    /// returns 0 rather than wrapping.
586    pub fn ops_eliminated(&self) -> usize {
587        self.ops_before.saturating_sub(self.ops_after)
588    }
589
590    pub fn bindings_dropped(&self) -> usize {
591        self.bindings_before.saturating_sub(self.bindings_after)
592    }
593
594    /// True iff the pipeline made no change at all  -  no ops eliminated,
595    /// no bindings dropped, no literals dropped, AND op count is
596    /// stable. The kernel was either already optimal or out of the
597    /// pipeline's reach. Useful for tooling that wants to skip emit
598    /// re-runs when nothing changed.
599    pub fn is_no_op(&self) -> bool {
600        self.ops_before == self.ops_after
601            && self.bindings_before == self.bindings_after
602            && self.literals_before == self.literals_after
603    }
604
605    /// Total off-graph data dropped (bindings + literals). Useful as
606    /// a single-number "how much cleanup did the tail of the pipeline
607    /// do?" signal.
608    pub fn off_graph_dropped(&self) -> usize {
609        self.bindings_dropped() + self.literals_before.saturating_sub(self.literals_after)
610    }
611
612    /// Merge another [`OptimizationStats`] into a running aggregate.
613    /// Adds counts and ORs the converged flag (any non-converged run
614    /// → aggregate is non-converged). `iterations` accumulates.
615    /// Useful for tooling that runs `run_all_with_stats` over a corpus
616    /// of N kernels and wants a single rolled-up summary.
617    pub fn merge(&mut self, other: OptimizationStats) {
618        self.ops_before = self.ops_before.saturating_add(other.ops_before);
619        self.ops_after = self.ops_after.saturating_add(other.ops_after);
620        self.bindings_before = self.bindings_before.saturating_add(other.bindings_before);
621        self.bindings_after = self.bindings_after.saturating_add(other.bindings_after);
622        self.literals_before = self.literals_before.saturating_add(other.literals_before);
623        self.literals_after = self.literals_after.saturating_add(other.literals_after);
624        self.iterations = self.iterations.saturating_add(other.iterations);
625        self.converged = self.converged && other.converged;
626    }
627
628    /// Identity element for [`Self::merge`]. Useful as the seed of a fold
629    /// over a corpus.
630    pub fn zero() -> Self {
631        OptimizationStats {
632            ops_before: 0,
633            ops_after: 0,
634            bindings_before: 0,
635            bindings_after: 0,
636            literals_before: 0,
637            literals_after: 0,
638            iterations: 0,
639            converged: true,
640        }
641    }
642
643    /// One-line human-readable summary suitable for log lines.
644    /// Format: `"ops X→Y (-N), bindings A→B (-M), iters K (converged|stopped)"`.
645    /// Mirrors the format_short pattern on the audit reports.
646    pub fn format_short(&self) -> String {
647        format!(
648            "ops {}→{} (-{}), bindings {}→{} (-{}), iters {} ({})",
649            self.ops_before,
650            self.ops_after,
651            self.ops_eliminated(),
652            self.bindings_before,
653            self.bindings_after,
654            self.bindings_dropped(),
655            self.iterations,
656            if self.converged {
657                "converged"
658            } else {
659                "stopped"
660            },
661        )
662    }
663}
664
665impl std::fmt::Display for OptimizationStats {
666    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
667        f.write_str(&self.format_short())
668    }
669}
670
671/// Like [`run_all`] but also returns [`OptimizationStats`] so the
672/// caller can surface what happened.
673#[must_use]
674pub fn run_all_with_stats(
675    desc: &crate::KernelDescriptor,
676) -> (crate::KernelDescriptor, OptimizationStats) {
677    let ops_before = desc.body.ops.len();
678    let bindings_before = desc.bindings.slots.len();
679    let literals_before = desc.body.literals.len();
680
681    let mut current = run_all_once(desc);
682    #[cfg(debug_assertions)]
683    debug_verify_after_rewrite(&current, "run_all_once (iter 1)");
684    let mut iterations = 1usize;
685    let mut current_hash = descriptor_hash(&current);
686    let mut converged = current_hash == descriptor_hash(desc) && current == *desc;
687    while !converged && iterations < RUN_ALL_MAX_ITERS {
688        let next = run_all_once(&current);
689        iterations += 1;
690        #[cfg(debug_assertions)]
691        debug_verify_after_rewrite(&next, &format!("run_all_once (iter {iterations})"));
692        let next_hash = descriptor_hash(&next);
693        // Fast path: hash mismatch → definitely changed, skip deep eq.
694        // Hash match → confirm with full equality to guard against collisions.
695        converged = next_hash == current_hash && next == current;
696        current = next;
697        current_hash = next_hash;
698    }
699
700    let stats = OptimizationStats {
701        ops_before,
702        ops_after: current.body.ops.len(),
703        bindings_before,
704        bindings_after: current.bindings.slots.len(),
705        literals_before,
706        literals_after: current.body.literals.len(),
707        iterations,
708        converged,
709    };
710    (current, stats)
711}
712
713/// Like [`run_all_with_dataflow_facts`] but also returns
714/// [`OptimizationStats`].
715#[must_use]
716pub fn run_all_with_dataflow_stats(
717    desc: &crate::KernelDescriptor,
718    alias_facts: &crate::analyses::alias_facts::AliasFactSet,
719    reaching_defs: &crate::analyses::reaching_def_facts::ReachingDefFactSet,
720) -> (crate::KernelDescriptor, OptimizationStats) {
721    let ops_before = desc.body.ops.len();
722    let bindings_before = desc.bindings.slots.len();
723    let literals_before = desc.body.literals.len();
724
725    let mut current = run_all_once_with_dataflow_facts(desc, alias_facts, reaching_defs);
726    let mut iterations = 1usize;
727    let mut current_hash = descriptor_hash(&current);
728    let mut converged = current_hash == descriptor_hash(desc) && current == *desc;
729    while !converged && iterations < RUN_ALL_MAX_ITERS {
730        let next = run_all_once_with_dataflow_facts(&current, alias_facts, reaching_defs);
731        iterations += 1;
732        let next_hash = descriptor_hash(&next);
733        converged = next_hash == current_hash && next == current;
734        current = next;
735        current_hash = next_hash;
736    }
737
738    let stats = OptimizationStats {
739        ops_before,
740        ops_after: current.body.ops.len(),
741        bindings_before,
742        bindings_after: current.bindings.slots.len(),
743        literals_before,
744        literals_after: current.body.literals.len(),
745        iterations,
746        converged,
747    };
748    (current, stats)
749}
750
751#[cfg(test)]
752mod tests {
753    use super::*;
754    use crate::{
755        BindingLayout, Dispatch, KernelBody, KernelDescriptor, KernelOp, KernelOpKind, LiteralValue,
756    };
757
758    #[test]
759    fn run_all_on_empty_kernel_returns_empty() {
760        let desc = KernelDescriptor {
761            id: "k".into(),
762            bindings: BindingLayout { slots: vec![] },
763            dispatch: Dispatch::new(1, 1, 1),
764            body: KernelBody {
765                ops: vec![],
766                child_bodies: vec![],
767                literals: vec![],
768            },
769        };
770        let out = run_all(&desc);
771        assert!(out.body.ops.is_empty());
772    }
773
774    #[test]
775    fn run_all_is_idempotent() {
776        let desc = KernelDescriptor {
777            id: "k".into(),
778            bindings: BindingLayout { slots: vec![] },
779            dispatch: Dispatch::new(1, 1, 1),
780            body: KernelBody {
781                ops: vec![
782                    KernelOp {
783                        kind: KernelOpKind::Literal,
784                        operands: vec![0],
785                        result: Some(0),
786                    },
787                    KernelOp {
788                        kind: KernelOpKind::Literal,
789                        operands: vec![0],
790                        result: Some(1),
791                    }, // dup
792                    KernelOp {
793                        kind: KernelOpKind::Literal,
794                        operands: vec![1],
795                        result: Some(2),
796                    }, // dead
797                ],
798                child_bodies: vec![],
799                literals: vec![LiteralValue::U32(5), LiteralValue::U32(99)],
800            },
801        };
802        let once = run_all(&desc);
803        let twice = run_all(&once);
804        assert_eq!(once.body.ops.len(), twice.body.ops.len());
805        assert_eq!(once.body.literals, twice.body.literals);
806    }
807
808    #[test]
809    fn run_all_collapses_kitchen_sink_kernel() {
810        // Inefficiency stack  -  every shipped pass should contribute:
811        //
812        //   r0 = Lit(0)            // literal pool idx 0 → U32(0) (zero,
813        //                           //   identity for Add, absorbing for Mul)
814        //   r1 = Lit(1)            // literal pool idx 1 → U32(1) (one,
815        //                           //   identity for Mul)
816        //   r2 = Lit(8)            // literal pool idx 2 → U32(8) (pow2)
817        //   r3 = Lit(7)            // literal pool idx 3 → U32(7) (varying)
818        //   r4 = Lit(0)            // duplicate of r0 → CSE target
819        //   r5 = Add(r3, r1+r0)    // post-fold rhs becomes Lit(1) →
820        //                           //   identity_elim won't fire (rhs is 1
821        //                           //   not 0, op is Add). But r5 has a
822        //                           //   useful value (= 7+1 = 8). Actually
823        //                           //   simpler: just Add(r3, r0) → r3
824        //                           //   (identity).
825        //   r6 = Mul(r3, r2)       // strength_reduce: × 8 → << 3
826        //   r7 = Mul(r3, r1)       // identity_elim: × 1 → r3
827        //   r8 = Mul(r3, r0)       // absorbing-zero: → r0 (i.e. 0)
828        //   StoreGlobal(buf, r0, r5)  // dead store (overwritten below)
829        //   StoreGlobal(buf, r0, r6)  // surviving store
830        //
831        // Expectations:
832        //   - Op count drops substantially.
833        //   - The dead first store is gone (dead_store).
834        //   - The final store's value-id is the strength-reduced shift
835        //     result (or its forwarded equivalent).
836        //   - At least one literal is dropped or deduped.
837        let desc = KernelDescriptor {
838            id: "kitchen_sink".into(),
839            bindings: BindingLayout {
840                slots: vec![crate::BindingSlot {
841                    slot: 0,
842                    element_type: vyre_foundation::ir::DataType::U32,
843                    element_count: None,
844                    memory_class: crate::MemoryClass::Global,
845                    visibility: crate::BindingVisibility::ReadWrite,
846                    name: "buf".into(),
847                }],
848            },
849            dispatch: Dispatch::new(1, 1, 1),
850            body: KernelBody {
851                ops: vec![
852                    KernelOp {
853                        kind: KernelOpKind::Literal,
854                        operands: vec![0],
855                        result: Some(0),
856                    },
857                    KernelOp {
858                        kind: KernelOpKind::Literal,
859                        operands: vec![1],
860                        result: Some(1),
861                    },
862                    KernelOp {
863                        kind: KernelOpKind::Literal,
864                        operands: vec![2],
865                        result: Some(2),
866                    },
867                    KernelOp {
868                        kind: KernelOpKind::Literal,
869                        operands: vec![3],
870                        result: Some(3),
871                    },
872                    KernelOp {
873                        kind: KernelOpKind::Literal,
874                        operands: vec![0],
875                        result: Some(4),
876                    },
877                    KernelOp {
878                        kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Add),
879                        operands: vec![3, 0],
880                        result: Some(5),
881                    },
882                    KernelOp {
883                        kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul),
884                        operands: vec![3, 2],
885                        result: Some(6),
886                    },
887                    KernelOp {
888                        kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul),
889                        operands: vec![3, 1],
890                        result: Some(7),
891                    },
892                    KernelOp {
893                        kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul),
894                        operands: vec![3, 0],
895                        result: Some(8),
896                    },
897                    KernelOp {
898                        kind: KernelOpKind::StoreGlobal,
899                        operands: vec![0, 0, 5],
900                        result: None,
901                    },
902                    KernelOp {
903                        kind: KernelOpKind::StoreGlobal,
904                        operands: vec![0, 0, 6],
905                        result: None,
906                    },
907                ],
908                child_bodies: vec![],
909                literals: vec![
910                    LiteralValue::U32(0),
911                    LiteralValue::U32(1),
912                    LiteralValue::U32(8),
913                    LiteralValue::U32(7),
914                ],
915            },
916        };
917
918        let before_op_count = desc.body.ops.len();
919        let out = run_all(&desc);
920
921        // Op count should drop  -  multiple ops eliminated, but the
922        // exact final count depends on pass interaction. Be specific
923        // about WHICH ops are gone:
924
925        // 1. dead_store: only one StoreGlobal should survive.
926        let store_count = out
927            .body
928            .ops
929            .iter()
930            .filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
931            .count();
932        assert_eq!(store_count, 1, "dead_store must drop the overwritten store");
933
934        // 2. descriptor_dce: r5 (Add identity), r7 (Mul identity), r8 (absorbing
935        //    zero) are all dead by id substitution → must be gone.
936        //    Also r4 (CSE-merged dup of r0).
937        //    r6 itself becomes a Shl after strength_reduce; that survives
938        //    as the value-operand of the surviving store.
939        let mul_count = out
940            .body
941            .ops
942            .iter()
943            .filter(|o| {
944                matches!(
945                    o.kind,
946                    KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul)
947                )
948            })
949            .count();
950        assert_eq!(mul_count, 0, "all 3 Mul ops should be eliminated");
951
952        let add_count = out
953            .body
954            .ops
955            .iter()
956            .filter(|o| {
957                matches!(
958                    o.kind,
959                    KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Add)
960                )
961            })
962            .count();
963        assert_eq!(add_count, 0, "Add(r3, 0) → r3 should be eliminated");
964
965        // 3. strength_reduce: Mul(_, 8) became Shl(_, 3). Both operands
966        //    of THIS Mul happened to be literals (r3=Lit(7), r2=Lit(8)),
967        //    so descriptor_const_fold runs after strength_reduce and folds the new
968        //    Shl(Lit, Lit) into a Lit(56). Either outcome is acceptable  -
969        //    what matters is "no Mul, no Add, no Sub". Already asserted.
970
971        // 4. Op count must drop by at least 5 (5 dead ops eliminated).
972        assert!(
973            out.body.ops.len() <= before_op_count - 5,
974            "expected op count to drop by ≥5 (was {before_op_count}, now {})",
975            out.body.ops.len()
976        );
977
978        // 5. Idempotence: running again is a no-op.
979        let twice = run_all(&out);
980        assert_eq!(out.body.ops.len(), twice.body.ops.len());
981    }
982
983    #[test]
984    fn run_all_forwards_then_drops_redundant_load() {
985        // Store(buf, 0, 7); r = Load(buf, 0); Store(buf, 0, r).
986        // After load_forwarding: Store(buf, 0, 7); Load(buf, 0); Store(buf, 0, 7).
987        // After dead_store: only the LAST store survives (the middle Load
988        // is dead by then; DCE drops it).
989        let desc = KernelDescriptor {
990            id: "stl".into(),
991            bindings: BindingLayout {
992                slots: vec![crate::BindingSlot {
993                    slot: 0,
994                    element_type: vyre_foundation::ir::DataType::U32,
995                    element_count: None,
996                    memory_class: crate::MemoryClass::Global,
997                    visibility: crate::BindingVisibility::ReadWrite,
998                    name: "buf".into(),
999                }],
1000            },
1001            dispatch: Dispatch::new(1, 1, 1),
1002            body: KernelBody {
1003                ops: vec![
1004                    KernelOp {
1005                        kind: KernelOpKind::Literal,
1006                        operands: vec![0],
1007                        result: Some(0),
1008                    }, // idx
1009                    KernelOp {
1010                        kind: KernelOpKind::Literal,
1011                        operands: vec![1],
1012                        result: Some(1),
1013                    }, // val
1014                    KernelOp {
1015                        kind: KernelOpKind::StoreGlobal,
1016                        operands: vec![0, 0, 1],
1017                        result: None,
1018                    },
1019                    KernelOp {
1020                        kind: KernelOpKind::LoadGlobal,
1021                        operands: vec![0, 0],
1022                        result: Some(2),
1023                    },
1024                    KernelOp {
1025                        kind: KernelOpKind::StoreGlobal,
1026                        operands: vec![0, 0, 2],
1027                        result: None,
1028                    },
1029                ],
1030                child_bodies: vec![],
1031                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
1032            },
1033        };
1034        let out = run_all(&desc);
1035
1036        let store_count = out
1037            .body
1038            .ops
1039            .iter()
1040            .filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
1041            .count();
1042        assert_eq!(
1043            store_count, 1,
1044            "dead_store must drop the redundant first store"
1045        );
1046
1047        let load_count = out
1048            .body
1049            .ops
1050            .iter()
1051            .filter(|o| matches!(o.kind, KernelOpKind::LoadGlobal))
1052            .count();
1053        assert_eq!(load_count, 0, "descriptor_dce must drop the redundant load");
1054    }
1055
1056    #[test]
1057    fn run_all_unrolls_then_simplifies() {
1058        // Constant-bound loop (count=2) whose body just stores Lit(0) and Lit(0).
1059        // After unroll: 4 ops in straight line.
1060        // After dead_store: only the last surviving store stays (since
1061        // both are at the same idx).
1062        let desc = KernelDescriptor {
1063            id: "loop_then_dse".into(),
1064            bindings: BindingLayout {
1065                slots: vec![crate::BindingSlot {
1066                    slot: 0,
1067                    element_type: vyre_foundation::ir::DataType::U32,
1068                    element_count: None,
1069                    memory_class: crate::MemoryClass::Global,
1070                    visibility: crate::BindingVisibility::ReadWrite,
1071                    name: "buf".into(),
1072                }],
1073            },
1074            dispatch: Dispatch::new(1, 1, 1),
1075            body: KernelBody {
1076                ops: vec![
1077                    // lo = 0, hi = 2, step-id = 2 (ignored), body-child-idx = 0
1078                    KernelOp {
1079                        kind: KernelOpKind::Literal,
1080                        operands: vec![0],
1081                        result: Some(0),
1082                    },
1083                    KernelOp {
1084                        kind: KernelOpKind::Literal,
1085                        operands: vec![1],
1086                        result: Some(1),
1087                    },
1088                    KernelOp {
1089                        kind: KernelOpKind::Literal,
1090                        operands: vec![0],
1091                        result: Some(2),
1092                    },
1093                    KernelOp {
1094                        kind: KernelOpKind::StructuredForLoop {
1095                            loop_var: std::sync::Arc::from("i"),
1096                        },
1097                        operands: vec![0, 1, 0],
1098                        result: None,
1099                    },
1100                ],
1101                child_bodies: vec![KernelBody {
1102                    ops: vec![
1103                        KernelOp {
1104                            kind: KernelOpKind::Literal,
1105                            operands: vec![0],
1106                            result: Some(0),
1107                        },
1108                        KernelOp {
1109                            kind: KernelOpKind::Literal,
1110                            operands: vec![1],
1111                            result: Some(1),
1112                        },
1113                        KernelOp {
1114                            kind: KernelOpKind::StoreGlobal,
1115                            operands: vec![0, 0, 1],
1116                            result: None,
1117                        },
1118                    ],
1119                    child_bodies: vec![],
1120                    literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
1121                }],
1122                literals: vec![LiteralValue::U32(0), LiteralValue::U32(2)],
1123            },
1124        };
1125        let out = run_all(&desc);
1126
1127        // Loop must be gone (unrolled).
1128        let loop_count = out
1129            .body
1130            .ops
1131            .iter()
1132            .filter(|o| matches!(o.kind, KernelOpKind::StructuredForLoop { .. }))
1133            .count();
1134        assert_eq!(loop_count, 0, "loop should be unrolled");
1135
1136        // After unrolling 2 iterations, the body's stores are inlined
1137        // with fresh result-ids per iteration. CSE eventually merges
1138        // structurally-equal Lit ops, but the current pipeline runs
1139        // CSE after dead_store, so dead_store sees stores at "different"
1140        // (textually distinct) idx-ids and conservatively keeps both.
1141        // What we CAN guarantee: at least one store survives, and the
1142        // loop is gone.
1143        let store_count = out
1144            .body
1145            .ops
1146            .iter()
1147            .filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
1148            .count();
1149        assert!(
1150            store_count >= 1,
1151            "at least one store from the unrolled body should survive"
1152        );
1153    }
1154
1155    #[test]
1156    fn run_all_with_stats_reports_op_reduction() {
1157        // Same kitchen-sink shape as run_all_collapses_kitchen_sink_kernel,
1158        // but assert the stats reflect what happened.
1159        let desc = KernelDescriptor {
1160            id: "stats".into(),
1161            bindings: BindingLayout {
1162                slots: vec![crate::BindingSlot {
1163                    slot: 0,
1164                    element_type: vyre_foundation::ir::DataType::U32,
1165                    element_count: None,
1166                    memory_class: crate::MemoryClass::Global,
1167                    visibility: crate::BindingVisibility::ReadWrite,
1168                    name: "buf".into(),
1169                }],
1170            },
1171            dispatch: Dispatch::new(1, 1, 1),
1172            body: KernelBody {
1173                ops: vec![
1174                    KernelOp {
1175                        kind: KernelOpKind::Literal,
1176                        operands: vec![0],
1177                        result: Some(0),
1178                    },
1179                    KernelOp {
1180                        kind: KernelOpKind::Literal,
1181                        operands: vec![1],
1182                        result: Some(1),
1183                    },
1184                    KernelOp {
1185                        kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Add),
1186                        operands: vec![1, 0],
1187                        result: Some(2),
1188                    }, // identity (Add x 0)
1189                    KernelOp {
1190                        kind: KernelOpKind::BinOpKind(vyre_foundation::ir::BinOp::Mul),
1191                        operands: vec![1, 0],
1192                        result: Some(3),
1193                    }, // absorbing zero (Mul x 0)
1194                    KernelOp {
1195                        kind: KernelOpKind::StoreGlobal,
1196                        operands: vec![0, 0, 1],
1197                        result: None,
1198                    },
1199                ],
1200                child_bodies: vec![],
1201                literals: vec![LiteralValue::U32(0), LiteralValue::U32(99)],
1202            },
1203        };
1204        let (_out, stats) = run_all_with_stats(&desc);
1205        assert!(
1206            stats.ops_eliminated() >= 2,
1207            "expected ≥2 ops eliminated, got {} ({} → {})",
1208            stats.ops_eliminated(),
1209            stats.ops_before,
1210            stats.ops_after
1211        );
1212        assert!(stats.iterations >= 1);
1213        assert!(stats.iterations <= RUN_ALL_MAX_ITERS);
1214        assert!(stats.converged, "pipeline must converge");
1215    }
1216
1217    #[test]
1218    fn optimization_stats_format_short_includes_all_fields() {
1219        let s = OptimizationStats {
1220            ops_before: 11,
1221            ops_after: 3,
1222            bindings_before: 3,
1223            bindings_after: 1,
1224            literals_before: 4,
1225            literals_after: 2,
1226            iterations: 2,
1227            converged: true,
1228        };
1229        let f = s.format_short();
1230        assert!(f.contains("ops 11→3 (-8)"));
1231        assert!(f.contains("bindings 3→1 (-2)"));
1232        assert!(f.contains("iters 2"));
1233        assert!(f.contains("converged"));
1234    }
1235
1236    #[test]
1237    fn optimization_stats_merge_is_associative() {
1238        let a = OptimizationStats {
1239            ops_before: 10,
1240            ops_after: 3,
1241            bindings_before: 2,
1242            bindings_after: 1,
1243            literals_before: 5,
1244            literals_after: 2,
1245            iterations: 2,
1246            converged: true,
1247        };
1248        let b = OptimizationStats {
1249            ops_before: 7,
1250            ops_after: 4,
1251            bindings_before: 1,
1252            bindings_after: 1,
1253            literals_before: 3,
1254            literals_after: 3,
1255            iterations: 1,
1256            converged: true,
1257        };
1258        let c = OptimizationStats {
1259            ops_before: 5,
1260            ops_after: 2,
1261            bindings_before: 1,
1262            bindings_after: 0,
1263            literals_before: 2,
1264            literals_after: 1,
1265            iterations: 1,
1266            converged: true,
1267        };
1268
1269        let mut left = a;
1270        left.merge(b);
1271        left.merge(c);
1272
1273        let mut bc = b;
1274        bc.merge(c);
1275        let mut right = a;
1276        right.merge(bc);
1277
1278        assert_eq!(left, right);
1279    }
1280
1281    #[test]
1282    fn optimization_stats_merge_aggregates() {
1283        let mut acc = OptimizationStats::zero();
1284        acc.merge(OptimizationStats {
1285            ops_before: 10,
1286            ops_after: 3,
1287            bindings_before: 2,
1288            bindings_after: 1,
1289            literals_before: 5,
1290            literals_after: 2,
1291            iterations: 2,
1292            converged: true,
1293        });
1294        acc.merge(OptimizationStats {
1295            ops_before: 7,
1296            ops_after: 4,
1297            bindings_before: 1,
1298            bindings_after: 1,
1299            literals_before: 3,
1300            literals_after: 3,
1301            iterations: 1,
1302            converged: false,
1303        });
1304        assert_eq!(acc.ops_before, 17);
1305        assert_eq!(acc.ops_after, 7);
1306        assert_eq!(acc.iterations, 3);
1307        assert!(!acc.converged); // ORed: any false → false
1308    }
1309
1310    #[test]
1311    fn optimization_stats_zero_is_identity() {
1312        let s = OptimizationStats {
1313            ops_before: 5,
1314            ops_after: 3,
1315            bindings_before: 1,
1316            bindings_after: 1,
1317            literals_before: 2,
1318            literals_after: 1,
1319            iterations: 2,
1320            converged: true,
1321        };
1322        let mut acc = OptimizationStats::zero();
1323        acc.merge(s);
1324        assert_eq!(acc, s);
1325    }
1326
1327    #[test]
1328    fn optimization_stats_format_short_marks_stopped() {
1329        let s = OptimizationStats {
1330            ops_before: 5,
1331            ops_after: 5,
1332            bindings_before: 1,
1333            bindings_after: 1,
1334            literals_before: 1,
1335            literals_after: 1,
1336            iterations: 4,
1337            converged: false,
1338        };
1339        assert!(s.format_short().contains("stopped"));
1340    }
1341
1342    #[test]
1343    fn run_all_with_stats_reports_no_change_when_already_optimal() {
1344        let desc = KernelDescriptor {
1345            id: "minimal".into(),
1346            bindings: BindingLayout {
1347                slots: vec![crate::BindingSlot {
1348                    slot: 0,
1349                    element_type: vyre_foundation::ir::DataType::U32,
1350                    element_count: None,
1351                    memory_class: crate::MemoryClass::Global,
1352                    visibility: crate::BindingVisibility::ReadWrite,
1353                    name: "buf".into(),
1354                }],
1355            },
1356            dispatch: Dispatch::new(1, 1, 1),
1357            body: KernelBody {
1358                ops: vec![
1359                    KernelOp {
1360                        kind: KernelOpKind::Literal,
1361                        operands: vec![0],
1362                        result: Some(0),
1363                    },
1364                    KernelOp {
1365                        kind: KernelOpKind::Literal,
1366                        operands: vec![1],
1367                        result: Some(1),
1368                    },
1369                    KernelOp {
1370                        kind: KernelOpKind::StoreGlobal,
1371                        operands: vec![0, 0, 1],
1372                        result: None,
1373                    },
1374                ],
1375                child_bodies: vec![],
1376                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
1377            },
1378        };
1379        let (_out, stats) = run_all_with_stats(&desc);
1380        assert_eq!(stats.ops_before, stats.ops_after);
1381        assert_eq!(stats.bindings_before, stats.bindings_after);
1382        assert_eq!(stats.iterations, 1);
1383        assert!(stats.converged);
1384    }
1385
1386    #[test]
1387    fn run_all_with_stats_reports_dropped_bindings() {
1388        let desc = KernelDescriptor {
1389            id: "with_unused".into(),
1390            bindings: BindingLayout {
1391                slots: vec![
1392                    crate::BindingSlot {
1393                        slot: 0,
1394                        element_type: vyre_foundation::ir::DataType::U32,
1395                        element_count: None,
1396                        memory_class: crate::MemoryClass::Global,
1397                        visibility: crate::BindingVisibility::ReadWrite,
1398                        name: "used".into(),
1399                    },
1400                    crate::BindingSlot {
1401                        slot: 9,
1402                        element_type: vyre_foundation::ir::DataType::U32,
1403                        element_count: None,
1404                        memory_class: crate::MemoryClass::Global,
1405                        // Soundness: drop_unused_bindings now retains
1406                        // WriteOnly/ReadWrite outputs (host dispatch
1407                        // contract). The "unused" candidate must be
1408                        // ReadOnly for the drop to fire.
1409                        visibility: crate::BindingVisibility::ReadOnly,
1410                        name: "unused".into(),
1411                    },
1412                ],
1413            },
1414            dispatch: Dispatch::new(1, 1, 1),
1415            body: KernelBody {
1416                ops: vec![
1417                    KernelOp {
1418                        kind: KernelOpKind::Literal,
1419                        operands: vec![0],
1420                        result: Some(0),
1421                    },
1422                    KernelOp {
1423                        kind: KernelOpKind::Literal,
1424                        operands: vec![1],
1425                        result: Some(1),
1426                    },
1427                    KernelOp {
1428                        kind: KernelOpKind::StoreGlobal,
1429                        operands: vec![0, 0, 1],
1430                        result: None,
1431                    },
1432                ],
1433                child_bodies: vec![],
1434                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
1435            },
1436        };
1437        let (_out, stats) = run_all_with_stats(&desc);
1438        // drop_unused_bindings now retains every declared binding
1439        // regardless of visibility: dropping a ReadOnly silently shifted
1440        // the host's positional input mapping (parser pipeline `haystack`
1441        // got DCE'd → wrong slot at dispatch). The rewrite is now a
1442        // no-op for binding count; the stats counters reflect that.
1443        assert_eq!(stats.bindings_before, 2);
1444        assert_eq!(stats.bindings_after, 2);
1445        assert_eq!(stats.bindings_dropped(), 0);
1446    }
1447}
1448
1449#[cfg(test)]
1450mod value_differential {
1451    //! Bootstrap-validated scalar SSA value differential for the arithmetic
1452    //! literal-chain combine rewrites. The per-rewrite contract tests assert op
1453    //! SHAPE (combined literal, operand ids); this asserts the COMPUTED u32
1454    //! value of `(gid ∘ a) ∘ b` is preserved for every input, the SSA analog
1455    //! of the foundation value-differential suites. It catches a wrong
1456    //! reassociation or an unguarded combined-literal overflow that a shape
1457    //! check cannot see.
1458    //!
1459    //! A minimal interpreter over the arithmetic `KernelOpKind` subset evaluates
1460    //! the descriptor before and after each rewrite across adversarial
1461    //! `GlobalInvocationId` seeds. The interpreter is validated against
1462    //! hand-computed values first (`evaluator_is_validated_on_known_values`), so
1463    //! a divergence is the rewrite's fault, not the harness's. It panics, never
1464    //! silently guesses, on any op or literal outside the validated scalar
1465    //! subset, so an unexpected lowering shape fails loudly instead of producing
1466    //! a false "values match".
1467
1468    use super::rhs_lit_chain::test_support::{
1469        binop, descriptor_with, empty_body, lit_u32, nonliteral_source,
1470    };
1471    use crate::{KernelDescriptor, KernelOp, KernelOpKind, LiteralValue};
1472    use std::collections::HashMap;
1473    use vyre_foundation::ir::BinOp;
1474
1475    /// vyre u32 BinOp semantics (wrapping add/sub/mul, floor div/rem, shift
1476    /// amount masked mod 32), the subset the arithmetic chains use. Matches the
1477    /// reference oracle's integer contract.
1478    fn apply_u32(op: BinOp, l: u32, r: u32) -> u32 {
1479        match op {
1480            BinOp::Add => l.wrapping_add(r),
1481            BinOp::Sub => l.wrapping_sub(r),
1482            BinOp::Mul => l.wrapping_mul(r),
1483            BinOp::Div => {
1484                if r == 0 {
1485                    0
1486                } else {
1487                    l / r
1488                }
1489            }
1490            BinOp::Mod => {
1491                if r == 0 {
1492                    0
1493                } else {
1494                    l % r
1495                }
1496            }
1497            BinOp::BitAnd => l & r,
1498            BinOp::BitOr => l | r,
1499            BinOp::BitXor => l ^ r,
1500            BinOp::Shl => l.wrapping_shl(r & 31),
1501            BinOp::Shr => l.wrapping_shr(r & 31),
1502            other => panic!("value_differential evaluator: unhandled BinOp {other:?}"),
1503        }
1504    }
1505
1506    /// Resolve the value at `id` by recursively evaluating its producer's
1507    /// operands. Order-independent: it does not assume where a rewrite inserted
1508    /// the freshly combined literal op relative to its consumer in `ops`.
1509    fn eval_id(
1510        id: u32,
1511        producers: &HashMap<u32, &KernelOp>,
1512        literals: &[LiteralValue],
1513        gid: u32,
1514        memo: &mut HashMap<u32, u32>,
1515    ) -> u32 {
1516        if let Some(v) = memo.get(&id) {
1517            return *v;
1518        }
1519        let op = producers
1520            .get(&id)
1521            .unwrap_or_else(|| panic!("value_differential: result id {id} has no producer"));
1522        let value = match &op.kind {
1523            KernelOpKind::GlobalInvocationId => gid,
1524            KernelOpKind::Literal => {
1525                match literals
1526                    .get(op.operands[0] as usize)
1527                    .expect("literal pool index in range")
1528                {
1529                    LiteralValue::U32(v) => *v,
1530                    other => panic!("value_differential: non-u32 literal {other:?}"),
1531                }
1532            }
1533            KernelOpKind::Copy => eval_id(op.operands[0], producers, literals, gid, memo),
1534            KernelOpKind::BinOpKind(b) => {
1535                let l = eval_id(op.operands[0], producers, literals, gid, memo);
1536                let r = eval_id(op.operands[1], producers, literals, gid, memo);
1537                apply_u32(*b, l, r)
1538            }
1539            other => panic!("value_differential evaluator: unhandled op kind {other:?}"),
1540        };
1541        memo.insert(id, value);
1542        value
1543    }
1544
1545    fn eval_scalar(desc: &KernelDescriptor, output_id: u32, gid: u32) -> u32 {
1546        let mut producers: HashMap<u32, &KernelOp> = HashMap::new();
1547        for op in &desc.body.ops {
1548            if let Some(rid) = op.result {
1549                producers.insert(rid, op);
1550            }
1551        }
1552        let mut memo = HashMap::new();
1553        eval_id(output_id, &producers, &desc.body.literals, gid, &mut memo)
1554    }
1555
1556    /// `result 4 = (gid ∘ a) ∘ b`; gid = `GlobalInvocationId` at result 0,
1557    /// a/b are rhs literals. Mirrors the per-rewrite combine fixtures so the
1558    /// rewrite actually fires (non-literal base, single-consumer intermediate).
1559    fn chain(op: BinOp, a: u32, b: u32) -> KernelDescriptor {
1560        let mut body = empty_body();
1561        nonliteral_source(&mut body, 0);
1562        lit_u32(&mut body, a, 1);
1563        binop(&mut body, op, 0, 1, 2);
1564        lit_u32(&mut body, b, 3);
1565        binop(&mut body, op, 2, 3, 4);
1566        descriptor_with("value_differential", body)
1567    }
1568
1569    const PROBES: &[u32] = &[
1570        0,
1571        1,
1572        2,
1573        3,
1574        7,
1575        8,
1576        255,
1577        256,
1578        0x0001_0000,
1579        0x5555_5555,
1580        0xAAAA_AAAA,
1581        0x8000_0000,
1582        0x7FFF_FFFF,
1583        0xFFFF_FFFF,
1584    ];
1585
1586    #[test]
1587    fn evaluator_is_validated_on_known_values() {
1588        assert_eq!(eval_scalar(&chain(BinOp::Add, 3, 5), 4, 10), 18); // (10+3)+5
1589        assert_eq!(eval_scalar(&chain(BinOp::Mul, 3, 5), 4, 4), 60); // (4*3)*5
1590        assert_eq!(eval_scalar(&chain(BinOp::Div, 3, 5), 4, 100), 6); // ⌊⌊100/3⌋/5⌋
1591        assert_eq!(eval_scalar(&chain(BinOp::Sub, 3, 5), 4, 20), 12); // (20-3)-5
1592    }
1593
1594    #[test]
1595    fn arithmetic_chain_combines_preserve_value_for_every_input() {
1596        type Rewrite = fn(&KernelDescriptor) -> KernelDescriptor;
1597        let cases: &[(&str, Rewrite, BinOp)] = &[
1598            (
1599                "add_combine",
1600                super::arithmetic_combine::add_combine::add_combine,
1601                BinOp::Add,
1602            ),
1603            (
1604                "mul_combine",
1605                super::arithmetic_combine::mul_combine::mul_combine,
1606                BinOp::Mul,
1607            ),
1608            ("sub_combine", super::sub_combine::sub_combine, BinOp::Sub),
1609            ("div_combine", super::div_combine::div_combine, BinOp::Div),
1610        ];
1611        // a=3, b=5: combined literal (8 / 15 / 8 / 15) neither wraps nor zeroes;
1612        // divisors are non-zero. The combine therefore fires (not the
1613        // overflow-guard decline path), so this exercises the rewritten op.
1614        for (name, rewrite, op) in cases {
1615            let original = chain(*op, 3, 5);
1616            let rewritten = rewrite(&original);
1617            for &gid in PROBES {
1618                let before = eval_scalar(&original, 4, gid);
1619                let after = eval_scalar(&rewritten, 4, gid);
1620                assert_eq!(
1621                    before, after,
1622                    "{name} changed the computed value of (gid {op:?} 3) {op:?} 5 \
1623                     at gid={gid:#010x}: {before} -> {after}"
1624                );
1625            }
1626        }
1627    }
1628}