use crate::ifds_gpu::{validate_ifds_problem, IfdsShape};
pub(crate) fn pack_seed_sets_for_gpu_into(
shape: IfdsShape,
seed_sets: &[&[(u32, u32, u32)]],
triples: &mut Vec<u32>,
offsets: &mut Vec<u32>,
) -> Result<u32, String> {
let plan = SeedPackWordPlan::for_seed_sets(seed_sets)?;
let mut max_per_query = 0u32;
for (query_index, seed_facts) in seed_sets.iter().enumerate() {
validate_ifds_problem(
"weir IFDS GPU seed scatter",
shape.num_procs,
shape.blocks_per_proc,
shape.facts_per_proc,
&[],
&[],
&[],
&[],
seed_facts,
)
.map_err(|error| format!("{error} Query index: {query_index}."))?;
let seed_len = u32::try_from(seed_facts.len()).map_err(|error| {
format!(
"weir IFDS GPU seed scatter query {query_index} has too many seed facts for u32 indexing: {error}. Fix: shard the IFDS seed batch before dispatch."
)
})?;
max_per_query = max_per_query.max(seed_len);
}
reserve_total(
triples,
plan.triple_words,
"IFDS GPU seed scatter triple word staging",
)?;
reserve_total(
offsets,
plan.offset_words,
"IFDS GPU seed scatter offset word staging",
)?;
let stable_triples_len = triples.len() == plan.triple_words;
let stable_offsets_len = offsets.len() == plan.offset_words;
if !stable_triples_len {
triples.clear();
}
if stable_offsets_len {
offsets[0] = 0;
} else {
offsets.clear();
offsets.push(0u32);
}
let mut triple_word_index = 0usize;
let mut previous_offset = 0u32;
for (query_index, seed_facts) in seed_sets.iter().enumerate() {
let seed_len = seed_facts.len() as u32;
for &(proc_id, block_id, fact_id) in *seed_facts {
if stable_triples_len {
triples[triple_word_index] = proc_id;
triples[triple_word_index + 1] = block_id;
triples[triple_word_index + 2] = fact_id;
triple_word_index += 3;
} else {
triples.extend_from_slice(&[proc_id, block_id, fact_id]);
}
}
let next = previous_offset
.checked_add(seed_len)
.ok_or_else(|| {
"weir IFDS GPU seed scatter offsets overflow u32 after validated seed packing. Fix: shard the IFDS seed batch before dispatch.".to_string()
})?;
if stable_offsets_len {
offsets[query_index + 1] = next;
} else {
offsets.push(next);
}
previous_offset = next;
}
Ok(max_per_query)
}
pub(crate) fn pack_seed_sets_for_gpu_for_capacity_into(
shape: IfdsShape,
query_capacity: usize,
seed_sets: &[&[(u32, u32, u32)]],
seed_triples: &mut Vec<u32>,
seed_offsets: &mut Vec<u32>,
) -> Result<u32, String> {
if seed_sets.is_empty() {
seed_triples.clear();
let offset_words = query_capacity.checked_add(1).ok_or_else(|| {
"weir IFDS resident seed scatter capacity offset count overflows host indexing. Fix: allocate smaller scratch or shard the batch."
.to_string()
})?;
crate::dispatch_decode::try_write_zero_words(
seed_offsets,
offset_words,
"IFDS resident seed scatter empty-offset staging",
)?;
return Ok(0);
}
if query_capacity == 0 {
return Err(
"weir IFDS resident seed scatter requires non-zero query capacity. Fix: allocate scratch before solving."
.to_string(),
);
}
if seed_sets.len() > query_capacity {
return Err(format!(
"weir IFDS resident seed scatter capacity is {query_capacity} queries but {} seed sets were requested. Fix: allocate larger scratch or shard the batch.",
seed_sets.len()
));
}
let max_seeds_per_query =
pack_seed_sets_for_gpu_into(shape, seed_sets, seed_triples, seed_offsets)?;
let final_seed_offset = seed_offsets.last().copied().ok_or_else(|| {
"weir IFDS resident seed scatter offsets were empty. Fix: pass at least one seed set."
.to_string()
})?;
let offset_words = query_capacity.checked_add(1).ok_or_else(|| {
"weir IFDS resident seed scatter capacity offset count overflows host indexing. Fix: allocate smaller scratch or shard the batch."
.to_string()
})?;
pad_seed_offsets_to_capacity(seed_offsets, offset_words, final_seed_offset)?;
Ok(max_seeds_per_query)
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct SeedPackWordPlan {
triple_words: usize,
offset_words: usize,
}
impl SeedPackWordPlan {
fn for_seed_sets(seed_sets: &[&[(u32, u32, u32)]]) -> Result<Self, String> {
let total_seed_facts = seed_sets.iter().try_fold(0usize, |total, seed_facts| {
total.checked_add(seed_facts.len()).ok_or_else(|| {
"weir IFDS GPU seed scatter seed fact count overflows host indexing. Fix: shard the IFDS seed batch before dispatch.".to_string()
})
})?;
let triple_words = total_seed_facts.checked_mul(3).ok_or_else(|| {
"weir IFDS GPU seed scatter triple word count overflows host indexing. Fix: shard the IFDS seed batch before dispatch.".to_string()
})?;
let offset_words = seed_sets.len().checked_add(1).ok_or_else(|| {
"weir IFDS GPU seed scatter offset word count overflows host indexing. Fix: shard the IFDS seed batch before dispatch.".to_string()
})?;
Ok(Self {
triple_words,
offset_words,
})
}
}
fn pad_seed_offsets_to_capacity(
seed_offsets: &mut Vec<u32>,
offset_words: usize,
final_seed_offset: u32,
) -> Result<(), String> {
crate::staging_reserve::reserve_vec(
seed_offsets,
offset_words,
"IFDS resident seed scatter padded offset staging",
)?;
if seed_offsets.len() < offset_words {
seed_offsets.extend(std::iter::repeat_n(
final_seed_offset,
offset_words - seed_offsets.len(),
));
} else {
seed_offsets.truncate(offset_words);
}
Ok(())
}
#[inline]
fn reserve_total<T>(values: &mut Vec<T>, total: usize, field: &'static str) -> Result<(), String> {
crate::staging_reserve::reserve_vec(values, total, field)
}
#[cfg(test)]
mod tests {
use super::{pack_seed_sets_for_gpu_for_capacity_into, pack_seed_sets_for_gpu_into};
use crate::ifds_gpu::IfdsShape;
#[test]
fn seed_pack_flattens_query_major_triples_and_offsets() {
let shape = IfdsShape {
num_procs: 2,
blocks_per_proc: 3,
facts_per_proc: 4,
edge_count: 0,
};
let first = [(0, 1, 2), (1, 2, 3)];
let second = [(1, 0, 0)];
let mut triples = Vec::new();
let mut offsets = Vec::new();
let max = pack_seed_sets_for_gpu_into(
shape,
&[first.as_slice(), second.as_slice()],
&mut triples,
&mut offsets,
)
.expect("valid seed triples must pack");
assert_eq!(max, 2);
assert_eq!(triples, [0, 1, 2, 1, 2, 3, 1, 0, 0]);
assert_eq!(offsets, [0, 2, 3]);
}
#[test]
fn seed_pack_capacity_pads_offsets_without_fake_triples() {
let shape = IfdsShape {
num_procs: 2,
blocks_per_proc: 3,
facts_per_proc: 4,
edge_count: 0,
};
let first = [(0, 1, 2)];
let mut triples = Vec::new();
let mut offsets = Vec::new();
let max = pack_seed_sets_for_gpu_for_capacity_into(
shape,
3,
&[first.as_slice()],
&mut triples,
&mut offsets,
)
.expect("capacity larger than query count must pad offsets");
assert_eq!(max, 1);
assert_eq!(triples, [0, 1, 2]);
assert_eq!(offsets, [0, 1, 1, 1]);
}
#[test]
fn seed_pack_validation_failure_preserves_existing_staging() {
let shape = IfdsShape {
num_procs: 1,
blocks_per_proc: 1,
facts_per_proc: 1,
edge_count: 0,
};
let valid_first = [(0, 0, 0)];
let invalid_second = [(0, 0, 9)];
let mut triples = vec![7, 8, 9];
let mut offsets = vec![0, 1];
let triples_capacity = triples.capacity();
let offsets_capacity = offsets.capacity();
let err = pack_seed_sets_for_gpu_into(
shape,
&[valid_first.as_slice(), invalid_second.as_slice()],
&mut triples,
&mut offsets,
)
.expect_err("invalid seed fact must reject the whole batch");
assert!(
err.contains("Query index: 1"),
"error should identify the invalid query, got {err}"
);
assert_eq!(triples, [7, 8, 9]);
assert_eq!(offsets, [0, 1]);
assert_eq!(triples.capacity(), triples_capacity);
assert_eq!(offsets.capacity(), offsets_capacity);
}
#[test]
fn seed_pack_source_has_no_resize_based_offset_padding() {
let source = include_str!("ifds_seed_pack.rs");
let production = source
.split("#[cfg(test)]")
.next()
.expect("IFDS seed pack production source must precede tests");
assert!(
!production.contains(".resize(")
&& production.contains("seed_offsets.extend(std::iter::repeat_n("),
"Fix: IFDS resident seed offset padding must extend after fallible reservation instead of resize-driven growth."
);
assert!(
production.contains("struct SeedPackWordPlan")
&& production.contains("fn for_seed_sets")
&& production.contains("fn pad_seed_offsets_to_capacity"),
"Fix: IFDS seed packing must keep sizing, reservation, and capacity padding modular."
);
assert!(
production
.find("for (query_index, seed_facts)")
.expect("validation pass must identify query indexes")
< production
.find("reserve_total(")
.expect("reservation must happen after validation"),
"Fix: IFDS seed packing must validate the full batch before mutating reusable staging buffers."
);
}
}