weirflow 0.1.0

//! Parity test: GPU IFDS solver matches reference oracle.

#![cfg(test)]

use vyre::DispatchConfig;
use vyre_driver_cuda::{CudaBackend, CudaBackendRegistration, CudaResidentBuffer};
use weir::ifds_gpu::{
    allocate_resident_ifds_parallel_scratch, allocate_resident_ifds_scratch,
    free_resident_ifds_parallel_scratch, free_resident_ifds_scratch,
    free_resident_prepared_ifds_csr, prepare_ifds_csr_borrowed_via, prepare_ifds_csr_resident_via,
    solve_borrowed_via, solve_resident_prepared_many_parallel_via,
    solve_resident_prepared_many_parallel_with_scratch_via, solve_resident_prepared_many_via,
    solve_resident_prepared_many_with_scratch_via, solve_resident_prepared_via,
    IfdsResidentDispatch,
};
use weir::ifds_resident_batch::ResidentIfdsBatch;
use weir::ifds_resident_direct_prepare::prepare_ifds_csr_resident_direct_via;
use weir::ifds_vyre_resident::VyreResidentIfds;
use weir::oracle::graph::ifds_solve as solve_cpu;

#[path = "ifds_gpu_solve_gpu_parity/borrowed_solver_contracts.rs"]
mod borrowed_solver_contracts;
#[path = "ifds_gpu_solve_gpu_parity/direct_and_facade_contracts.rs"]
mod direct_and_facade_contracts;
#[path = "ifds_gpu_solve_gpu_parity/parallel_batch_contracts.rs"]
mod parallel_batch_contracts;
#[path = "ifds_gpu_solve_gpu_parity/resident_prepared_contracts.rs"]
mod resident_prepared_contracts;
#[path = "ifds_gpu_solve_gpu_parity/scratch_telemetry_contracts.rs"]
mod scratch_telemetry_contracts;

fn run_dispatch(
    backend: &CudaBackend,
) -> impl Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String> + '_ {
    move |program, inputs, grid_override| {
        let mut config = DispatchConfig::default();
        config.grid_override = grid_override;
        backend
            .dispatch_borrowed(program, inputs, &config)
            .map_err(|err| err.to_string())
    }
}

fn live_dispatcher() -> CudaBackend {
    CudaBackend::acquire()
        .expect("CudaBackend::acquire failed on a host that must have an NVIDIA GPU.")
}

struct CudaResidentIfds<'a> {
    backend: &'a CudaBackend,
}

impl IfdsResidentDispatch for CudaResidentIfds<'_> {
    type Resource = CudaResidentBuffer;

    fn resident_backend_id(&self) -> &'static str {
        "weir_cuda_resident_ifds"
    }

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

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

    fn upload_resident_many(&self, uploads: &[(&Self::Resource, &[u8])]) -> Result<(), String> {
        let concrete: Vec<(CudaResidentBuffer, &[u8])> = uploads
            .iter()
            .map(|(resource, bytes)| (**resource, *bytes))
            .collect();
        self.backend
            .upload_resident_many(&concrete)
            .map_err(|error| error.to_string())
    }

    fn upload_resident_at_many(
        &self,
        uploads: &[(&Self::Resource, usize, &[u8])],
    ) -> Result<(), String> {
        let concrete: Vec<(CudaResidentBuffer, usize, &[u8])> = uploads
            .iter()
            .map(|(resource, dst_offset_bytes, bytes)| (**resource, *dst_offset_bytes, *bytes))
            .collect();
        self.backend
            .upload_resident_at_many(&concrete)
            .map_err(|error| error.to_string())
    }

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

    fn download_resident_into(
        &self,
        resource: &Self::Resource,
        output: &mut Vec<u8>,
    ) -> Result<(), String> {
        self.backend
            .download_resident_into(*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> {
        self.backend
            .download_resident_range(*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> {
        self.backend
            .download_resident_range_into(*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> {
        let concrete = ranges
            .iter()
            .map(|(resource, byte_offset, byte_len)| (**resource, *byte_offset, *byte_len))
            .collect::<Vec<_>>();
        self.backend
            .download_resident_ranges_into(&concrete, outputs)
            .map_err(|error| error.to_string())
    }

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

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