use super::{DirectIfdsKey, DirectResidentIfdsBatch};
use crate::ifds_gpu::{free_resident_prepared_ifds_csr, IfdsResidentDispatch};
impl<R> DirectResidentIfdsBatch<R>
where
R: Clone,
{
pub fn free_all<D>(&mut self, dispatch: &D) -> Result<(), String>
where
D: IfdsResidentDispatch<Resource = R>,
{
let solver_result = self.solver.free_all(dispatch);
let mut first_error = solver_result.err();
for (_, entry) in self.entries.drain() {
self.retained_bytes = self
.retained_bytes
.checked_sub(entry.retained_bytes)
.ok_or_else(|| {
"weir direct resident IFDS retained byte accounting underflowed while freeing all entries. Fix: rebuild the direct resident batch facade before reuse."
.to_string()
})?;
if let Err(error) = free_resident_prepared_ifds_csr(dispatch, entry.graph) {
if first_error.is_none() {
first_error = Some(error);
}
}
}
self.lru.clear();
if let Some(error) = first_error {
Err(error)
} else {
Ok(())
}
}
pub(crate) fn touch_key(&mut self, key: &DirectIfdsKey) -> Result<(), String> {
let last_seen = self.next_cache_tick()?;
if let Some(entry) = self.entries.get_mut(key) {
entry.last_seen = last_seen;
}
self.record_lru(key.clone(), last_seen)
}
pub(crate) fn evict_until_room<D>(
&mut self,
dispatch: &D,
incoming_bytes: usize,
) -> Result<(), String>
where
D: IfdsResidentDispatch<Resource = R>,
{
loop {
let entry_over_limit = self
.max_entries
.is_some_and(|max_entries| self.entries.len() >= max_entries);
let byte_over_limit = self.max_retained_bytes.is_some_and(|max_retained_bytes| {
self.retained_bytes
.checked_add(incoming_bytes)
.is_none_or(|total| total > max_retained_bytes)
});
if !entry_over_limit && !byte_over_limit {
return Ok(());
}
let Some(key) = self.pop_lru_key() else {
return Err(
"weir direct resident IFDS batch cache order lost an entry. Fix: rebuild the direct resident batch facade.".to_string(),
);
};
let Some(entry) = self.entries.remove(&key) else {
continue;
};
self.retained_bytes = self
.retained_bytes
.checked_sub(entry.retained_bytes)
.ok_or_else(|| {
"weir direct resident IFDS retained byte accounting underflowed during eviction. Fix: rebuild the direct resident batch facade before reuse."
.to_string()
})?;
free_resident_prepared_ifds_csr(dispatch, entry.graph)?;
self.evictions = self
.evictions
.checked_add(1)
.ok_or_else(|| {
"weir direct resident IFDS eviction counter overflowed u64. Fix: rebuild the direct resident batch facade before reuse."
.to_string()
})?;
}
}
pub(crate) fn reserve_cache_admission_slot(&mut self) -> Result<(), String> {
if self.entries.len() >= self.entries.capacity() {
self.entries.try_reserve(1).map_err(|error| {
format!(
"weir direct resident IFDS cache could not reserve one entry slot before resident graph prepare: {error}. Fix: lower direct resident cache pressure or shard the IFDS workload."
)
})?;
}
Ok(())
}
fn pop_lru_key(&mut self) -> Option<DirectIfdsKey> {
let entries = &self.entries;
self.lru.pop_valid(
|key| entries.get(key).map(|entry| entry.last_seen),
|| entries.keys().next().cloned(),
)
}
fn record_lru(&mut self, key: DirectIfdsKey, last_seen: u64) -> Result<(), String> {
self.lru.record(
key,
last_seen,
self.entries
.iter()
.map(|(key, entry)| (key.clone(), entry.last_seen)),
"weir direct resident IFDS",
)
}
fn next_cache_tick(&mut self) -> Result<u64, String> {
self.clock = self
.clock
.checked_add(1)
.ok_or_else(|| {
"weir direct resident IFDS logical clock overflowed u64. Fix: rebuild the direct resident batch facade before reuse."
.to_string()
})?;
Ok(self.clock)
}
#[cfg(test)]
pub(super) fn lru_len_for_tests(&self) -> usize {
self.lru.len()
}
}