weirflow 0.1.0

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

use super::{summarize_function_with_count, SummarizeFunctionScratch};
use crate::dispatch_input_cache::OwnedDispatchInputs;
use crate::soundness::{DataflowEvidence, Soundness};

pub fn summarize_function_evidence_via<F>(
    dispatch: &F,
    fn_ast_in: &[u32],
    callgraph_in: &[u32],
    cached_summary_in: &[u32],
    bit_count: u32,
) -> Result<DataflowEvidence<Vec<u32>>, String>
where
    F: Fn(&vyre::ir::Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    summarize_function_via(
        dispatch,
        fn_ast_in,
        callgraph_in,
        cached_summary_in,
        bit_count,
    )
    .map(|value| DataflowEvidence::new(value, Soundness::MayOver))
}

pub fn summarize_function_via<F>(
    dispatch: &F,
    fn_ast_in: &[u32],
    callgraph_in: &[u32],
    cached_summary_in: &[u32],
    bit_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();
    summarize_function_borrowed_via(
        &|program, inputs, grid| owned_inputs.dispatch(program, inputs, grid, dispatch),
        fn_ast_in,
        callgraph_in,
        cached_summary_in,
        bit_count,
    )
}

/// GPU function-summary projection using borrowed dispatch inputs.
///
/// This avoids constructing an owned input vector for backends that already
/// accept borrowed byte slices.
pub fn summarize_function_borrowed_via<F>(
    dispatch: &F,
    fn_ast_in: &[u32],
    callgraph_in: &[u32],
    cached_summary_in: &[u32],
    bit_count: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let mut scratch = SummarizeFunctionScratch::default();
    summarize_function_borrowed_with_scratch_via(
        dispatch,
        fn_ast_in,
        callgraph_in,
        cached_summary_in,
        bit_count,
        &mut scratch,
    )
}

/// GPU function-summary projection using borrowed dispatch inputs and
/// caller-owned staging/output scratch.
pub fn summarize_function_borrowed_with_scratch_via<F>(
    dispatch: &F,
    fn_ast_in: &[u32],
    callgraph_in: &[u32],
    cached_summary_in: &[u32],
    bit_count: u32,
    scratch: &mut SummarizeFunctionScratch,
) -> 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("summarize_function result", bit_count)?
            .max(1);
    let mut result = crate::staging_reserve::reserved_vec(words, "summary result word")?;
    summarize_function_borrowed_into_result_with_scratch_via(
        dispatch,
        fn_ast_in,
        callgraph_in,
        cached_summary_in,
        bit_count,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// GPU function-summary projection using caller-owned staging/output scratch
/// and caller-owned decoded result storage.
pub fn summarize_function_borrowed_into_result_with_scratch_via<F>(
    dispatch: &F,
    fn_ast_in: &[u32],
    callgraph_in: &[u32],
    cached_summary_in: &[u32],
    bit_count: u32,
    scratch: &mut SummarizeFunctionScratch,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    summarize_function_staged_into_result_via(
        &|program, inputs, grid, outputs| {
            let result = dispatch(program, inputs, grid)?;
            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
            Ok(())
        },
        fn_ast_in,
        callgraph_in,
        cached_summary_in,
        bit_count,
        &mut scratch.outputs,
        &mut scratch.ast_bytes,
        &mut scratch.callgraph_bytes,
        &mut scratch.cached_summary_bytes,
        &mut scratch.out_bytes,
        result,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "summarize_function output",
        result,
        bit_count,
    )
}

/// GPU function-summary projection using borrowed dispatch inputs and
/// caller-owned output storage.
pub fn summarize_function_borrowed_into_via<F>(
    dispatch: &F,
    fn_ast_in: &[u32],
    callgraph_in: &[u32],
    cached_summary_in: &[u32],
    bit_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 semantic_words = crate::dispatch_decode::bitset_word_capacity(
        "summarize_function semantic width",
        bit_count,
    )?;
    let words = semantic_words.max(1);
    let byte_capacity = words
        .checked_mul(std::mem::size_of::<u32>())
        .ok_or_else(|| {
            "summarize_function byte capacity overflows host usize. Fix: split the summary bitset before GPU dispatch."
                .to_string()
        })?;
    let mut ast_bytes = crate::staging_reserve::reserved_vec(byte_capacity, "summary AST byte")?;
    let mut callgraph_bytes =
        crate::staging_reserve::reserved_vec(byte_capacity, "summary callgraph byte")?;
    let mut cached_summary_bytes =
        crate::staging_reserve::reserved_vec(byte_capacity, "summary cached byte")?;
    let mut out_bytes = crate::staging_reserve::reserved_vec(byte_capacity, "summary output byte")?;
    let mut result = crate::staging_reserve::reserved_vec(words, "summary result word")?;
    summarize_function_staged_into_result_via(
        dispatch,
        fn_ast_in,
        callgraph_in,
        cached_summary_in,
        bit_count,
        outputs,
        &mut ast_bytes,
        &mut callgraph_bytes,
        &mut cached_summary_bytes,
        &mut out_bytes,
        &mut result,
    )?;
    Ok(result)
}

fn summarize_function_staged_into_result_via<F>(
    dispatch: &F,
    fn_ast_in: &[u32],
    callgraph_in: &[u32],
    cached_summary_in: &[u32],
    bit_count: u32,
    outputs: &mut Vec<Vec<u8>>,
    ast_bytes: &mut Vec<u8>,
    callgraph_bytes: &mut Vec<u8>,
    cached_summary_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 semantic_words = crate::dispatch_decode::bitset_word_capacity(
        "summarize_function semantic width",
        bit_count,
    )?;
    let words = semantic_words.max(1);
    let program = summarize_function_with_count("ast", "cg", "cached", "out", bit_count);
    crate::dispatch_decode::pack_exact_u32_slots_into(
        fn_ast_in,
        semantic_words,
        "summarize_function ast",
        ast_bytes,
    )?;
    crate::dispatch_decode::pack_exact_u32_slots_into(
        callgraph_in,
        semantic_words,
        "summarize_function callgraph",
        callgraph_bytes,
    )?;
    crate::dispatch_decode::pack_exact_u32_slots_into(
        cached_summary_in,
        semantic_words,
        "summarize_function cached",
        cached_summary_bytes,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "summarize_function ast",
        fn_ast_in,
        bit_count,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "summarize_function callgraph",
        callgraph_in,
        bit_count,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "summarize_function cached",
        cached_summary_in,
        bit_count,
    )?;
    crate::dispatch_decode::pad_dispatch_min_words(ast_bytes, words)?;
    crate::dispatch_decode::pad_dispatch_min_words(callgraph_bytes, words)?;
    crate::dispatch_decode::pad_dispatch_min_words(cached_summary_bytes, words)?;
    let out_byte_len = words.checked_mul(4).ok_or_else(|| {
        "summarize_function output byte length overflows host usize. Fix: split the summary bitset before GPU dispatch.".to_string()
    })?;
    let dispatch_words = u32::try_from(words)
        .map_err(|_| "summarize_function word count exceeds u32 dispatch metadata. Fix: split the summary bitset before GPU dispatch.".to_string())?
        .max(1);
    crate::dispatch_decode::try_write_zero_bytes(
        out_bytes,
        out_byte_len,
        "summarize_function output zero staging",
    )?;
    let inputs = [
        ast_bytes.as_slice(),
        callgraph_bytes.as_slice(),
        cached_summary_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,
        "summarize_function",
        "summary",
        words,
        result,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "summarize_function output",
        result,
        bit_count,
    )
}