#![allow(clippy::too_many_arguments)]
use crate::ifds_frontier_decode::ifds_encoded_frontier_nodes_into;
use crate::ifds_gpu::ifds_gpu_step;
use crate::ifds_resident_types::PreparedIfdsCsr;
use crate::ifds_seed_frontiers::seed_frontier_words_into;
use vyre_foundation::ir::Program;
pub(super) fn solve_prepared_borrowed_with_parts_via<F>(
dispatch: &F,
prepared: &PreparedIfdsCsr,
seed_facts: &[(u32, u32, u32)],
max_iterations: u32,
frontier: &mut Vec<u32>,
frontier_b: &mut Vec<u8>,
next: &mut Vec<u32>,
outputs: &mut Vec<Vec<u8>>,
) -> Result<Vec<u32>, String>
where
F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
let mut result = crate::staging_reserve::reserved_vec(
crate::dispatch_decode::u32_to_usize(
prepared.node_count,
"weir IFDS prepared solve node count",
)?,
"IFDS borrowed solve result node",
)?;
solve_prepared_borrowed_with_parts_into_via(
dispatch,
prepared,
seed_facts,
max_iterations,
frontier,
frontier_b,
next,
outputs,
&mut result,
)?;
Ok(result)
}
pub(super) fn solve_prepared_borrowed_with_parts_into_via<F>(
dispatch: &F,
prepared: &PreparedIfdsCsr,
seed_facts: &[(u32, u32, u32)],
max_iterations: u32,
frontier: &mut Vec<u32>,
frontier_b: &mut Vec<u8>,
next: &mut Vec<u32>,
outputs: &mut Vec<Vec<u8>>,
result: &mut Vec<u32>,
) -> Result<(), String>
where
F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
crate::dispatch_decode::require_positive_iterations(
"weir IFDS prepared GPU solve",
max_iterations,
)?;
let has_frontier = seed_frontier_words_into(
prepared.shape,
prepared.node_count,
prepared.words,
seed_facts,
frontier,
)?;
if !has_frontier {
result.clear();
return Ok(());
}
let program = ifds_gpu_step(prepared.shape, "fin", "fout")?;
for _ in 0..max_iterations {
crate::dispatch_decode::try_pack_u32_into(frontier, frontier_b, "weir u32 byte staging")?;
let inputs: [&[u8]; 7] = [
&prepared.pg_nodes_b,
&prepared.row_ptr_b,
&prepared.col_idx_b,
&prepared.pg_edge_kind_mask_b,
&prepared.pg_node_tags_b,
frontier_b.as_slice(),
frontier_b.as_slice(),
];
dispatch(
&program,
&inputs,
Some([prepared.node_count.max(1), 1, 1]),
outputs,
)?;
crate::dispatch_decode::unpack_only_exact_u32_into(
outputs,
"ifds solve_via",
"frontier",
prepared.words,
next,
)?;
crate::dispatch_decode::require_bitset_tail_clear(
"ifds solve_via output",
next,
prepared.node_count,
)?;
if next.as_slice() == frontier.as_slice() {
result.clear();
ifds_encoded_frontier_nodes_into(next, prepared.node_count, prepared.shape, result)?;
return Ok(());
}
std::mem::swap(frontier, next);
}
Err(format!(
"IFDS solve_via did not converge within {max_iterations} iterations"
))
}