weir/ifds_resident_types/

dispatch.rs

use smallvec::SmallVec;
use vyre_foundation::ir::Program;

/// Backend-neutral resident dispatch hooks for IFDS.
///
/// Weir owns this trait so production dataflow APIs do not depend on concrete
/// driver crates. A Vyre/CUDA adapter can implement it by forwarding to
/// `VyreBackend::{allocate_resident, upload_resident, dispatch_resident_timed,
/// download_resident, free_resident}`.
pub trait IfdsResidentDispatch {
    /// Opaque backend-owned resident handle.
    type Resource: Clone + PartialEq;

    /// Stable backend/device identity for resident resource ownership.
    fn resident_backend_id(&self) -> &'static str;

    /// Stable backend implementation version for resident cache compatibility.
    ///
    /// Implementers backed by `vyre::VyreBackend` should forward
    /// `vyre::VyreBackend::version`. The default keeps existing non-Vyre
    /// test dispatchers source-compatible while preventing accidental reuse
    /// across dispatchers that opt into versioned identity.
    fn resident_backend_version(&self) -> &'static str {
        "unspecified"
    }

    /// Allocate a resident buffer of exactly `byte_len` bytes.
    fn allocate_resident(&self, byte_len: usize) -> Result<Self::Resource, String>;

    /// Upload bytes into a resident buffer.
    fn upload_resident(&self, resource: &Self::Resource, bytes: &[u8]) -> Result<(), String>;

    /// Upload several resident buffers as one logical staging operation.
    ///
    /// Implementers must provide the batching boundary explicitly. A default
    /// per-buffer loop hides stream synchronization and defeats the resident
    /// seed-scatter path's upload amortization contract.
    fn upload_resident_many(&self, uploads: &[(&Self::Resource, &[u8])]) -> Result<(), String>;

    /// Upload byte ranges into resident buffers as one logical staging
    /// operation.
    ///
    /// Backends with true ranged resident uploads should override this method.
    /// The default only supports zero-offset uploads and forwards them through
    /// [`Self::upload_resident_many`], preserving existing test dispatchers
    /// while allowing CUDA-backed paths to avoid padding small seed batches to
    /// the full retained capacity.
    fn upload_resident_at_many(
        &self,
        uploads: &[(&Self::Resource, usize, &[u8])],
    ) -> Result<(), String> {
        if uploads.iter().any(|(_, offset, _)| *offset != 0) {
            return Err(
                "weir IFDS resident dispatch does not implement ranged uploads. Fix: provide upload_resident_at_many for this backend or use zero-offset uploads only."
                    .to_string(),
            );
        }
        let mut zero_offset_uploads = SmallVec::<[(&Self::Resource, &[u8]); 8]>::new();
        zero_offset_uploads.try_reserve(uploads.len()).map_err(|error| {
            format!(
                "weir IFDS resident dispatch could not reserve {} zero-offset upload tuple(s): {error}. Fix: shard the resident upload batch before dispatch.",
                uploads.len()
            )
        })?;
        for &(resource, _, bytes) in uploads {
            zero_offset_uploads.push((resource, bytes));
        }
        self.upload_resident_many(&zero_offset_uploads)
    }

    /// Download a resident buffer.
    fn download_resident(&self, resource: &Self::Resource) -> Result<Vec<u8>, String>;

    /// Download a resident buffer into caller-owned host storage.
    fn download_resident_into(
        &self,
        resource: &Self::Resource,
        output: &mut Vec<u8>,
    ) -> Result<(), String>;

    /// Download one byte range from a resident buffer.
    ///
    /// Implementers must provide a real ranged transfer. A default full-buffer
    /// download would turn every convergence poll into an O(frontier) copy and
    /// hide the exact production regression this resident API exists to avoid.
    fn download_resident_range(
        &self,
        resource: &Self::Resource,
        byte_offset: usize,
        byte_len: usize,
    ) -> Result<Vec<u8>, String>;

    /// Download one byte range from a resident buffer into caller-owned host
    /// storage.
    fn download_resident_range_into(
        &self,
        resource: &Self::Resource,
        byte_offset: usize,
        byte_len: usize,
        output: &mut Vec<u8>,
    ) -> Result<(), String>;

    /// Download several resident byte ranges into caller-owned host buffers as
    /// one logical readback operation.
    fn download_resident_ranges_into(
        &self,
        ranges: &[(&Self::Resource, usize, usize)],
        outputs: &mut [&mut Vec<u8>],
    ) -> Result<(), String> {
        if ranges.len() != outputs.len() {
            return Err(format!(
                "weir IFDS resident ranged batch download expected matching range/output counts but got {} range(s) and {} output(s). Fix: pass one caller-owned output Vec per readback range.",
                ranges.len(),
                outputs.len()
            ));
        }
        for ((resource, byte_offset, byte_len), output) in ranges.iter().zip(outputs.iter_mut()) {
            self.download_resident_range_into(resource, *byte_offset, *byte_len, output)?;
        }
        Ok(())
    }

    /// Free a resident buffer.
    fn free_resident(&self, resource: Self::Resource) -> Result<(), String>;

    /// Dispatch `program` with resident buffers in binding order.
    fn dispatch_resident(
        &self,
        program: &Program,
        resources: &[Self::Resource],
        grid_override: Option<[u32; 3]>,
    ) -> Result<(), String>;

    /// Dispatch an ordered sequence of resident-buffer programs and read
    /// selected resident byte ranges into caller-owned host buffers.
    ///
    /// The default preserves correctness through `dispatch_resident` plus
    /// `download_resident_ranges_into`. CUDA-backed adapters override this so
    /// a fixed-point window and final compact readback are enqueued behind one
    /// device stream fence instead of synchronizing once per iteration.
    #[allow(clippy::type_complexity)]
    fn dispatch_resident_sequence_read_ranges_into(
        &self,
        steps: &[(&Program, &[Self::Resource], Option<[u32; 3]>)],
        read_ranges: &[(&Self::Resource, usize, usize)],
        outputs: &mut [&mut Vec<u8>],
    ) -> Result<(), String> {
        for (program, resources, grid_override) in steps {
            self.dispatch_resident(program, resources, *grid_override)?;
        }
        self.download_resident_ranges_into(read_ranges, outputs)
    }

    /// Dispatch a resident sequence prefix, repeat a resident sub-sequence, and
    /// read selected resident byte ranges into caller-owned host buffers.
    ///
    /// The default materializes no extra host sequence vector; it loops the
    /// repeated kernel group directly. CUDA-backed adapters can override this
    /// to forward repeat structure to the backend, avoiding per-iteration
    /// sequence construction and repeated launch preparation.
    #[allow(clippy::type_complexity)]
    fn dispatch_resident_repeated_sequence_read_ranges_into(
        &self,
        prefix_steps: &[(&Program, &[Self::Resource], Option<[u32; 3]>)],
        repeated_steps: &[(&Program, &[Self::Resource], Option<[u32; 3]>)],
        repeat_count: u32,
        read_ranges: &[(&Self::Resource, usize, usize)],
        outputs: &mut [&mut Vec<u8>],
    ) -> Result<(), String> {
        for (program, resources, grid_override) in prefix_steps {
            self.dispatch_resident(program, resources, *grid_override)?;
        }
        for _ in 0..repeat_count {
            for (program, resources, grid_override) in repeated_steps {
                self.dispatch_resident(program, resources, *grid_override)?;
            }
        }
        self.download_resident_ranges_into(read_ranges, outputs)
    }
}