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