#![allow(clippy::too_many_arguments)]
use crate::ifds_gpu::{
allocate_resident_ifds_parallel_scratch_with_capacities, free_resident_ifds_parallel_scratch,
solve_resident_prepared_many_parallel_with_scratch_and_host_into_via, IfdsResidentDispatch,
ResidentIfdsParallelHostScratch, ResidentIfdsParallelScratch, ResidentPreparedIfdsCsr,
};
use std::mem;
pub(super) fn solve_resident_many_with_retained_scratch_into<D, R>(
dispatch: &D,
resident: &ResidentPreparedIfdsCsr<R>,
parallel_scratch: &mut Option<ResidentIfdsParallelScratch<R>>,
scratch_allocations: &mut u64,
scratch_reuses: &mut u64,
scratch_frees: &mut u64,
max_parallel_scratch_bytes: Option<usize>,
host_scratch: &mut ResidentIfdsParallelHostScratch,
seed_sets: &[&[(u32, u32, u32)]],
max_iterations: u32,
results: &mut Vec<Vec<u32>>,
) -> Result<(), String>
where
D: IfdsResidentDispatch<Resource = R>,
R: Clone,
{
let max_iterations_usize = usize::try_from(max_iterations).map_err(|error| {
format!(
"weir resident IFDS batch max_iterations does not fit usize: {error}. Fix: use a smaller iteration budget."
)
})?;
let total_seed_facts = seed_sets.iter().try_fold(0usize, |total, seed_set| {
total.checked_add(seed_set.len()).ok_or_else(|| {
"weir resident IFDS batch seed count overflowed usize. Fix: shard the IFDS seed batch."
.to_string()
})
})?;
if total_seed_facts == 0 {
crate::staging_reserve::resize_result_rows(
results,
seed_sets.len(),
"resident IFDS retained-scratch result row",
)?;
crate::staging_reserve::clear_result_rows(results);
return Ok(());
}
require_parallel_scratch_budget(
resident.frontier_words(),
seed_sets.len(),
max_iterations_usize,
total_seed_facts,
max_parallel_scratch_bytes,
)?;
let scratch_fits = parallel_scratch.as_ref().is_some_and(|scratch| {
scratch.words_per_query() == resident.frontier_words()
&& scratch.max_queries() >= seed_sets.len()
&& scratch.changed_iteration_capacity() >= max_iterations_usize
&& scratch.max_seed_facts() >= total_seed_facts
});
if !scratch_fits {
if let Some(old_scratch) = parallel_scratch.take() {
free_resident_ifds_parallel_scratch(dispatch, old_scratch)?;
*scratch_frees = increment_counter(*scratch_frees)?;
}
*parallel_scratch = Some(allocate_resident_ifds_parallel_scratch_with_capacities(
dispatch,
resident,
seed_sets.len(),
max_iterations_usize,
total_seed_facts,
)?);
*scratch_allocations = increment_counter(*scratch_allocations)?;
} else {
*scratch_reuses = increment_counter(*scratch_reuses)?;
}
let scratch = parallel_scratch.as_ref().ok_or_else(|| {
"weir resident IFDS batch failed to retain parallel scratch after allocation".to_string()
})?;
solve_resident_prepared_many_parallel_with_scratch_and_host_into_via(
dispatch,
resident,
scratch,
host_scratch,
seed_sets,
max_iterations,
results,
)
}
fn increment_counter(counter: u64) -> Result<u64, String> {
counter.checked_add(1).ok_or_else(|| {
"weir resident IFDS retained scratch counter overflowed u64. Fix: recreate retained scratch before continuing an unbounded solve stream."
.to_string()
})
}
pub(super) fn solve_resident_one_with_retained_scratch_into<D, R>(
dispatch: &D,
resident: &ResidentPreparedIfdsCsr<R>,
parallel_scratch: &mut Option<ResidentIfdsParallelScratch<R>>,
scratch_allocations: &mut u64,
scratch_reuses: &mut u64,
scratch_frees: &mut u64,
max_parallel_scratch_bytes: Option<usize>,
host_scratch: &mut ResidentIfdsParallelHostScratch,
seed_facts: &[(u32, u32, u32)],
max_iterations: u32,
single_result_scratch: &mut Vec<Vec<u32>>,
result: &mut Vec<u32>,
label: &str,
) -> Result<(), String>
where
D: IfdsResidentDispatch<Resource = R>,
R: Clone,
{
crate::staging_reserve::resize_result_rows(
single_result_scratch,
1,
"resident IFDS single-result scratch row",
)?;
result.clear();
mem::swap(result, &mut single_result_scratch[0]);
let solve_result = solve_resident_many_with_retained_scratch_into(
dispatch,
resident,
parallel_scratch,
scratch_allocations,
scratch_reuses,
scratch_frees,
max_parallel_scratch_bytes,
host_scratch,
&[seed_facts],
max_iterations,
single_result_scratch,
);
let route_result = match solve_result {
Ok(()) if single_result_scratch.len() == 1 => Ok(()),
Ok(()) => Err(format!(
"{label} produced {} result slot(s) for one seed set. Fix: resident IFDS batch result routing is corrupt.",
single_result_scratch.len()
)),
Err(error) => Err(error),
};
if let Some(slot) = single_result_scratch.get_mut(0) {
mem::swap(result, slot);
}
route_result
}
pub(super) fn require_parallel_scratch_budget(
words_per_query: usize,
query_count: usize,
_max_iterations: usize,
total_seed_facts: usize,
max_parallel_scratch_bytes: Option<usize>,
) -> Result<(), String> {
let Some(max_parallel_scratch_bytes) = max_parallel_scratch_bytes else {
return Ok(());
};
let frontier_bytes = words_per_query
.checked_mul(query_count)
.and_then(|words| words.checked_mul(std::mem::size_of::<u32>()))
.ok_or_else(|| {
"weir resident IFDS parallel scratch frontier byte budget overflowed. Fix: shard the IFDS query batch before dispatch.".to_string()
})?;
let changed_bytes = std::mem::size_of::<u32>();
let seed_triples_bytes = total_seed_facts
.checked_mul(3)
.and_then(|words| words.max(1).checked_mul(std::mem::size_of::<u32>()))
.ok_or_else(|| {
"weir resident IFDS parallel scratch seed triple byte budget overflowed. Fix: shard the seed batch before dispatch.".to_string()
})?;
let seed_offsets_bytes = query_count
.checked_add(1)
.and_then(|words| words.checked_mul(std::mem::size_of::<u32>()))
.ok_or_else(|| {
"weir resident IFDS parallel scratch seed offset byte budget overflowed. Fix: shard the query batch before dispatch.".to_string()
})?;
let required = frontier_bytes
.checked_add(changed_bytes)
.and_then(|bytes| bytes.checked_add(seed_triples_bytes))
.and_then(|bytes| bytes.checked_add(seed_offsets_bytes))
.ok_or_else(|| {
"weir resident IFDS parallel scratch total byte budget overflowed. Fix: shard the IFDS query batch before dispatch.".to_string()
})?;
if required > max_parallel_scratch_bytes {
return Err(format!(
"weir resident IFDS parallel scratch requires {required} bytes but the configured device-memory budget allows {max_parallel_scratch_bytes}. Fix: increase max_parallel_scratch_bytes, shard seed queries, or lower max_iterations before dispatch."
));
}
Ok(())
}
pub(super) fn take_single_result_slot(
results: &mut [Vec<u32>],
label: &str,
) -> Result<Vec<u32>, String> {
if results.len() != 1 {
return Err(format!(
"{label} produced {} result slot(s) for one seed set. Fix: resident IFDS batch result routing is corrupt.",
results.len()
));
}
Ok(mem::take(&mut results[0]))
}