use super::{
graph_hash, DirectIfdsEntry, DirectIfdsKey, DirectResidentIfdsBatch, DirectResidentIfdsGraph,
DirectResidentIfdsGraphKey,
};
use crate::ifds_gpu::{IfdsResidentDispatch, IfdsShape, ResidentPreparedIfdsCsr};
use std::cell::{Cell, RefCell};
mod identity_contracts;
mod lru_accounting_contracts;
mod solve_reuse_contracts;
mod validation_contracts;
struct FreeOnlyDispatch {
freed: RefCell<Vec<u64>>,
}
impl IfdsResidentDispatch for FreeOnlyDispatch {
type Resource = u64;
fn resident_backend_id(&self) -> &'static str {
"weir_test_direct_ifds_lru"
}
fn allocate_resident(&self, _byte_len: usize) -> Result<Self::Resource, String> {
unreachable!("LRU test only frees preloaded resident entries")
}
fn upload_resident(&self, _resource: &Self::Resource, _bytes: &[u8]) -> Result<(), String> {
unreachable!("LRU test only frees preloaded resident entries")
}
fn upload_resident_many(&self, _uploads: &[(&Self::Resource, &[u8])]) -> Result<(), String> {
unreachable!("LRU test only frees preloaded resident entries")
}
fn download_resident(&self, _resource: &Self::Resource) -> Result<Vec<u8>, String> {
unreachable!("LRU test never downloads resident entries")
}
fn download_resident_into(
&self,
_resource: &Self::Resource,
_output: &mut Vec<u8>,
) -> Result<(), String> {
unreachable!("LRU test never downloads resident entries")
}
fn download_resident_range(
&self,
_resource: &Self::Resource,
_byte_offset: usize,
_byte_len: usize,
) -> Result<Vec<u8>, String> {
unreachable!("LRU test never downloads resident entries")
}
fn download_resident_range_into(
&self,
_resource: &Self::Resource,
_byte_offset: usize,
_byte_len: usize,
_output: &mut Vec<u8>,
) -> Result<(), String> {
unreachable!("LRU test never downloads resident entries")
}
fn free_resident(&self, resource: Self::Resource) -> Result<(), String> {
self.freed.borrow_mut().push(resource);
Ok(())
}
fn dispatch_resident(
&self,
_program: &vyre::ir::Program,
_resources: &[Self::Resource],
_grid_override: Option<[u32; 3]>,
) -> Result<(), String> {
unreachable!("LRU test never dispatches resident entries")
}
}
fn key(graph_hash: u64) -> DirectIfdsKey {
DirectIfdsKey {
backend_id: "weir_test_direct_ifds_lru",
backend_version: "unspecified",
num_procs: 1,
blocks_per_proc: 1,
facts_per_proc: 1,
graph_hash,
}
}
fn resident(first_resource: u64) -> ResidentPreparedIfdsCsr<u64> {
ResidentPreparedIfdsCsr::from_parts(
IfdsShape {
num_procs: 1,
blocks_per_proc: 1,
facts_per_proc: 1,
edge_count: 0,
},
1,
1,
0,
first_resource,
first_resource + 1,
first_resource + 2,
first_resource + 3,
first_resource + 4,
)
}
struct DirectSolveDispatch {
buffers: RefCell<Vec<Vec<u8>>>,
dispatches: Cell<u32>,
allocate_calls: Cell<u32>,
upload_many_calls: Cell<u32>,
freed: RefCell<Vec<usize>>,
}
impl DirectSolveDispatch {
fn new() -> Self {
Self {
buffers: RefCell::new(Vec::new()),
dispatches: Cell::new(0),
allocate_calls: Cell::new(0),
upload_many_calls: Cell::new(0),
freed: RefCell::new(Vec::new()),
}
}
}
impl IfdsResidentDispatch for DirectSolveDispatch {
type Resource = usize;
fn resident_backend_id(&self) -> &'static str {
"weir_test_direct_ifds_solve"
}
fn allocate_resident(&self, byte_len: usize) -> Result<Self::Resource, String> {
self.allocate_calls.set(self.allocate_calls.get() + 1);
let mut buffers = self.buffers.borrow_mut();
let resource = buffers.len();
buffers.push(vec![0; byte_len]);
Ok(resource)
}
fn upload_resident(&self, resource: &Self::Resource, bytes: &[u8]) -> Result<(), String> {
self.buffers.borrow_mut()[*resource] = bytes.to_vec();
Ok(())
}
fn upload_resident_many(&self, uploads: &[(&Self::Resource, &[u8])]) -> Result<(), String> {
self.upload_many_calls.set(self.upload_many_calls.get() + 1);
for &(resource, bytes) in uploads {
self.upload_resident(resource, bytes)?;
}
Ok(())
}
fn download_resident(&self, resource: &Self::Resource) -> Result<Vec<u8>, String> {
Ok(self.buffers.borrow()[*resource].clone())
}
fn download_resident_into(
&self,
resource: &Self::Resource,
output: &mut Vec<u8>,
) -> Result<(), String> {
output.clear();
output.extend_from_slice(&self.buffers.borrow()[*resource]);
Ok(())
}
fn download_resident_range(
&self,
resource: &Self::Resource,
byte_offset: usize,
byte_len: usize,
) -> Result<Vec<u8>, String> {
let buffer = self.buffers.borrow()[*resource].clone();
let end = byte_offset
.checked_add(byte_len)
.ok_or_else(|| "fake direct resident range overflow".to_string())?;
Ok(buffer[byte_offset..end].to_vec())
}
fn download_resident_range_into(
&self,
resource: &Self::Resource,
byte_offset: usize,
byte_len: usize,
output: &mut Vec<u8>,
) -> Result<(), String> {
output.clear();
output.extend_from_slice(&self.download_resident_range(resource, byte_offset, byte_len)?);
Ok(())
}
fn free_resident(&self, resource: Self::Resource) -> Result<(), String> {
self.freed.borrow_mut().push(resource);
Ok(())
}
fn dispatch_resident(
&self,
_program: &vyre::ir::Program,
resources: &[Self::Resource],
_grid_override: Option<[u32; 3]>,
) -> Result<(), String> {
match resources.len() {
1 => {
self.buffers.borrow_mut()[resources[0]].fill(0);
Ok(())
}
3 => {
let mut buffers = self.buffers.borrow_mut();
let seed_words = crate::dispatch_decode::unpack_exact_u32(
&buffers[resources[0]],
buffers[resources[0]].len() / std::mem::size_of::<u32>(),
"fake direct IFDS seed triples",
)?;
let offsets = crate::dispatch_decode::unpack_exact_u32(
&buffers[resources[1]],
buffers[resources[1]].len() / std::mem::size_of::<u32>(),
"fake direct IFDS seed offsets",
)?;
let query_count = offsets.len().saturating_sub(1);
if query_count == 0 {
return Ok(());
}
let frontier_word_count = buffers[resources[2]].len() / std::mem::size_of::<u32>();
let words_per_query = frontier_word_count / query_count;
let mut frontiers = crate::dispatch_decode::unpack_exact_u32(
&buffers[resources[2]],
frontier_word_count,
"fake direct IFDS seed frontiers",
)?;
for query in 0..query_count {
let start = offsets[query] as usize;
let end = offsets[query + 1] as usize;
for seed_index in start..end {
let triple = seed_index * 3;
if triple + 2 >= seed_words.len() {
continue;
}
let dense = seed_words[triple + 1].saturating_add(seed_words[triple + 2]);
let word = dense / 32;
let index = query
.checked_mul(words_per_query)
.and_then(|base| base.checked_add(word as usize));
if let Some(slot) = index.and_then(|idx| frontiers.get_mut(idx)) {
*slot |= 1u32 << (dense % 32);
}
}
}
buffers[resources[2]] = crate::dispatch_decode::pack_u32(&frontiers);
Ok(())
}
4 => Ok(()),
7 => {
self.dispatches.set(self.dispatches.get() + 1);
let changed_resource = resources[6];
let mut buffers = self.buffers.borrow_mut();
if buffers[changed_resource].len() < std::mem::size_of::<u32>() {
buffers[changed_resource].resize(std::mem::size_of::<u32>(), 0);
}
buffers[changed_resource][..std::mem::size_of::<u32>()]
.copy_from_slice(&0u32.to_le_bytes());
Ok(())
}
17 => {
let col_len_resource = resources[16];
self.buffers.borrow_mut()[col_len_resource] = 1u32.to_le_bytes().to_vec();
Ok(())
}
len => Err(format!(
"unexpected fake direct resident dispatch resource count {len}"
)),
}
}
}