1use vyre_foundation::ir::{BufferAccess, BufferDecl, Program};
20
21type IfdsIntraRule = (u32, u32, u32);
22type IfdsInterRule = (u32, u32, u32, u32);
23type IfdsFactRule = (u32, u32, u32);
24type ParsedIfdsRules = (
25 Vec<IfdsIntraRule>,
26 Vec<IfdsInterRule>,
27 Vec<IfdsFactRule>,
28 Vec<IfdsFactRule>,
29);
30
31#[must_use]
43pub fn declared_dispatch_outputs(program: &Program) -> Vec<&BufferDecl> {
44 program
45 .buffers()
46 .iter()
47 .filter(|decl| {
48 matches!(
49 decl.access(),
50 BufferAccess::ReadWrite | BufferAccess::WriteOnly
51 )
52 })
53 .collect()
54}
55
56pub struct ResidentDispatchStep<'a> {
58 pub program: &'a Program,
60 pub handle_ids: &'a [u64],
62 pub grid_override: Option<[u32; 3]>,
64}
65
66#[derive(Clone, Copy, Debug, Eq, PartialEq)]
68pub struct ResidentReadRange {
69 pub handle_id: u64,
71 pub byte_offset: usize,
73 pub byte_len: usize,
75}
76
77#[derive(Debug)]
86pub struct ResidentStaticBufferSet {
87 pub handles: Vec<u64>,
90 pub cache_hit: bool,
92 pub retained_by_dispatcher: bool,
94}
95
96#[derive(Debug)]
100pub enum DispatchError {
101 Rejected(String),
104 BadInputs(String),
107 BackendError(String),
110}
111
112impl std::fmt::Display for DispatchError {
113 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114 match self {
115 Self::Rejected(msg) => write!(f, "dispatcher rejected program: {msg}"),
116 Self::BadInputs(msg) => write!(f, "dispatcher input mismatch: {msg}"),
117 Self::BackendError(msg) => write!(f, "dispatcher backend error: {msg}"),
118 }
119 }
120}
121
122impl std::error::Error for DispatchError {}
123
124pub trait OptimizerDispatcher {
131 fn dispatch(
142 &self,
143 program: &Program,
144 inputs: &[Vec<u8>],
145 grid_override: Option<[u32; 3]>,
146 ) -> Result<Vec<Vec<u8>>, DispatchError>;
147
148 fn supports_persistent(&self) -> bool {
154 false
155 }
156
157 fn device_feature_cache_key(&self) -> u64 {
165 0
166 }
167
168 fn alloc_resident(&self, _byte_len: usize) -> Result<u64, DispatchError> {
171 Err(DispatchError::Rejected(
172 "Fix: this dispatcher does not implement the persistent path; \
173 use `dispatch` instead, or wire the resident-buffer methods."
174 .to_string(),
175 ))
176 }
177
178 fn alloc_resident_many(&self, byte_lens: &[usize]) -> Result<Vec<u64>, DispatchError> {
181 let mut handles = Vec::new();
182 handles.try_reserve(byte_lens.len()).map_err(|error| {
183 DispatchError::BackendError(format!(
184 "Fix: reserve resident handle group before allocation; requested {} buffer(s): {error}.",
185 byte_lens.len()
186 ))
187 })?;
188 for (index, &byte_len) in byte_lens.iter().enumerate() {
189 match self.alloc_resident(byte_len) {
190 Ok(handle) => handles.push(handle),
191 Err(error) => {
192 let allocation_error = error.to_string();
193 if let Err(free_error) = free_resident_handles(
194 self,
195 &handles,
196 "resident grouped allocation rollback",
197 ) {
198 return Err(DispatchError::BackendError(format!(
199 "Fix: resident grouped allocation failed at buffer {index} after {} partial allocation(s): {allocation_error}; rollback also failed: {free_error}.",
200 handles.len()
201 )));
202 }
203 return Err(error);
204 }
205 }
206 }
207 Ok(handles)
208 }
209
210 fn upload_resident(&self, _handle: u64, _bytes: &[u8]) -> Result<(), DispatchError> {
212 Err(DispatchError::Rejected(
213 "Fix: dispatcher does not implement upload_resident.".to_string(),
214 ))
215 }
216
217 fn upload_resident_many(&self, uploads: &[(u64, &[u8])]) -> Result<(), DispatchError> {
219 for &(handle, bytes) in uploads {
220 self.upload_resident(handle, bytes)?;
221 }
222 Ok(())
223 }
224
225 fn acquire_resident_static_uploads(
233 &self,
234 _cache_domain: u64,
235 payloads: &[&[u8]],
236 ) -> Result<ResidentStaticBufferSet, DispatchError> {
237 let mut byte_lens = Vec::new();
238 byte_lens.try_reserve(payloads.len()).map_err(|error| {
239 DispatchError::BackendError(format!(
240 "Fix: reserve resident static byte lengths before upload; requested {} payload(s): {error}.",
241 payloads.len()
242 ))
243 })?;
244 for payload in payloads {
245 byte_lens.push(payload.len());
246 }
247 let handles = self.alloc_resident_many(&byte_lens)?;
248
249 let mut uploads = Vec::new();
250 uploads.try_reserve(payloads.len()).map_err(|error| {
251 DispatchError::BackendError(format!(
252 "Fix: reserve resident static upload storage before upload; requested {} payload(s): {error}.",
253 payloads.len()
254 ))
255 })?;
256 for (&handle, &payload) in handles.iter().zip(payloads.iter()) {
257 uploads.push((handle, payload));
258 }
259
260 if let Err(error) = self.upload_resident_many(&uploads) {
261 let upload_error = error.to_string();
262 if let Err(free_error) =
263 free_resident_handles(self, &handles, "resident static upload rollback")
264 {
265 return Err(DispatchError::BackendError(format!(
266 "Fix: resident static upload failed after allocating {} buffer(s): {upload_error}; rollback also failed: {free_error}.",
267 handles.len()
268 )));
269 }
270 return Err(error);
271 }
272
273 Ok(ResidentStaticBufferSet {
274 handles,
275 cache_hit: false,
276 retained_by_dispatcher: false,
277 })
278 }
279
280 fn release_resident_static_uploads(
283 &self,
284 set: ResidentStaticBufferSet,
285 ) -> Result<(), DispatchError> {
286 if set.retained_by_dispatcher {
287 return Ok(());
288 }
289 for handle in set.handles {
290 self.free_resident(handle)?;
291 }
292 Ok(())
293 }
294
295 fn read_resident(&self, _handle: u64) -> Result<Vec<u8>, DispatchError> {
297 Err(DispatchError::Rejected(
298 "Fix: dispatcher does not implement read_resident.".to_string(),
299 ))
300 }
301
302 fn read_resident_many(&self, handles: &[u64]) -> Result<Vec<Vec<u8>>, DispatchError> {
304 handles
305 .iter()
306 .map(|&handle| self.read_resident(handle))
307 .collect()
308 }
309
310 fn read_resident_ranges(
312 &self,
313 ranges: &[ResidentReadRange],
314 ) -> Result<Vec<Vec<u8>>, DispatchError> {
315 let mut outputs = Vec::new();
316 self.read_resident_ranges_into(ranges, &mut outputs)?;
317 Ok(outputs)
318 }
319
320 fn read_resident_ranges_into(
323 &self,
324 ranges: &[ResidentReadRange],
325 outputs: &mut Vec<Vec<u8>>,
326 ) -> Result<(), DispatchError> {
327 let mut unique_handles = Vec::new();
328 unique_handles.try_reserve(ranges.len()).map_err(|error| {
329 DispatchError::BackendError(format!(
330 "Fix: reserve resident ranged-read handle dedupe storage before dispatch; requested {} range(s): {error}.",
331 ranges.len()
332 ))
333 })?;
334 let mut range_handle_indices = Vec::new();
335 range_handle_indices
336 .try_reserve(ranges.len())
337 .map_err(|error| {
338 DispatchError::BackendError(format!(
339 "Fix: reserve resident ranged-read index storage before dispatch; requested {} range(s): {error}.",
340 ranges.len()
341 ))
342 })?;
343 for range in ranges {
344 if let Some(index) = unique_handles
345 .iter()
346 .position(|&handle| handle == range.handle_id)
347 {
348 range_handle_indices.push(index);
349 } else {
350 let index = unique_handles.len();
351 unique_handles.push(range.handle_id);
352 range_handle_indices.push(index);
353 }
354 }
355 let full_buffers = self.read_resident_many(&unique_handles)?;
356 if full_buffers.len() != unique_handles.len() {
357 return Err(DispatchError::BackendError(format!(
358 "Fix: resident ranged-read batch returned {} buffer(s) for {} unique handle(s).",
359 full_buffers.len(),
360 unique_handles.len()
361 )));
362 }
363 if outputs.len() < ranges.len() {
364 outputs
365 .try_reserve(ranges.len() - outputs.len())
366 .map_err(|error| {
367 DispatchError::BackendError(format!(
368 "Fix: reserve resident ranged-read output storage before dispatch; requested {} range(s): {error}.",
369 ranges.len()
370 ))
371 })?;
372 outputs.resize_with(ranges.len(), Vec::new);
373 } else {
374 outputs.truncate(ranges.len());
375 }
376 for ((range, &buffer_index), output) in ranges
377 .iter()
378 .zip(range_handle_indices.iter())
379 .zip(outputs.iter_mut())
380 {
381 let full = full_buffers.get(buffer_index).ok_or_else(|| {
382 DispatchError::BackendError(format!(
383 "Fix: resident ranged-read handle index {buffer_index} missing from {} readback buffer(s).",
384 full_buffers.len()
385 ))
386 })?;
387 let end = range
388 .byte_offset
389 .checked_add(range.byte_len)
390 .ok_or_else(|| {
391 DispatchError::BadInputs(format!(
392 "Fix: resident read range for handle {} overflows usize at offset {} len {}.",
393 range.handle_id, range.byte_offset, range.byte_len
394 ))
395 })?;
396 if end > full.len() {
397 return Err(DispatchError::BadInputs(format!(
398 "Fix: resident read range for handle {} requested bytes [{}..{}) but buffer readback has {} bytes.",
399 range.handle_id,
400 range.byte_offset,
401 end,
402 full.len()
403 )));
404 }
405 output.clear();
406 output.extend_from_slice(&full[range.byte_offset..end]);
407 }
408 Ok(())
409 }
410
411 fn free_resident(&self, _handle: u64) -> Result<(), DispatchError> {
413 Err(DispatchError::Rejected(
414 "Fix: dispatcher does not implement free_resident.".to_string(),
415 ))
416 }
417
418 fn dispatch_resident(
423 &self,
424 _program: &Program,
425 _handles: &[u64],
426 _grid_override: Option<[u32; 3]>,
427 ) -> Result<(), DispatchError> {
428 Err(DispatchError::Rejected(
429 "Fix: dispatcher does not implement dispatch_resident.".to_string(),
430 ))
431 }
432
433 fn dispatch_resident_sequence(
439 &self,
440 steps: &[ResidentDispatchStep<'_>],
441 ) -> Result<(), DispatchError> {
442 for step in steps {
443 self.dispatch_resident(step.program, step.handle_ids, step.grid_override)?;
444 }
445 Ok(())
446 }
447
448 fn dispatch_resident_sequence_read_many(
455 &self,
456 steps: &[ResidentDispatchStep<'_>],
457 read_handles: &[u64],
458 ) -> Result<Vec<Vec<u8>>, DispatchError> {
459 self.dispatch_resident_sequence(steps)?;
460 self.read_resident_many(read_handles)
461 }
462
463 fn dispatch_resident_sequence_read_ranges(
465 &self,
466 steps: &[ResidentDispatchStep<'_>],
467 read_ranges: &[ResidentReadRange],
468 ) -> Result<Vec<Vec<u8>>, DispatchError> {
469 self.dispatch_resident_sequence(steps)?;
470 self.read_resident_ranges(read_ranges)
471 }
472
473 fn upload_resident_many_sequence_read_many(
480 &self,
481 uploads: &[(u64, &[u8])],
482 steps: &[ResidentDispatchStep<'_>],
483 read_handles: &[u64],
484 ) -> Result<Vec<Vec<u8>>, DispatchError> {
485 self.upload_resident_many(uploads)?;
486 self.dispatch_resident_sequence_read_many(steps, read_handles)
487 }
488
489 fn upload_resident_many_sequence_read_ranges(
492 &self,
493 uploads: &[(u64, &[u8])],
494 steps: &[ResidentDispatchStep<'_>],
495 read_ranges: &[ResidentReadRange],
496 ) -> Result<Vec<Vec<u8>>, DispatchError> {
497 self.upload_resident_many(uploads)?;
498 self.dispatch_resident_sequence_read_ranges(steps, read_ranges)
499 }
500
501 fn upload_resident_many_sequence_read_many_into(
504 &self,
505 uploads: &[(u64, &[u8])],
506 steps: &[ResidentDispatchStep<'_>],
507 read_handles: &[u64],
508 outputs: &mut Vec<Vec<u8>>,
509 ) -> Result<(), DispatchError> {
510 let readbacks =
511 self.upload_resident_many_sequence_read_many(uploads, steps, read_handles)?;
512 if outputs.len() < readbacks.len() {
513 outputs.resize_with(readbacks.len(), Vec::new);
514 } else {
515 outputs.truncate(readbacks.len());
516 }
517 for (slot, readback) in outputs.iter_mut().zip(readbacks) {
518 slot.clear();
519 slot.extend_from_slice(&readback);
520 }
521 Ok(())
522 }
523
524 fn clear_upload_resident_many_sequence_read_many_into(
533 &self,
534 clears: &[(u64, usize)],
535 uploads: &[(u64, &[u8])],
536 steps: &[ResidentDispatchStep<'_>],
537 read_handles: &[u64],
538 outputs: &mut Vec<Vec<u8>>,
539 ) -> Result<(), DispatchError> {
540 if clears.is_empty() {
541 return self.upload_resident_many_sequence_read_many_into(
542 uploads,
543 steps,
544 read_handles,
545 outputs,
546 );
547 }
548 let mut fills = Vec::new();
549 fills.try_reserve(clears.len()).map_err(|error| {
550 DispatchError::BackendError(format!(
551 "Fix: reserve resident clear fill descriptors before dispatch; requested {} clear(s): {error}.",
552 clears.len()
553 ))
554 })?;
555 for &(handle, byte_len) in clears {
556 fills.push((handle, byte_len, 0));
557 }
558 self.fill_upload_resident_many_sequence_read_many_into(
559 &fills,
560 uploads,
561 steps,
562 read_handles,
563 outputs,
564 )
565 }
566
567 fn fill_upload_resident_many_sequence_read_many_into(
571 &self,
572 fills: &[(u64, usize, u8)],
573 uploads: &[(u64, &[u8])],
574 steps: &[ResidentDispatchStep<'_>],
575 read_handles: &[u64],
576 outputs: &mut Vec<Vec<u8>>,
577 ) -> Result<(), DispatchError> {
578 if fills.is_empty() {
579 return self.upload_resident_many_sequence_read_many_into(
580 uploads,
581 steps,
582 read_handles,
583 outputs,
584 );
585 }
586
587 with_staged_fill_uploads(
588 fills,
589 uploads,
590 "resident fill payloads",
591 "resident fill/upload payloads",
592 |combined_uploads| {
593 self.upload_resident_many_sequence_read_many_into(
594 combined_uploads,
595 steps,
596 read_handles,
597 outputs,
598 )
599 },
600 )
601 }
602
603 fn fill_upload_resident_many_sequence_read_ranges_into(
607 &self,
608 fills: &[(u64, usize, u8)],
609 uploads: &[(u64, &[u8])],
610 steps: &[ResidentDispatchStep<'_>],
611 read_ranges: &[ResidentReadRange],
612 outputs: &mut Vec<Vec<u8>>,
613 ) -> Result<(), DispatchError> {
614 if fills.is_empty() {
615 return self.upload_resident_many_sequence_read_ranges_into(
616 uploads,
617 steps,
618 read_ranges,
619 outputs,
620 );
621 }
622
623 with_staged_fill_uploads(
624 fills,
625 uploads,
626 "resident range-fill payloads",
627 "resident range-fill/upload payloads",
628 |combined_uploads| {
629 self.upload_resident_many_sequence_read_ranges_into(
630 combined_uploads,
631 steps,
632 read_ranges,
633 outputs,
634 )
635 },
636 )
637 }
638
639 fn upload_resident_many_sequence_read_ranges_into(
642 &self,
643 uploads: &[(u64, &[u8])],
644 steps: &[ResidentDispatchStep<'_>],
645 read_ranges: &[ResidentReadRange],
646 outputs: &mut Vec<Vec<u8>>,
647 ) -> Result<(), DispatchError> {
648 self.upload_resident_many(uploads)?;
649 self.dispatch_resident_sequence(steps)?;
650 self.read_resident_ranges_into(read_ranges, outputs)
651 }
652}
653
654fn free_resident_handles<D: OptimizerDispatcher + ?Sized>(
655 dispatcher: &D,
656 handles: &[u64],
657 context: &str,
658) -> Result<(), DispatchError> {
659 for (index, &handle) in handles.iter().enumerate() {
660 dispatcher.free_resident(handle).map_err(|error| {
661 DispatchError::BackendError(format!(
662 "Fix: {context} failed to free resident handle {handle} at index {index}: {error}."
663 ))
664 })?;
665 }
666 Ok(())
667}
668
669fn with_staged_fill_uploads<R>(
670 fills: &[(u64, usize, u8)],
671 uploads: &[(u64, &[u8])],
672 fill_context: &'static str,
673 combined_context: &'static str,
674 run: impl FnOnce(&[(u64, &[u8])]) -> Result<R, DispatchError>,
675) -> Result<R, DispatchError> {
676 let mut fill_payloads = Vec::new();
677 fill_payloads.try_reserve(fills.len()).map_err(|error| {
678 DispatchError::BackendError(format!(
679 "Fix: reserve {fill_context} before dispatch; requested {} fill(s): {error}.",
680 fills.len()
681 ))
682 })?;
683 for &(_, byte_len, value) in fills {
684 fill_payloads.push(vec![value; byte_len]);
685 }
686
687 let mut combined_uploads = Vec::new();
688 combined_uploads
689 .try_reserve(fills.len() + uploads.len())
690 .map_err(|error| {
691 DispatchError::BackendError(format!(
692 "Fix: reserve {combined_context} before dispatch; requested {} fill(s) and {} upload(s): {error}.",
693 fills.len(),
694 uploads.len()
695 ))
696 })?;
697 for ((handle, _, _), fill) in fills.iter().zip(fill_payloads.iter()) {
698 combined_uploads.push((*handle, fill.as_slice()));
699 }
700 combined_uploads.extend_from_slice(uploads);
701
702 run(&combined_uploads)
703}
704
705#[cfg(any(test, feature = "cpu-parity"))]
706pub mod oracle {
707 use super::{DispatchError, OptimizerDispatcher, ParsedIfdsRules};
726 use vyre_foundation::ir::Program;
727
728 pub struct CpuOracleDispatcher;
732
733 impl CpuOracleDispatcher {
734 #[must_use]
737 pub fn new() -> Self {
738 Self
739 }
740 }
741
742 impl Default for CpuOracleDispatcher {
743 fn default() -> Self {
744 Self::new()
745 }
746 }
747
748 impl OptimizerDispatcher for CpuOracleDispatcher {
749 fn dispatch(
750 &self,
751 program: &Program,
752 inputs: &[Vec<u8>],
753 _grid_override: Option<[u32; 3]>,
754 ) -> Result<Vec<Vec<u8>>, DispatchError> {
755 let generator = top_level_region_generator(program).ok_or_else(|| {
759 DispatchError::Rejected(
760 "Fix: oracle dispatcher only accepts canonical \
761 graph-primitive Programs whose entry is a single \
762 wrapping Region with a generator id."
763 .to_string(),
764 )
765 })?;
766
767 match generator {
768 vyre_primitives::graph::persistent_bfs::OP_ID => {
769 persistent_bfs_oracle(program, inputs)
770 }
771 crate::optimizer::dce_program::OP_ID => persistent_bfs_oracle(program, inputs),
772 vyre_primitives::graph::exploded::OP_ID => {
773 exploded_ifds_csr_oracle(program, inputs)
774 }
775 other => Err(DispatchError::Rejected(format!(
776 "Fix: oracle dispatcher does not recognize generator \
777 `{other}`. Wire the oracle for this primitive or \
778 dispatch through the production backend."
779 ))),
780 }
781 }
782 }
783
784 fn top_level_region_generator(program: &Program) -> Option<&str> {
785 match program.entry() {
786 [vyre_foundation::ir::Node::Region { generator, .. }] => Some(generator.as_str()),
787 _ => None,
788 }
789 }
790
791 fn persistent_bfs_oracle(
792 program: &Program,
793 inputs: &[Vec<u8>],
794 ) -> Result<Vec<Vec<u8>>, DispatchError> {
795 if inputs.len() < 6 {
807 return Err(DispatchError::BadInputs(format!(
808 "Fix: persistent_bfs oracle expects ≥ 6 input buffers, got {}",
809 inputs.len()
810 )));
811 }
812 let nodes = crate::hardware::dispatch_buffers::read_u32s(&inputs[0]);
813 let edge_offsets = crate::hardware::dispatch_buffers::read_u32s(&inputs[1]);
814 let edge_targets_raw = crate::hardware::dispatch_buffers::read_u32s(&inputs[2]);
815 let edge_kind_mask_raw = crate::hardware::dispatch_buffers::read_u32s(&inputs[3]);
816 let _node_tags = crate::hardware::dispatch_buffers::read_u32s(&inputs[4]);
817 let frontier_in = crate::hardware::dispatch_buffers::read_u32s(&inputs[5]);
818
819 let node_count = nodes.len() as u32;
823
824 let max_iters = node_count.max(1);
829
830 let allow_mask = u32::MAX;
831 let edge_count = declared_edge_count(&edge_offsets)?;
832 let edge_targets = trim_padded_edge_buffer("edge_targets", &edge_targets_raw, edge_count)?;
833 let edge_kind_mask =
834 trim_padded_edge_buffer("edge_kind_mask", &edge_kind_mask_raw, edge_count)?;
835
836 let (frontier_out, convergence) =
837 vyre_primitives::graph::persistent_bfs::try_cpu_ref_converged(
838 node_count,
839 &edge_offsets,
840 edge_targets,
841 edge_kind_mask,
842 &frontier_in,
843 allow_mask,
844 max_iters,
845 )
846 .map_err(DispatchError::BadInputs)?;
847
848 let declared = super::declared_dispatch_outputs(program);
856 let mut outputs = Vec::with_capacity(declared.len());
857 for decl in declared {
858 let words = match decl.name() {
859 "frontier_out" => frontier_out.clone(),
860 "changed" => changed_words_for(decl.count(), &convergence)?,
861 "converged" => vec![u32::from(convergence.converged)],
862 other => {
863 return Err(DispatchError::Rejected(format!(
864 concat!(
865 "Fix: persistent_bfs oracle has no value for declared output ",
866 "buffer `{other}`. Teach the oracle to produce it or stop ",
867 "declaring it; returning a short output list would silently ",
868 "shift every later output index."
869 ),
870 other = other
871 )))
872 }
873 };
874 outputs.push(u32_buffer_to_bytes(&words));
875 }
876 Ok(outputs)
877 }
878
879 fn changed_words_for(
897 count: u32,
898 convergence: &vyre_primitives::graph::persistent_bfs::PersistentBfsConvergence,
899 ) -> Result<Vec<u32>, DispatchError> {
900 let per_iteration = u32::from(!convergence.converged);
901 match count {
902 1 => Ok(vec![per_iteration]),
903 2 => Ok(vec![per_iteration, convergence.changed]),
904 other => Err(DispatchError::Rejected(format!(
905 concat!(
906 "Fix: persistent_bfs oracle understands a `changed` buffer of 1 ",
907 "element (per-iteration) or 2 (per-iteration plus sticky), not ",
908 "{other}. Declare one of those or teach the oracle what the extra ",
909 "slots mean."
910 ),
911 other = other
912 ))),
913 }
914 }
915
916 #[cfg(test)]
917 #[allow(clippy::items_after_test_module)]
918 mod changed_words_tests {
919 use super::changed_words_for;
920 use vyre_primitives::graph::persistent_bfs::PersistentBfsConvergence;
921
922 fn convergence(changed: u32, converged: bool) -> PersistentBfsConvergence {
923 PersistentBfsConvergence {
924 changed,
925 converged,
926 stop_iter: 0,
927 }
928 }
929
930 #[test]
935 fn a_single_changed_word_reports_the_last_iterations_progress() {
936 assert_eq!(
937 changed_words_for(1, &convergence(1, true)).expect("count 1 is supported"),
938 vec![0],
939 "a converged run exits on a zero compare, so the per-iteration flag is 0"
940 );
941 assert_eq!(
942 changed_words_for(1, &convergence(1, false)).expect("count 1 is supported"),
943 vec![1],
944 "an exhausted budget means the final iteration still grew the frontier"
945 );
946 }
947
948 #[test]
952 fn two_changed_words_keep_the_per_iteration_and_sticky_flags_distinct() {
953 assert_eq!(
954 changed_words_for(2, &convergence(1, true)).expect("count 2 is supported"),
955 vec![0, 1],
956 "slot 0 is the final compare, slot 1 latched because earlier steps grew"
957 );
958 assert_eq!(
959 changed_words_for(2, &convergence(0, true)).expect("count 2 is supported"),
960 vec![0, 0],
961 "a traversal that never grew converges immediately with nothing latched"
962 );
963 }
964
965 #[test]
968 fn an_unrecognized_changed_count_is_refused() {
969 let err = changed_words_for(3, &convergence(1, true))
970 .expect_err("an unknown changed layout must not be guessed");
971 assert!(
972 format!("{err}").contains("not 3"),
973 "the refusal must name the count it saw, got: {err}"
974 );
975 }
976 }
977
978 fn exploded_ifds_csr_oracle(
979 program: &Program,
980 inputs: &[Vec<u8>],
981 ) -> Result<Vec<Vec<u8>>, DispatchError> {
982 if inputs.len() != 18 {
983 return Err(DispatchError::BadInputs(format!(
984 "Fix: exploded IFDS oracle expected 18 input buffers, got {}.",
985 inputs.len()
986 )));
987 }
988
989 let key = vyre_primitives::graph::exploded::ifds_program_cache_key_from_program(program)
990 .map_err(DispatchError::BackendError)?;
991 let (intra_edges, inter_edges, flow_gen, flow_kill) = parse_ifds_rule_inputs(&key, inputs)?;
992
993 let (row_ptr, col_idx) = vyre_primitives::graph::exploded::build_cpu_reference(
994 key.num_procs,
995 key.blocks_per_proc,
996 key.facts_per_proc,
997 &intra_edges,
998 &inter_edges,
999 &flow_gen,
1000 &flow_kill,
1001 );
1002
1003 let col_len = u32::try_from(col_idx.len()).map_err(|error| {
1004 DispatchError::BackendError(format!(
1005 "Fix: exploded IFDS oracle col_idx length does not fit u32: {error}."
1006 ))
1007 })?;
1008 let col_idx_words = program
1009 .buffer("col_idx")
1010 .map(|buffer| buffer.count() as usize)
1011 .unwrap_or(1);
1012 let mut col_idx_padded = vec![0u32; col_idx_words];
1013 if col_idx.len() > col_idx_words {
1014 return Err(DispatchError::BackendError(format!(
1015 "Fix: exploded IFDS oracle emitted {} columns but program allocates {col_idx_words}."
1016 ,
1017 col_idx.len()
1018 )));
1019 }
1020 col_idx_padded[..col_idx.len()].copy_from_slice(&col_idx);
1021
1022 let row_cursor_words = program
1023 .buffer("row_cursor")
1024 .map(|buffer| buffer.count() as usize)
1025 .unwrap_or(1);
1026 let row_cursor = vec![0u32; row_cursor_words];
1027
1028 Ok(vec![
1029 u32_buffer_to_bytes(&row_ptr),
1030 u32_buffer_to_bytes(&row_cursor),
1031 u32_buffer_to_bytes(&col_idx_padded),
1032 u32_buffer_to_bytes(&[col_len]),
1033 ])
1034 }
1035
1036 fn parse_ifds_rule_inputs(
1037 key: &vyre_primitives::graph::exploded::IfdsCsrProgramCacheKey,
1038 inputs: &[Vec<u8>],
1039 ) -> Result<ParsedIfdsRules, DispatchError> {
1040 let intra_proc = crate::hardware::dispatch_buffers::read_u32s(&inputs[0]);
1041 let intra_src_block = crate::hardware::dispatch_buffers::read_u32s(&inputs[1]);
1042 let intra_dst_block = crate::hardware::dispatch_buffers::read_u32s(&inputs[2]);
1043 let inter_src_proc = crate::hardware::dispatch_buffers::read_u32s(&inputs[3]);
1044 let inter_src_block = crate::hardware::dispatch_buffers::read_u32s(&inputs[4]);
1045 let inter_dst_proc = crate::hardware::dispatch_buffers::read_u32s(&inputs[5]);
1046 let inter_dst_block = crate::hardware::dispatch_buffers::read_u32s(&inputs[6]);
1047 let gen_proc = crate::hardware::dispatch_buffers::read_u32s(&inputs[7]);
1048 let gen_block = crate::hardware::dispatch_buffers::read_u32s(&inputs[8]);
1049 let gen_fact = crate::hardware::dispatch_buffers::read_u32s(&inputs[9]);
1050 let kill_proc = crate::hardware::dispatch_buffers::read_u32s(&inputs[10]);
1051 let kill_block = crate::hardware::dispatch_buffers::read_u32s(&inputs[11]);
1052 let kill_fact = crate::hardware::dispatch_buffers::read_u32s(&inputs[12]);
1053
1054 let intra_edges = read_ifds_triples(
1055 "intra",
1056 key.intra_count,
1057 &intra_proc,
1058 &intra_src_block,
1059 &intra_dst_block,
1060 )?;
1061 let inter_edges = read_ifds_quads(
1062 "inter",
1063 key.inter_count,
1064 &inter_src_proc,
1065 &inter_src_block,
1066 &inter_dst_proc,
1067 &inter_dst_block,
1068 )?;
1069 let flow_gen = read_ifds_triples("GEN", key.gen_count, &gen_proc, &gen_block, &gen_fact)?;
1070 let flow_kill =
1071 read_ifds_triples("KILL", key.kill_count, &kill_proc, &kill_block, &kill_fact)?;
1072
1073 Ok((intra_edges, inter_edges, flow_gen, flow_kill))
1074 }
1075
1076 fn read_ifds_triples(
1077 kind: &str,
1078 count: u32,
1079 proc: &[u32],
1080 a: &[u32],
1081 b: &[u32],
1082 ) -> Result<Vec<(u32, u32, u32)>, DispatchError> {
1083 let count = count as usize;
1084 for (name, column) in [("proc", proc), ("a", a), ("b", b)] {
1085 if column.len() < count {
1086 return Err(DispatchError::BadInputs(format!(
1087 "Fix: exploded IFDS oracle {kind} {name} column has {} word(s), expected {count}."
1088 ,
1089 column.len()
1090 )));
1091 }
1092 }
1093 Ok((0..count)
1094 .map(|index| (proc[index], a[index], b[index]))
1095 .collect())
1096 }
1097
1098 fn read_ifds_quads(
1099 kind: &str,
1100 count: u32,
1101 a: &[u32],
1102 b: &[u32],
1103 c: &[u32],
1104 d: &[u32],
1105 ) -> Result<Vec<(u32, u32, u32, u32)>, DispatchError> {
1106 let count = count as usize;
1107 for (name, column) in [
1108 ("src_proc", a),
1109 ("src_block", b),
1110 ("dst_proc", c),
1111 ("dst_block", d),
1112 ] {
1113 if column.len() < count {
1114 return Err(DispatchError::BadInputs(format!(
1115 "Fix: exploded IFDS oracle {kind} {name} column has {} word(s), expected {count}."
1116 ,
1117 column.len()
1118 )));
1119 }
1120 }
1121 Ok((0..count)
1122 .map(|index| (a[index], b[index], c[index], d[index]))
1123 .collect())
1124 }
1125
1126 fn declared_edge_count(edge_offsets: &[u32]) -> Result<usize, DispatchError> {
1127 edge_offsets
1128 .last()
1129 .copied()
1130 .map(|edge_count| edge_count as usize)
1131 .ok_or_else(|| {
1132 DispatchError::BadInputs(
1133 "Fix: persistent_bfs oracle requires a CSR offset sentinel.".to_string(),
1134 )
1135 })
1136 }
1137
1138 fn trim_padded_edge_buffer<'a>(
1139 name: &str,
1140 buffer: &'a [u32],
1141 edge_count: usize,
1142 ) -> Result<&'a [u32], DispatchError> {
1143 if buffer.len() < edge_count {
1144 return Err(DispatchError::BadInputs(format!(
1145 "Fix: persistent_bfs oracle {name} has {} words but CSR declares {edge_count} edges.",
1146 buffer.len()
1147 )));
1148 }
1149 Ok(&buffer[..edge_count])
1150 }
1151
1152 fn u32_buffer_to_bytes(words: &[u32]) -> Vec<u8> {
1153 vyre_primitives::wire::pack_u32_slice(words)
1154 }
1155}
1156
1157#[cfg(test)]
1158mod tests {
1159 use super::*;
1160 use std::cell::{Cell, RefCell};
1161
1162 struct RangedReadDispatcher {
1163 buffers: Vec<(u64, Vec<u8>)>,
1164 read_calls: Cell<usize>,
1165 batched_handles: RefCell<Vec<u64>>,
1166 }
1167
1168 impl OptimizerDispatcher for RangedReadDispatcher {
1169 fn dispatch(
1170 &self,
1171 _program: &Program,
1172 _inputs: &[Vec<u8>],
1173 _grid_override: Option<[u32; 3]>,
1174 ) -> Result<Vec<Vec<u8>>, DispatchError> {
1175 Err(DispatchError::Rejected(
1176 "Fix: ranged-read test dispatcher does not implement dispatch.".to_string(),
1177 ))
1178 }
1179
1180 fn read_resident(&self, handle: u64) -> Result<Vec<u8>, DispatchError> {
1181 self.read_calls.set(self.read_calls.get() + 1);
1182 self.buffers
1183 .iter()
1184 .find(|(candidate, _)| *candidate == handle)
1185 .map(|(_, bytes)| bytes.clone())
1186 .ok_or_else(|| {
1187 DispatchError::BadInputs(format!(
1188 "Fix: test dispatcher missing resident handle {handle}."
1189 ))
1190 })
1191 }
1192
1193 fn read_resident_many(&self, handles: &[u64]) -> Result<Vec<Vec<u8>>, DispatchError> {
1194 self.batched_handles.borrow_mut().extend_from_slice(handles);
1195 handles
1196 .iter()
1197 .map(|&handle| self.read_resident(handle))
1198 .collect()
1199 }
1200 }
1201
1202 struct FailingAllocDispatcher {
1203 next_handle: Cell<u64>,
1204 fail_at_call: usize,
1205 allocations: RefCell<Vec<usize>>,
1206 freed: RefCell<Vec<u64>>,
1207 }
1208
1209 impl FailingAllocDispatcher {
1210 fn new(first_handle: u64, fail_at_call: usize) -> Self {
1211 Self {
1212 next_handle: Cell::new(first_handle),
1213 fail_at_call,
1214 allocations: RefCell::new(Vec::new()),
1215 freed: RefCell::new(Vec::new()),
1216 }
1217 }
1218 }
1219
1220 impl OptimizerDispatcher for FailingAllocDispatcher {
1221 fn dispatch(
1222 &self,
1223 _program: &Program,
1224 _inputs: &[Vec<u8>],
1225 _grid_override: Option<[u32; 3]>,
1226 ) -> Result<Vec<Vec<u8>>, DispatchError> {
1227 Err(DispatchError::Rejected(
1228 "Fix: failing allocation test dispatcher does not implement dispatch.".to_string(),
1229 ))
1230 }
1231
1232 fn alloc_resident(&self, byte_len: usize) -> Result<u64, DispatchError> {
1233 let call = self.allocations.borrow().len();
1234 self.allocations.borrow_mut().push(byte_len);
1235 if call == self.fail_at_call {
1236 return Err(DispatchError::BackendError(
1237 "Fix: injected optimizer resident allocation failure".to_string(),
1238 ));
1239 }
1240 let handle = self.next_handle.get();
1241 self.next_handle.set(handle + 1);
1242 Ok(handle)
1243 }
1244
1245 fn free_resident(&self, handle: u64) -> Result<(), DispatchError> {
1246 self.freed.borrow_mut().push(handle);
1247 Ok(())
1248 }
1249 }
1250
1251 #[test]
1252 fn generated_fill_upload_staging_preserves_fill_then_upload_order() {
1253 let host_payload = [0xA5_u8, 0x5A];
1254 let mut staged = Vec::new();
1255
1256 with_staged_fill_uploads(
1257 &[(7, 3, 0x11), (9, 2, 0x22)],
1258 &[(13, host_payload.as_slice())],
1259 "test fill payloads",
1260 "test combined uploads",
1261 |uploads| {
1262 for &(handle, bytes) in uploads {
1263 staged.push((handle, bytes.to_vec()));
1264 }
1265 Ok(())
1266 },
1267 )
1268 .expect("Fix: shared resident fill staging should succeed");
1269
1270 assert_eq!(
1271 staged,
1272 vec![
1273 (7, vec![0x11, 0x11, 0x11]),
1274 (9, vec![0x22, 0x22]),
1275 (13, host_payload.to_vec()),
1276 ],
1277 "resident fill staging must preserve device-fill uploads before caller uploads"
1278 );
1279 }
1280
1281 #[test]
1282 fn resident_grouped_allocation_rolls_back_partial_handles() {
1283 let dispatcher = FailingAllocDispatcher::new(90, 2);
1284
1285 let err = dispatcher
1286 .alloc_resident_many(&[4, 8, 12])
1287 .expect_err("Fix: injected grouped allocation failure should surface");
1288
1289 assert!(
1290 matches!(err, DispatchError::BackendError(message) if message.contains("injected optimizer resident allocation failure"))
1291 );
1292 assert_eq!(dispatcher.allocations.borrow().as_slice(), &[4, 8, 12]);
1293 assert_eq!(
1294 dispatcher.freed.borrow().as_slice(),
1295 &[90, 91],
1296 "Fix: grouped resident allocation must free every prior handle on failure."
1297 );
1298 }
1299
1300 #[test]
1301 fn ranged_readback_deduplicates_full_buffer_reads_by_handle() {
1302 let dispatcher = RangedReadDispatcher {
1303 buffers: vec![(7, (0u8..32).collect()), (9, (100u8..132).collect())],
1304 read_calls: Cell::new(0),
1305 batched_handles: RefCell::new(Vec::new()),
1306 };
1307
1308 let outputs = dispatcher
1309 .read_resident_ranges(&[
1310 ResidentReadRange {
1311 handle_id: 7,
1312 byte_offset: 4,
1313 byte_len: 4,
1314 },
1315 ResidentReadRange {
1316 handle_id: 9,
1317 byte_offset: 2,
1318 byte_len: 3,
1319 },
1320 ResidentReadRange {
1321 handle_id: 7,
1322 byte_offset: 12,
1323 byte_len: 5,
1324 },
1325 ])
1326 .expect("Fix: ranged readback must succeed for in-bounds dedup keys; return Err on overlap violations - deduplicated ranged readback must succeed");
1327
1328 assert_eq!(
1329 outputs,
1330 vec![
1331 vec![4, 5, 6, 7],
1332 vec![102, 103, 104],
1333 vec![12, 13, 14, 15, 16]
1334 ]
1335 );
1336 assert_eq!(
1337 dispatcher.read_calls.get(),
1338 2,
1339 "Fix: default ranged readback must read each unique resident handle once, not once per range."
1340 );
1341 assert_eq!(
1342 dispatcher.batched_handles.borrow().as_slice(),
1343 &[7, 9],
1344 "Fix: default ranged readback must preserve first-seen handle order for batched backend overrides."
1345 );
1346 }
1347
1348 #[test]
1349 fn generated_ranged_readbacks_deduplicate_handles_without_reordering_ranges() {
1350 let dispatcher = RangedReadDispatcher {
1351 buffers: (0..8u64)
1352 .map(|handle| {
1353 (
1354 handle,
1355 (0..64u8)
1356 .map(|byte| byte.wrapping_add((handle as u8).wrapping_mul(17)))
1357 .collect::<Vec<_>>(),
1358 )
1359 })
1360 .collect(),
1361 read_calls: Cell::new(0),
1362 batched_handles: RefCell::new(Vec::new()),
1363 };
1364 let ranges = (0..2048usize)
1365 .map(|case| ResidentReadRange {
1366 handle_id: ((case.wrapping_mul(5).wrapping_add(case / 11)) % 8) as u64,
1367 byte_offset: (case.wrapping_mul(7)) % 48,
1368 byte_len: (case % 16) + 1,
1369 })
1370 .collect::<Vec<_>>();
1371
1372 let outputs = dispatcher
1373 .read_resident_ranges(&ranges)
1374 .expect("Fix: generated matrix fixtures must stay in-bounds; fix fixture or return Err - generated ranged readback matrix must succeed");
1375
1376 assert_eq!(outputs.len(), ranges.len());
1377 for (range, output) in ranges.iter().zip(outputs.iter()) {
1378 let full = dispatcher
1379 .buffers
1380 .iter()
1381 .find(|(handle, _)| *handle == range.handle_id)
1382 .map(|(_, bytes)| bytes.as_slice())
1383 .expect("Fix: replace expect with fallible API or document caller precondition; panic only on programmer error - generated range uses known handle");
1384 assert_eq!(
1385 output.as_slice(),
1386 &full[range.byte_offset..range.byte_offset + range.byte_len],
1387 "generated range must preserve caller range order and byte-exact slices"
1388 );
1389 }
1390 assert_eq!(
1391 dispatcher.read_calls.get(),
1392 8,
1393 "Fix: generated ranged readback matrix must issue one full read per unique handle."
1394 );
1395 }
1396
1397 #[test]
1401 fn declared_dispatch_outputs_are_the_writable_storage_buffers_in_order() {
1402 use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType};
1403
1404 let program = Program::wrapped(
1405 vec![
1406 BufferDecl::storage("ro_in", 0, BufferAccess::ReadOnly, DataType::U32),
1407 BufferDecl::storage("frontier_out", 1, BufferAccess::ReadWrite, DataType::U32)
1408 .with_count(4),
1409 BufferDecl::workgroup("wg_scratch", 256, DataType::U32),
1410 BufferDecl::storage("changed", 2, BufferAccess::ReadWrite, DataType::U32)
1411 .with_count(1),
1412 BufferDecl::storage("sink", 3, BufferAccess::WriteOnly, DataType::U32)
1413 .with_count(1),
1414 ],
1415 [1, 1, 1],
1416 Vec::new(),
1417 );
1418
1419 let names: Vec<&str> = declared_dispatch_outputs(&program)
1420 .iter()
1421 .map(|decl| decl.name())
1422 .collect();
1423 assert_eq!(names, vec!["frontier_out", "changed", "sink"]);
1424 }
1425
1426 #[test]
1430 fn declared_dispatch_outputs_is_empty_when_nothing_is_writable() {
1431 use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType};
1432
1433 let program = Program::wrapped(
1434 vec![
1435 BufferDecl::storage("ro_in", 0, BufferAccess::ReadOnly, DataType::U32),
1436 BufferDecl::workgroup("wg_scratch", 64, DataType::U32),
1437 ],
1438 [1, 1, 1],
1439 Vec::new(),
1440 );
1441 assert!(declared_dispatch_outputs(&program).is_empty());
1442 }
1443}