use super::*;
use crate::ifds_gpu::IfdsResidentDispatch;
use crate::ifds_gpu::{prepare_ifds_csr_borrowed_with_scratch_via, IfdsPrepareScratch};
use crate::resident_cache_identity::{
ResidentGraphCacheDomain, ResidentGraphCacheIdentity, ResidentGraphCacheMissReason,
};
use std::cell::{Cell, RefCell};
use vyre::ir::Program;
struct FakeResidentDispatch {
backend_id: &'static str,
backend_version: &'static str,
next: Cell<u64>,
upload_many_calls: Cell<u32>,
uploads: RefCell<Vec<(u64, usize)>>,
freed: RefCell<Vec<u64>>,
}
impl FakeResidentDispatch {
fn new() -> Self {
Self::new_with_backend_id("weir_test_ifds_resident_cache_fake")
}
fn new_with_backend_id(backend_id: &'static str) -> Self {
Self::new_with_backend_identity(backend_id, "test-v1")
}
fn new_with_backend_identity(backend_id: &'static str, backend_version: &'static str) -> Self {
Self {
backend_id,
backend_version,
next: Cell::new(1),
upload_many_calls: Cell::new(0),
uploads: RefCell::new(Vec::new()),
freed: RefCell::new(Vec::new()),
}
}
}
impl IfdsResidentDispatch for FakeResidentDispatch {
type Resource = u64;
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> {
let next = self.next.get();
self.next.set(next + 1);
Ok(next)
}
fn upload_resident(&self, resource: &Self::Resource, bytes: &[u8]) -> Result<(), String> {
self.uploads.borrow_mut().push((*resource, bytes.len()));
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> {
Err("fake IFDS cache dispatch does not download resident buffers".to_string())
}
fn download_resident_into(
&self,
_resource: &Self::Resource,
_output: &mut Vec<u8>,
) -> Result<(), String> {
Err("fake IFDS cache dispatch does not download resident buffers".to_string())
}
fn download_resident_range(
&self,
_resource: &Self::Resource,
_byte_offset: usize,
_byte_len: usize,
) -> Result<Vec<u8>, String> {
Err("fake IFDS cache dispatch does not range-download resident buffers".to_string())
}
fn download_resident_range_into(
&self,
_resource: &Self::Resource,
_byte_offset: usize,
_byte_len: usize,
_output: &mut Vec<u8>,
) -> Result<(), String> {
Err("fake IFDS cache dispatch does not range-download resident buffers".to_string())
}
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> {
Err("fake IFDS cache dispatch does not execute resident programs".to_string())
}
}
fn prepared_fixture_with_col_idx(col_idx_word: u32) -> crate::ifds_gpu::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] {
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(&col_idx_word.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,
2,
1,
&[(0, 0, 1)],
&[],
&[],
&[],
&mut scratch,
)
.expect("fixture IFDS CSR must prepare")
}
fn prepared_fixture() -> crate::ifds_gpu::PreparedIfdsCsr {
prepared_fixture_with_col_idx(1)
}
fn ifds_identity(
dispatch: &FakeResidentDispatch,
prepared: &crate::ifds_gpu::PreparedIfdsCsr,
) -> ResidentGraphCacheIdentity {
ResidentGraphCacheIdentity::ifds_csr(
dispatch.resident_backend_id(),
dispatch.resident_backend_version(),
prepared.stable_layout_hash(),
prepared.node_count(),
prepared.shape().edge_count,
u32::try_from(prepared.frontier_words()).expect("fixture frontier words must fit u32"),
)
}
#[test]
fn resident_ifds_csr_cache_reuses_equivalent_prepared_layout() {
let prepared = prepared_fixture();
let dispatch = FakeResidentDispatch::new();
let mut cache = ResidentIfdsCsrCache::new();
{
let first = cache
.get_or_upload(&dispatch, &prepared)
.expect("first IFDS resident lookup must upload");
assert_eq!(first.node_count(), 2);
assert_eq!(first.edge_count(), prepared.shape().edge_count);
assert_eq!(first.frontier_words(), 1);
assert_eq!(first.stable_layout_hash(), prepared.stable_layout_hash());
}
{
let second = cache
.get_or_upload(&dispatch, &prepared)
.expect("second IFDS resident lookup must reuse cached CSR");
assert_eq!(second.node_count(), 2);
assert_eq!(second.frontier_words(), 1);
}
assert_eq!(cache.len(), 1);
assert_eq!(
cache.stats(),
ResidentIfdsCsrCacheStats {
hits: 1,
misses: 1,
resident_uploads: 1,
resident_upload_bytes: prepared.retained_graph_bytes() as u64,
resident_avoided_upload_bytes: prepared.retained_graph_bytes() as u64,
evictions: 0,
retained_bytes: prepared.retained_graph_bytes(),
entries: 1,
}
);
assert_eq!(
cache.stats().resident_graph_reuse_telemetry(),
vyre::ResidentGraphReuseTelemetry::from_counters(
1,
1,
prepared.retained_graph_bytes() as u64,
prepared.retained_graph_bytes() as u64
)
);
assert_eq!(dispatch.upload_many_calls.get(), 1);
assert_eq!(dispatch.uploads.borrow().len(), 4);
cache
.free_all(&dispatch)
.expect("IFDS resident cache must free retained CSR resources");
assert!(cache.is_empty());
assert_eq!(cache.stats().retained_bytes, 0);
assert_eq!(dispatch.freed.borrow().len(), 4);
}
#[test]
fn resident_ifds_csr_cache_exposes_shared_identity_and_miss_reason() {
let prepared = prepared_fixture();
let dispatch = FakeResidentDispatch::new();
let mut cache = ResidentIfdsCsrCache::new();
let requested = ifds_identity(&dispatch, &prepared);
assert_eq!(
cache.miss_reason_for_identity(&requested),
ResidentGraphCacheMissReason::EmptyCache
);
cache
.get_or_upload(&dispatch, &prepared)
.expect("first IFDS resident lookup must upload");
assert_eq!(
cache.last_miss_reason(),
Some(ResidentGraphCacheMissReason::EmptyCache)
);
let identities = cache.resident_identities();
assert_eq!(identities.len(), 1);
assert_eq!(identities[0], requested);
assert_eq!(identities[0].domain, ResidentGraphCacheDomain::IfdsCsr);
assert_eq!(identities[0].backend_id, dispatch.resident_backend_id());
assert_eq!(
identities[0].backend_version,
dispatch.resident_backend_version()
);
assert_eq!(identities[0].node_count, prepared.node_count());
assert_eq!(identities[0].edge_count, prepared.shape().edge_count);
assert_eq!(identities[0].frontier_words, 1);
assert_ne!(identities[0].stable_digest64(), 0);
let same_shape_other_layout = prepared_fixture_with_col_idx(0);
assert_eq!(
cache.miss_reason_for_identity(&ifds_identity(&dispatch, &same_shape_other_layout)),
ResidentGraphCacheMissReason::LayoutChanged
);
let other_backend =
FakeResidentDispatch::new_with_backend_id("weir_test_ifds_cache_identity_other_backend");
assert_eq!(
cache.miss_reason_for_identity(&ifds_identity(&other_backend, &prepared)),
ResidentGraphCacheMissReason::BackendChanged
);
let shape_collision = ResidentGraphCacheIdentity::ifds_csr(
dispatch.resident_backend_id(),
dispatch.resident_backend_version(),
prepared.stable_layout_hash(),
prepared.node_count(),
prepared.shape().edge_count,
2,
);
assert_eq!(
cache.miss_reason_for_identity(&shape_collision),
ResidentGraphCacheMissReason::ShapeChanged
);
cache
.get_or_upload(&dispatch, &prepared)
.expect("second IFDS resident lookup must hit");
assert_eq!(cache.last_miss_reason(), None);
}
#[test]
fn resident_ifds_csr_cache_does_not_share_handles_across_backends() {
let prepared = prepared_fixture();
let first_backend = FakeResidentDispatch::new_with_backend_id("weir_test_ifds_cache_gpu_a");
let second_backend = FakeResidentDispatch::new_with_backend_id("weir_test_ifds_cache_gpu_b");
let mut cache = ResidentIfdsCsrCache::new();
cache
.get_or_upload(&first_backend, &prepared)
.expect("first backend lookup must upload CSR");
cache
.get_or_upload(&second_backend, &prepared)
.expect("same layout on another backend must upload its own CSR");
assert_eq!(cache.len(), 2);
assert_eq!(
cache.stats(),
ResidentIfdsCsrCacheStats {
hits: 0,
misses: 2,
resident_uploads: 2,
resident_upload_bytes: (prepared.retained_graph_bytes() * 2) as u64,
resident_avoided_upload_bytes: 0,
evictions: 0,
retained_bytes: prepared.retained_graph_bytes() * 2,
entries: 2,
}
);
assert_eq!(first_backend.upload_many_calls.get(), 1);
assert_eq!(second_backend.upload_many_calls.get(), 1);
cache
.free_all(&first_backend)
.expect("freeing mixed fake handles should drain cache in tests");
}
#[test]
fn resident_ifds_csr_cache_separates_backend_versions() {
let prepared = prepared_fixture();
let first_backend = FakeResidentDispatch::new_with_backend_identity(
"weir_test_ifds_cache_gpu_versioned",
"test-v1",
);
let second_backend = FakeResidentDispatch::new_with_backend_identity(
"weir_test_ifds_cache_gpu_versioned",
"test-v2",
);
let mut cache = ResidentIfdsCsrCache::new();
cache
.get_or_upload(&first_backend, &prepared)
.expect("first backend version must upload CSR");
cache
.get_or_upload(&second_backend, &prepared)
.expect("changed backend version must upload a distinct CSR");
assert_eq!(cache.len(), 2);
assert_eq!(
cache.stats(),
ResidentIfdsCsrCacheStats {
hits: 0,
misses: 2,
resident_uploads: 2,
resident_upload_bytes: (prepared.retained_graph_bytes() * 2) as u64,
resident_avoided_upload_bytes: 0,
evictions: 0,
retained_bytes: prepared.retained_graph_bytes() * 2,
entries: 2,
}
);
assert_eq!(first_backend.upload_many_calls.get(), 1);
assert_eq!(second_backend.upload_many_calls.get(), 1);
}
#[test]
fn resident_ifds_csr_cache_evicts_least_recently_used_layout() {
let first = prepared_fixture_with_col_idx(1);
let second = prepared_fixture_with_col_idx(0);
let dispatch = FakeResidentDispatch::new();
let mut cache = ResidentIfdsCsrCache::with_max_retained_bytes(first.retained_graph_bytes());
cache
.get_or_upload(&dispatch, &first)
.expect("first IFDS resident lookup must upload");
cache
.get_or_upload(&dispatch, &second)
.expect("second IFDS resident lookup must evict first layout");
cache
.get_or_upload(&dispatch, &first)
.expect("third IFDS resident lookup must evict second layout");
assert_eq!(cache.len(), 1);
assert_eq!(
cache.stats(),
ResidentIfdsCsrCacheStats {
hits: 0,
misses: 3,
resident_uploads: 3,
resident_upload_bytes: (first.retained_graph_bytes()
+ second.retained_graph_bytes()
+ first.retained_graph_bytes()) as u64,
resident_avoided_upload_bytes: 0,
evictions: 2,
retained_bytes: first.retained_graph_bytes(),
entries: 1,
}
);
assert_eq!(dispatch.upload_many_calls.get(), 3);
assert_eq!(dispatch.freed.borrow().len(), 8);
cache
.free_all(&dispatch)
.expect("IFDS resident cache must free final retained CSR resources");
assert_eq!(dispatch.freed.borrow().len(), 12);
assert_eq!(cache.stats().retained_bytes, 0);
}
#[test]
fn resident_ifds_csr_cache_lru_heap_compacts_hot_hits() {
let prepared = prepared_fixture();
let dispatch = FakeResidentDispatch::new();
let mut cache = ResidentIfdsCsrCache::new();
cache
.get_or_upload(&dispatch, &prepared)
.expect("first IFDS resident lookup must upload");
for _ in 0..160 {
cache
.get_or_upload(&dispatch, &prepared)
.expect("hot IFDS resident cache hit must succeed");
}
assert_eq!(cache.len(), 1);
assert_eq!(cache.stats().hits, 160);
let stale_limit = cache
.len()
.checked_mul(4)
.and_then(|value| value.checked_add(32))
.expect("test IFDS resident cache stale LRU limit must fit usize");
assert!(
cache.lru_len_for_tests() <= stale_limit,
"Fix: IFDS resident CSR cache LRU metadata must compact stale hit records instead of growing with every access."
);
cache
.free_all(&dispatch)
.expect("IFDS resident cache must free retained CSR resources");
assert_eq!(cache.lru_len_for_tests(), 0);
}
#[test]
fn resident_ifds_csr_cache_accounting_uses_checked_arithmetic() {
let source = include_str!("../cache.rs");
for forbidden in [
"retained_bytes.saturating_add",
"retained_bytes.saturating_sub",
"resident_upload_bytes\n .saturating_add",
"resident_avoided_upload_bytes\n .saturating_add",
"self.lru.len() <= self.entries.len().saturating_mul(4).saturating_add(32)",
"BinaryHeap::with_capacity(self.entries.len())",
concat!(".", "expect("),
concat!("panic", "!("),
concat!("unimplemented", "!("),
concat!("todo", "!("),
] {
assert!(
!source.contains(forbidden),
"Fix: IFDS resident CSR cache accounting must use checked arithmetic, not saturating math that hides budget/accounting corruption: {forbidden}"
);
}
}