weirflow 0.1.0

GPU-first dataflow analysis primitives for Vyre and Santh compiler pipelines.
Documentation
//! Reusable backend-resident frontier resources for fixed-point hot loops.

/// Reusable resident frontier resources for fixed-point hot loops.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct FixedPointResidentFrontierScratch {
    frontier_in: Option<vyre::backend::Resource>,
    frontier_out: Option<vyre::backend::Resource>,
    changed: Option<vyre::backend::Resource>,
    byte_len: usize,
}

impl FixedPointResidentFrontierScratch {
    /// Ensure the resident frontier pair exists with at least `byte_len` bytes.
    pub fn ensure_pair_refs(
        &mut self,
        backend: &dyn vyre::VyreBackend,
        byte_len: usize,
    ) -> Result<(&vyre::backend::Resource, &vyre::backend::Resource), String> {
        if self.byte_len < byte_len || self.frontier_in.is_none() || self.frontier_out.is_none() {
            self.clear(backend).map_err(|error| error.to_string())?;
            let frontier_in = backend
                .allocate_resident(byte_len)
                .map_err(|error| error.to_string())?;
            let frontier_out = match backend.allocate_resident(byte_len) {
                Ok(resource) => resource,
                Err(error) => {
                    let _ = backend.free_resident(frontier_in);
                    return Err(error.to_string());
                }
            };
            self.frontier_in = Some(frontier_in);
            self.frontier_out = Some(frontier_out);
            self.byte_len = byte_len;
        }
        Ok((
            self.frontier_in.as_ref().ok_or_else(|| {
                "weir resident fixed-point frontier input resource missing after ensure. Fix: clear and recreate the resident frontier scratch before dispatch."
                    .to_string()
            })?,
            self.frontier_out.as_ref().ok_or_else(|| {
                "weir resident fixed-point frontier output resource missing after ensure. Fix: clear and recreate the resident frontier scratch before dispatch."
                    .to_string()
            })?,
        ))
    }

    /// Ensure an in-place resident frontier and one u32 changed flag exist.
    pub fn ensure_inplace_refs(
        &mut self,
        backend: &dyn vyre::VyreBackend,
        byte_len: usize,
    ) -> Result<(&vyre::backend::Resource, &vyre::backend::Resource), String> {
        let mut allocated_frontier_this_call = false;
        if self.byte_len < byte_len || self.frontier_in.is_none() {
            self.clear(backend).map_err(|error| error.to_string())?;
            let frontier_in = backend
                .allocate_resident(byte_len)
                .map_err(|error| error.to_string())?;
            self.frontier_in = Some(frontier_in);
            self.byte_len = byte_len;
            allocated_frontier_this_call = true;
        }
        if self.changed.is_none() {
            match backend.allocate_resident(std::mem::size_of::<u32>()) {
                Ok(changed) => self.changed = Some(changed),
                Err(error) => {
                    if allocated_frontier_this_call {
                        let _ = self.clear(backend);
                    }
                    return Err(error.to_string());
                }
            }
        }
        Ok((
            self.frontier_in.as_ref().ok_or_else(|| {
                "weir resident fixed-point in-place frontier resource missing after ensure. Fix: clear and recreate the resident frontier scratch before dispatch."
                    .to_string()
            })?,
            self.changed.as_ref().ok_or_else(|| {
                "weir resident fixed-point changed-flag resource missing after ensure. Fix: clear and recreate the resident frontier scratch before dispatch."
                    .to_string()
            })?,
        ))
    }

    /// Ensure the resident frontier pair and return owned handles for tests.
    #[cfg(test)]
    pub fn ensure_pair(
        &mut self,
        backend: &dyn vyre::VyreBackend,
        byte_len: usize,
    ) -> Result<(vyre::backend::Resource, vyre::backend::Resource), String> {
        let (frontier_in, frontier_out) = self.ensure_pair_refs(backend, byte_len)?;
        Ok((frontier_in.clone(), frontier_out.clone()))
    }

    /// Current resident frontier byte capacity.
    #[must_use]
    pub fn byte_len(&self) -> usize {
        self.byte_len
    }

    /// Whether both resident frontier resources are currently allocated.
    #[must_use]
    pub fn is_allocated(&self) -> bool {
        self.frontier_in.is_some() && self.frontier_out.is_some()
    }

    /// Free resident frontier resources, preserving the scratch object itself.
    pub fn clear(&mut self, backend: &dyn vyre::VyreBackend) -> Result<(), vyre::BackendError> {
        let mut first_error = None;
        if let Some(resource) = self.frontier_in.take() {
            if let Err(error) = backend.free_resident(resource) {
                first_error.get_or_insert(error);
            }
        }
        if let Some(resource) = self.frontier_out.take() {
            if let Err(error) = backend.free_resident(resource) {
                first_error.get_or_insert(error);
            }
        }
        if let Some(resource) = self.changed.take() {
            if let Err(error) = backend.free_resident(resource) {
                first_error.get_or_insert(error);
            }
        }
        self.byte_len = 0;
        if let Some(error) = first_error {
            Err(error)
        } else {
            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::sync::Mutex;
    use vyre::backend::{private, BackendError, DispatchConfig, Resource};

    use super::*;

    struct FakeBackend {
        next: AtomicU64,
        allocations: Mutex<Vec<usize>>,
        freed: Mutex<Vec<u64>>,
        fail_byte_len: Option<usize>,
    }

    impl FakeBackend {
        fn new() -> Self {
            Self {
                next: AtomicU64::new(1),
                allocations: Mutex::new(Vec::new()),
                freed: Mutex::new(Vec::new()),
                fail_byte_len: None,
            }
        }

        fn failing_on(byte_len: usize) -> Self {
            Self {
                next: AtomicU64::new(1),
                allocations: Mutex::new(Vec::new()),
                freed: Mutex::new(Vec::new()),
                fail_byte_len: Some(byte_len),
            }
        }
    }

    impl private::Sealed for FakeBackend {}

    impl vyre::VyreBackend for FakeBackend {
        fn id(&self) -> &'static str {
            "weir_test_frontier_scratch"
        }

        fn supported_ops(&self) -> &HashSet<vyre::ir::OpId> {
            static OPS: std::sync::OnceLock<HashSet<vyre::ir::OpId>> = std::sync::OnceLock::new();
            OPS.get_or_init(HashSet::new)
        }

        fn dispatch(
            &self,
            _program: &vyre::ir::Program,
            _inputs: &[Vec<u8>],
            _config: &DispatchConfig,
        ) -> Result<Vec<Vec<u8>>, BackendError> {
            unreachable!("frontier scratch tests do not dispatch")
        }

        fn allocate_resident(&self, byte_len: usize) -> Result<Resource, BackendError> {
            self.allocations.lock().unwrap().push(byte_len);
            if self.fail_byte_len == Some(byte_len) {
                return Err(BackendError::InvalidProgram {
                    fix: format!("fake allocation failure for {byte_len} bytes"),
                });
            }
            Ok(Resource::Resident(
                self.next.fetch_add(1, Ordering::Relaxed),
            ))
        }

        fn free_resident(&self, resource: Resource) -> Result<(), BackendError> {
            let Resource::Resident(handle) = resource else {
                panic!("fake backend only frees resident resources");
            };
            self.freed.lock().unwrap().push(handle);
            Ok(())
        }
    }

    #[test]
    fn resident_frontier_scratch_reuses_larger_capacity_for_smaller_runs() {
        let backend = FakeBackend::new();
        let mut scratch = FixedPointResidentFrontierScratch::default();

        let first = scratch
            .ensure_pair(&backend, 128)
            .expect("first frontier allocation must succeed");
        let second = scratch
            .ensure_pair(&backend, 32)
            .expect("smaller frontier must reuse larger allocation");

        assert_eq!(first, second);
        assert_eq!(scratch.byte_len(), 128);
        assert_eq!(backend.allocations.lock().unwrap().as_slice(), &[128, 128]);
        assert!(
            backend.freed.lock().unwrap().is_empty(),
            "Fix: smaller resident frontier requests must not free and reallocate larger reusable slots."
        );
    }

    #[test]
    fn resident_inplace_frontier_rolls_back_new_frontier_when_changed_flag_fails() {
        let backend = FakeBackend::failing_on(std::mem::size_of::<u32>());
        let mut scratch = FixedPointResidentFrontierScratch::default();

        let error = scratch
            .ensure_inplace_refs(&backend, 128)
            .expect_err("changed-flag allocation failure must fail the in-place frontier ensure");

        assert!(error.contains("fake allocation failure"));
        assert_eq!(
            backend.allocations.lock().unwrap().as_slice(),
            &[128, std::mem::size_of::<u32>()],
            "Fix: in-place frontier ensure must attempt frontier then changed flag exactly once."
        );
        assert_eq!(
            backend.freed.lock().unwrap().as_slice(),
            &[1],
            "Fix: newly allocated frontier must be released when changed flag allocation fails."
        );
        assert_eq!(scratch.byte_len(), 0);
        assert!(!scratch.is_allocated());
    }

    #[test]
    fn resident_frontier_scratch_production_ensure_has_no_panic_path() {
        let source = include_str!("fixed_point_resident_frontier.rs");
        let production = source
            .split("#[cfg(test)]")
            .next()
            .expect("production frontier source must precede tests");

        for forbidden in [
            concat!(".", "expect("),
            concat!("panic", "!("),
            concat!("unimplemented", "!("),
            concat!("todo", "!("),
        ] {
            assert!(
                !production.contains(forbidden),
                "resident frontier production ensure path must return typed errors instead of using {forbidden}"
            );
        }
    }
}