#![allow(clippy::too_many_arguments)]
use super::program_cache::ensure_resident_parallel_step_programs;
use super::result_slots::resize_result_slots;
use crate::ifds_frontier_decode::ifds_encoded_frontier_nodes_from_le_bytes_into;
use crate::ifds_resident_alloc::{
allocate_resident_ifds_parallel_scratch_with_capacities, free_resident_ifds_parallel_scratch,
};
use crate::ifds_resident_types::{
IfdsResidentDispatch, ResidentIfdsParallelHostScratch, ResidentIfdsParallelScratch,
ResidentPreparedIfdsCsr,
};
use crate::ifds_seed_frontiers::scatter_seed_frontiers_resident_staged;
use crate::ifds_seed_pack::pack_seed_sets_for_gpu_for_capacity_into;
pub fn solve_resident_prepared_many_parallel_via<D>(
dispatch: &D,
prepared: &ResidentPreparedIfdsCsr<D::Resource>,
seed_sets: &[&[(u32, u32, u32)]],
max_iterations: u32,
) -> Result<Vec<Vec<u32>>, String>
where
D: IfdsResidentDispatch,
{
if seed_sets.is_empty() {
return Ok(Vec::new());
}
crate::dispatch_decode::require_positive_iterations(
"weir IFDS resident parallel batch GPU solve",
max_iterations,
)?;
if seed_sets.iter().all(|seed_facts| seed_facts.is_empty()) {
return crate::staging_reserve::reserved_result_rows(
seed_sets.len(),
"resident IFDS empty parallel result row",
);
}
let scratch = allocate_resident_ifds_parallel_scratch_with_capacities(
dispatch,
prepared,
seed_sets.len(),
usize::try_from(max_iterations).map_err(|error| {
format!(
"weir IFDS resident parallel max_iterations does not fit usize: {error}. Fix: use a smaller iteration budget."
)
})?,
seed_sets.iter().map(|seed_set| seed_set.len()).sum(),
)?;
let result = solve_resident_prepared_many_parallel_with_scratch_via(
dispatch,
prepared,
&scratch,
seed_sets,
max_iterations,
);
let free_result = free_resident_ifds_parallel_scratch(dispatch, scratch);
match (result, free_result) {
(Ok(results), Ok(())) => Ok(results),
(Err(error), _) => Err(error),
(Ok(_), Err(error)) => Err(error),
}
}
pub fn solve_resident_prepared_many_parallel_with_scratch_via<D>(
dispatch: &D,
prepared: &ResidentPreparedIfdsCsr<D::Resource>,
scratch: &ResidentIfdsParallelScratch<D::Resource>,
seed_sets: &[&[(u32, u32, u32)]],
max_iterations: u32,
) -> Result<Vec<Vec<u32>>, String>
where
D: IfdsResidentDispatch,
{
let mut host_scratch = ResidentIfdsParallelHostScratch::default();
solve_resident_prepared_many_parallel_with_scratch_and_host_via(
dispatch,
prepared,
scratch,
&mut host_scratch,
seed_sets,
max_iterations,
)
}
pub fn solve_resident_prepared_many_parallel_with_scratch_and_host_via<D>(
dispatch: &D,
prepared: &ResidentPreparedIfdsCsr<D::Resource>,
scratch: &ResidentIfdsParallelScratch<D::Resource>,
host_scratch: &mut ResidentIfdsParallelHostScratch,
seed_sets: &[&[(u32, u32, u32)]],
max_iterations: u32,
) -> Result<Vec<Vec<u32>>, String>
where
D: IfdsResidentDispatch,
{
let mut results =
crate::staging_reserve::reserved_vec(seed_sets.len(), "resident IFDS parallel result row")?;
solve_resident_prepared_many_parallel_with_scratch_and_host_into_via(
dispatch,
prepared,
scratch,
host_scratch,
seed_sets,
max_iterations,
&mut results,
)?;
Ok(results)
}
pub fn solve_resident_prepared_many_parallel_with_scratch_and_host_into_via<D>(
dispatch: &D,
prepared: &ResidentPreparedIfdsCsr<D::Resource>,
scratch: &ResidentIfdsParallelScratch<D::Resource>,
host_scratch: &mut ResidentIfdsParallelHostScratch,
seed_sets: &[&[(u32, u32, u32)]],
max_iterations: u32,
results: &mut Vec<Vec<u32>>,
) -> Result<(), String>
where
D: IfdsResidentDispatch,
{
let _caller_owned_scratch_contract = (
&host_scratch.seed_triples,
&host_scratch.seed_offsets,
&host_scratch.changed_bytes,
&host_scratch.frontier_bytes,
);
solve_resident_prepared_many_parallel_with_scratch_and_host_into_mode(
dispatch,
prepared,
scratch,
host_scratch,
seed_sets,
max_iterations,
results,
false,
)
}
fn solve_resident_prepared_many_parallel_with_scratch_and_host_into_mode<D>(
dispatch: &D,
prepared: &ResidentPreparedIfdsCsr<D::Resource>,
scratch: &ResidentIfdsParallelScratch<D::Resource>,
host_scratch: &mut ResidentIfdsParallelHostScratch,
seed_sets: &[&[(u32, u32, u32)]],
max_iterations: u32,
results: &mut Vec<Vec<u32>>,
poll_each_iteration: bool,
) -> Result<(), String>
where
D: IfdsResidentDispatch,
{
if seed_sets.is_empty() {
results.clear();
return Ok(());
}
crate::dispatch_decode::require_positive_iterations(
"weir IFDS resident parallel batch GPU solve",
max_iterations,
)?;
if seed_sets.iter().all(|seed_facts| seed_facts.is_empty()) {
resize_result_slots(results, seed_sets.len())?;
for result in results.iter_mut() {
result.clear();
}
return Ok(());
}
if scratch.words_per_query != prepared.words {
return Err(format!(
"weir IFDS resident parallel scratch has {} frontier words per query but prepared CSR requires {}. Fix: allocate parallel scratch from the same resident IFDS CSR before solving.",
scratch.words_per_query, prepared.words
));
}
if seed_sets.len() > scratch.max_queries {
return Err(format!(
"weir IFDS resident parallel scratch supports {} queries but solve requested {}. Fix: allocate larger parallel scratch or shard the query batch.",
scratch.max_queries,
seed_sets.len()
));
}
let max_iterations_usize = usize::try_from(max_iterations).map_err(|error| {
format!(
"weir IFDS resident parallel max_iterations does not fit usize: {error}. Fix: use a smaller iteration budget."
)
})?;
if max_iterations_usize > scratch.changed_iteration_capacity {
return Err(format!(
"weir IFDS resident parallel scratch is configured for {} iteration(s) but solve requested {max_iterations}. Fix: allocate scratch with allocate_resident_ifds_parallel_scratch_with_iterations for the intended budget.",
scratch.changed_iteration_capacity
));
}
let active_query_count = seed_sets.len();
let query_count = u32::try_from(active_query_count).map_err(|error| {
format!(
"weir IFDS resident parallel active query count does not fit u32: {error}. Fix: shard the query batch."
)
})?;
let scratch_query_capacity = u32::try_from(scratch.max_queries).map_err(|error| {
format!(
"weir IFDS resident parallel scratch query capacity does not fit u32: {error}. Fix: allocate smaller parallel scratch or shard the query batch."
)
})?;
let changed_slot_capacity = u32::try_from(scratch.changed_iteration_capacity).map_err(|error| {
format!(
"weir IFDS resident parallel iteration capacity does not fit u32: {error}. Fix: allocate a smaller iteration budget."
)
})?;
let scratch_words = scratch
.words_per_query
.checked_mul(active_query_count)
.ok_or_else(|| {
"weir IFDS resident parallel active frontier words overflow usize. Fix: shard the query batch.".to_string()
})?;
let seed_triples = &mut host_scratch.seed_triples;
let seed_offsets = &mut host_scratch.seed_offsets;
let max_seeds_per_query = pack_seed_sets_for_gpu_for_capacity_into(
prepared.shape,
active_query_count,
seed_sets,
seed_triples,
seed_offsets,
)?;
if seed_triples.len() / 3 > scratch.max_seed_facts {
return Err(format!(
"weir IFDS resident parallel scratch stages {} seed fact(s) but solve requested {}. Fix: allocate scratch with allocate_resident_ifds_parallel_scratch_with_capacities for the intended seed batch.",
scratch.max_seed_facts,
seed_triples.len() / 3
));
}
let seed_triples_bytes = &mut host_scratch.seed_triples_bytes;
let seed_offsets_bytes = &mut host_scratch.seed_offsets_bytes;
scatter_seed_frontiers_resident_staged(
dispatch,
prepared.words,
active_query_count,
seed_triples,
seed_offsets,
max_seeds_per_query,
&scratch.frontiers,
&scratch.seed_triples,
&scratch.seed_offsets,
scratch.seed_triples_byte_len,
scratch.seed_offsets_byte_len,
&scratch.seed_clear_program,
&scratch.seed_scatter_program,
&scratch.seed_scatter_resources,
seed_triples_bytes,
seed_offsets_bytes,
)?;
let resources = [
prepared.pg_nodes.clone(),
prepared.row_ptr.clone(),
prepared.col_idx.clone(),
prepared.pg_edge_kind_mask.clone(),
prepared.pg_node_tags.clone(),
scratch.frontiers.clone(),
scratch.changed.clone(),
];
ensure_resident_parallel_step_programs(
prepared,
host_scratch,
scratch_query_capacity,
changed_slot_capacity,
max_iterations_usize,
)?;
let changed_bytes = &mut host_scratch.changed_bytes;
let frontier_bytes = &mut host_scratch.frontier_bytes;
let clear_changed_resources = std::slice::from_ref(&scratch.changed);
let repeated_steps = [
(
&scratch.clear_changed_program,
clear_changed_resources,
Some([1, 1, 1]),
),
(
&host_scratch.step_programs[0],
resources.as_slice(),
Some([prepared.node_count.max(1), query_count, 1]),
),
];
let changed_flag_bytes = std::mem::size_of::<u32>();
let expected_frontier_bytes =
scratch_words
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| {
"weir IFDS resident parallel final frontier byte length overflows usize. Fix: shard the IFDS batch before solving.".to_string()
})?;
if poll_each_iteration {
let mut converged = false;
for iteration in 0..max_iterations {
dispatch
.dispatch_resident_sequence_read_ranges_into(
&repeated_steps,
&[(&scratch.changed, 0, changed_flag_bytes)],
&mut [&mut *changed_bytes],
)
.map_err(|error| {
format!(
"weir IFDS resident parallel batch poll changed flag at iteration {iteration} failed: {error}"
)
})?;
let changed_word = crate::dispatch_decode::unpack_exact_u32_scalar(
changed_bytes,
"ifds resident parallel solve changed flag",
)?;
if changed_word == 0 {
converged = true;
break;
}
}
if !converged {
return Err(format!(
"IFDS resident parallel batch solve did not converge within {max_iterations} iterations"
));
}
dispatch
.download_resident_range_into(&scratch.frontiers, 0, expected_frontier_bytes, frontier_bytes)
.map_err(|error| {
format!(
"weir IFDS resident parallel batch download final frontier matrix failed: {error}"
)
})?;
} else {
dispatch
.dispatch_resident_repeated_sequence_read_ranges_into(
&[],
&repeated_steps,
max_iterations,
&[
(&scratch.changed, 0, changed_flag_bytes),
(&scratch.frontiers, 0, expected_frontier_bytes),
],
&mut [&mut *changed_bytes, &mut *frontier_bytes],
)
.map_err(|error| {
format!(
"weir IFDS resident parallel batch download final changed flag and frontier matrix failed: {error}"
)
})?;
let changed_word = crate::dispatch_decode::unpack_exact_u32_scalar(
changed_bytes,
"ifds resident parallel solve changed flag",
)?;
if changed_word != 0 {
return Err(format!(
"IFDS resident parallel batch solve did not converge within {max_iterations} iterations"
));
}
}
if frontier_bytes.len() != expected_frontier_bytes {
return Err(format!(
"ifds resident parallel batch final frontier matrix has {} bytes, expected {expected_frontier_bytes}. Fix: backend returned malformed resident frontier storage.",
frontier_bytes.len()
));
}
let query_frontier_bytes =
prepared
.words
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| {
"weir IFDS resident parallel query frontier byte length overflows usize. Fix: shard the IFDS graph before solving.".to_string()
})?;
resize_result_slots(results, seed_sets.len())?;
for (query_index, encoded) in results.iter_mut().enumerate().take(seed_sets.len()) {
let start = query_index
.checked_mul(prepared.words)
.and_then(|word| word.checked_mul(std::mem::size_of::<u32>()))
.ok_or_else(|| {
"weir IFDS resident parallel frontier slice byte offset overflows usize. Fix: shard the query batch before solving.".to_string()
})?;
let end = start
.checked_add(query_frontier_bytes)
.ok_or_else(|| {
"weir IFDS resident parallel frontier slice byte end overflows usize. Fix: shard the query batch before solving.".to_string()
})?;
encoded.clear();
ifds_encoded_frontier_nodes_from_le_bytes_into(
&frontier_bytes[start..end],
prepared.words,
prepared.node_count,
prepared.shape,
encoded,
)?;
}
Ok(())
}