weirflow 0.1.0

GPU-first dataflow analysis primitives for Vyre and Santh compiler pipelines.
Documentation
#[cfg(any(test, feature = "cpu-parity"))]
use super::dominator_step;
#[cfg(any(test, feature = "cpu-parity"))]
use crate::dispatch_input_cache::OwnedDispatchInputs;

#[cfg(any(test, feature = "cpu-parity"))]
pub(crate) fn compute_via<F>(
    dispatch: &F,
    node_count: u32,
    succ: &[Vec<u32>],
    entry: u32,
    max_iterations: 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();
    let mut outputs = Vec::new();
    let out = compute_borrowed_into_via(
        &|program, inputs, grid, outputs| {
            let result = owned_inputs.dispatch(program, inputs, grid, dispatch)?;
            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
            Ok(())
        },
        node_count,
        succ,
        entry,
        max_iterations,
        &mut outputs,
    )?;
    Ok(out)
}

#[cfg(any(test, feature = "cpu-parity"))]
pub(crate) fn compute_borrowed_into_via<F>(
    dispatch: &F,
    node_count: u32,
    succ: &[Vec<u32>],
    entry: u32,
    max_iterations: 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 nodes = usize::try_from(node_count).map_err(|_| {
        format!("dominators compute_via node_count={node_count} cannot fit usize. Fix: shard the CFG before GPU dispatch.")
    })?;
    if node_count == 0 {
        if !succ.is_empty() {
            return Err(format!(
                "dominators compute_via received {} successor rows for node_count=0. Fix: pass an empty successor table for empty CFGs.",
                succ.len()
            ));
        }
        return Ok(Vec::new());
    }
    crate::dispatch_decode::require_positive_iterations("dominators compute_via", max_iterations)?;
    if succ.len() != nodes {
        return Err(format!(
            "dominators compute_via received {} successor rows for node_count={node_count}. Fix: pass exactly one successor row per CFG node.",
            succ.len()
        ));
    }
    if entry >= node_count {
        return Err(format!(
            "dominators compute_via received entry={entry} outside node_count={node_count}. Fix: pass a valid CFG entry node."
        ));
    }
    let words = bitset_words(node_count)?;
    let matrix_words = nodes.checked_mul(words).ok_or_else(|| {
        format!("dominators compute_via node_count={node_count} overflows dominator matrix word count. Fix: shard the CFG before GPU dispatch.")
    })?;
    let (pred_offsets, pred_targets) = build_predecessor_csr(node_count, succ)?;
    let edge_count = u32::try_from(pred_targets.len()).map_err(|_| {
        format!("dominators compute_via edge count {} exceeds u32 dispatch metadata. Fix: shard the CFG before GPU dispatch.", pred_targets.len())
    })?;
    let program = dominator_step(
        node_count,
        edge_count,
        entry,
        "pred_offsets",
        "pred_targets",
        "dom_current",
        "dom_next",
    )?;
    let mut pred_offset_bytes = Vec::new();
    crate::dispatch_decode::try_pack_u32_into(
        &pred_offsets,
        &mut pred_offset_bytes,
        "dominators predecessor offset byte",
    )?;
    let mut pred_target_bytes = Vec::new();
    crate::dispatch_decode::try_pack_u32_into(
        &pred_targets,
        &mut pred_target_bytes,
        "dominators predecessor target byte",
    )?;
    let mut current = initial_matrix(node_count, entry, words, matrix_words)?;
    let mut next = Vec::new();
    crate::dispatch_decode::try_write_zero_words(
        &mut next,
        matrix_words,
        "dominators next matrix word",
    )?;
    let matrix_byte_len = matrix_words
        .checked_mul(std::mem::size_of::<u32>())
        .ok_or_else(|| {
            "dominators compute_via matrix byte length overflows host usize. Fix: shard the CFG before GPU dispatch."
                .to_string()
        })?;
    let mut current_bytes =
        crate::staging_reserve::reserved_vec(matrix_byte_len, "dominators current matrix byte")?;
    let mut next_bytes = Vec::new();
    crate::dispatch_decode::try_write_zero_bytes(
        &mut next_bytes,
        matrix_byte_len,
        "dominators next matrix byte",
    )?;
    let grid_words = u32::try_from(matrix_words).map_err(|_| {
        format!("dominators compute_via matrix word count {matrix_words} exceeds u32 dispatch metadata. Fix: shard the CFG before GPU dispatch.")
    })?;
    for _ in 0..max_iterations {
        crate::dispatch_decode::try_pack_u32_into(&current, &mut current_bytes, "weir u32 byte staging")?;
        let inputs = [
            pred_offset_bytes.as_slice(),
            pred_target_bytes.as_slice(),
            current_bytes.as_slice(),
            next_bytes.as_slice(),
        ];
        dispatch(&program, &inputs, Some([grid_words, 1, 1]), outputs)?;
        crate::dispatch_decode::unpack_only_exact_u32_into(
            &outputs,
            "dominators compute_via",
            "matrix",
            matrix_words,
            &mut next,
        )?;
        require_matrix_tail_clear(&next, words, node_count)?;
        if next == current {
            return Ok(current);
        }
        std::mem::swap(&mut current, &mut next);
        next.fill(0);
    }
    Err(format!(
        "dominators compute_via did not converge within max_iterations={max_iterations}. Fix: increase the fixpoint budget or split the CFG before dispatch."
    ))
}

#[cfg(any(test, feature = "cpu-parity"))]
fn bitset_words(node_count: u32) -> Result<usize, String> {
    usize::try_from(crate::graph_layout::LinearDomain::new(node_count).bitset_words()).map_err(
        |_| {
            format!(
                "dominators compute_via node_count={node_count} overflows host bitset word count. Fix: shard the CFG before GPU dispatch."
            )
        },
    )
}

#[cfg(any(test, feature = "cpu-parity"))]
fn initial_matrix(
    node_count: u32,
    entry: u32,
    words: usize,
    matrix_words: usize,
) -> Result<Vec<u32>, String> {
    let nodes = usize::try_from(node_count).map_err(|source| {
        format!("dominators initial matrix node_count cannot fit host usize: {source}. Fix: shard the CFG before GPU dispatch.")
    })?;
    let entry_idx = usize::try_from(entry).map_err(|source| {
        format!("dominators initial matrix entry cannot fit host usize: {source}. Fix: shard the CFG before GPU dispatch.")
    })?;
    let all_mask = all_node_mask(node_count, words)?;
    let mut matrix = Vec::new();
    crate::dispatch_decode::try_write_zero_words(
        &mut matrix,
        matrix_words,
        "dominators initial matrix word",
    )?;
    for node in 0..nodes {
        let row = &mut matrix[node * words..(node + 1) * words];
        if node == entry_idx {
            row[entry_idx / 32] = 1u32 << (entry_idx % 32);
        } else {
            row.copy_from_slice(&all_mask);
        }
    }
    Ok(matrix)
}

#[cfg(any(test, feature = "cpu-parity"))]
fn all_node_mask(node_count: u32, words: usize) -> Result<Vec<u32>, String> {
    let mut mask = Vec::new();
    crate::dispatch_decode::try_write_zero_words(
        &mut mask,
        words,
        "dominators all-node mask word",
    )?;
    let node_count = usize::try_from(node_count).map_err(|source| {
        format!("dominators all-node mask node_count cannot fit host usize: {source}. Fix: shard the CFG before GPU dispatch.")
    })?;
    for node in 0..node_count {
        mask[node / 32] |= 1u32 << (node % 32);
    }
    Ok(mask)
}

#[cfg(any(test, feature = "cpu-parity"))]
fn build_predecessor_csr(
    node_count: u32,
    succ: &[Vec<u32>],
) -> Result<(Vec<u32>, Vec<u32>), String> {
    let nodes = usize::try_from(node_count).map_err(|_| {
        format!("dominators compute_via node_count={node_count} cannot fit usize. Fix: shard the CFG before GPU dispatch.")
    })?;
    let mut counts = crate::staging_reserve::reserved_vec(nodes, "dominators predecessor count")?;
    counts.extend(std::iter::repeat_n(0usize, nodes));
    let mut edge_count = 0usize;
    for (src, targets) in succ.iter().enumerate() {
        edge_count = edge_count.checked_add(targets.len()).ok_or_else(|| {
            "dominators compute_via edge count overflows host usize. Fix: shard the CFG before GPU dispatch.".to_string()
        })?;
        for &dst in targets {
            let dst_idx = usize::try_from(dst).map_err(|_| {
                format!("dominators compute_via destination node {dst} cannot fit usize. Fix: shard the CFG before GPU dispatch.")
            })?;
            if dst_idx >= nodes {
                return Err(format!(
                    "dominators compute_via received edge {src}->{dst} outside node_count={node_count}. Fix: validate CFG edge endpoints before dominator construction."
                ));
            }
            counts[dst_idx] = counts[dst_idx].checked_add(1).ok_or_else(|| {
                format!("dominators compute_via predecessor count for node {dst} overflows host usize. Fix: shard the CFG before GPU dispatch.")
            })?;
        }
    }

    let offset_len = nodes.checked_add(1).ok_or_else(|| {
        "dominators compute_via predecessor offset length overflows host usize. Fix: shard the CFG before GPU dispatch."
            .to_string()
    })?;
    let mut offsets = Vec::new();
    crate::dispatch_decode::try_write_zero_words(
        &mut offsets,
        offset_len,
        "dominators predecessor offset word",
    )?;
    let mut running = 0usize;
    for (node, count) in counts.iter().enumerate() {
        offsets[node] = u32::try_from(running).map_err(|_| {
            format!("dominators compute_via predecessor offset {running} exceeds u32 dispatch metadata. Fix: shard the CFG before GPU dispatch.")
        })?;
        running = running.checked_add(*count).ok_or_else(|| {
            "dominators compute_via predecessor offsets overflow host usize. Fix: shard the CFG before GPU dispatch.".to_string()
        })?;
    }
    offsets[nodes] = u32::try_from(running).map_err(|_| {
        format!("dominators compute_via predecessor edge count {running} exceeds u32 dispatch metadata. Fix: shard the CFG before GPU dispatch.")
    })?;
    let mut cursor = crate::staging_reserve::reserved_vec(nodes, "dominators predecessor cursor")?;
    for &value in &offsets[..nodes] {
        cursor.push(usize::try_from(value).map_err(|_| {
            format!("dominators compute_via predecessor cursor offset {value} cannot fit usize. Fix: shard the CFG before GPU dispatch.")
        })?);
    }
    let mut targets = Vec::new();
    crate::dispatch_decode::try_write_zero_words(
        &mut targets,
        edge_count,
        "dominators predecessor target word",
    )?;
    for (src, succs) in succ.iter().enumerate() {
        let src = u32::try_from(src).map_err(|_| {
            format!("dominators compute_via source node {src} exceeds u32 dispatch metadata. Fix: shard the CFG before GPU dispatch.")
        })?;
        for &dst in succs {
            let dst_idx = usize::try_from(dst).map_err(|_| {
                format!("dominators compute_via destination node {dst} cannot fit usize. Fix: shard the CFG before GPU dispatch.")
            })?;
            let slot = cursor[dst_idx];
            targets[slot] = src;
            cursor[dst_idx] += 1;
        }
    }
    Ok((offsets, targets))
}

#[cfg(test)]
mod source_contract_tests {
    #[test]
    fn dominator_dispatch_source_has_no_resize_based_predecessor_count_staging() {
        let source = include_str!("dispatch.rs");
        assert!(
            !source.contains(concat!("counts", ".resize("))
                && source.contains("counts.extend(std::iter::repeat_n(0usize, nodes))"),
            "Fix: dominator predecessor-count staging must extend after fallible reservation instead of resize-driven growth."
        );
    }
}

#[cfg(any(test, feature = "cpu-parity"))]
fn require_matrix_tail_clear(matrix: &[u32], words: usize, node_count: u32) -> Result<(), String> {
    for row in matrix.chunks(words) {
        crate::dispatch_decode::require_bitset_tail_clear(
            "dominators compute_via output",
            row,
            node_count,
        )?;
    }
    Ok(())
}