use crate::ifds_resident_types::{
IfdsResidentDispatch, ResidentIfdsParallelScratch, ResidentPreparedIfdsCsr,
};
use crate::ifds_seed_frontiers::{ifds_clear_frontiers_program, ifds_seed_frontiers_program};
const DEFAULT_PARALLEL_ITERATION_BUDGET: usize = 64;
const PARALLEL_CHANGED_FLAG_WORDS: usize = 1;
pub fn allocate_resident_ifds_parallel_scratch<D>(
dispatch: &D,
prepared: &ResidentPreparedIfdsCsr<D::Resource>,
max_queries: usize,
) -> Result<ResidentIfdsParallelScratch<D::Resource>, String>
where
D: IfdsResidentDispatch,
{
allocate_resident_ifds_parallel_scratch_with_iterations(
dispatch,
prepared,
max_queries,
DEFAULT_PARALLEL_ITERATION_BUDGET,
)
}
pub fn allocate_resident_ifds_parallel_scratch_with_iterations<D>(
dispatch: &D,
prepared: &ResidentPreparedIfdsCsr<D::Resource>,
max_queries: usize,
max_iterations: usize,
) -> Result<ResidentIfdsParallelScratch<D::Resource>, String>
where
D: IfdsResidentDispatch,
{
allocate_resident_ifds_parallel_scratch_with_capacities(
dispatch,
prepared,
max_queries,
max_iterations,
max_queries,
)
}
pub fn allocate_resident_ifds_parallel_scratch_with_capacities<D>(
dispatch: &D,
prepared: &ResidentPreparedIfdsCsr<D::Resource>,
max_queries: usize,
max_iterations: usize,
max_seed_facts: usize,
) -> Result<ResidentIfdsParallelScratch<D::Resource>, String>
where
D: IfdsResidentDispatch,
{
if max_queries == 0 {
return Err(
"weir IFDS resident parallel scratch requires max_queries > 0. Fix: allocate scratch only for non-empty query batches."
.to_string(),
);
}
if max_iterations == 0 {
return Err(
"weir IFDS resident parallel scratch requires max_iterations > 0. Fix: allocate scratch for the intended solve budget."
.to_string(),
);
}
let total_words = prepared.words.checked_mul(max_queries).ok_or_else(|| {
"weir IFDS resident parallel scratch frontier words overflow usize. Fix: shard the query batch before dispatch.".to_string()
})?;
let frontier_byte_len = total_words
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| {
"weir IFDS resident parallel scratch frontier bytes overflow usize. Fix: shard the query batch before dispatch.".to_string()
})?;
let changed_byte_len = PARALLEL_CHANGED_FLAG_WORDS
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| {
"weir IFDS resident parallel scratch changed-flag bytes overflow usize. Fix: use a smaller iteration budget.".to_string()
})?;
let seed_triples_words = max_seed_facts.checked_mul(3).ok_or_else(|| {
"weir IFDS resident parallel scratch seed triple word count overflows usize. Fix: shard the seed batch.".to_string()
})?;
let seed_triples_byte_len = seed_triples_words
.max(1)
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| {
"weir IFDS resident parallel scratch seed triple bytes overflow usize. Fix: shard the seed batch.".to_string()
})?;
let seed_offsets_byte_len = max_queries
.checked_add(1)
.and_then(|words| words.checked_mul(std::mem::size_of::<u32>()))
.ok_or_else(|| {
"weir IFDS resident parallel scratch seed offset bytes overflow usize. Fix: shard the query batch.".to_string()
})?;
let retained_byte_len = frontier_byte_len
.checked_add(changed_byte_len)
.and_then(|bytes| bytes.checked_add(seed_triples_byte_len))
.and_then(|bytes| bytes.checked_add(seed_offsets_byte_len))
.ok_or_else(|| {
"weir IFDS resident parallel scratch retained byte accounting overflowed usize. Fix: shard the query batch before dispatch.".to_string()
})?;
let frontiers = dispatch
.allocate_resident(frontier_byte_len)
.map_err(|error| {
format!("weir IFDS resident allocate parallel scratch frontier matrix failed: {error}")
})?;
let changed = match dispatch.allocate_resident(changed_byte_len) {
Ok(resource) => resource,
Err(error) => {
let _ = super::free_unique_resident_resources(dispatch, [("frontiers", frontiers)]);
return Err(format!(
"weir IFDS resident allocate parallel scratch changed flag failed: {error}"
));
}
};
let seed_triples = match dispatch.allocate_resident(seed_triples_byte_len) {
Ok(resource) => resource,
Err(error) => {
let _ = super::free_unique_resident_resources(
dispatch,
[("changed", changed), ("frontiers", frontiers)],
);
return Err(format!(
"weir IFDS resident allocate parallel scratch seed triples failed: {error}"
));
}
};
let seed_offsets = match dispatch.allocate_resident(seed_offsets_byte_len) {
Ok(resource) => resource,
Err(error) => {
let _ = super::free_unique_resident_resources(
dispatch,
[
("seed_triples", seed_triples),
("changed", changed),
("frontiers", frontiers),
],
);
return Err(format!(
"weir IFDS resident allocate parallel scratch seed offsets failed: {error}"
));
}
};
let seed_clear_words = u32::try_from(prepared.words).map_err(|error| {
let _ = super::free_unique_resident_resources(
dispatch,
[
("seed_offsets", seed_offsets.clone()),
("seed_triples", seed_triples.clone()),
("changed", changed.clone()),
("frontiers", frontiers.clone()),
],
);
format!(
"weir IFDS resident parallel scratch seed clear frontier word count does not fit u32: {error}. Fix: shard the IFDS problem before solving."
)
})?;
let query_count = u32::try_from(max_queries).map_err(|error| {
let _ = super::free_unique_resident_resources(
dispatch,
[
("seed_offsets", seed_offsets.clone()),
("seed_triples", seed_triples.clone()),
("changed", changed.clone()),
("frontiers", frontiers.clone()),
],
);
format!(
"weir IFDS resident parallel scratch query capacity does not fit u32: {error}. Fix: allocate smaller parallel scratch or shard the query batch."
)
})?;
let seed_capacity = u32::try_from(max_seed_facts).map_err(|error| {
let _ = super::free_unique_resident_resources(
dispatch,
[
("seed_offsets", seed_offsets.clone()),
("seed_triples", seed_triples.clone()),
("changed", changed.clone()),
("frontiers", frontiers.clone()),
],
);
format!(
"weir IFDS resident parallel scratch seed capacity does not fit u32: {error}. Fix: shard the seed batch."
)
})?;
let seed_clear_program = ifds_clear_frontiers_program(seed_clear_words, query_count)?;
let seed_scatter_program = ifds_seed_frontiers_program(
prepared.shape,
seed_clear_words,
query_count,
seed_capacity,
seed_capacity,
)?;
let clear_changed_program =
ifds_clear_frontiers_program(PARALLEL_CHANGED_FLAG_WORDS as u32, 1)?;
Ok(ResidentIfdsParallelScratch {
words_per_query: prepared.words,
max_queries,
max_seed_facts,
changed_iteration_capacity: max_iterations,
frontier_byte_len,
changed_byte_len,
seed_triples_byte_len,
seed_offsets_byte_len,
retained_byte_len,
seed_scatter_resources: [
seed_triples.clone(),
seed_offsets.clone(),
frontiers.clone(),
],
frontiers,
changed,
seed_triples,
seed_offsets,
seed_clear_program,
seed_scatter_program,
clear_changed_program,
})
}
pub fn free_resident_ifds_parallel_scratch<D>(
dispatch: &D,
scratch: ResidentIfdsParallelScratch<D::Resource>,
) -> Result<(), String>
where
D: IfdsResidentDispatch,
{
let ResidentIfdsParallelScratch {
frontiers,
changed,
seed_triples,
seed_offsets,
..
} = scratch;
super::free_unique_resident_resources(
dispatch,
[
("parallel scratch seed offsets", seed_offsets),
("parallel scratch seed triples", seed_triples),
("parallel scratch frontier matrix", frontiers),
("parallel scratch changed flag", changed),
],
)
}