weir/loop_sum/

dispatch.rs

#![allow(clippy::too_many_arguments)]

use super::{try_loop_summarize_with_count, LoopSummarizeScratch};
use crate::dispatch_input_cache::OwnedDispatchInputs;

pub fn loop_summarize_via<F>(
    dispatch: &F,
    cfg_in: &[u32],
    ranges_in: &[u32],
    var_count: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let owned_inputs = OwnedDispatchInputs::new();
    loop_summarize_borrowed_via(
        &|program, inputs, grid| owned_inputs.dispatch(program, inputs, grid, dispatch),
        cfg_in,
        ranges_in,
        var_count,
    )
}

/// GPU loop summarization using borrowed dispatch inputs.
///
/// Prefer this allocation-aware entrypoint when the caller can provide a
/// borrowed-dispatch backend adapter.
pub fn loop_summarize_borrowed_via<F>(
    dispatch: &F,
    cfg_in: &[u32],
    ranges_in: &[u32],
    var_count: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let mut scratch = LoopSummarizeScratch::default();
    loop_summarize_borrowed_with_scratch_via(dispatch, cfg_in, ranges_in, var_count, &mut scratch)
}

/// GPU loop summarization using borrowed dispatch inputs and caller-owned
/// staging/output scratch.
pub fn loop_summarize_borrowed_with_scratch_via<F>(
    dispatch: &F,
    cfg_in: &[u32],
    ranges_in: &[u32],
    var_count: u32,
    scratch: &mut LoopSummarizeScratch,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let mut result = crate::staging_reserve::reserved_vec(
        loop_summary_slots(var_count)?,
        "loop summary result slot",
    )?;
    loop_summarize_borrowed_into_result_with_scratch_via(
        dispatch,
        cfg_in,
        ranges_in,
        var_count,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// GPU loop summarization using caller-owned staging/output scratch and
/// caller-owned decoded result storage.
pub fn loop_summarize_borrowed_into_result_with_scratch_via<F>(
    dispatch: &F,
    cfg_in: &[u32],
    ranges_in: &[u32],
    var_count: u32,
    scratch: &mut LoopSummarizeScratch,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    loop_summarize_staged_into_result_via(
        &|program, inputs, grid, outputs| {
            let result = dispatch(program, inputs, grid)?;
            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
            Ok(())
        },
        cfg_in,
        ranges_in,
        var_count,
        &mut scratch.outputs,
        &mut scratch.cfg_bytes,
        &mut scratch.ranges_bytes,
        &mut scratch.summary_bytes,
        result,
    )
}

/// GPU loop summarization using borrowed dispatch inputs and caller-owned
/// output storage.
pub fn loop_summarize_borrowed_into_via<F>(
    dispatch: &F,
    cfg_in: &[u32],
    ranges_in: &[u32],
    var_count: u32,
    outputs: &mut Vec<Vec<u8>>,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let slots = loop_summary_slots(var_count)?;
    let byte_len = slots.checked_mul(4).ok_or_else(|| {
        "loop_summarize output byte length overflows host usize. Fix: shard loop summarization before GPU dispatch.".to_string()
    })?;
    let mut cfg_bytes = crate::staging_reserve::reserved_vec(byte_len, "loop summary CFG byte")?;
    let mut ranges_bytes =
        crate::staging_reserve::reserved_vec(byte_len, "loop summary range byte")?;
    let mut summary_bytes =
        crate::staging_reserve::reserved_vec(byte_len, "loop summary output byte")?;
    let mut result = crate::staging_reserve::reserved_vec(slots, "loop summary result slot")?;
    loop_summarize_staged_into_result_via(
        dispatch,
        cfg_in,
        ranges_in,
        var_count,
        outputs,
        &mut cfg_bytes,
        &mut ranges_bytes,
        &mut summary_bytes,
        &mut result,
    )?;
    Ok(result)
}

fn loop_summary_slots(var_count: u32) -> Result<usize, String> {
    crate::dispatch_decode::u32_to_usize(var_count, "loop_summarize variable count")?
        .checked_mul(2)
        .ok_or_else(|| {
        format!(
            "loop_summarize var_count={var_count} overflows host slot indexing. Fix: shard loop summarization before GPU dispatch."
        )
    })
}

fn loop_summarize_staged_into_result_via<F>(
    dispatch: &F,
    cfg_in: &[u32],
    ranges_in: &[u32],
    var_count: u32,
    outputs: &mut Vec<Vec<u8>>,
    cfg_bytes: &mut Vec<u8>,
    ranges_bytes: &mut Vec<u8>,
    summary_bytes: &mut Vec<u8>,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let slots = loop_summary_slots(var_count)?;
    let program = try_loop_summarize_with_count("cfg", "ranges", "summary", var_count)?;
    crate::dispatch_decode::pack_exact_u32_slots_into(
        cfg_in,
        slots,
        "loop_summarize cfg",
        cfg_bytes,
    )?;
    crate::dispatch_decode::pack_exact_u32_slots_into(
        ranges_in,
        slots,
        "loop_summarize ranges",
        ranges_bytes,
    )?;
    let summary_byte_len = slots.checked_mul(4).ok_or_else(|| {
        "loop_summarize output byte length overflows host usize. Fix: shard loop summarization before GPU dispatch.".to_string()
    })?;
    crate::dispatch_decode::try_write_zero_bytes(
        summary_bytes,
        summary_byte_len,
        "loop_summarize output zero staging",
    )?;
    let inputs = [
        cfg_bytes.as_slice(),
        ranges_bytes.as_slice(),
        summary_bytes.as_slice(),
    ];
    dispatch(&program, &inputs, Some([var_count.max(1), 1, 1]), outputs)?;
    crate::dispatch_decode::unpack_only_exact_u32_into(
        outputs,
        "loop_summarize",
        "summary",
        slots,
        result,
    )
}