vyre-lower 0.6.4

//! B6  -  promote a recognized matmul fragment shape into a single
//! [`KernelOpKind::MatrixMma`] op.
//!
//! ## What this catches
//!
//! The promotable shape is a tight, contiguous descriptor-body
//! sub-sequence with the operand contract for the
//! `M16N8K16/F16/F16/F32` MatrixMma fragment:
//!
//!   - 4 ops producing `a0..a3` (the A fragment lanes  -  typically
//!     `LoadShared` / `LoadGlobal` of the A tile);
//!   - 2 ops producing `b0..b1` (the B fragment lanes);
//!   - 4 `Fma` ops producing `c0..c3` accumulators in order, where
//!     each `Fma`'s third operand (`c_in`) is one of four pre-loaded
//!     C accumulators threaded through the chain.
//!
//! The pre-loaded `c_init0..c_init3` ops are NOT consumed by the
//! rewrite (they remain in the descriptor)  -  only the `Fma` chain
//! collapses into one `MatrixMma`. The MatrixMma's operand vector
//! is `[a0,a1,a2,a3, b0,b1, c0,c1,c2,c3]` and its 4 result ids start
//! at the new fresh result base.
//!
//! ## What this DOES NOT catch
//!
//! Generic 2D matmul-tile loops (the `for k in 0..K { c += a[k] * b[k] }`
//! shape that compiles to a serial Fma chain over a runtime-bounded
//! loop). Those need cross-iteration dataflow + tile-shape inference
//! and are an open follow-up. v0.4.1 ships the contiguous-fragment
//! detector so a kernel that explicitly lays out the M16N8K16 shape
//! (e.g. a hand-written tile primitive) gets the MatrixMma
//! lowering for free.
//!
//! ## Why this is safe to land
//!
//! The MatrixMma op is already verified and emitted (PTX side ships
//! `mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32` on sm_70+ per
//! B6 row text). This rewrite only adds the source-side promoter so
//! the existing emit path activates without explicit MatrixMma in
//! the input descriptor.

use std::collections::BTreeMap;

use super::literal::ResultAllocator;
use crate::descriptor::{
    KernelBody, KernelDescriptor, KernelOp, KernelOpKind, MatrixMmaElement, MatrixMmaLayout,
    MatrixMmaShape,
};
use crate::operand_semantics::remap_body_result_ids;

const MATMUL_TILE_LEN: usize = 4; // 4 Fma ops produce c0..c3
const A_FRAGMENT_LEN: usize = 4;
const B_FRAGMENT_LEN: usize = 2;

/// One structured loop whose body contains a promotable matmul fragment.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MatmulTileLoopPlan {
    /// Path to the body containing the loop. Empty path means the root body.
    pub body_path: Vec<usize>,
    /// Op index of the `StructuredForLoop` within `body_path`.
    pub loop_op_index: usize,
    /// Child body index referenced by the loop.
    pub child_body_index: u32,
    /// Op index inside the child body where the promotable FMA block begins.
    pub fma_start_index: usize,
}

/// Infer structured-loop tensor-core opportunities without rewriting.
///
/// Backends and diagnostics use this to prove generic matmul tile loops are
/// no longer invisible: a `StructuredForLoop` whose child body contains the
/// canonical four-FMA M16N8K16 fragment is reported with a stable body path.
#[must_use]
pub fn infer_matmul_tile_loops(desc: &KernelDescriptor) -> Vec<MatmulTileLoopPlan> {
    let mut plans = Vec::new();
    infer_body(&desc.body, &mut Vec::new(), &mut plans);
    plans
}

/// Apply the matmul-fragment promoter recursively. Returns the
/// transformed descriptor with every recognized fragment chain
/// collapsed into a single `MatrixMma`.
#[must_use]
pub fn matmul_promote(desc: &KernelDescriptor) -> KernelDescriptor {
    let mut out = desc.clone();
    let mut allocator = ResultAllocator::for_body_tree(&desc.body);
    // Each promotion collapses 4 Fma ops (which produced the 4 accumulator
    // result ids) into one MatrixMma carrying a FRESH 4-wide result block. The
    // old accumulator ids vanish, so every downstream consumer of them must be
    // re-pointed at the matching MatrixMma lane or it would dangle on an id no
    // op produces. `id_map` records old-accumulator -> new-lane for the whole
    // tree; the final pass re-points consumers in one sound, operand-aware walk.
    let mut id_map = BTreeMap::new();
    out.body = promote_body(&out.body, &mut allocator, &mut id_map);
    if !id_map.is_empty() {
        out.body = remap_body_result_ids(&out.body, &id_map);
    }
    out
}

fn promote_body(
    body: &KernelBody,
    allocator: &mut ResultAllocator,
    id_map: &mut BTreeMap<u32, u32>,
) -> KernelBody {
    let mut new_ops: Vec<KernelOp> = Vec::with_capacity(body.ops.len());

    let ops = &body.ops;
    let mut i = 0;
    while i < ops.len() {
        if let Some((promoted, advance, remap)) = try_promote_at(ops, i, allocator) {
            new_ops.push(promoted);
            id_map.extend(remap);
            i += advance;
        } else {
            new_ops.push(ops[i].clone());
            i += 1;
        }
    }

    let new_children: Vec<KernelBody> = body
        .child_bodies
        .iter()
        .map(|c| promote_body(c, allocator, id_map))
        .collect();

    KernelBody {
        ops: new_ops,
        child_bodies: new_children,
        literals: body.literals.clone(),
    }
}

fn infer_body(body: &KernelBody, path: &mut Vec<usize>, plans: &mut Vec<MatmulTileLoopPlan>) {
    for (op_index, op) in body.ops.iter().enumerate() {
        if matches!(op.kind, KernelOpKind::StructuredForLoop { .. }) {
            let Some(child_body_index) = op.operands.get(2).copied() else {
                continue;
            };
            let Some(child) = body.child_bodies.get(child_body_index as usize) else {
                continue;
            };
            for fma_start_index in promotable_fma_starts(&child.ops) {
                plans.push(MatmulTileLoopPlan {
                    body_path: path.clone(),
                    loop_op_index: op_index,
                    child_body_index,
                    fma_start_index,
                });
            }
        }
    }
    for (child_index, child) in body.child_bodies.iter().enumerate() {
        path.push(child_index);
        infer_body(child, path, plans);
        path.pop();
    }
}

fn promotable_fma_starts(ops: &[KernelOp]) -> Vec<usize> {
    let mut starts = Vec::new();
    for i in 0..ops.len() {
        if match_fragment_at(ops, i).is_some() {
            starts.push(i);
        }
    }
    starts
}

/// If `ops[i..]` begins with the promotable fragment shape, return the
/// synthesized `MatrixMma` op + how many ops to advance past in the
/// source. `None` means no match.
fn try_promote_at(
    ops: &[KernelOp],
    i: usize,
    allocator: &mut ResultAllocator,
) -> Option<(KernelOp, usize, [(u32, u32); MATMUL_TILE_LEN])> {
    let FragmentMatch {
        a_ids,
        b_unique,
        c_ids,
        c_out,
    } = match_fragment_at(ops, i)?;

    let mut operands = Vec::with_capacity(A_FRAGMENT_LEN + B_FRAGMENT_LEN + MATMUL_TILE_LEN);
    operands.extend_from_slice(&a_ids);
    operands.extend_from_slice(&b_unique);
    operands.extend_from_slice(&c_ids);

    let result_base = allocator.fresh_block(MATMUL_TILE_LEN as u32);

    let promoted = KernelOp {
        kind: KernelOpKind::MatrixMma {
            shape: MatrixMmaShape::M16N8K16,
            a_layout: MatrixMmaLayout::RowMajor,
            b_layout: MatrixMmaLayout::ColMajor,
            a_type: MatrixMmaElement::F16,
            b_type: MatrixMmaElement::F16,
            accum_type: MatrixMmaElement::F32,
        },
        operands,
        result: Some(result_base),
    };

    // Output lane k of the MatrixMma is the updated accumulator that Fma k
    // produced, so old c_out[k] is replaced by result_base + k. Consumers of
    // the dropped Fma results are re-pointed via this map.
    let remap = [
        (c_out[0], result_base),
        (c_out[1], result_base + 1),
        (c_out[2], result_base + 2),
        (c_out[3], result_base + 3),
    ];

    Some((promoted, MATMUL_TILE_LEN, remap))
}

struct FragmentMatch {
    a_ids: [u32; A_FRAGMENT_LEN],
    b_unique: [u32; B_FRAGMENT_LEN],
    c_ids: [u32; MATMUL_TILE_LEN],
    /// The 4 Fma result ids (the accumulator outputs) the MatrixMma replaces.
    c_out: [u32; MATMUL_TILE_LEN],
}

fn match_fragment_at(ops: &[KernelOp], i: usize) -> Option<FragmentMatch> {
    let needed = MATMUL_TILE_LEN; // we look at the 4 Fma ops
    if i + needed > ops.len() {
        return None;
    }
    // The MatrixMma operand vector requires 4 a + 2 b + 4 c  -  the
    // 4 c values come from the prior cumulative state. We look for a
    // contiguous block of 4 `Fma` ops whose first operand cycles over
    // 4 unique a ids and whose second cycles over 2 unique b ids.

    let block = &ops[i..i + needed];
    if !block.iter().all(|op| matches!(op.kind, KernelOpKind::Fma)) {
        return None;
    }
    if !block
        .iter()
        .all(|op| op.operands.len() == 3 && op.result.is_some())
    {
        return None;
    }
    // Each Fma's c-in is a distinct id (the 4 pre-loaded accumulators).
    let c_ids = [
        block[0].operands[2],
        block[1].operands[2],
        block[2].operands[2],
        block[3].operands[2],
    ];
    if !all_distinct(&c_ids) {
        return None;
    }
    // The 4 Fma RESULT ids (accumulator outputs). The block was already
    // checked to have a result on every op; capture them so consumers can be
    // re-pointed to the MatrixMma after the Fmas are dropped.
    let c_out = [
        block[0].result?,
        block[1].result?,
        block[2].result?,
        block[3].result?,
    ];
    // a operand cycles 4-wide; b cycles 2-wide.
    let a_ids = [
        block[0].operands[0],
        block[1].operands[0],
        block[2].operands[0],
        block[3].operands[0],
    ];
    let b_ids = [
        block[0].operands[1],
        block[1].operands[1],
        block[2].operands[1],
        block[3].operands[1],
    ];
    if !all_distinct(&a_ids) {
        return None;
    }
    // Two of the b ids must repeat; the other two are the unique b
    // fragment lanes. Reject if more than 2 unique.
    let mut b_unique: Vec<u32> = b_ids.to_vec();
    b_unique.sort_unstable();
    b_unique.dedup();
    if b_unique.len() != B_FRAGMENT_LEN {
        return None;
    }
    Some(FragmentMatch {
        a_ids,
        b_unique: [b_unique[0], b_unique[1]],
        c_ids,
        c_out,
    })
}

fn all_distinct(ids: &[u32]) -> bool {
    let mut sorted = ids.to_vec();
    sorted.sort_unstable();
    sorted.windows(2).all(|w| w[0] != w[1])
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::descriptor::{
        BindingLayout, Dispatch, KernelBody, KernelDescriptor, KernelOp, KernelOpKind,
    };

    fn empty_desc(body: KernelBody) -> KernelDescriptor {
        KernelDescriptor {
            id: "k".into(),
            bindings: BindingLayout { slots: vec![] },
            dispatch: Dispatch::new(1, 1, 1),
            body,
        }
    }

    fn lit(result: u32) -> KernelOp {
        KernelOp {
            kind: KernelOpKind::Literal,
            operands: vec![0],
            result: Some(result),
        }
    }

    fn fma(a: u32, b: u32, c: u32, result: u32) -> KernelOp {
        KernelOp {
            kind: KernelOpKind::Fma,
            operands: vec![a, b, c],
            result: Some(result),
        }
    }

    fn store_global(slot: u32, idx: u32, val: u32) -> KernelOp {
        KernelOp {
            kind: KernelOpKind::StoreGlobal,
            operands: vec![slot, idx, val],
            result: None,
        }
    }

    #[test]
    fn empty_body_is_unchanged() {
        let desc = empty_desc(KernelBody {
            ops: vec![],
            child_bodies: vec![],
            literals: vec![],
        });
        let out = matmul_promote(&desc);
        assert!(out.body.ops.is_empty());
    }

    #[test]
    fn non_matmul_body_is_unchanged() {
        let desc = empty_desc(KernelBody {
            ops: vec![lit(0), lit(1), lit(2)],
            child_bodies: vec![],
            literals: vec![],
        });
        let out = matmul_promote(&desc);
        assert_eq!(out.body.ops.len(), 3);
        assert!(out
            .body
            .ops
            .iter()
            .all(|op| matches!(op.kind, KernelOpKind::Literal)));
    }

    #[test]
    fn four_fma_with_correct_shape_promotes_to_matrix_mma() {
        // Pre-load 4 a + 2 b + 4 c accumulators, then 4 Fmas:
        //   Fma(a0, b0, c0) -> r10
        //   Fma(a1, b1, c1) -> r11
        //   Fma(a2, b0, c2) -> r12  (b cycles)
        //   Fma(a3, b1, c3) -> r13
        let prelude = vec![
            lit(0), // a0
            lit(1), // a1
            lit(2), // a2
            lit(3), // a3
            lit(4), // b0
            lit(5), // b1
            lit(6), // c0
            lit(7), // c1
            lit(8), // c2
            lit(9), // c3
        ];
        let fmas = vec![
            fma(0, 4, 6, 10),
            fma(1, 5, 7, 11),
            fma(2, 4, 8, 12),
            fma(3, 5, 9, 13),
        ];
        let mut ops = prelude;
        ops.extend(fmas);
        let desc = empty_desc(KernelBody {
            ops,
            child_bodies: vec![],
            literals: vec![],
        });
        let out = matmul_promote(&desc);
        // The 10 prelude Literals stay; the 4 Fmas collapse to 1 MatrixMma.
        assert_eq!(out.body.ops.len(), 11);
        let mma = out.body.ops.last().unwrap();
        assert!(matches!(mma.kind, KernelOpKind::MatrixMma { .. }));
        // operand layout: [a0..a3, b0..b1, c0..c3]
        assert_eq!(mma.operands, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
    }

    #[test]
    fn promotion_keeps_accumulator_consumers_wired() {
        // A real tiled matmul consumes the 4 accumulator outputs downstream
        // (store to the output tile / loop carry). Collapsing the 4 Fmas to a
        // MatrixMma allocates a fresh result block, so the dropped Fma result
        // ids must not be left dangling in their consumers: every SSA id still
        // referenced after promotion must have a producer.
        let prelude = (0..10).map(lit).collect::<Vec<_>>();
        let fmas = vec![
            fma(0, 4, 6, 10),
            fma(1, 5, 7, 11),
            fma(2, 4, 8, 12),
            fma(3, 5, 9, 13),
        ];
        // Downstream consumers: store each accumulator result to a buffer.
        let stores = vec![
            store_global(0, 0, 10),
            store_global(0, 0, 11),
            store_global(0, 0, 12),
            store_global(0, 0, 13),
        ];
        let mut ops = prelude;
        ops.extend(fmas);
        ops.extend(stores);
        let desc = empty_desc(KernelBody {
            ops,
            child_bodies: vec![],
            literals: vec![],
        });
        let out = matmul_promote(&desc);

        // The promotion must have fired: 4 Fmas collapsed into 1 MatrixMma.
        assert_eq!(
            out.body
                .ops
                .iter()
                .filter(|op| matches!(op.kind, KernelOpKind::Fma))
                .count(),
            0,
            "all four Fmas collapse"
        );
        assert_eq!(
            out.body
                .ops
                .iter()
                .filter(|op| matches!(op.kind, KernelOpKind::MatrixMma { .. }))
                .count(),
            1,
            "into exactly one MatrixMma"
        );

        // No-dangling contract: every value still consumed by a StoreGlobal
        // must be produced by some op in the body.
        let produced: Vec<u32> = out.body.ops.iter().flat_map(|op| op.result_ids()).collect();
        let store_values: Vec<u32> = out
            .body
            .ops
            .iter()
            .filter(|op| matches!(op.kind, KernelOpKind::StoreGlobal))
            .map(|op| op.operands[2])
            .collect();
        assert_eq!(store_values.len(), 4, "all four stores survive");
        for value in &store_values {
            assert!(
                produced.contains(value),
                "store value {value} dangles: promotion dropped the Fma that \
                 produced it without re-pointing the consumer to the MatrixMma"
            );
        }
    }

    #[test]
    fn fma_chain_with_repeated_a_does_not_promote() {
        // a id repeats  -  not a valid M16N8K16 fragment.
        let prelude = (0..10).map(lit).collect::<Vec<_>>();
        let fmas = vec![
            fma(0, 4, 6, 10),
            fma(0, 5, 7, 11), // repeated a
            fma(2, 4, 8, 12),
            fma(3, 5, 9, 13),
        ];
        let mut ops = prelude;
        ops.extend(fmas);
        let desc = empty_desc(KernelBody {
            ops,
            child_bodies: vec![],
            literals: vec![],
        });
        let out = matmul_promote(&desc);
        // No promotion: original 14 ops remain.
        assert_eq!(out.body.ops.len(), 14);
        assert!(!out
            .body
            .ops
            .iter()
            .any(|op| matches!(op.kind, KernelOpKind::MatrixMma { .. })));
    }

    #[test]
    fn fma_chain_with_three_unique_b_does_not_promote() {
        // 3 unique b ids violates the M16N8K16 contract (B fragment is 2 lanes).
        let prelude = (0..11).map(lit).collect::<Vec<_>>();
        let fmas = vec![
            fma(0, 4, 6, 11),
            fma(1, 5, 7, 12),
            fma(2, 10, 8, 13), // third unique b
            fma(3, 5, 9, 14),
        ];
        let mut ops = prelude;
        ops.extend(fmas);
        let desc = empty_desc(KernelBody {
            ops,
            child_bodies: vec![],
            literals: vec![],
        });
        let out = matmul_promote(&desc);
        assert!(!out
            .body
            .ops
            .iter()
            .any(|op| matches!(op.kind, KernelOpKind::MatrixMma { .. })));
    }

    #[test]
    fn matmul_promote_recurses_into_child_bodies() {
        let prelude = (0..10).map(lit).collect::<Vec<_>>();
        let fmas = vec![
            fma(0, 4, 6, 10),
            fma(1, 5, 7, 11),
            fma(2, 4, 8, 12),
            fma(3, 5, 9, 13),
        ];
        let mut child_ops = prelude;
        child_ops.extend(fmas);
        let child = KernelBody {
            ops: child_ops,
            child_bodies: vec![],
            literals: vec![],
        };
        let desc = empty_desc(KernelBody {
            ops: vec![lit(0)],
            child_bodies: vec![child],
            literals: vec![],
        });
        let out = matmul_promote(&desc);
        let promoted_child = &out.body.child_bodies[0];
        assert!(promoted_child
            .ops
            .iter()
            .any(|op| matches!(op.kind, KernelOpKind::MatrixMma { .. })));
    }

    #[test]
    fn matmul_tile_loop_inference_finds_fma_fragment_inside_structured_loop() {
        let prelude = (0..10).map(lit).collect::<Vec<_>>();
        let fmas = vec![
            fma(0, 4, 6, 10),
            fma(1, 5, 7, 11),
            fma(2, 4, 8, 12),
            fma(3, 5, 9, 13),
        ];

        let mut child_ops = prelude;
        child_ops.extend(fmas);
        let child = KernelBody {
            ops: child_ops,
            child_bodies: vec![],
            literals: vec![],
        };
        let desc = empty_desc(KernelBody {
            ops: vec![
                lit(20),
                lit(21),
                KernelOp {
                    kind: KernelOpKind::StructuredForLoop {
                        loop_var: "k".into(),
                    },
                    operands: vec![20, 21, 0],
                    result: None,
                },
            ],
            child_bodies: vec![child],
            literals: vec![],
        });
        let plans = infer_matmul_tile_loops(&desc);
        assert_eq!(
            plans,
            vec![MatmulTileLoopPlan {
                body_path: vec![],
                loop_op_index: 2,
                child_body_index: 0,
                fma_start_index: 10,
            }]
        );
    }

    #[test]
    fn matmul_promote_is_idempotent() {
        let prelude = (0..10).map(lit).collect::<Vec<_>>();
        let fmas = vec![
            fma(0, 4, 6, 10),
            fma(1, 5, 7, 11),
            fma(2, 4, 8, 12),
            fma(3, 5, 9, 13),
        ];
        let mut ops = prelude;
        ops.extend(fmas);
        let desc = empty_desc(KernelBody {
            ops,
            child_bodies: vec![],
            literals: vec![],
        });
        let once = matmul_promote(&desc);
        let twice = matmul_promote(&once);
        assert_eq!(once.body.ops.len(), twice.body.ops.len());
    }
}