weirflow 0.1.0

GPU-first dataflow analysis primitives for Vyre and Santh compiler pipelines.
Documentation
#![allow(clippy::too_many_arguments)]

use super::{callgraph_build_with_count, CallgraphBuildScratch};
use crate::dispatch_input_cache::OwnedDispatchInputs;

pub fn callgraph_build_via<F>(
    dispatch: &F,
    direct_edges: &[u32],
    indirect_sites: &[u32],
    pts_closure: &[u32],
    node_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();
    callgraph_build_borrowed_via(
        &|program, inputs, grid| owned_inputs.dispatch(program, inputs, grid, dispatch),
        direct_edges,
        indirect_sites,
        pts_closure,
        node_count,
    )
}

/// GPU callgraph projection using borrowed dispatch inputs.
///
/// Prefer this over [`callgraph_build_via`] for backend adapters that support
/// borrowed input slices.
pub fn callgraph_build_borrowed_via<F>(
    dispatch: &F,
    direct_edges: &[u32],
    indirect_sites: &[u32],
    pts_closure: &[u32],
    node_count: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let mut scratch = CallgraphBuildScratch::default();
    callgraph_build_borrowed_with_scratch_via(
        dispatch,
        direct_edges,
        indirect_sites,
        pts_closure,
        node_count,
        &mut scratch,
    )
}

/// GPU callgraph projection using borrowed dispatch inputs and caller-owned
/// staging/output scratch.
pub fn callgraph_build_borrowed_with_scratch_via<F>(
    dispatch: &F,
    direct_edges: &[u32],
    indirect_sites: &[u32],
    pts_closure: &[u32],
    node_count: u32,
    scratch: &mut CallgraphBuildScratch,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let words = crate::dispatch_decode::bitset_word_capacity("callgraph_build result", node_count)?;
    let mut result = crate::staging_reserve::reserved_vec(words, "callgraph result word")?;
    callgraph_build_borrowed_into_result_with_scratch_via(
        dispatch,
        direct_edges,
        indirect_sites,
        pts_closure,
        node_count,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// GPU callgraph projection using caller-owned staging/output scratch and
/// caller-owned decoded result storage.
pub fn callgraph_build_borrowed_into_result_with_scratch_via<F>(
    dispatch: &F,
    direct_edges: &[u32],
    indirect_sites: &[u32],
    pts_closure: &[u32],
    node_count: u32,
    scratch: &mut CallgraphBuildScratch,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    callgraph_build_staged_into_result_via(
        &|program, inputs, grid, outputs| {
            let result = dispatch(program, inputs, grid)?;
            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
            Ok(())
        },
        direct_edges,
        indirect_sites,
        pts_closure,
        node_count,
        &mut scratch.outputs,
        &mut scratch.direct_edges_bytes,
        &mut scratch.indirect_sites_bytes,
        &mut scratch.pts_closure_bytes,
        &mut scratch.out_bytes,
        result,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear("callgraph_build output", result, node_count)
}

/// GPU callgraph projection using borrowed dispatch inputs and caller-owned
/// output storage.
pub fn callgraph_build_borrowed_into_via<F>(
    dispatch: &F,
    direct_edges: &[u32],
    indirect_sites: &[u32],
    pts_closure: &[u32],
    node_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 words =
        crate::dispatch_decode::bitset_word_capacity("callgraph_build semantic width", node_count)?;
    let byte_capacity = words
        .checked_mul(std::mem::size_of::<u32>())
        .ok_or_else(|| {
            "callgraph_build byte capacity overflows host usize. Fix: split the callgraph bitset before GPU dispatch."
                .to_string()
        })?;
    let mut direct_edges_bytes =
        crate::staging_reserve::reserved_vec(byte_capacity, "callgraph direct-edge byte")?;
    let mut indirect_sites_bytes =
        crate::staging_reserve::reserved_vec(byte_capacity, "callgraph indirect-site byte")?;
    let mut pts_closure_bytes =
        crate::staging_reserve::reserved_vec(byte_capacity, "callgraph points-to byte")?;
    let mut out_bytes =
        crate::staging_reserve::reserved_vec(byte_capacity, "callgraph output byte")?;
    let mut result = crate::staging_reserve::reserved_vec(words, "callgraph result word")?;
    callgraph_build_staged_into_result_via(
        dispatch,
        direct_edges,
        indirect_sites,
        pts_closure,
        node_count,
        outputs,
        &mut direct_edges_bytes,
        &mut indirect_sites_bytes,
        &mut pts_closure_bytes,
        &mut out_bytes,
        &mut result,
    )?;
    Ok(result)
}

fn callgraph_build_staged_into_result_via<F>(
    dispatch: &F,
    direct_edges: &[u32],
    indirect_sites: &[u32],
    pts_closure: &[u32],
    node_count: u32,
    outputs: &mut Vec<Vec<u8>>,
    direct_edges_bytes: &mut Vec<u8>,
    indirect_sites_bytes: &mut Vec<u8>,
    pts_closure_bytes: &mut Vec<u8>,
    out_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 words =
        crate::dispatch_decode::bitset_word_capacity("callgraph_build semantic width", node_count)?;
    let program = callgraph_build_with_count("direct", "indirect", "pts", "out", node_count);
    crate::dispatch_decode::pack_exact_u32_slots_into(
        direct_edges,
        words,
        "callgraph_build direct",
        direct_edges_bytes,
    )?;
    crate::dispatch_decode::pack_exact_u32_slots_into(
        indirect_sites,
        words,
        "callgraph_build indirect",
        indirect_sites_bytes,
    )?;
    crate::dispatch_decode::pack_exact_u32_slots_into(
        pts_closure,
        words,
        "callgraph_build pts",
        pts_closure_bytes,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "callgraph_build direct",
        direct_edges,
        node_count,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "callgraph_build indirect",
        indirect_sites,
        node_count,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "callgraph_build pts",
        pts_closure,
        node_count,
    )?;
    let out_byte_len = words.checked_mul(4).ok_or_else(|| {
        "callgraph_build output byte length overflows host usize. Fix: split the callgraph bitset before GPU dispatch.".to_string()
    })?;
    let dispatch_words = u32::try_from(words)
        .map_err(|_| "callgraph_build word count exceeds u32 dispatch metadata. Fix: split the callgraph bitset before GPU dispatch.".to_string())?
        .max(1);
    crate::dispatch_decode::try_write_zero_bytes(
        out_bytes,
        out_byte_len,
        "callgraph_build output zero staging",
    )?;
    let inputs = [
        direct_edges_bytes.as_slice(),
        indirect_sites_bytes.as_slice(),
        pts_closure_bytes.as_slice(),
        out_bytes.as_slice(),
    ];
    dispatch(&program, &inputs, Some([dispatch_words, 1, 1]), outputs)?;
    crate::dispatch_decode::unpack_only_exact_u32_into(
        outputs,
        "callgraph_build",
        "callgraph",
        words,
        result,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear("callgraph_build output", result, node_count)
}