vyre-lower 0.4.1

//! Tail-mask store predication for power-of-two dispatch coercion.
//!
//! `vyre-driver::program_walks::coerce_to_pow2_with_tail_mask` decides
//! when a launch may be rounded up to a uniform power-of-two element
//! count. This descriptor rewrite is the lowering-side safety half: it
//! wraps global/shared stores whose index is a global lane id in
//! `if index < logical_element_count`.

use vyre_foundation::ir::BinOp;

use crate::{KernelBody, KernelDescriptor, KernelOp, KernelOpKind, LiteralValue};

/// Predicate eligible stores with `index < logical_element_count`.
///
/// This pass is intentionally opt-in: callers run it only when they
/// actually round the dispatch element count above the logical output
/// element count. If `logical_element_count == 0`, the descriptor is
/// returned unchanged because dispatching zero logical elements should
/// be handled by the launch planner.
#[must_use]
pub fn apply_tail_mask(desc: &KernelDescriptor, logical_element_count: u32) -> KernelDescriptor {
    if logical_element_count == 0 {
        return desc.clone();
    }
    let mut out = desc.clone();
    let mut next_result = next_result_id(&out.body);
    apply_to_body(&mut out.body, logical_element_count, &mut next_result);
    out
}

fn apply_to_body(body: &mut KernelBody, logical_element_count: u32, next_result: &mut u32) {
    for child in &mut body.child_bodies {
        apply_to_body(child, logical_element_count, next_result);
    }

    let original_ops = std::mem::take(&mut body.ops);
    let mut rewritten = Vec::with_capacity(original_ops.len());
    for op in original_ops {
        if let Some(index_id) = tail_mask_index(&op, &rewritten) {
            let limit_lit_idx = body.literals.len() as u32;
            body.literals.push(LiteralValue::U32(logical_element_count));

            let limit_id = fresh_result(next_result);
            rewritten.push(KernelOp {
                kind: KernelOpKind::Literal,
                operands: vec![limit_lit_idx],
                result: Some(limit_id),
            });

            let cond_id = fresh_result(next_result);
            rewritten.push(KernelOp {
                kind: KernelOpKind::BinOpKind(BinOp::Lt),
                operands: vec![index_id, limit_id],
                result: Some(cond_id),
            });

            let child_idx = body.child_bodies.len() as u32;
            body.child_bodies.push(KernelBody {
                ops: vec![op],
                child_bodies: Vec::new(),
                literals: Vec::new(),
            });
            rewritten.push(KernelOp {
                kind: KernelOpKind::StructuredIfThen,
                operands: vec![cond_id, child_idx],
                result: None,
            });
        } else {
            rewritten.push(op);
        }
    }
    body.ops = rewritten;
}

fn tail_mask_index(op: &KernelOp, earlier_ops: &[KernelOp]) -> Option<u32> {
    if !matches!(
        op.kind,
        KernelOpKind::StoreGlobal | KernelOpKind::StoreShared
    ) {
        return None;
    }
    let index_id = op.operands.get(1).copied()?;
    let producer = earlier_ops
        .iter()
        .find(|candidate| candidate.result == Some(index_id))?;
    if matches!(producer.kind, KernelOpKind::GlobalInvocationId)
        && producer.operands.first().copied().unwrap_or(0) == 0
    {
        Some(index_id)
    } else {
        None
    }
}

fn next_result_id(body: &KernelBody) -> u32 {
    fn walk(body: &KernelBody, max_seen: &mut u32) {
        for op in &body.ops {
            for result in op.result_ids() {
                *max_seen = (*max_seen).max(result.saturating_add(1));
            }
        }
        for child in &body.child_bodies {
            walk(child, max_seen);
        }
    }
    let mut next = 0;
    walk(body, &mut next);
    next
}

fn fresh_result(next: &mut u32) -> u32 {
    let id = *next;
    *next = next.saturating_add(1);
    id
}

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

    fn desc_with_store_index(index_kind: KernelOpKind) -> KernelDescriptor {
        KernelDescriptor {
            id: "tail-mask-test".to_string(),
            bindings: BindingLayout {
                slots: vec![BindingSlot {
                    slot: 0,
                    element_type: DataType::U32,
                    element_count: Some(100),
                    memory_class: MemoryClass::Global,
                    visibility: BindingVisibility::WriteOnly,
                    name: "out".to_string(),
                }],
            },
            dispatch: Dispatch::new(64, 1, 1),
            body: KernelBody {
                ops: vec![
                    KernelOp {
                        kind: index_kind,
                        operands: vec![0],
                        result: Some(0),
                    },
                    KernelOp {
                        kind: KernelOpKind::Literal,
                        operands: vec![0],
                        result: Some(1),
                    },
                    KernelOp {
                        kind: KernelOpKind::StoreGlobal,
                        operands: vec![0, 0, 1],
                        result: None,
                    },
                ],
                child_bodies: Vec::new(),
                literals: vec![LiteralValue::U32(7)],
            },
        }
    }

    #[test]
    fn masks_global_lane_store() {
        let desc = desc_with_store_index(KernelOpKind::GlobalInvocationId);
        let out = apply_tail_mask(&desc, 100);
        crate::verify::verify(&out)
            .expect("Fix: tail-mask rewrite must preserve descriptor invariants");
        assert_eq!(out.body.child_bodies.len(), 1);
        assert!(matches!(
            out.body.ops.last().map(|op| &op.kind),
            Some(KernelOpKind::StructuredIfThen)
        ));
        assert!(matches!(
            out.body.ops[out.body.ops.len() - 2].kind,
            KernelOpKind::BinOpKind(BinOp::Lt)
        ));
        assert_eq!(out.body.literals.last(), Some(&LiteralValue::U32(100)));
        assert!(matches!(
            out.body.child_bodies[0].ops[0].kind,
            KernelOpKind::StoreGlobal
        ));
    }

    #[test]
    fn leaves_non_global_lane_store_unmasked() {
        let desc = desc_with_store_index(KernelOpKind::LocalInvocationId);
        let out = apply_tail_mask(&desc, 100);
        assert_eq!(out, desc);
    }

    #[test]
    fn zero_logical_count_is_noop() {
        let desc = desc_with_store_index(KernelOpKind::GlobalInvocationId);
        let out = apply_tail_mask(&desc, 0);
        assert_eq!(out, desc);
    }

    #[test]
    fn generated_result_ids_follow_existing_ids() {
        let desc = desc_with_store_index(KernelOpKind::GlobalInvocationId);
        let out = apply_tail_mask(&desc, 100);
        let lit = &out.body.ops[out.body.ops.len() - 3];
        let cond = &out.body.ops[out.body.ops.len() - 2];
        assert_eq!(lit.result, Some(2));
        assert_eq!(cond.result, Some(3));
    }
}