vyre-lower 0.4.1

//! Dead-store elimination.
//!
//! When two `StoreGlobal` (or `StoreShared`) ops in sequence write to
//! the same binding-slot AND the same index-operand-id, AND no
//! intervening op reads that location (or any aliased location), the
//! first store is dead — its value gets immediately overwritten and
//! never observed.
//!
//! Phase-1 conservative rules:
//! - Same binding slot.
//! - Same index operand id (textual equality — no commutative-index
//!   normalization, no alias analysis across bindings).
//! - No intervening LoadGlobal/LoadShared from the same binding.
//! - No intervening structured-control-flow op (an if/loop body might
//!   read the location; phase 1 stays out of those).
//! - No intervening Barrier or Atomic on the same slot.
//!
//! Phase 2 (follow-up): cross-arm DSE in megakernels via points-to
//! analysis, alias-aware between separate slots, etc.

use crate::{KernelBody, KernelDescriptor, KernelOp, KernelOpKind};
use rustc_hash::FxHashMap;

#[must_use]
pub fn dead_store(desc: &KernelDescriptor) -> KernelDescriptor {
    let mut out = desc.clone();
    out.body = dead_store_body(out.body);
    out
}

fn dead_store_body(mut body: KernelBody) -> KernelBody {
    let mut keep = vec![true; body.ops.len()];
    let mut pending_stores = FxHashMap::<StoreKey, usize>::default();

    for (op_index, op) in body.ops.iter().enumerate() {
        match store_key(op) {
            Some(key) => {
                if let Some(previous_store) = pending_stores.insert(key, op_index) {
                    keep[previous_store] = false;
                }
            }
            None => invalidate_pending_stores(&mut pending_stores, &op.kind),
        }
    }

    let old_ops = std::mem::take(&mut body.ops);
    body.ops = old_ops
        .into_iter()
        .enumerate()
        .filter_map(|(i, op)| keep[i].then_some(op))
        .collect();

    body.child_bodies = body.child_bodies.into_iter().map(dead_store_body).collect();

    body
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct StoreKey {
    space: StoreSpace,
    slot: u32,
    index: u32,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum StoreSpace {
    Global,
    Shared,
}

fn store_key(op: &KernelOp) -> Option<StoreKey> {
    let space = match op.kind {
        KernelOpKind::StoreGlobal => StoreSpace::Global,
        KernelOpKind::StoreShared => StoreSpace::Shared,
        _ => return None,
    };
    Some(StoreKey {
        space,
        slot: *op.operands.first()?,
        index: *op.operands.get(1)?,
    })
}

fn invalidate_pending_stores(pending_stores: &mut FxHashMap<StoreKey, usize>, kind: &KernelOpKind) {
    match invalidated_store_space(kind) {
        Invalidation::None => {}
        Invalidation::Only(space) => {
            pending_stores.retain(|key, _| key.space != space);
        }
        Invalidation::All => pending_stores.clear(),
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Invalidation {
    None,
    Only(StoreSpace),
    All,
}

fn invalidated_store_space(kind: &KernelOpKind) -> Invalidation {
    use KernelOpKind::*;
    match kind {
        // Reads from the same slot make the previous store live.
        // Conservative: any load from the same address space kills the
        // dead-store removal opportunity for that address space.
        LoadGlobal => Invalidation::Only(StoreSpace::Global),
        LoadShared => Invalidation::Only(StoreSpace::Shared),
        // Atomics on the same slot also use/may-use the value.
        Atomic { .. } => Invalidation::All,
        // Structured control flow may read.
        StructuredIfThen
        | StructuredIfThenElse
        | StructuredForLoop { .. }
        | StructuredBlock
        | Region { .. } => Invalidation::All,
        // Barriers and async serialize visibility.
        Barrier { .. } | AsyncLoad { .. } | AsyncStore { .. } | AsyncWait { .. } => {
            Invalidation::All
        }
        // Trap/Resume/Return end the kernel; preceding store is observable
        // by the trap handler, so don't elide.
        Trap { .. } | Resume { .. } | Return => Invalidation::All,
        // Calls: opaque side effects.
        Call { .. } | OpaqueExpr { .. } | OpaqueNode { .. } => Invalidation::All,
        _ => Invalidation::None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        BindingLayout, BindingSlot, BindingVisibility, Dispatch, KernelBody, KernelDescriptor,
        KernelOp, KernelOpKind, LiteralValue, MemoryClass,
    };
    use vyre_foundation::ir::DataType;

    fn out_binding() -> BindingSlot {
        BindingSlot {
            slot: 0,
            element_type: DataType::U32,
            element_count: None,
            memory_class: MemoryClass::Global,
            visibility: BindingVisibility::WriteOnly,
            name: "out".into(),
        }
    }

    #[test]
    fn dead_store_on_empty_kernel_is_noop() {
        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 = dead_store(&desc);
        assert!(out.body.ops.is_empty());
    }

    #[test]
    fn two_back_to_back_stores_to_same_address_keep_only_last() {
        let desc = KernelDescriptor {
            id: "double_store".into(),
            bindings: BindingLayout {
                slots: vec![out_binding()],
            },
            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::StoreGlobal,
                        operands: vec![0, 0, 1], // store value 1 at index 0
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2], // store value 2 at index 0 — kills the prev store
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(7),
                    LiteralValue::U32(99),
                ],
            },
        };
        let out = dead_store(&desc);
        // Started with 5 ops; first store is dead → 4 ops survive.
        assert_eq!(out.body.ops.len(), 4);
        // Confirm the surviving store stores value 2.
        let store = out
            .body
            .ops
            .iter()
            .find(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
            .unwrap();
        assert_eq!(store.operands, vec![0, 0, 2]);
    }

    #[test]
    fn intervening_load_keeps_first_store_alive() {
        let mut binding = out_binding();
        binding.visibility = BindingVisibility::ReadWrite;
        let desc = KernelDescriptor {
            id: "store_load_store".into(),
            bindings: BindingLayout {
                slots: vec![binding],
            },
            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::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(3),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(7),
                    LiteralValue::U32(99),
                ],
            },
        };
        let out = dead_store(&desc);
        // The intervening load reads what the first store wrote → both stores stay.
        let store_count = out
            .body
            .ops
            .iter()
            .filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
            .count();
        assert_eq!(store_count, 2);
    }

    #[test]
    fn stores_to_different_indices_both_survive() {
        let desc = KernelDescriptor {
            id: "two_indices".into(),
            bindings: BindingLayout {
                slots: vec![out_binding()],
            },
            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::StoreGlobal,
                        operands: vec![0, 0, 1], // index = result-id 0
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 1, 2], // index = result-id 1 — DIFFERENT
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(1),
                    LiteralValue::U32(99),
                ],
            },
        };
        let out = dead_store(&desc);
        let store_count = out
            .body
            .ops
            .iter()
            .filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
            .count();
        assert_eq!(store_count, 2);
    }

    #[test]
    fn three_consecutive_stores_keep_only_last() {
        let desc = KernelDescriptor {
            id: "triple".into(),
            bindings: BindingLayout {
                slots: vec![out_binding()],
            },
            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::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 3],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(7),
                    LiteralValue::U32(8),
                    LiteralValue::U32(9),
                ],
            },
        };
        let out = dead_store(&desc);
        let store_count = out
            .body
            .ops
            .iter()
            .filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
            .count();
        assert_eq!(store_count, 1);
    }

    #[test]
    fn interleaved_same_address_stores_keep_only_last_per_address() {
        let desc = KernelDescriptor {
            id: "interleaved".into(),
            bindings: BindingLayout {
                slots: vec![out_binding()],
            },
            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::StoreGlobal,
                        operands: vec![0, 0, 2],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 1, 2],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 3],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(1),
                    LiteralValue::U32(7),
                    LiteralValue::U32(9),
                ],
            },
        };
        let out = dead_store(&desc);
        let stores = out
            .body
            .ops
            .iter()
            .filter(|op| matches!(op.kind, KernelOpKind::StoreGlobal))
            .collect::<Vec<_>>();
        assert_eq!(stores.len(), 2);
        assert_eq!(stores[0].operands, vec![0, 1, 2]);
        assert_eq!(stores[1].operands, vec![0, 0, 3]);
    }

    #[test]
    fn intervening_barrier_keeps_first_store_alive() {
        let desc = KernelDescriptor {
            id: "barrier_between".into(),
            bindings: BindingLayout {
                slots: vec![out_binding()],
            },
            dispatch: Dispatch::new(64, 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::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::Barrier {
                            ordering:
                                vyre_foundation::runtime::memory_model::MemoryOrdering::SeqCst,
                        },
                        operands: vec![],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(7),
                    LiteralValue::U32(99),
                ],
            },
        };
        let out = dead_store(&desc);
        // Barrier serializes visibility — first store is observable by other threads.
        let store_count = out
            .body
            .ops
            .iter()
            .filter(|o| matches!(o.kind, KernelOpKind::StoreGlobal))
            .count();
        assert_eq!(store_count, 2);
    }

    #[test]
    fn dead_store_is_idempotent() {
        let desc = KernelDescriptor {
            id: "k".into(),
            bindings: BindingLayout {
                slots: vec![out_binding()],
            },
            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::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(7),
                    LiteralValue::U32(99),
                ],
            },
        };
        let once = dead_store(&desc);
        let twice = dead_store(&once);
        assert_eq!(once.body.ops.len(), twice.body.ops.len());
    }
}