1use std::collections::HashSet;
15use std::sync::Arc;
16
17use vyre_foundation::ir::model::expr::Ident;
18use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
19use vyre_foundation::MemoryOrdering;
20
21use crate::bitset::bitset_words;
22use crate::bitset::frontier::{
23 frontier_absorb_new_bits_body_prefixed_with_flag, frontier_tail_mask,
24};
25use crate::graph::program_graph::{
26 ProgramGraphShape, BINDING_PRIMITIVE_START, NAME_EDGE_KIND_MASK, NAME_EDGE_OFFSETS,
27 NAME_EDGE_TARGETS,
28};
29
30pub const OP_ID: &str = "vyre-primitives::graph::reachable_program";
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub struct UnknownNode {
37 pub index: usize,
39 pub node: u32,
41 pub node_count: u32,
43}
44
45impl std::fmt::Display for UnknownNode {
46 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47 write!(
48 f,
49 "reachable: edges[{}] references node {} but node_count = {}. \
50 Fix: callers must deduplicate and bounds-check edges before \
51 calling this primitive.",
52 self.index, self.node, self.node_count
53 )
54 }
55}
56
57impl std::error::Error for UnknownNode {}
58
59#[derive(Debug, Clone, PartialEq, Eq)]
62pub enum ReachableError {
63 UnknownNode(UnknownNode),
65 Allocation(String),
67}
68
69impl std::fmt::Display for ReachableError {
70 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71 match self {
72 Self::UnknownNode(error) => error.fmt(f),
73 Self::Allocation(message) => f.write_str(message),
74 }
75 }
76}
77
78impl std::error::Error for ReachableError {}
79
80impl From<UnknownNode> for ReachableError {
81 fn from(error: UnknownNode) -> Self {
82 Self::UnknownNode(error)
83 }
84}
85
86pub fn reachable(
97 node_count: u32,
98 edges: &[(u32, u32)],
99 sources: &[u32],
100) -> Result<HashSet<u32>, UnknownNode> {
101 match try_reachable(node_count, edges, sources) {
102 Ok(result) => Ok(result),
103 Err(ReachableError::UnknownNode(error)) => Err(error),
104 Err(ReachableError::Allocation(_)) => Ok(HashSet::new()),
105 }
106}
107
108pub fn try_reachable(
114 node_count: u32,
115 edges: &[(u32, u32)],
116 sources: &[u32],
117) -> Result<HashSet<u32>, ReachableError> {
118 const NONE: usize = usize::MAX;
119
120 let n = node_count as usize;
121 for (index, &(from, to)) in edges.iter().enumerate() {
122 if (from as usize) >= n {
123 return Err(ReachableError::UnknownNode(UnknownNode {
124 index,
125 node: from,
126 node_count,
127 }));
128 }
129 if (to as usize) >= n {
130 return Err(ReachableError::UnknownNode(UnknownNode {
131 index,
132 node: to,
133 node_count,
134 }));
135 }
136 }
137 let mut head: Vec<usize> = Vec::new();
138 crate::graph::scratch::reserve_graph_items(
139 &mut head,
140 n,
141 "reachable CPU oracle",
142 "adjacency heads",
143 )
144 .map_err(ReachableError::Allocation)?;
145 head.resize(n, NONE);
146 let mut to_nodes: Vec<u32> = Vec::new();
147 crate::graph::scratch::reserve_graph_items(
148 &mut to_nodes,
149 edges.len(),
150 "reachable CPU oracle",
151 "adjacency destinations",
152 )
153 .map_err(ReachableError::Allocation)?;
154 let mut next_edges: Vec<usize> = Vec::new();
155 crate::graph::scratch::reserve_graph_items(
156 &mut next_edges,
157 edges.len(),
158 "reachable CPU oracle",
159 "adjacency next links",
160 )
161 .map_err(ReachableError::Allocation)?;
162 for &(from, to) in edges {
163 let edge_index = to_nodes.len();
164 to_nodes.push(to);
165 next_edges.push(head[from as usize]);
166 head[from as usize] = edge_index;
167 }
168 let mut visited: Vec<bool> = Vec::new();
169 crate::graph::scratch::reserve_graph_items(
170 &mut visited,
171 n,
172 "reachable CPU oracle",
173 "visited bitmap",
174 )
175 .map_err(ReachableError::Allocation)?;
176 visited.resize(n, false);
177 let mut out_of_range_sources: Vec<u32> = Vec::new();
178 crate::graph::scratch::reserve_graph_items(
179 &mut out_of_range_sources,
180 sources.len(),
181 "reachable CPU oracle",
182 "out-of-range source list",
183 )
184 .map_err(ReachableError::Allocation)?;
185 let mut stack: Vec<u32> = Vec::new();
186 crate::graph::scratch::reserve_graph_items(
187 &mut stack,
188 sources.len(),
189 "reachable CPU oracle",
190 "DFS stack",
191 )
192 .map_err(ReachableError::Allocation)?;
193 stack.extend_from_slice(sources);
194 while let Some(v) = stack.pop() {
195 let idx = v as usize;
196 if idx >= n {
197 out_of_range_sources.push(v);
198 continue;
199 }
200 if visited[idx] {
201 continue;
202 }
203 visited[idx] = true;
204 let mut edge = head[idx];
205 while edge != NONE {
206 let next = to_nodes[edge];
207 if !visited[next as usize] {
208 stack.push(next);
209 }
210 edge = next_edges[edge];
211 }
212 }
213 let result_capacity = visited
214 .iter()
215 .filter(|&&is_visited| is_visited)
216 .count()
217 .saturating_add(out_of_range_sources.len());
218 let mut result = HashSet::new();
219 result.try_reserve(result_capacity).map_err(|error| {
220 ReachableError::Allocation(format!(
221 "Fix: reachable CPU oracle could not reserve {result_capacity} result nodes: {error}"
222 ))
223 })?;
224 for (idx, is_visited) in visited.into_iter().enumerate() {
225 if is_visited {
226 result.insert(idx as u32);
227 }
228 }
229 result.extend(out_of_range_sources);
230 Ok(result)
231}
232
233#[must_use]
257pub fn reachable_program(
258 node_count: u32,
259 edge_count: u32,
260 sources_buf: &str,
261 reach_out: &str,
262 max_iters: u32,
263) -> Program {
264 let shape = ProgramGraphShape::new(node_count, edge_count);
265 let words = bitset_words(node_count);
266 let frontier_a = "reach_frontier_a";
267 let frontier_b = "reach_frontier_b";
268 let active_flag_idx = words;
269 let Some(frontier_b_storage_words) = words.checked_add(1) else {
270 return crate::invalid_output_program(
271 OP_ID,
272 reach_out,
273 DataType::U32,
274 "Fix: reachable_program active-flag scratch word overflows u32.".to_string(),
275 );
276 };
277
278 let Some(iter_nodes) = (max_iters as usize).checked_mul(8) else {
279 return crate::invalid_output_program(
280 OP_ID,
281 reach_out,
282 DataType::U32,
283 "Fix: reachable_program max_iters*8 overflows usize.".to_string(),
284 );
285 };
286 let Some(node_capacity) = iter_nodes.checked_add(4) else {
287 return crate::invalid_output_program(
288 OP_ID,
289 reach_out,
290 DataType::U32,
291 "Fix: reachable_program node capacity overflows usize.".to_string(),
292 );
293 };
294 let mut entry: Vec<Node> = Vec::new();
295 if let Err(error) = entry.try_reserve(node_capacity) {
296 return crate::invalid_output_program(
297 OP_ID,
298 reach_out,
299 DataType::U32,
300 format!("Fix: reachable_program could not reserve {node_capacity} IR nodes: {error}"),
301 );
302 }
303 let lane = Expr::gid_x();
304
305 entry.push(Node::if_then(
306 Expr::lt(lane.clone(), Expr::u32(words)),
307 vec![
308 Node::store(
309 reach_out,
310 lane.clone(),
311 Expr::load(sources_buf, lane.clone()),
312 ),
313 Node::store(frontier_a, lane.clone(), Expr::u32(0)),
314 Node::store(frontier_b, lane.clone(), Expr::u32(0)),
315 ],
316 ));
317 entry.push(Node::if_then(
318 Expr::eq(lane.clone(), Expr::u32(0)),
319 vec![Node::store(
320 frontier_b,
321 Expr::u32(active_flag_idx),
322 Expr::u32(0),
323 )],
324 ));
325 if max_iters > 0 {
326 entry.push(reachable_wave_barrier(node_count));
327 }
328
329 for i in 0..max_iters {
330 let current_wave = if i == 0 { sources_buf } else { frontier_b };
331 let active_var = format!("iter_{i}_active");
332 let active_expr = if i == 0 {
333 Expr::u32(1)
334 } else {
335 Expr::load(frontier_b, Expr::u32(active_flag_idx))
336 };
337 let active_cond = Expr::ne(Expr::var(active_var.as_str()), Expr::u32(0));
338 entry.push(Node::let_bind(active_var.as_str(), active_expr));
339 entry.push(Node::if_then(
340 Expr::lt(lane.clone(), Expr::u32(words)),
341 vec![Node::store(frontier_a, lane.clone(), Expr::u32(0))],
342 ));
343 entry.push(reachable_wave_barrier(node_count));
344 entry.push(Node::if_then(
345 active_cond.clone(),
346 vec![reachable_forward_wave_node(
347 shape,
348 current_wave,
349 frontier_a,
350 &format!("iter_{i}_expand"),
351 )],
352 ));
353 entry.push(reachable_wave_barrier(node_count));
354 entry.push(Node::if_then(
355 Expr::eq(lane.clone(), Expr::u32(0)),
356 vec![Node::store(
357 frontier_b,
358 Expr::u32(active_flag_idx),
359 Expr::u32(0),
360 )],
361 ));
362 entry.push(reachable_wave_barrier(node_count));
363 entry.extend(frontier_absorb_new_bits_body_prefixed_with_flag(
364 reach_out,
365 frontier_a,
366 frontier_b,
367 None,
368 Some((frontier_b, Expr::u32(active_flag_idx))),
369 words,
370 frontier_tail_mask(node_count),
371 &format!("iter_{i}_absorb"),
372 ));
373 if i + 1 < max_iters {
374 entry.push(reachable_wave_barrier(node_count));
375 }
376 }
377
378 let storage_words = words.max(1);
379 let mut buffers = shape.read_only_buffers();
380 buffers.push(
381 BufferDecl::storage(
382 sources_buf,
383 BINDING_PRIMITIVE_START,
384 BufferAccess::ReadOnly,
385 DataType::U32,
386 )
387 .with_count(storage_words),
388 );
389 buffers.push(
390 BufferDecl::storage(
391 reach_out,
392 BINDING_PRIMITIVE_START + 1,
393 BufferAccess::ReadWrite,
394 DataType::U32,
395 )
396 .with_count(storage_words),
397 );
398 buffers.push(
399 BufferDecl::storage(
400 frontier_a,
401 BINDING_PRIMITIVE_START + 2,
402 BufferAccess::ReadWrite,
403 DataType::U32,
404 )
405 .with_count(storage_words),
406 );
407 buffers.push(
408 BufferDecl::storage(
409 frontier_b,
410 BINDING_PRIMITIVE_START + 3,
411 BufferAccess::ReadWrite,
412 DataType::U32,
413 )
414 .with_count(frontier_b_storage_words),
415 );
416
417 Program::wrapped(
418 buffers,
419 [256, 1, 1],
420 vec![Node::Region {
421 generator: Ident::from(OP_ID),
422 source_region: None,
423 body: Arc::new(entry),
424 }],
425 )
426}
427
428fn reachable_wave_barrier(node_count: u32) -> Node {
429 if node_count <= 256 {
430 Node::barrier()
431 } else {
432 Node::barrier_with_ordering(MemoryOrdering::GridSync)
433 }
434}
435
436fn reachable_forward_wave_node(
437 shape: ProgramGraphShape,
438 frontier_in: &str,
439 frontier_out: &str,
440 local_prefix: &str,
441) -> Node {
442 let local = |name: &str| -> String { format!("{local_prefix}_{name}") };
443 let lane = Expr::gid_x();
444 let word_idx = local("word_idx");
445 let bit_mask = local("bit_mask");
446 let src_word = local("src_word");
447 let edge_start = local("edge_start");
448 let edge_end = local("edge_end");
449 let edge_iter = local("edge");
450 let kind_mask = local("kind_mask");
451 let dst = local("dst");
452 let dst_word_idx = local("dst_word_idx");
453 let dst_bit = local("dst_bit");
454 let previous = local("_prev");
455
456 Node::if_then(
457 Expr::lt(lane.clone(), Expr::u32(shape.node_count)),
458 vec![
459 Node::let_bind(word_idx.as_str(), Expr::shr(lane.clone(), Expr::u32(5))),
460 Node::let_bind(
461 bit_mask.as_str(),
462 Expr::shl(Expr::u32(1), Expr::bitand(lane.clone(), Expr::u32(31))),
463 ),
464 Node::let_bind(
465 src_word.as_str(),
466 Expr::load(frontier_in, Expr::var(word_idx.as_str())),
467 ),
468 Node::if_then(
469 Expr::ne(
470 Expr::bitand(Expr::var(src_word.as_str()), Expr::var(bit_mask.as_str())),
471 Expr::u32(0),
472 ),
473 vec![
474 Node::let_bind(
475 edge_start.as_str(),
476 Expr::load(NAME_EDGE_OFFSETS, lane.clone()),
477 ),
478 Node::let_bind(
479 edge_end.as_str(),
480 Expr::load(NAME_EDGE_OFFSETS, Expr::add(lane.clone(), Expr::u32(1))),
481 ),
482 Node::loop_for(
483 edge_iter.as_str(),
484 Expr::var(edge_start.as_str()),
485 Expr::var(edge_end.as_str()),
486 vec![
487 Node::let_bind(
488 kind_mask.as_str(),
489 Expr::load(NAME_EDGE_KIND_MASK, Expr::var(edge_iter.as_str())),
490 ),
491 Node::if_then(
492 Expr::ne(Expr::var(kind_mask.as_str()), Expr::u32(0)),
493 vec![
494 Node::let_bind(
495 dst.as_str(),
496 Expr::load(
497 NAME_EDGE_TARGETS,
498 Expr::var(edge_iter.as_str()),
499 ),
500 ),
501 Node::if_then(
502 Expr::lt(
503 Expr::var(dst.as_str()),
504 Expr::u32(shape.node_count),
505 ),
506 vec![
507 Node::let_bind(
508 dst_word_idx.as_str(),
509 Expr::shr(Expr::var(dst.as_str()), Expr::u32(5)),
510 ),
511 Node::let_bind(
512 dst_bit.as_str(),
513 Expr::shl(
514 Expr::u32(1),
515 Expr::bitand(
516 Expr::var(dst.as_str()),
517 Expr::u32(31),
518 ),
519 ),
520 ),
521 Node::let_bind(
522 previous.as_str(),
523 Expr::atomic_or(
524 frontier_out,
525 Expr::var(dst_word_idx.as_str()),
526 Expr::var(dst_bit.as_str()),
527 ),
528 ),
529 ],
530 ),
531 ],
532 ),
533 ],
534 ),
535 ],
536 ),
537 ],
538 )
539}
540
541#[cfg(test)]
542mod tests {
543 use super::*;
544
545 fn hs(items: &[u32]) -> HashSet<u32> {
546 items.iter().copied().collect()
547 }
548
549 #[test]
550 fn generated_try_reachable_matches_legacy_reachable() {
551 for node_count in 1u32..=64 {
552 for seed in 0u32..64 {
553 let edges: Vec<(u32, u32)> = (0..node_count)
554 .filter_map(|node| {
555 let step = (seed % 7) + 1;
556 let dst = node.saturating_add(step);
557 (dst < node_count).then_some((node, dst))
558 })
559 .collect();
560 let sources = [seed % node_count, node_count + seed];
561 let fallible = try_reachable(node_count, &edges, &sources).unwrap();
562 let legacy = reachable(node_count, &edges, &sources).unwrap();
563 assert_eq!(fallible, legacy);
564 assert!(fallible.contains(&(node_count + seed)));
565 }
566 }
567 }
568
569 #[test]
570 fn empty_sources_reach_nothing() {
571 let got = reachable(3, &[(0, 1), (1, 2)], &[]).unwrap();
572 assert!(got.is_empty());
573 }
574
575 #[test]
576 fn single_source_reaches_chain() {
577 let got = reachable(3, &[(0, 1), (1, 2)], &[0]).unwrap();
578 assert_eq!(got, hs(&[0, 1, 2]));
579 }
580
581 #[test]
582 fn cycle_terminates() {
583 let got = reachable(2, &[(0, 1), (1, 0)], &[0]).unwrap();
585 assert_eq!(got, hs(&[0, 1]));
586 }
587
588 #[test]
589 fn disconnected_source_not_included() {
590 let got = reachable(4, &[(0, 1), (2, 3)], &[0]).unwrap();
591 assert_eq!(got, hs(&[0, 1]));
592 assert!(!got.contains(&2));
593 assert!(!got.contains(&3));
594 }
595
596 #[test]
597 fn unknown_source_is_noop() {
598 let got = reachable(2, &[(0, 1)], &[7]).unwrap();
602 assert_eq!(got, hs(&[7]));
603 }
604
605 #[test]
606 fn out_of_range_edge_is_reported_not_silently_dropped() {
607 let err = reachable(3, &[(0, 1), (5, 1)], &[0]).unwrap_err();
610 assert_eq!(err.index, 1);
611 assert_eq!(err.node, 5);
612 assert_eq!(err.node_count, 3);
613 }
614
615 #[test]
616 fn reachable_program_smoke() {
617 let program = reachable_program(4, 4, "sources", "reach", 2);
620 assert!(!program.is_explicit_noop());
621 assert!(!program.buffers().is_empty());
622 assert!(!program.entry().is_empty());
623 assert_eq!(program.workgroup_size(), [256, 1, 1]);
624
625 let names: Vec<&str> = program.buffers().iter().map(|b| b.name()).collect();
628 assert!(names.contains(&"pg_edge_offsets"));
629 assert!(names.contains(&"pg_edge_targets"));
630 assert!(names.contains(&"sources"));
631 assert!(names.contains(&"reach"));
632 assert!(names.contains(&"reach_frontier_a"));
633 assert!(names.contains(&"reach_frontier_b"));
634 let frontier_b = program
635 .buffers()
636 .iter()
637 .find(|buffer| buffer.name() == "reach_frontier_b")
638 .expect("Fix: reachable wavefront scratch must be declared.");
639 assert_eq!(
640 frontier_b.count(),
641 bitset_words(4) + 1,
642 "Fix: reach_frontier_b must reserve one extra word for the converged-wave flag."
643 );
644 }
645
646 #[test]
647 fn reachable_program_zero_iters_seeds_only() {
648 let program = reachable_program(4, 4, "sources", "reach", 0);
651 assert!(!program.is_explicit_noop());
652 assert!(!program.buffers().is_empty());
653 }
654
655 #[test]
656 fn generated_wavefront_depth_limited_reachability_matches_scalar_reference() {
657 for seed in 0..10_000_u32 {
658 let mut state = mix32(seed ^ 0xA11C_E5E7);
659 let node_count = 1 + (state % 96);
660 state = mix32(state);
661 let edge_budget = state % (node_count * 3);
662 let mut edges = Vec::new();
663 for edge_idx in 0..edge_budget {
664 state = mix32(state ^ edge_idx.wrapping_mul(0x9E37_79B9));
665 let from = state % node_count;
666 state = mix32(state.rotate_left(7));
667 let to = match edge_idx % 11 {
668 0 => from,
669 1 => (from + 1) % node_count,
670 2 => node_count - 1,
671 _ => state % node_count,
672 };
673 edges.push((from, to));
674 }
675 let source_count = 1 + (mix32(state ^ 0x5150_ACE5) % 4);
676 let mut sources = Vec::new();
677 for idx in 0..source_count {
678 state = mix32(state ^ idx.wrapping_mul(0x85EB_CA6B));
679 sources.push(state % node_count);
680 }
681 let max_iters = mix32(state ^ 0xD47A_F10D) % (node_count.min(16) + 1);
682
683 let wave = depth_limited_wavefront(node_count, &edges, &sources, max_iters);
684 let scalar = depth_limited_scalar(node_count, &edges, &sources, max_iters);
685
686 assert_eq!(
687 wave, scalar,
688 "seed={seed} node_count={node_count} max_iters={max_iters}"
689 );
690 }
691 }
692
693 fn depth_limited_wavefront(
694 node_count: u32,
695 edges: &[(u32, u32)],
696 sources: &[u32],
697 max_iters: u32,
698 ) -> HashSet<u32> {
699 let mut visited = HashSet::new();
700 let mut current = HashSet::new();
701 for &source in sources {
702 if source < node_count && visited.insert(source) {
703 current.insert(source);
704 }
705 }
706
707 for _ in 0..max_iters {
708 let mut next = HashSet::new();
709 for &(from, to) in edges {
710 if from < node_count
711 && to < node_count
712 && current.contains(&from)
713 && visited.insert(to)
714 {
715 next.insert(to);
716 }
717 }
718 current = next;
719 }
720 visited
721 }
722
723 fn depth_limited_scalar(
724 node_count: u32,
725 edges: &[(u32, u32)],
726 sources: &[u32],
727 max_iters: u32,
728 ) -> HashSet<u32> {
729 let mut min_depth = vec![u32::MAX; node_count as usize];
730 let mut queue = std::collections::VecDeque::new();
731 for &source in sources {
732 if source < node_count && min_depth[source as usize] > 0 {
733 min_depth[source as usize] = 0;
734 queue.push_back(source);
735 }
736 }
737 while let Some(node) = queue.pop_front() {
738 let depth = min_depth[node as usize];
739 if depth >= max_iters {
740 continue;
741 }
742 let next_depth = depth + 1;
743 for &(from, to) in edges {
744 if from == node && to < node_count && next_depth < min_depth[to as usize] {
745 min_depth[to as usize] = next_depth;
746 queue.push_back(to);
747 }
748 }
749 }
750 min_depth
751 .into_iter()
752 .enumerate()
753 .filter_map(|(node, depth)| (depth <= max_iters).then_some(node as u32))
754 .collect()
755 }
756
757 fn mix32(mut value: u32) -> u32 {
758 value ^= value >> 16;
759 value = value.wrapping_mul(0x7FEB_352D);
760 value ^= value >> 15;
761 value = value.wrapping_mul(0x846C_A68B);
762 value ^ (value >> 16)
763 }
764}