weirflow 0.1.0

//! Adapter from the common Vyre resident backend contract to Weir IFDS.
//!
//! IFDS owns its resident dispatch trait so Weir does not depend on concrete
//! driver crates. This module bridges that trait to any `VyreBackend` that
//! implements resident allocation, batched upload, resident dispatch, ranged
//! download, and free through the common Vyre contract.

use crate::ifds_gpu::IfdsResidentDispatch;

/// Weir IFDS resident dispatcher backed by a generic Vyre backend.
#[derive(Clone, Copy, Debug)]
pub struct VyreResidentIfds<'a, B: vyre::VyreBackend + ?Sized> {
    backend: &'a B,
}

impl<'a, B: vyre::VyreBackend + ?Sized> VyreResidentIfds<'a, B> {
    /// Create an IFDS resident dispatcher over `backend`.
    #[must_use]
    pub const fn new(backend: &'a B) -> Self {
        Self { backend }
    }

    /// Underlying Vyre backend.
    #[must_use]
    pub const fn backend(&self) -> &'a B {
        self.backend
    }
}

impl<B: vyre::VyreBackend + ?Sized> IfdsResidentDispatch for VyreResidentIfds<'_, B> {
    type Resource = vyre::backend::Resource;

    fn resident_backend_id(&self) -> &'static str {
        self.backend.id()
    }

    fn resident_backend_version(&self) -> &'static str {
        self.backend.version()
    }

    fn allocate_resident(&self, byte_len: usize) -> Result<Self::Resource, String> {
        vyre::VyreBackend::allocate_resident(self.backend, byte_len)
            .map_err(|error| error.to_string())
    }

    fn upload_resident(&self, resource: &Self::Resource, bytes: &[u8]) -> Result<(), String> {
        vyre::VyreBackend::upload_resident(self.backend, resource, bytes)
            .map_err(|error| error.to_string())
    }

    fn upload_resident_many(&self, uploads: &[(&Self::Resource, &[u8])]) -> Result<(), String> {
        vyre::VyreBackend::upload_resident_many(self.backend, uploads)
            .map_err(|error| error.to_string())
    }

    fn upload_resident_at_many(
        &self,
        uploads: &[(&Self::Resource, usize, &[u8])],
    ) -> Result<(), String> {
        vyre::VyreBackend::upload_resident_at_many(self.backend, uploads)
            .map_err(|error| error.to_string())
    }

    fn download_resident(&self, resource: &Self::Resource) -> Result<Vec<u8>, String> {
        vyre::VyreBackend::download_resident(self.backend, resource)
            .map_err(|error| error.to_string())
    }

    fn download_resident_into(
        &self,
        resource: &Self::Resource,
        output: &mut Vec<u8>,
    ) -> Result<(), String> {
        vyre::VyreBackend::download_resident_into(self.backend, resource, output)
            .map_err(|error| error.to_string())
    }

    fn download_resident_range(
        &self,
        resource: &Self::Resource,
        byte_offset: usize,
        byte_len: usize,
    ) -> Result<Vec<u8>, String> {
        vyre::VyreBackend::download_resident_range(self.backend, resource, byte_offset, byte_len)
            .map_err(|error| error.to_string())
    }

    fn download_resident_range_into(
        &self,
        resource: &Self::Resource,
        byte_offset: usize,
        byte_len: usize,
        output: &mut Vec<u8>,
    ) -> Result<(), String> {
        vyre::VyreBackend::download_resident_range_into(
            self.backend,
            resource,
            byte_offset,
            byte_len,
            output,
        )
        .map_err(|error| error.to_string())
    }

    fn download_resident_ranges_into(
        &self,
        ranges: &[(&Self::Resource, usize, usize)],
        outputs: &mut [&mut Vec<u8>],
    ) -> Result<(), String> {
        vyre::VyreBackend::download_resident_ranges_into(self.backend, ranges, outputs)
            .map_err(|error| error.to_string())
    }

    fn free_resident(&self, resource: Self::Resource) -> Result<(), String> {
        vyre::VyreBackend::free_resident(self.backend, resource).map_err(|error| error.to_string())
    }

    fn dispatch_resident(
        &self,
        program: &vyre::Program,
        resources: &[Self::Resource],
        grid_override: Option<[u32; 3]>,
    ) -> Result<(), String> {
        let mut config = vyre::DispatchConfig::default();
        config.grid_override = grid_override;
        vyre::VyreBackend::dispatch_resident_timed(self.backend, program, resources, &config)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    fn dispatch_resident_sequence_read_ranges_into(
        &self,
        steps: &[(&vyre::Program, &[Self::Resource], Option<[u32; 3]>)],
        read_ranges: &[(&Self::Resource, usize, usize)],
        outputs: &mut [&mut Vec<u8>],
    ) -> Result<(), String> {
        let resident_steps = resident_steps_from(steps)?;
        let resident_read_ranges = resident_read_ranges_from(read_ranges)?;
        vyre::VyreBackend::dispatch_resident_sequence_read_ranges_into(
            self.backend,
            &resident_steps,
            &resident_read_ranges,
            outputs,
        )
        .map_err(|error| error.to_string())
    }

    fn dispatch_resident_repeated_sequence_read_ranges_into(
        &self,
        prefix_steps: &[(&vyre::Program, &[Self::Resource], Option<[u32; 3]>)],
        repeated_steps: &[(&vyre::Program, &[Self::Resource], Option<[u32; 3]>)],
        repeat_count: u32,
        read_ranges: &[(&Self::Resource, usize, usize)],
        outputs: &mut [&mut Vec<u8>],
    ) -> Result<(), String> {
        let resident_prefix_steps = resident_steps_from(prefix_steps)?;
        let resident_repeated_steps = resident_steps_from(repeated_steps)?;
        let resident_read_ranges = resident_read_ranges_from(read_ranges)?;
        vyre::VyreBackend::dispatch_resident_repeated_sequence_read_ranges_into(
            self.backend,
            &resident_prefix_steps,
            &resident_repeated_steps,
            repeat_count,
            &resident_read_ranges,
            outputs,
        )
        .map_err(|error| error.to_string())
    }
}

#[allow(clippy::type_complexity)]
fn resident_steps_from<'a>(
    steps: &'a [(
        &'a vyre::Program,
        &'a [vyre::backend::Resource],
        Option<[u32; 3]>,
    )],
) -> Result<Vec<vyre::backend::ResidentDispatchStep<'a>>, String> {
    let mut resident_steps =
        crate::staging_reserve::reserved_vec(steps.len(), "Vyre resident IFDS dispatch step")?;
    for &(program, resources, grid_override) in steps {
        resident_steps.push(vyre::backend::ResidentDispatchStep {
            program,
            resources,
            grid_override,
            workgroup_override: None,
        });
    }
    Ok(resident_steps)
}

fn resident_read_ranges_from<'a>(
    read_ranges: &'a [(&'a vyre::backend::Resource, usize, usize)],
) -> Result<Vec<vyre::backend::ResidentReadRange<'a>>, String> {
    let mut resident_read_ranges =
        crate::staging_reserve::reserved_vec(read_ranges.len(), "Vyre resident IFDS read range")?;
    for &(resource, byte_offset, byte_len) in read_ranges {
        resident_read_ranges.push(vyre::backend::ResidentReadRange {
            resource,
            byte_offset,
            byte_len,
        });
    }
    Ok(resident_read_ranges)
}

#[cfg(test)]
#[path = "ifds_vyre_resident_tests.rs"]
mod tests;