use std::sync::Arc;
use crate::ifds_gpu::{validate_ifds_problem, IfdsShape};
use crate::ifds_gpu_bytes::padded_pack_u32_into;
use crate::ifds_resident_types::IfdsResidentDispatch;
use vyre_foundation::ir::model::expr::Ident;
use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
const IFDS_SEED_TRIPLES: &str = "ifds_seed_triples";
const IFDS_SEED_OFFSETS: &str = "ifds_seed_offsets";
const IFDS_FRONTIERS: &str = "frontiers";
const IFDS_SEED_CLEAR_OP: &str = "weir::ifds_gpu::clear_frontiers";
const IFDS_SEED_SCATTER_OP: &str = "weir::ifds_gpu::scatter_seed_frontiers";
fn set_frontier_bit(bits: &mut [u32], dense: u32) -> Result<(), String> {
let word = crate::dispatch_decode::u32_to_usize(dense / 32, "IFDS seed frontier word")?;
let bit = dense & 31;
let word_count = bits.len();
let slot = bits.get_mut(word).ok_or_else(|| {
format!(
"weir IFDS seed frontier dense node {dense} maps to word {word}, outside {word_count} frontier words. Fix: size seed frontier storage from the IFDS node count before scatter."
)
})?;
*slot |= 1u32 << bit;
Ok(())
}
pub(crate) fn seed_frontier_words_into(
shape: IfdsShape,
node_count: u32,
words: usize,
seed_facts: &[(u32, u32, u32)],
frontier: &mut Vec<u32>,
) -> Result<bool, String> {
validate_ifds_problem(
"weir IFDS frontier seed",
shape.num_procs,
shape.blocks_per_proc,
shape.facts_per_proc,
&[],
&[],
&[],
&[],
seed_facts,
)?;
if seed_facts.is_empty() {
frontier.clear();
return Ok(false);
}
crate::dispatch_decode::try_write_zero_words(frontier, words, "weir zero word staging")?;
let mut has_frontier = false;
for &(proc_id, block_id, fact_id) in seed_facts {
let proc_base = proc_id
.checked_mul(shape.blocks_per_proc)
.and_then(|value| value.checked_mul(shape.facts_per_proc))
.ok_or_else(|| {
"weir IFDS seed frontier proc*blocks*facts overflowed u32. Fix: shard the IFDS problem before seed scatter."
.to_string()
})?;
let block_base = block_id.checked_mul(shape.facts_per_proc).ok_or_else(|| {
"weir IFDS seed frontier block*facts overflowed u32. Fix: shard the IFDS problem before seed scatter."
.to_string()
})?;
let dense = proc_base
.checked_add(block_base)
.and_then(|value| value.checked_add(fact_id))
.ok_or_else(|| {
"weir IFDS seed frontier dense id overflowed u32. Fix: shard the IFDS problem before seed scatter."
.to_string()
})?;
if dense < node_count {
set_frontier_bit(frontier, dense)?;
has_frontier = true;
}
}
Ok(has_frontier)
}
pub(crate) fn ifds_clear_frontiers_program(
words: u32,
query_count: u32,
) -> Result<Program, String> {
let frontier_words = words
.checked_mul(query_count)
.ok_or_else(|| {
"weir IFDS seed clear frontier matrix word count overflowed u32. Fix: shard the query batch before resident seed scatter."
.to_string()
})?;
let word = Expr::InvocationId { axis: 0 };
let query = Expr::InvocationId { axis: 1 };
Ok(Program::wrapped(
vec![
BufferDecl::storage(IFDS_FRONTIERS, 0, BufferAccess::ReadWrite, DataType::U32)
.with_count(frontier_words.max(1)),
],
[1, 1, 1],
vec![Node::Region {
generator: Ident::from(IFDS_SEED_CLEAR_OP),
source_region: None,
body: Arc::new(vec![Node::if_then(
Expr::lt(word.clone(), Expr::u32(words)),
vec![Node::store(
IFDS_FRONTIERS,
Expr::add(Expr::mul(query, Expr::u32(words)), word),
Expr::u32(0),
)],
)]),
}],
))
}
pub(crate) fn ifds_seed_frontiers_program(
shape: IfdsShape,
words: u32,
query_count: u32,
seed_count: u32,
max_seeds_per_query: u32,
) -> Result<Program, String> {
if !shape.fits() {
return Err(format!(
"weir IFDS seed scatter dimensions exceed 32-bit exploded-node encoding: procs={} blocks={} facts={}. Fix: shard the IFDS problem before resident seed scatter.",
shape.num_procs, shape.blocks_per_proc, shape.facts_per_proc
));
}
let block_fact_span = shape
.blocks_per_proc
.checked_mul(shape.facts_per_proc)
.ok_or_else(|| {
"weir IFDS seed scatter blocks*facts overflowed u32. Fix: shard the IFDS problem before resident seed scatter."
.to_string()
})?;
let frontier_words = words
.checked_mul(query_count)
.ok_or_else(|| {
"weir IFDS seed scatter frontier matrix word count overflowed u32. Fix: shard the query batch before resident seed scatter."
.to_string()
})?;
let seed_words = seed_count
.checked_mul(3)
.ok_or_else(|| {
"weir IFDS seed scatter triple word count overflowed u32. Fix: shard the seed batch before resident seed scatter."
.to_string()
})?;
let seed_offset_words = query_count.checked_add(1).ok_or_else(|| {
"weir IFDS seed scatter offset word count overflowed u32. Fix: shard the query batch before resident seed scatter."
.to_string()
})?;
let local_seed = Expr::InvocationId { axis: 0 };
let query = Expr::InvocationId { axis: 1 };
let seed_idx = Expr::add(
Expr::load(IFDS_SEED_OFFSETS, query.clone()),
local_seed.clone(),
);
let seed_base = Expr::mul(Expr::var("seed_idx"), Expr::u32(3));
let dense = Expr::add(
Expr::add(
Expr::mul(
Expr::load(IFDS_SEED_TRIPLES, seed_base.clone()),
Expr::u32(block_fact_span),
),
Expr::mul(
Expr::load(
IFDS_SEED_TRIPLES,
Expr::add(seed_base.clone(), Expr::u32(1)),
),
Expr::u32(shape.facts_per_proc),
),
),
Expr::load(IFDS_SEED_TRIPLES, Expr::add(seed_base, Expr::u32(2))),
);
let body = vec![
Node::let_bind("seed_start", Expr::load(IFDS_SEED_OFFSETS, query.clone())),
Node::let_bind(
"seed_end",
Expr::load(IFDS_SEED_OFFSETS, Expr::add(query.clone(), Expr::u32(1))),
),
Node::if_then(
Expr::lt(
local_seed.clone(),
Expr::sub(Expr::var("seed_end"), Expr::var("seed_start")),
),
vec![
Node::let_bind("seed_idx", seed_idx),
Node::let_bind("dense", dense),
Node::let_bind(
"frontier_word",
Expr::add(
Expr::mul(query, Expr::u32(words)),
Expr::shr(Expr::var("dense"), Expr::u32(5)),
),
),
Node::let_bind(
"frontier_bit",
Expr::shl(
Expr::u32(1),
Expr::bitand(Expr::var("dense"), Expr::u32(31)),
),
),
Node::let_bind(
"_seed_or",
Expr::atomic_or(
IFDS_FRONTIERS,
Expr::var("frontier_word"),
Expr::var("frontier_bit"),
),
),
],
),
];
Ok(Program::wrapped(
vec![
BufferDecl::storage(IFDS_SEED_TRIPLES, 0, BufferAccess::ReadOnly, DataType::U32)
.with_count(seed_words.max(1)),
BufferDecl::storage(IFDS_SEED_OFFSETS, 1, BufferAccess::ReadOnly, DataType::U32)
.with_count(seed_offset_words),
BufferDecl::storage(IFDS_FRONTIERS, 2, BufferAccess::ReadWrite, DataType::U32)
.with_count(frontier_words.max(1)),
],
[1, 1, 1],
vec![Node::Region {
generator: Ident::from(IFDS_SEED_SCATTER_OP),
source_region: None,
body: Arc::new(vec![Node::if_then(
Expr::lt(local_seed.clone(), Expr::u32(max_seeds_per_query)),
body,
)]),
}],
))
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn scatter_seed_frontiers_resident_staged<D>(
dispatch: &D,
words: usize,
query_capacity: usize,
seed_triples: &[u32],
seed_offsets: &[u32],
max_seeds_per_query: u32,
frontiers: &D::Resource,
seed_triples_resource: &D::Resource,
seed_offsets_resource: &D::Resource,
seed_triples_byte_len: usize,
seed_offsets_byte_len: usize,
clear_program: &Program,
scatter_program: &Program,
scatter_resources: &[D::Resource],
seed_triples_bytes: &mut Vec<u8>,
seed_offsets_bytes: &mut Vec<u8>,
) -> Result<(), String>
where
D: IfdsResidentDispatch,
{
let query_count = u32::try_from(query_capacity).map_err(|error| {
format!(
"weir IFDS resident seed scatter query capacity does not fit u32: {error}. Fix: allocate smaller scratch or shard the batch."
)
})?;
let words_u32 = u32::try_from(words).map_err(|error| {
format!(
"weir IFDS resident seed scatter frontier word count does not fit u32: {error}. Fix: shard the IFDS problem before solving."
)
})?;
padded_pack_u32_into(seed_triples, seed_triples_bytes)?;
crate::dispatch_decode::try_pack_u32_into(
seed_offsets,
seed_offsets_bytes,
"IFDS resident seed offset byte staging",
)?;
if seed_triples_bytes.len() > seed_triples_byte_len {
return Err(format!(
"weir IFDS resident seed scatter triple staging buffer holds {seed_triples_byte_len} bytes but solve needs {}. Fix: allocate larger seed scratch or shard the seed batch.",
seed_triples_bytes.len()
));
}
if seed_offsets_bytes.len() > seed_offsets_byte_len {
return Err(format!(
"weir IFDS resident seed scatter offset staging buffer holds {seed_offsets_byte_len} bytes but solve needs {}. Fix: allocate larger query scratch or shard the batch.",
seed_offsets_bytes.len()
));
}
crate::dispatch_decode::try_pad_zero_bytes_to_len(
seed_offsets_bytes,
seed_offsets_byte_len,
"IFDS resident seed offset padding",
)?;
dispatch
.upload_resident_at_many(&[
(seed_triples_resource, 0, seed_triples_bytes.as_slice()),
(seed_offsets_resource, 0, seed_offsets_bytes.as_slice()),
])
.map_err(|error| format!("weir IFDS resident seed scatter batch upload failed: {error}"))?;
dispatch.dispatch_resident(
clear_program,
std::slice::from_ref(frontiers),
Some([words_u32.max(1), query_count, 1]),
)?;
dispatch.dispatch_resident(
scatter_program,
scatter_resources,
Some([max_seeds_per_query.max(1), query_count, 1]),
)
}
#[cfg(test)]
mod tests {
use super::{
ifds_clear_frontiers_program, ifds_seed_frontiers_program,
scatter_seed_frontiers_resident_staged,
};
use crate::ifds_gpu::IfdsShape;
use crate::ifds_resident_types::IfdsResidentDispatch;
use std::cell::RefCell;
use vyre_foundation::ir::Program;
struct UploadLenDispatch {
upload_lengths: RefCell<Vec<usize>>,
}
impl IfdsResidentDispatch for UploadLenDispatch {
type Resource = u32;
fn resident_backend_id(&self) -> &'static str {
"weir_test_seed_scatter_upload_lengths"
}
fn allocate_resident(&self, _byte_len: usize) -> Result<Self::Resource, String> {
unreachable!("seed scatter upload length test does not allocate")
}
fn upload_resident(&self, _resource: &Self::Resource, bytes: &[u8]) -> Result<(), String> {
self.upload_lengths.borrow_mut().push(bytes.len());
Ok(())
}
fn upload_resident_many(&self, uploads: &[(&Self::Resource, &[u8])]) -> Result<(), String> {
for &(_, bytes) in uploads {
self.upload_lengths.borrow_mut().push(bytes.len());
}
Ok(())
}
fn download_resident(&self, _resource: &Self::Resource) -> Result<Vec<u8>, String> {
unreachable!("seed scatter upload length test does not download")
}
fn download_resident_into(
&self,
_resource: &Self::Resource,
_output: &mut Vec<u8>,
) -> Result<(), String> {
unreachable!("seed scatter upload length test does not download")
}
fn download_resident_range(
&self,
_resource: &Self::Resource,
_byte_offset: usize,
_byte_len: usize,
) -> Result<Vec<u8>, String> {
unreachable!("seed scatter upload length test does not download")
}
fn download_resident_range_into(
&self,
_resource: &Self::Resource,
_byte_offset: usize,
_byte_len: usize,
_output: &mut Vec<u8>,
) -> Result<(), String> {
unreachable!("seed scatter upload length test does not download")
}
fn free_resident(&self, _resource: Self::Resource) -> Result<(), String> {
unreachable!("seed scatter upload length test does not free")
}
fn dispatch_resident(
&self,
_program: &Program,
_resources: &[Self::Resource],
_grid_override: Option<[u32; 3]>,
) -> Result<(), String> {
Ok(())
}
}
#[test]
fn resident_seed_scatter_uploads_only_packed_seed_triples() {
let dispatch = UploadLenDispatch {
upload_lengths: RefCell::new(Vec::new()),
};
let shape = IfdsShape {
num_procs: 1,
blocks_per_proc: 8,
facts_per_proc: 4,
edge_count: 0,
};
let clear_program = ifds_clear_frontiers_program(1, 2)
.expect("clear frontier program must fit test dimensions");
let scatter_program = ifds_seed_frontiers_program(shape, 1, 2, 1, 1)
.expect("seed scatter program must fit test dimensions");
let mut seed_triples_bytes = Vec::new();
let mut seed_offsets_bytes = Vec::new();
scatter_seed_frontiers_resident_staged(
&dispatch,
1,
2,
&[0, 1, 2],
&[0, 1, 1],
1,
&0,
&1,
&2,
4096,
12,
&clear_program,
&scatter_program,
&[1, 2, 0],
&mut seed_triples_bytes,
&mut seed_offsets_bytes,
)
.expect("seed scatter staging must accept packed seed triples below resident capacity");
assert_eq!(
dispatch.upload_lengths.borrow().as_slice(),
&[12, 12],
"resident seed scatter must upload packed seed triples, not the full retained seed capacity"
);
}
#[test]
fn seed_frontier_program_builders_reject_dimension_overflow_without_panic() {
let clear_error = ifds_clear_frontiers_program(u32::MAX, 2)
.expect_err("clear frontier word matrix overflow must return an error");
assert!(
clear_error.contains("frontier matrix word count overflowed u32")
&& clear_error.contains("Fix: shard the query batch"),
"clear frontier overflow diagnostic must name the failed matrix and operator action"
);
let shape = IfdsShape {
num_procs: u32::MAX,
blocks_per_proc: u32::MAX,
facts_per_proc: u32::MAX,
edge_count: 0,
};
let scatter_error = ifds_seed_frontiers_program(shape, 1, 1, 1, 1)
.expect_err("invalid seed scatter shape must return an error");
assert!(
scatter_error.contains("exceed 32-bit exploded-node encoding")
&& scatter_error.contains("Fix: shard the IFDS problem"),
"seed scatter shape diagnostic must name the encoding bound and operator action"
);
}
}