use crate::ifds_csr_split::split_columns_into;
use crate::ifds_frontier_decode::try_ifds_words;
use crate::ifds_gpu::{validate_ifds_problem, IfdsShape};
use crate::ifds_gpu_bytes::{
ensure_ifds_byte_slots, write_ifds_padded_u32_bytes, write_ifds_zero_bytes,
};
use crate::ifds_resident_types::{IfdsPrepareScratch, PreparedIfdsCsr};
use vyre_foundation::ir::Program;
use vyre_primitives::graph::exploded::{build_ifds_csr_program, validate_ifds_csr_layout};
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_exploded_csr_borrowed_into_via<F>(
dispatch: &F,
num_procs: u32,
blocks_per_proc: u32,
facts_per_proc: u32,
intra_edges: &[(u32, u32, u32)],
inter_edges: &[(u32, u32, u32, u32)],
flow_gen: &[(u32, u32, u32)],
flow_kill: &[(u32, u32, u32)],
scratch: &mut IfdsPrepareScratch,
) -> Result<(Vec<u32>, Vec<u32>), String>
where
F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
let intra_count = u32::try_from(intra_edges.len()).map_err(|_| {
format!(
"weir IFDS GPU CSR build received {} intra edges, exceeding u32 index space. Fix: shard the IFDS graph before dispatch.",
intra_edges.len()
)
})?;
let inter_count = u32::try_from(inter_edges.len()).map_err(|_| {
format!(
"weir IFDS GPU CSR build received {} inter edges, exceeding u32 index space. Fix: shard the IFDS graph before dispatch.",
inter_edges.len()
)
})?;
let gen_count = u32::try_from(flow_gen.len()).map_err(|_| {
format!(
"weir IFDS GPU CSR build received {} GEN entries, exceeding u32 index space. Fix: shard the IFDS graph before dispatch.",
flow_gen.len()
)
})?;
let kill_count = u32::try_from(flow_kill.len()).map_err(|_| {
format!(
"weir IFDS GPU CSR build received {} KILL entries, exceeding u32 index space. Fix: shard the IFDS graph before dispatch.",
flow_kill.len()
)
})?;
let layout = validate_ifds_csr_layout(
num_procs,
blocks_per_proc,
facts_per_proc,
intra_count,
inter_count,
gen_count,
)
.map_err(|error| format!("weir IFDS GPU CSR build invalid layout: {error}"))?;
split_columns_into(
intra_edges,
&mut [
&mut scratch.intra_proc,
&mut scratch.intra_src_block,
&mut scratch.intra_dst_block,
],
&[
"IFDS CSR split triple first column",
"IFDS CSR split triple second column",
"IFDS CSR split triple third column",
],
)?;
split_columns_into(
inter_edges,
&mut [
&mut scratch.inter_src_proc,
&mut scratch.inter_src_block,
&mut scratch.inter_dst_proc,
&mut scratch.inter_dst_block,
],
&[
"IFDS CSR split quad first column",
"IFDS CSR split quad second column",
"IFDS CSR split quad third column",
"IFDS CSR split quad fourth column",
],
)?;
split_columns_into(
flow_gen,
&mut [
&mut scratch.gen_proc,
&mut scratch.gen_block,
&mut scratch.gen_fact,
],
&[
"IFDS CSR split triple first column",
"IFDS CSR split triple second column",
"IFDS CSR split triple third column",
],
)?;
split_columns_into(
flow_kill,
&mut [
&mut scratch.kill_proc,
&mut scratch.kill_block,
&mut scratch.kill_fact,
],
&[
"IFDS CSR split triple first column",
"IFDS CSR split triple second column",
"IFDS CSR split triple third column",
],
)?;
let program = build_ifds_csr_program(
num_procs,
blocks_per_proc,
facts_per_proc,
intra_count,
inter_count,
gen_count,
kill_count,
layout.max_col_count,
);
ensure_ifds_byte_slots(&mut scratch.inputs, 17)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[0], &scratch.intra_proc)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[1], &scratch.intra_src_block)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[2], &scratch.intra_dst_block)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[3], &scratch.inter_src_proc)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[4], &scratch.inter_src_block)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[5], &scratch.inter_dst_proc)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[6], &scratch.inter_dst_block)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[7], &scratch.gen_proc)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[8], &scratch.gen_block)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[9], &scratch.gen_fact)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[10], &scratch.kill_proc)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[11], &scratch.kill_block)?;
write_ifds_padded_u32_bytes(&mut scratch.inputs[12], &scratch.kill_fact)?;
write_ifds_zero_bytes(
&mut scratch.inputs[13],
layout
.row_words
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| {
"weir IFDS row_ptr byte length overflows usize. Fix: shard the IFDS graph before dispatch.".to_string()
})?,
)?;
write_ifds_zero_bytes(
&mut scratch.inputs[14],
layout
.row_cursor_words
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| {
"weir IFDS row_cursor byte length overflows usize. Fix: shard the IFDS graph before dispatch.".to_string()
})?,
)?;
write_ifds_zero_bytes(
&mut scratch.inputs[15],
layout
.col_buffer_words
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| {
"weir IFDS col_idx byte length overflows usize. Fix: shard the IFDS graph before dispatch.".to_string()
})?,
)?;
write_ifds_zero_bytes(&mut scratch.inputs[16], 4)?;
let input_refs = [
scratch.inputs[0].as_slice(),
scratch.inputs[1].as_slice(),
scratch.inputs[2].as_slice(),
scratch.inputs[3].as_slice(),
scratch.inputs[4].as_slice(),
scratch.inputs[5].as_slice(),
scratch.inputs[6].as_slice(),
scratch.inputs[7].as_slice(),
scratch.inputs[8].as_slice(),
scratch.inputs[9].as_slice(),
scratch.inputs[10].as_slice(),
scratch.inputs[11].as_slice(),
scratch.inputs[12].as_slice(),
scratch.inputs[13].as_slice(),
scratch.inputs[14].as_slice(),
scratch.inputs[15].as_slice(),
scratch.inputs[16].as_slice(),
];
dispatch(&program, &input_refs, Some([1, 1, 1]), &mut scratch.outputs)?;
crate::dispatch_decode::require_output_count(
&scratch.outputs,
"weir IFDS GPU CSR build",
4,
"row_ptr,row_cursor,col_idx,col_len",
)?;
let row_ptr = crate::dispatch_decode::unpack_exact_u32(
&scratch.outputs[0],
layout.row_words,
"weir IFDS row_ptr",
)?;
let col_len =
crate::dispatch_decode::unpack_exact_u32(&scratch.outputs[3], 1, "weir IFDS col_len")?[0];
if col_len > layout.max_col_count {
return Err(format!(
"weir IFDS GPU CSR build reported col_len {col_len} above allocated maximum {}. Fix: repair IFDS row counting.",
layout.max_col_count
));
}
let mut col_idx = crate::dispatch_decode::unpack_exact_u32(
&scratch.outputs[2],
layout.col_buffer_words,
"weir IFDS col_idx",
)?;
col_idx.truncate(crate::dispatch_decode::u32_to_usize(
col_len,
"weir IFDS reported CSR column count",
)?);
crate::dispatch_decode::require_csr_offsets_targets(
"weir IFDS GPU CSR build output",
layout.total_nodes,
&row_ptr,
&col_idx,
)?;
Ok((row_ptr, col_idx))
}
#[allow(clippy::too_many_arguments)]
pub fn prepare_ifds_csr_borrowed_via<F>(
dispatch: &F,
num_procs: u32,
blocks_per_proc: u32,
facts_per_proc: u32,
intra_edges: &[(u32, u32, u32)],
inter_edges: &[(u32, u32, u32, u32)],
flow_gen: &[(u32, u32, u32)],
flow_kill: &[(u32, u32, u32)],
) -> Result<PreparedIfdsCsr, String>
where
F: Fn(&Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
let mut scratch = IfdsPrepareScratch::default();
prepare_ifds_csr_borrowed_with_scratch_via(
&|program, inputs, grid, outputs| {
let result = dispatch(program, inputs, grid)?;
crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
Ok(())
},
num_procs,
blocks_per_proc,
facts_per_proc,
intra_edges,
inter_edges,
flow_gen,
flow_kill,
&mut scratch,
)
}
#[allow(clippy::too_many_arguments)]
pub fn prepare_ifds_csr_borrowed_with_scratch_via<F>(
dispatch: &F,
num_procs: u32,
blocks_per_proc: u32,
facts_per_proc: u32,
intra_edges: &[(u32, u32, u32)],
inter_edges: &[(u32, u32, u32, u32)],
flow_gen: &[(u32, u32, u32)],
flow_kill: &[(u32, u32, u32)],
scratch: &mut IfdsPrepareScratch,
) -> Result<PreparedIfdsCsr, String>
where
F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
validate_ifds_problem(
"weir IFDS GPU CSR prepare",
num_procs,
blocks_per_proc,
facts_per_proc,
intra_edges,
inter_edges,
flow_gen,
flow_kill,
&[],
)?;
let pre_edge_count = intra_edges
.len()
.checked_add(inter_edges.len())
.and_then(|count| count.checked_add(flow_gen.len()))
.and_then(|count| count.checked_add(flow_kill.len()))
.ok_or_else(|| {
"weir IFDS CSR prepare edge count overflows usize. Fix: shard the IFDS problem before dispatch.".to_string()
})?;
let pre_edge_count = u32::try_from(pre_edge_count).map_err(|error| {
format!(
"weir IFDS CSR prepare edge count does not fit u32: {error}. Fix: shard the IFDS problem or reduce the encoded domain before dispatch."
)
})?;
let pre_shape = IfdsShape {
num_procs,
blocks_per_proc,
facts_per_proc,
edge_count: pre_edge_count,
};
if !pre_shape.fits() {
return Err(format!(
"weir IFDS CSR prepare dimensions exceed 32-bit exploded-node encoding: procs={num_procs} blocks={blocks_per_proc} facts={facts_per_proc}. Fix: shard the IFDS problem or reduce the encoded domain before dispatch."
));
}
let (row_ptr, col_idx) = build_exploded_csr_borrowed_into_via(
dispatch,
num_procs,
blocks_per_proc,
facts_per_proc,
intra_edges,
inter_edges,
flow_gen,
flow_kill,
scratch,
)?;
let shape = IfdsShape {
edge_count: u32::try_from(col_idx.len()).map_err(|error| {
format!(
"weir IFDS CSR prepare built CSR has too many edges for u32 encoding: {error}. Fix: shard the IFDS problem before dispatch."
)
})?,
..pre_shape
};
let node_count = shape.node_count()?;
let words = try_ifds_words(node_count)?;
Ok(PreparedIfdsCsr {
shape,
node_count,
words,
pg_nodes_b: crate::dispatch_decode::pack_repeated_u32(
0,
crate::dispatch_decode::u32_to_usize(node_count, "weir IFDS node count")?,
)?,
row_ptr_b: crate::dispatch_decode::try_pack_u32(&row_ptr, "weir IFDS CSR row_ptr bytes")?,
col_idx_b: crate::dispatch_decode::try_pack_u32(&col_idx, "weir IFDS CSR col_idx bytes")?,
pg_edge_kind_mask_b: crate::dispatch_decode::pack_repeated_u32(1, col_idx.len())?,
pg_node_tags_b: crate::dispatch_decode::pack_repeated_u32(
0,
crate::dispatch_decode::u32_to_usize(node_count, "weir IFDS node count")?,
)?,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn csr_prepare_truncates_overallocated_col_idx_without_copying_tail() {
let mut scratch = IfdsPrepareScratch::default();
let dispatch = |_: &Program,
_: &[&[u8]],
_: Option<[u32; 3]>,
outputs: &mut Vec<Vec<u8>>|
-> Result<(), String> {
outputs.clear();
outputs.push(crate::dispatch_decode::pack_u32(&[0, 1, 1]));
outputs.push(crate::dispatch_decode::pack_u32(&[0, 0]));
outputs.push(crate::dispatch_decode::pack_u32(&[1, 0]));
outputs.push(crate::dispatch_decode::pack_u32(&[1]));
Ok(())
};
let (row_ptr, col_idx) = build_exploded_csr_borrowed_into_via(
&dispatch,
1,
2,
1,
&[(0, 0, 1)],
&[],
&[(0, 0, 0)],
&[],
&mut scratch,
)
.expect("fixture CSR output must decode");
assert_eq!(row_ptr, vec![0, 1, 1]);
assert_eq!(col_idx, vec![1]);
}
}