weirflow 0.1.0

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

use crate::ifds_frontier_decode::ifds_encoded_frontier_nodes_into;
use crate::ifds_gpu::ifds_gpu_step;
use crate::ifds_resident_types::PreparedIfdsCsr;
use crate::ifds_seed_frontiers::seed_frontier_words_into;
use vyre_foundation::ir::Program;

pub(super) fn solve_prepared_borrowed_with_parts_via<F>(
    dispatch: &F,
    prepared: &PreparedIfdsCsr,
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
    frontier: &mut Vec<u32>,
    frontier_b: &mut Vec<u8>,
    next: &mut Vec<u32>,
    outputs: &mut Vec<Vec<u8>>,
) -> Result<Vec<u32>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let mut result = crate::staging_reserve::reserved_vec(
        crate::dispatch_decode::u32_to_usize(
            prepared.node_count,
            "weir IFDS prepared solve node count",
        )?,
        "IFDS borrowed solve result node",
    )?;
    solve_prepared_borrowed_with_parts_into_via(
        dispatch,
        prepared,
        seed_facts,
        max_iterations,
        frontier,
        frontier_b,
        next,
        outputs,
        &mut result,
    )?;
    Ok(result)
}

pub(super) fn solve_prepared_borrowed_with_parts_into_via<F>(
    dispatch: &F,
    prepared: &PreparedIfdsCsr,
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
    frontier: &mut Vec<u32>,
    frontier_b: &mut Vec<u8>,
    next: &mut Vec<u32>,
    outputs: &mut Vec<Vec<u8>>,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    crate::dispatch_decode::require_positive_iterations(
        "weir IFDS prepared GPU solve",
        max_iterations,
    )?;
    let has_frontier = seed_frontier_words_into(
        prepared.shape,
        prepared.node_count,
        prepared.words,
        seed_facts,
        frontier,
    )?;
    if !has_frontier {
        result.clear();
        return Ok(());
    }
    let program = ifds_gpu_step(prepared.shape, "fin", "fout")?;
    for _ in 0..max_iterations {
        crate::dispatch_decode::try_pack_u32_into(frontier, frontier_b, "weir u32 byte staging")?;
        let inputs: [&[u8]; 7] = [
            &prepared.pg_nodes_b,
            &prepared.row_ptr_b,
            &prepared.col_idx_b,
            &prepared.pg_edge_kind_mask_b,
            &prepared.pg_node_tags_b,
            frontier_b.as_slice(),
            frontier_b.as_slice(),
        ];
        dispatch(
            &program,
            &inputs,
            Some([prepared.node_count.max(1), 1, 1]),
            outputs,
        )?;
        crate::dispatch_decode::unpack_only_exact_u32_into(
            outputs,
            "ifds solve_via",
            "frontier",
            prepared.words,
            next,
        )?;
        crate::dispatch_decode::require_bitset_tail_clear(
            "ifds solve_via output",
            next,
            prepared.node_count,
        )?;
        if next.as_slice() == frontier.as_slice() {
            result.clear();
            ifds_encoded_frontier_nodes_into(next, prepared.node_count, prepared.shape, result)?;
            return Ok(());
        }
        std::mem::swap(frontier, next);
    }
    Err(format!(
        "IFDS solve_via did not converge within {max_iterations} iterations"
    ))
}