vyre-lower 0.4.1

//! Store-to-load and load-to-load forwarding.
//!
//! When a `Load(slot, idx)` follows a `Store(slot, idx, val)` (or an
//! earlier `Load(slot, idx) -> r`) with no intervening op that could
//! invalidate `(slot, idx)`, the second load's result-id is rewritten
//! to point at the prior value-id. The Load op itself is left in place;
//! `descriptor_dce` removes it if it becomes unused.
//!
//! ## Invalidation
//!
//! Per-slot tracking. Any of these intervening ops invalidate ALL
//! entries for the relevant slot (or all slots, when scope is unknown):
//! - `Store*` to the same slot at a different index — may alias the
//!   tracked index, so the cached value is no longer guaranteed.
//! - `Atomic` — read-modify-writes can change the value.
//! - `Barrier` — re-publishes other threads' writes.
//! - `Async{Load,Store,Wait}` — staged writes may land asynchronously.
//! - Structured control flow (`If`/`ForLoop`/`Block`/`Region`) — body
//!   could write anywhere.
//! - `Trap`/`Resume`/`Return` — terminator.
//! - `Call`/`OpaqueExpr`/`OpaqueNode` — opaque side effects.
//!
//! ## Why two flavors at once
//!
//! Both store-to-load (`store; load → forward`) and load-to-load
//! (`load r1; load r2 → r2 := r1`) share the same per-slot cache; doing
//! one without the other would leave easy redundancy on the table.

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

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

fn load_forwarding_body(mut body: KernelBody) -> KernelBody {
    // Per-slot map: slot_id → (idx_id → val_id).
    let mut cache: FxHashMap<u32, FxHashMap<u32, u32>> = FxHashMap::default();
    let mut id_remap: FxHashMap<u32, u32> = FxHashMap::default();

    for op in &body.ops {
        match &op.kind {
            KernelOpKind::StoreGlobal | KernelOpKind::StoreShared => {
                if op.operands.len() < 3 {
                    continue;
                }
                let slot = op.operands[0];
                let idx = resolve(op.operands[1], &id_remap);
                let val = resolve(op.operands[2], &id_remap);
                // Conservative: stores at a different index in the same
                // slot might alias the tracked one. Wipe other entries
                // for this slot, then install (idx → val).
                let entry = cache.entry(slot).or_default();
                entry.clear();
                entry.insert(idx, val);
            }
            KernelOpKind::LoadGlobal | KernelOpKind::LoadShared => {
                if op.operands.len() < 2 {
                    continue;
                }
                let slot = op.operands[0];
                let idx = resolve(op.operands[1], &id_remap);
                let Some(load_result) = op.result else {
                    continue;
                };
                let entry = cache.entry(slot).or_default();
                if let Some(&forwarded) = entry.get(&idx) {
                    // Forward: rewrite this load's result-id refs to point
                    // at the cached value.
                    id_remap.insert(load_result, forwarded);
                    // Don't update the cache — the cached value is the
                    // canonical id; this load is now redundant.
                } else {
                    // First time seeing (slot, idx) at this point —
                    // remember THIS load as the canonical for any later
                    // duplicate.
                    entry.insert(idx, load_result);
                }
            }
            KernelOpKind::LoadConstant => {
                // Constants are immutable — perfectly safe to forward.
                if op.operands.len() < 2 {
                    continue;
                }
                let slot = op.operands[0];
                let idx = resolve(op.operands[1], &id_remap);
                let Some(load_result) = op.result else {
                    continue;
                };
                let entry = cache.entry(slot).or_default();
                if let Some(&forwarded) = entry.get(&idx) {
                    id_remap.insert(load_result, forwarded);
                } else {
                    entry.insert(idx, load_result);
                }
            }
            KernelOpKind::Atomic { .. } => {
                // Atomic touches one slot; conservatively wipe just that slot.
                if let Some(&slot) = op.operands.first() {
                    cache.remove(&slot);
                }
            }
            KernelOpKind::Barrier { .. } => {
                // Barrier republishes other threads' writes — wipe everything.
                cache.clear();
            }
            KernelOpKind::AsyncLoad { .. }
            | KernelOpKind::AsyncStore { .. }
            | KernelOpKind::AsyncWait { .. } => {
                cache.clear();
            }
            KernelOpKind::StructuredIfThen
            | KernelOpKind::StructuredIfThenElse
            | KernelOpKind::StructuredForLoop { .. }
            | KernelOpKind::StructuredBlock
            | KernelOpKind::Region { .. } => {
                cache.clear();
            }
            KernelOpKind::Trap { .. } | KernelOpKind::Resume { .. } | KernelOpKind::Return => {
                cache.clear();
            }
            KernelOpKind::Call { .. }
            | KernelOpKind::OpaqueExpr { .. }
            | KernelOpKind::OpaqueNode { .. } => {
                cache.clear();
            }
            // Pure ops — no memory effect.
            _ => {}
        }
    }

    if !id_remap.is_empty() {
        for op in &mut body.ops {
            for pos in 0..op.operands.len() {
                if operand_is_result_reference(&op.kind, pos) {
                    if let Some(&new) = id_remap.get(&op.operands[pos]) {
                        op.operands[pos] = new;
                    }
                }
            }
        }
    }

    body.child_bodies = body
        .child_bodies
        .into_iter()
        .map(load_forwarding_body)
        .collect();
    body
}

fn resolve(id: u32, remap: &FxHashMap<u32, u32>) -> u32 {
    let mut cur = id;
    let mut hops = 0usize;
    while let Some(&next) = remap.get(&cur) {
        if next == cur {
            break;
        }
        cur = next;
        hops += 1;
        if hops > remap.len() + 1 {
            break;
        }
    }
    cur
}

/// Mirror — must stay in sync with the others.
fn operand_is_result_reference(kind: &KernelOpKind, pos: usize) -> bool {
    use KernelOpKind::*;
    match kind {
        Literal => false,
        LocalInvocationId | GlobalInvocationId | WorkgroupId => false,
        SubgroupLocalId | SubgroupSize | LoopIndex { .. } => false,
        LoopCarrier { .. } | LoopCarrierEnd { .. } => pos == 0,
        LoopCarrierFinal { .. } => false,
        LoadGlobal | LoadShared | LoadConstant => pos != 0,
        BufferLength => false,
        StoreGlobal | StoreShared => pos != 0,
        BinOpKind(_) | UnOpKind(_) | Fma | MatrixMma { .. } | Select | Cast { .. } => true,
        Atomic { .. } => pos != 0,
        SubgroupBallot | SubgroupShuffle | SubgroupAdd => true,
        StructuredIfThen | StructuredIfThenElse => pos == 0,
        StructuredForLoop { .. } => pos != 2,
        StructuredBlock | Region { .. } => false,
        Return | Barrier { .. } => false,
        AsyncLoad { .. } | AsyncStore { .. } => pos >= 2,
        AsyncWait { .. } => false,
        Trap { .. } => pos == 0,
        Resume { .. } => false,
        IndirectDispatch { .. } => false,
        Call { .. } => true,
        OpaqueExpr { .. } | OpaqueNode { .. } => true,
    }
}

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

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

    #[test]
    fn store_then_load_forwards_value() {
        // r0=Lit(0) (idx), r1=Lit(7) (val), Store(buf, r0, r1), r2=Load(buf, r0),
        // Store(buf, r0, r2). The Load should forward r2 := r1.
        let desc = KernelDescriptor {
            id: "stl".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(2),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2], // value should now resolve to r1
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
            },
        };
        let out = load_forwarding(&desc);
        // The trailing store's value-operand was 2 → should now be 1.
        assert_eq!(out.body.ops[4].operands, vec![0, 0, 1]);
    }

    #[test]
    fn load_then_load_same_idx_forwards() {
        let desc = KernelDescriptor {
            id: "ll".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(2),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(0)],
            },
        };
        let out = load_forwarding(&desc);
        // Store value should resolve to r1 (the first load), not r2.
        assert_eq!(out.body.ops[3].operands, vec![0, 0, 1]);
    }

    #[test]
    fn cache_keys_resolve_through_prior_load_forwarding() {
        let desc = KernelDescriptor {
            id: "resolved_idx".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            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(3),
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(2),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 1, 3],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 2],
                        result: Some(4),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 1, 4],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(9)],
            },
        };

        let out = load_forwarding(&desc);
        assert_eq!(
            out.body.ops[6].operands,
            vec![0, 1, 3],
            "load at a remapped-equivalent index should forward from the earlier store"
        );
    }

    #[test]
    fn intervening_store_to_different_idx_invalidates_cache() {
        let desc = KernelDescriptor {
            id: "intervening".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    }, // idx_a
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    }, // val
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![2],
                        result: Some(2),
                    }, // idx_b (different)
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![3],
                        result: Some(3),
                    }, // val_b
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 1], // store val at idx_a
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 2, 3], // store val_b at idx_b — may alias idx_a
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0], // load idx_a
                        result: Some(4),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 4],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(7),
                    LiteralValue::U32(99),
                    LiteralValue::U32(8),
                ],
            },
        };
        let out = load_forwarding(&desc);
        // The second store wiped the cache for slot 0 then installed (idx_b -> val_b).
        // The load(idx_a) is no longer in cache → no forwarding → trailing store keeps r4.
        assert_eq!(out.body.ops[7].operands, vec![0, 0, 4]);
    }

    #[test]
    fn barrier_invalidates_cache() {
        let desc = KernelDescriptor {
            id: "barrier".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            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::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::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(2),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
            },
        };
        let out = load_forwarding(&desc);
        // Barrier wiped cache; trailing store still references r2.
        assert_eq!(out.body.ops[5].operands, vec![0, 0, 2]);
    }

    #[test]
    fn different_slots_dont_interfere() {
        let mut s1 = rw_slot();
        s1.slot = 1;
        s1.name = "buf2".into();
        let desc = KernelDescriptor {
            id: "twoslots".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot(), s1],
            },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    }, // idx
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    }, // val_a
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![2],
                        result: Some(2),
                    }, // val_b
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 1], // slot 0
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![1, 0, 2], // slot 1 — different slot, doesn't invalidate slot 0
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0], // slot 0, idx 0 — should forward to val_a (r1)
                        result: Some(3),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 3],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![
                    LiteralValue::U32(0),
                    LiteralValue::U32(7),
                    LiteralValue::U32(99),
                ],
            },
        };
        let out = load_forwarding(&desc);
        assert_eq!(out.body.ops[6].operands, vec![0, 0, 1]);
    }

    #[test]
    fn atomic_invalidates_only_target_slot() {
        let desc = KernelDescriptor {
            id: "atomic_slot".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::Atomic {
                            op: vyre_foundation::ir::AtomicOp::Add,
                            ordering:
                                vyre_foundation::runtime::memory_model::MemoryOrdering::SeqCst,
                        },
                        operands: vec![0, 0, 1], // slot 0
                        result: Some(2),
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(3),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 3],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
            },
        };
        let out = load_forwarding(&desc);
        // Atomic on slot 0 invalidated cache; load can't forward.
        assert_eq!(out.body.ops[5].operands, vec![0, 0, 3]);
    }

    #[test]
    fn structured_if_invalidates_cache() {
        let desc = KernelDescriptor {
            id: "if_invalidates".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            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,
                    },
                    // condition operand
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(2),
                    },
                    KernelOp {
                        kind: KernelOpKind::StructuredIfThen,
                        operands: vec![2, 0], // cond=r2, body-child-idx=0
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(3),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 3],
                        result: None,
                    },
                ],
                child_bodies: vec![KernelBody {
                    ops: vec![],
                    child_bodies: vec![],
                    literals: vec![],
                }],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
            },
        };
        let out = load_forwarding(&desc);
        assert_eq!(out.body.ops[6].operands, vec![0, 0, 3]);
    }

    #[test]
    fn nothing_to_forward_is_noop() {
        let desc = KernelDescriptor {
            id: "noop".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
            },
        };
        let out = load_forwarding(&desc);
        assert_eq!(out.body.ops.len(), 3);
        assert_eq!(out.body.ops[2].operands, vec![0, 0, 1]);
    }

    #[test]
    fn loadconstant_is_forwardable_across_arbitrary_ops() {
        // LoadConstant points at immutable memory — even a Barrier shouldn't
        // wipe the cached entry. But our impl does (it's slot-keyed and
        // Barrier wipes everything). This test just documents current
        // conservative behavior.
        let desc = KernelDescriptor {
            id: "lc".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadConstant,
                        operands: vec![0, 0],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadConstant,
                        operands: vec![0, 0], // same (slot, idx) — should forward
                        result: Some(2),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(0)],
            },
        };
        let out = load_forwarding(&desc);
        assert_eq!(out.body.ops[3].operands, vec![0, 0, 1]);
    }

    #[test]
    fn idempotent() {
        let desc = KernelDescriptor {
            id: "idemp".into(),
            bindings: BindingLayout {
                slots: vec![rw_slot()],
            },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![1],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                    KernelOp {
                        kind: KernelOpKind::LoadGlobal,
                        operands: vec![0, 0],
                        result: Some(2),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 2],
                        result: None,
                    },
                ],
                child_bodies: vec![],
                literals: vec![LiteralValue::U32(0), LiteralValue::U32(7)],
            },
        };
        let once = load_forwarding(&desc);
        let twice = load_forwarding(&once);
        assert_eq!(once.body.ops, twice.body.ops);
    }

    #[test]
    fn empty_kernel_is_noop() {
        let desc = KernelDescriptor {
            id: "empty".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(1, 1, 1),
            body: KernelBody {
                ops: vec![],
                child_bodies: vec![],
                literals: vec![],
            },
        };
        let out = load_forwarding(&desc);
        assert!(out.body.ops.is_empty());
    }
}