weirflow 0.1.0

GPU-first dataflow analysis primitives for Vyre and Santh compiler pipelines.
Documentation
use std::cell::RefCell;

use vyre::ir::Program;

const DEFAULT_OWNED_INPUT_SLOT_CAPACITY: usize = 8;
pub(crate) const FIXPOINT_FRONTIER_INPUT_SLOT: usize = 5;
pub(crate) const FIXPOINT_FRONTIER_INPUT_SLOTS: &[usize] = &[FIXPOINT_FRONTIER_INPUT_SLOT];

pub(crate) struct OwnedDispatchInputs {
    buffers: RefCell<Vec<Vec<u8>>>,
}

impl OwnedDispatchInputs {
    pub(crate) fn new() -> Self {
        Self {
            buffers: RefCell::new(Vec::new()),
        }
    }

    pub(crate) fn dispatch<F>(
        &self,
        program: &Program,
        inputs: &[&[u8]],
        grid: Option<[u32; 3]>,
        dispatch: &F,
    ) -> Result<Vec<Vec<u8>>, String>
    where
        F: Fn(&Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
    {
        self.with_owned(inputs, |owned| dispatch(program, owned, grid))
    }

    pub(crate) fn dispatch_refreshing_slots<F>(
        &self,
        program: &Program,
        inputs: &[&[u8]],
        grid: Option<[u32; 3]>,
        refresh_slots: &[usize],
        dispatch: &F,
    ) -> Result<Vec<Vec<u8>>, String>
    where
        F: Fn(&Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
    {
        self.with_owned_refreshing_slots(inputs, refresh_slots, |owned| {
            dispatch(program, owned, grid)
        })
    }

    pub(crate) fn dispatch_refreshing_frontier<F>(
        &self,
        program: &Program,
        inputs: &[&[u8]],
        grid: Option<[u32; 3]>,
        dispatch: &F,
    ) -> Result<Vec<Vec<u8>>, String>
    where
        F: Fn(&Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
    {
        self.dispatch_refreshing_slots(
            program,
            inputs,
            grid,
            FIXPOINT_FRONTIER_INPUT_SLOTS,
            dispatch,
        )
    }

    pub(crate) fn with_owned<R, F>(&self, inputs: &[&[u8]], call: F) -> Result<R, String>
    where
        F: FnOnce(&[Vec<u8>]) -> Result<R, String>,
    {
        let mut buffers = self.buffers.borrow_mut();
        refresh_all(&mut buffers, inputs)?;
        call(&buffers)
    }

    pub(crate) fn with_owned_refreshing_slots<R, F>(
        &self,
        inputs: &[&[u8]],
        refresh_slots: &[usize],
        call: F,
    ) -> Result<R, String>
    where
        F: FnOnce(&[Vec<u8>]) -> Result<R, String>,
    {
        let mut buffers = self.buffers.borrow_mut();
        refresh_slots_only(&mut buffers, inputs, refresh_slots)?;
        call(&buffers)
    }
}

fn refresh_all(buffers: &mut Vec<Vec<u8>>, inputs: &[&[u8]]) -> Result<(), String> {
    if buffers.len() < inputs.len() {
        let reserved_slots = inputs.len().max(DEFAULT_OWNED_INPUT_SLOT_CAPACITY);
        crate::staging_reserve::reserve_vec(buffers, reserved_slots, "owned dispatch input slot")?;
        buffers.extend(std::iter::repeat_with(Vec::new).take(inputs.len() - buffers.len()));
    } else if buffers.len() > inputs.len() {
        buffers.truncate(inputs.len());
    }
    for (slot, input) in buffers.iter_mut().zip(inputs.iter().copied()) {
        refresh_slot(slot, input)?;
    }
    Ok(())
}

fn refresh_slots_only(
    buffers: &mut Vec<Vec<u8>>,
    inputs: &[&[u8]],
    refresh_slots: &[usize],
) -> Result<(), String> {
    if buffers.len() != inputs.len() {
        return refresh_all(buffers, inputs);
    }
    for &idx in refresh_slots {
        let Some(slot) = buffers.get_mut(idx) else {
            return Err(format!(
                "owned dispatch input refresh slot {idx} is out of bounds for {} input(s). Fix: keep borrowed and owned ABI adapters aligned.",
                buffers.len()
            ));
        };
        let Some(input) = inputs.get(idx) else {
            return Err(format!(
                "borrowed dispatch input refresh slot {idx} is out of bounds for {} input(s). Fix: keep borrowed and owned ABI adapters aligned.",
                inputs.len()
            ));
        };
        refresh_slot(slot, input)?;
    }
    Ok(())
}

#[inline]
fn refresh_slot(slot: &mut Vec<u8>, input: &[u8]) -> Result<(), String> {
    if slot.len() == input.len() {
        slot.copy_from_slice(input);
        return Ok(());
    }
    crate::staging_reserve::reserve_vec(slot, input.len(), "owned dispatch input byte")?;
    slot.clear();
    slot.extend_from_slice(input);
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::OwnedDispatchInputs;

    #[test]
    fn full_refresh_reuses_outer_and_inner_slots() {
        let cache = OwnedDispatchInputs::new();
        let first_a = [1_u8, 2, 3, 4];
        let first_b = [5_u8, 6, 7, 8];
        let second_a = [9_u8, 10];
        let second_b = [11_u8, 12];
        let mut outer_ptr = std::ptr::null();
        let mut first_inner_ptr = std::ptr::null();
        let mut second_inner_ptr = std::ptr::null();

        cache
            .with_owned(&[&first_a, &first_b], |owned| {
                outer_ptr = owned.as_ptr();
                first_inner_ptr = owned[0].as_ptr();
                second_inner_ptr = owned[1].as_ptr();
                Ok(())
            })
            .unwrap();

        cache
            .with_owned(&[&second_a, &second_b], |owned| {
                assert_eq!(owned.as_ptr(), outer_ptr);
                assert_eq!(owned[0].as_ptr(), first_inner_ptr);
                assert_eq!(owned[1].as_ptr(), second_inner_ptr);
                assert_eq!(owned[0], second_a);
                assert_eq!(owned[1], second_b);
                Ok(())
            })
            .unwrap();
    }

    #[test]
    fn full_refresh_growing_arity_preserves_existing_slots() {
        let cache = OwnedDispatchInputs::new();
        let first_a = [1_u8, 2, 3, 4];
        let first_b = [5_u8, 6, 7, 8];
        let second_a = [9_u8, 10, 11, 12];
        let second_b = [13_u8, 14, 15, 16];
        let second_c = [17_u8, 18, 19, 20];
        let mut outer_ptr = std::ptr::null();
        let mut first_inner_ptr = std::ptr::null();
        let mut second_inner_ptr = std::ptr::null();

        cache
            .with_owned(&[&first_a, &first_b], |owned| {
                outer_ptr = owned.as_ptr();
                first_inner_ptr = owned[0].as_ptr();
                second_inner_ptr = owned[1].as_ptr();
                Ok(())
            })
            .unwrap();

        cache
            .with_owned(&[&second_a, &second_b, &second_c], |owned| {
                assert_eq!(owned.as_ptr(), outer_ptr);
                assert_eq!(owned[0].as_ptr(), first_inner_ptr);
                assert_eq!(owned[1].as_ptr(), second_inner_ptr);
                assert_eq!(owned[0], second_a);
                assert_eq!(owned[1], second_b);
                assert_eq!(owned[2], second_c);
                Ok(())
            })
            .unwrap();
    }

    #[test]
    fn hot_slot_refresh_preserves_invariant_inputs() {
        let cache = OwnedDispatchInputs::new();
        let invariant = [1_u8, 2, 3, 4];
        let mutable_a = [5_u8, 6, 7, 8];
        let mutable_b = [9_u8, 10, 11, 12];
        let updated_a = [13_u8, 14];
        let updated_b = [15_u8, 16];
        let mut invariant_ptr = std::ptr::null();
        let mut first_hot_ptr = std::ptr::null();
        let mut second_hot_ptr = std::ptr::null();

        cache
            .with_owned_refreshing_slots(&[&invariant, &mutable_a, &mutable_b], &[1, 2], |owned| {
                invariant_ptr = owned[0].as_ptr();
                first_hot_ptr = owned[1].as_ptr();
                second_hot_ptr = owned[2].as_ptr();
                Ok(())
            })
            .unwrap();

        cache
            .with_owned_refreshing_slots(
                &[&[99_u8][..], &updated_a, &updated_b],
                &[1, 2],
                |owned| {
                    assert_eq!(owned[0].as_ptr(), invariant_ptr);
                    assert_eq!(owned[1].as_ptr(), first_hot_ptr);
                    assert_eq!(owned[2].as_ptr(), second_hot_ptr);
                    assert_eq!(owned[0], invariant);
                    assert_eq!(owned[1], updated_a);
                    assert_eq!(owned[2], updated_b);
                    Ok(())
                },
            )
            .unwrap();
    }

    #[test]
    fn hot_slot_refresh_rejects_bad_slot_index() {
        let cache = OwnedDispatchInputs::new();
        let input = [1_u8, 2, 3, 4];
        cache
            .with_owned_refreshing_slots(&[&input], &[0], |_| Ok(()))
            .unwrap();
        let err = cache
            .with_owned_refreshing_slots(&[&input], &[1], |_| Ok(()))
            .unwrap_err();
        assert!(err.contains("out of bounds"));
        assert!(err.contains("Fix:"));
    }

    #[test]
    fn owned_dispatch_input_cache_source_has_no_resize_based_slot_growth() {
        let source = include_str!("dispatch_input_cache.rs");
        let production = source
            .split("#[cfg(test)]")
            .next()
            .expect("dispatch input cache production source must precede tests");
        assert!(
            !production.contains(".resize_with(")
                && production.contains("buffers.extend(std::iter::repeat_with(Vec::new)"),
            "Fix: owned dispatch input slot staging must extend after fallible reserve instead of resize-driven growth."
        );
    }
}