use super::*;
use std::cell::{Cell, RefCell};
struct FakeResidentDispatch {
buffers: RefCell<Vec<Vec<u8>>>,
dispatches: Cell<u32>,
sequence_read_ranges_calls: Cell<u32>,
downloads: RefCell<Vec<usize>>,
converge_after_dispatches: u32,
allocate_calls: Cell<u32>,
upload_calls: Cell<u32>,
upload_many_calls: Cell<u32>,
uploaded_lengths: RefCell<Vec<usize>>,
fail_upload_call: Cell<Option<u32>>,
freed: RefCell<Vec<usize>>,
}
impl FakeResidentDispatch {
fn new_with_program_graph_buffers(converge_after_dispatches: u32) -> Self {
Self {
buffers: RefCell::new(vec![Vec::new(); 5]),
dispatches: Cell::new(0),
sequence_read_ranges_calls: Cell::new(0),
downloads: RefCell::new(Vec::new()),
converge_after_dispatches,
allocate_calls: Cell::new(0),
upload_calls: Cell::new(0),
upload_many_calls: Cell::new(0),
uploaded_lengths: RefCell::new(Vec::new()),
fail_upload_call: Cell::new(None),
freed: RefCell::new(Vec::new()),
}
}
fn failing_on_upload_call(converge_after_dispatches: u32, call: u32) -> Self {
let dispatch = Self::new_with_program_graph_buffers(converge_after_dispatches);
dispatch.fail_upload_call.set(Some(call));
dispatch
}
}
impl IfdsResidentDispatch for FakeResidentDispatch {
type Resource = usize;
fn resident_backend_id(&self) -> &'static str {
"weir_test_ifds_gpu_fake_resident"
}
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> {
let call = self.upload_calls.get() + 1;
self.upload_calls.set(call);
if self.fail_upload_call.get() == Some(call) {
return Err(format!("injected upload failure on call {call}"));
}
self.uploaded_lengths.borrow_mut().push(bytes.len());
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> {
let buffer = self.buffers.borrow()[*resource].clone();
self.downloads.borrow_mut().push(buffer.len());
Ok(buffer)
}
fn download_resident_into(
&self,
resource: &Self::Resource,
output: &mut Vec<u8>,
) -> Result<(), String> {
let buffer = self.buffers.borrow();
output.clear();
output.extend_from_slice(&buffer[*resource]);
self.downloads.borrow_mut().push(output.len());
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())?;
self.downloads.borrow_mut().push(byte_len);
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();
let end = byte_offset
.checked_add(byte_len)
.ok_or_else(|| "fake resident range overflow".to_string())?;
output.clear();
output.extend_from_slice(&buffer[*resource][byte_offset..end]);
self.downloads.borrow_mut().push(byte_len);
Ok(())
}
fn free_resident(&self, _resource: Self::Resource) -> Result<(), String> {
self.freed.borrow_mut().push(_resource);
let mut buffers = self.buffers.borrow_mut();
if _resource >= 7 && _resource + 1 == buffers.len() {
buffers.pop();
}
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);
let call = self.dispatches.get();
self.dispatches.set(call + 1);
let changed_resource = resources[6];
let changed_value = if call + 1 < self.converge_after_dispatches {
1u32
} else {
0u32
};
let mut buffers = self.buffers.borrow_mut();
let byte_offset = usize::try_from(call)
.unwrap_or(usize::MAX)
.saturating_mul(std::mem::size_of::<u32>());
let byte_end = byte_offset.saturating_add(std::mem::size_of::<u32>());
if byte_end <= buffers[changed_resource].len() {
buffers[changed_resource][byte_offset..byte_end]
.copy_from_slice(&changed_value.to_le_bytes());
} else if buffers[changed_resource].len() == std::mem::size_of::<u32>() {
buffers[changed_resource].copy_from_slice(&changed_value.to_le_bytes());
}
Ok(())
}
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> {
self.sequence_read_ranges_calls
.set(self.sequence_read_ranges_calls.get() + 1);
for (program, resources, grid_override) in steps {
self.dispatch_resident(program, resources, *grid_override)?;
}
self.download_resident_ranges_into(read_ranges, outputs)
}
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> {
self.sequence_read_ranges_calls
.set(self.sequence_read_ranges_calls.get() + 1);
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)
}
}
#[path = "ifds_resident_solve_tests/single_solver_contracts.rs"]
mod single_solver_contracts;
#[path = "ifds_resident_solve_tests/scratch_contracts.rs"]
mod scratch_contracts;
#[path = "ifds_resident_solve_tests/parallel_batch_contracts.rs"]
mod parallel_batch_contracts;
#[path = "ifds_resident_solve_tests/csr_upload_contracts.rs"]
mod csr_upload_contracts;