use super::*;
use crate::ifds_gpu::{
prepare_ifds_csr_borrowed_with_scratch_via, IfdsPrepareScratch, IfdsResidentDispatch,
PreparedIfdsCsr,
};
use std::cell::{Cell, RefCell};
use vyre::ir::Program;
struct FakeResidentDispatch {
buffers: RefCell<Vec<Vec<u8>>>,
dispatches: Cell<u32>,
allocate_calls: Cell<u32>,
upload_many_calls: Cell<u32>,
freed: RefCell<Vec<usize>>,
}
impl FakeResidentDispatch {
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 FakeResidentDispatch {
type Resource = usize;
fn resident_backend_id(&self) -> &'static str {
"weir_test_ifds_resident_batch_fake"
}
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 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> {
let buffer = self.buffers.borrow()[*resource].clone();
let end = byte_offset
.checked_add(byte_len)
.ok_or_else(|| "fake resident range overflow".to_string())?;
output.clear();
output.extend_from_slice(&buffer[byte_offset..end]);
Ok(())
}
fn free_resident(&self, resource: Self::Resource) -> Result<(), String> {
self.freed.borrow_mut().push(resource);
Ok(())
}
fn dispatch_resident(
&self,
_program: &Program,
resources: &[Self::Resource],
_grid_override: Option<[u32; 3]>,
) -> Result<(), String> {
if resources.len() == 1 {
self.buffers.borrow_mut()[resources[0]].fill(0);
return Ok(());
}
if resources.len() == 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 IFDS resident seed triples",
)?;
let offsets = crate::dispatch_decode::unpack_exact_u32(
&buffers[resources[1]],
buffers[resources[1]].len() / std::mem::size_of::<u32>(),
"fake IFDS resident 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 IFDS resident 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);
return Ok(());
}
assert_eq!(resources.len(), 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(())
}
}
fn prepared_fixture() -> PreparedIfdsCsr {
let dispatch =
|_: &Program, inputs: &[&[u8]], grid: Option<[u32; 3]>, outputs: &mut Vec<Vec<u8>>| {
assert_eq!(inputs.len(), 17);
assert_eq!(grid, Some([1, 1, 1]));
outputs.clear();
outputs.resize_with(4, Vec::new);
for word in [0u32, 1, 1, 1, 1] {
outputs[0].extend_from_slice(&word.to_le_bytes());
}
for word in [0u32, 0] {
outputs[1].extend_from_slice(&word.to_le_bytes());
}
outputs[2].extend_from_slice(&1u32.to_le_bytes());
outputs[3].extend_from_slice(&1u32.to_le_bytes());
Ok(())
};
let mut scratch = IfdsPrepareScratch::default();
prepare_ifds_csr_borrowed_with_scratch_via(
&dispatch,
1,
4,
1,
&[(0, 0, 1)],
&[],
&[],
&[],
&mut scratch,
)
.expect("fixture IFDS CSR must prepare")
}
mod cache_stats_contracts;
mod budget_contracts;
mod result_reuse_contracts;
mod accounting_error_contracts;