vyre-lower 0.4.1

//! Dead-op detection on `KernelDescriptor`.
//!
//! A KernelOp is dead when:
//! 1. It produces a result (`op.result.is_some()`)
//! 2. No other op in the same kernel body (or its child bodies) reads
//!    that result
//! 3. The op has no side effects (i.e. not a Store, Barrier, AtomicXxx,
//!    AsyncXxx, Trap/Resume, IndirectDispatch, Return)
//!
//! Phase 1: detection only — returns a list of op-indices flagged as
//! dead. Phase 2: a real DCE rewrite that strips them from the
//! descriptor (defers to vyre-opt's optimizer pipeline).

use crate::{KernelBody, KernelDescriptor, KernelOpKind};
use rustc_hash::FxHashSet;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct DeadOpReport {
    pub kernel_id: String,
    /// Op-indices (into `KernelBody.ops`) of detected dead ops.
    pub dead_op_indices: Vec<usize>,
    /// Total ops in the body (for context — `dead_count / total` gives
    /// the dead-code ratio).
    pub total_op_count: u32,
}

impl DeadOpReport {
    #[must_use]
    pub fn dead_count(&self) -> usize {
        self.dead_op_indices.len()
    }

    #[must_use]
    pub fn dead_ratio(&self) -> f32 {
        if self.total_op_count == 0 {
            0.0
        } else {
            self.dead_op_indices.len() as f32 / self.total_op_count as f32
        }
    }
}

#[must_use]
pub fn analyze(desc: &KernelDescriptor) -> DeadOpReport {
    // First pass: collect every result-id that any op produces. This
    // is the universe of valid "result references."
    let mut produced = FxHashSet::<u32>::with_capacity_and_hasher(
        count_ops(&desc.body) as usize,
        Default::default(),
    );
    walk_collect_produced(&desc.body, &mut produced);

    // Second pass: collect every operand that IS a produced result-id
    // (filtering out binding-slots, pool-indices, body-indices, axis
    // numbers, and any other operand that lives in a different
    // namespace from the result-id space).
    let mut referenced =
        FxHashSet::<u32>::with_capacity_and_hasher(produced.len(), Default::default());
    walk_collect_result_references(&desc.body, &produced, &mut referenced);

    // Third pass: find ops that produce a result not in `referenced`,
    // are pure (no side effects), and have a result id.
    let mut dead = Vec::new();
    walk_find_dead(&desc.body, &referenced, &mut dead, 0);

    DeadOpReport {
        kernel_id: desc.id.clone(),
        dead_op_indices: dead,
        total_op_count: count_ops(&desc.body),
    }
}

fn walk_collect_produced(body: &KernelBody, produced: &mut FxHashSet<u32>) {
    for op in &body.ops {
        for result in op.result_ids() {
            produced.insert(result);
        }
    }
    for child in &body.child_bodies {
        walk_collect_produced(child, produced);
    }
}

fn walk_collect_result_references(
    body: &KernelBody,
    produced: &FxHashSet<u32>,
    referenced: &mut FxHashSet<u32>,
) {
    for op in &body.ops {
        for (pos, operand_id) in op.operands.iter().enumerate() {
            if !operand_is_result_reference(&op.kind, pos) {
                continue;
            }
            if produced.contains(operand_id) {
                referenced.insert(*operand_id);
            }
        }
    }
    for child in &body.child_bodies {
        walk_collect_result_references(child, produced, referenced);
    }
}

/// Per-`KernelOpKind` operand classification: does the operand at
/// position `pos` reference an earlier result-id, or is it some other
/// kind of index (binding slot, literal-pool index, body-child index,
/// axis number)?
fn operand_is_result_reference(kind: &KernelOpKind, pos: usize) -> bool {
    use KernelOpKind::*;
    match kind {
        // Operand 0 is a pool index, not a result-id.
        Literal => false,
        // Operand 0 is the axis number for builtin id ops.
        LocalInvocationId | GlobalInvocationId | WorkgroupId => false,
        // No operands.
        SubgroupLocalId | SubgroupSize | LoopIndex { .. } => false,
        LoopCarrier { .. } | LoopCarrierEnd { .. } => pos == 0,
        LoopCarrierFinal { .. } => false,
                // Operand 0 is a binding-slot; rest are result-ids.
        LoadGlobal | LoadShared | LoadConstant => pos != 0,
        // Operand 0 is binding-slot; operands 1.. are result-ids.
        BufferLength => false,
        // Operand 0 is binding-slot; operands 1, 2 are result-ids.
        StoreGlobal | StoreShared => pos != 0,
        // All operands are result-ids.
        BinOpKind(_) | UnOpKind(_) | Fma | MatrixMma { .. } | Select | Cast { .. } => true,
        // Atomic: operand 0 is binding-slot; rest are result-ids.
        Atomic { .. } => pos != 0,
        // Subgroup ops: all operands are result-ids.
        SubgroupBallot | SubgroupShuffle | SubgroupAdd => true,
        // Operand 0 is condition (result-id); operand 1 (and 2 for
        // IfThenElse) are body-child indices, not result-ids.
        StructuredIfThen | StructuredIfThenElse => pos == 0,
        // Operands 0, 1 are lo/hi (result-ids); operand 2 is body index.
        StructuredForLoop { .. } => pos != 2,
        // Operand 0 is body-child index, not a result-id.
        StructuredBlock | Region { .. } => false,
        // No operands or all metadata.
        Return | Barrier { .. } => false,
        // Async load/store: operands 0, 1 are binding-slots; operands 2, 3 are result-ids.
        AsyncLoad { .. } | AsyncStore { .. } => pos >= 2,
        AsyncWait { .. } => false,
        // Trap: operand 0 is the address (result-id).
        Trap { .. } => pos == 0,
        Resume { .. } => false,
        // Operand 0 is binding-slot.
        IndirectDispatch { .. } => false,
        // Call: every operand is a result-id (call args).
        Call { .. } => true,
        // Conservative: assume all operands are result-id refs for
        // unknown extension ops, so DCE doesn't accidentally elide
        // their dependencies.
        OpaqueExpr { .. } | OpaqueNode { .. } => true,
    }
}

fn count_ops(body: &KernelBody) -> u32 {
    let mut total: u32 = body.ops.len() as u32;
    for child in &body.child_bodies {
        total = total.saturating_add(count_ops(child));
    }
    total
}

fn walk_find_dead(
    body: &KernelBody,
    referenced: &FxHashSet<u32>,
    dead: &mut Vec<usize>,
    op_index_offset: usize,
) {
    for (local_idx, op) in body.ops.iter().enumerate() {
        let op_index = op_index_offset + local_idx;
        if op.result.is_some() {
            if op.result_ids().all(|result| !referenced.contains(&result)) && is_pure(&op.kind) {
                dead.push(op_index);
            }
        }
    }
    for child in &body.child_bodies {
        walk_find_dead(child, referenced, dead, op_index_offset + body.ops.len());
    }
}

fn is_pure(kind: &KernelOpKind) -> bool {
    !matches!(
        kind,
        KernelOpKind::StoreGlobal
            | KernelOpKind::StoreShared
            | KernelOpKind::LoopCarrierEnd { .. }
            | KernelOpKind::LoopCarrierFinal { .. }
            | KernelOpKind::Barrier { .. }
            | KernelOpKind::Atomic { .. }
            | KernelOpKind::AsyncLoad { .. }
            | KernelOpKind::AsyncStore { .. }
            | KernelOpKind::AsyncWait { .. }
            | KernelOpKind::Trap { .. }
            | KernelOpKind::Resume { .. }
            | KernelOpKind::IndirectDispatch { .. }
            | KernelOpKind::Return
            | KernelOpKind::StructuredIfThen
            | KernelOpKind::StructuredIfThenElse
            | KernelOpKind::StructuredForLoop { .. }
            | KernelOpKind::StructuredBlock
            | KernelOpKind::Region { .. }
            | KernelOpKind::Call { .. }
            | KernelOpKind::OpaqueExpr { .. }
            | KernelOpKind::OpaqueNode { .. }
    )
}

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

    #[test]
    fn empty_kernel_has_no_dead_ops() {
        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 r = analyze(&desc);
        assert!(r.dead_op_indices.is_empty());
        assert_eq!(r.total_op_count, 0);
        assert!((r.dead_ratio() - 0.0).abs() < 1e-6);
    }

    #[test]
    fn unused_literal_is_dead() {
        // Two literals; one is used in a store, the other is dead.
        let desc = KernelDescriptor {
            id: "k".into(),
            bindings: BindingLayout {
                slots: vec![BindingSlot {
                    slot: 0,
                    element_type: DataType::U32,
                    element_count: None,
                    memory_class: MemoryClass::Global,
                    visibility: BindingVisibility::WriteOnly,
                    name: "out".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),
                    }, // dead
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(7),
                    LiteralValue::U32(99),
                ],
            },
        };
        let r = analyze(&desc);
        assert_eq!(r.dead_op_indices.len(), 1);
        assert_eq!(r.dead_op_indices[0], 2); // the third op
        assert_eq!(r.dead_count(), 1);
    }

    #[test]
    fn store_is_never_dead_even_with_no_result() {
        // Store has no result by definition; should NOT be flagged as dead.
        let desc = KernelDescriptor {
            id: "k".into(),
            bindings: BindingLayout {
                slots: vec![BindingSlot {
                    slot: 0,
                    element_type: DataType::U32,
                    element_count: None,
                    memory_class: MemoryClass::Global,
                    visibility: BindingVisibility::WriteOnly,
                    name: "out".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 r = analyze(&desc);
        assert!(r.dead_op_indices.is_empty());
    }

    #[test]
    fn unused_arithmetic_op_is_dead() {
        // tid; lit; add(tid, lit) → never used.
        let desc = KernelDescriptor {
            id: "k".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(64, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::LocalInvocationId,
                        operands: vec![],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::BinOpKind(BinOp::Add),
                        operands: vec![0, 1],
                        result: Some(2),
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(5)],
            },
        };
        let r = analyze(&desc);
        // op at index 2 (the Add) is dead — its result 2 is never used
        // (and operands 0 and 1 ARE used by it, so they're alive).
        // op at index 0 (tid) is dead too — result 0 is only used by
        // the dead Add, so transitively dead.
        // op at index 1 (lit) is dead too — same reason.
        // Phase-1 conservative: only flag direct deads (an op whose
        // result is unreferenced) — the chain-DCE is phase 2.
        assert_eq!(r.dead_op_indices.len(), 1);
        assert_eq!(r.dead_op_indices[0], 2);
    }

    #[test]
    fn dead_ratio_computed_correctly() {
        // 4 ops, 1 dead → ratio 0.25.
        let r = DeadOpReport {
            kernel_id: "k".into(),
            dead_op_indices: vec![2],
            total_op_count: 4,
        };
        assert!((r.dead_ratio() - 0.25).abs() < 1e-5);
    }

    #[test]
    fn return_is_never_dead() {
        let desc = KernelDescriptor {
            id: "k".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![KernelOp {
                    kind: KernelOpKind::Return,
                    operands: vec![],
                    result: None,
                }],
                child_bodies: vec![],
                literals: vec![],
            },
        };
        let r = analyze(&desc);
        assert!(r.dead_op_indices.is_empty());
    }

    #[test]
    fn barrier_is_never_dead() {
        let desc = KernelDescriptor {
            id: "k".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(64, 1, 1),
            body: KernelBody {
                ops: vec![KernelOp {
                    kind: KernelOpKind::Barrier {
                        ordering: vyre_foundation::runtime::memory_model::MemoryOrdering::SeqCst,
                    },
                    operands: vec![],
                    result: None,
                }],
                child_bodies: vec![],
                literals: vec![],
            },
        };
        let r = analyze(&desc);
        assert!(r.dead_op_indices.is_empty());
    }
}