Skip to main content

vyre_self_substrate/optimizer/
dispatcher.rs

1//! Dispatcher trait  -  the seam between the self-hosted optimizer and
2//! a backend that can actually run vyre Programs.
3//!
4//! The optimizer encodes the user's Program into ProgramGraph buffers,
5//! builds a vyre Program that does the analysis (e.g. `persistent_bfs`),
6//! and asks an `OptimizerDispatcher` to run that analysis Program. The
7//! returned bytes drive the rewrite.
8//!
9//! `vyre-self-substrate` cannot depend on a concrete backend  -  it sits
10//! below the driver layer. The trait inverts that dependency: the
11//! orchestrator code stays in self-substrate, and a backend crate
12//! (e.g. `vyre-driver-wgpu` or a runtime wrapper) provides the impl.
13//!
14//! Test code in this crate uses `oracle::CpuOracleDispatcher` so the
15//! encoder can be proven sound against the existing primitive oracles
16//! before any GPU backend is wired. The CPU oracle is gated to tests
17//! only  -  it is never on a production code path.
18
19use 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/// The buffers an [`OptimizerDispatcher`] returns, in declared order.
32///
33/// The trait contract is "the declared outputs in the same canonical order", which
34/// means every writable storage buffer. Workgroup scratch is never a dispatch
35/// output, so it is excluded.
36///
37/// This is the single owner of that rule. A dispatcher that hardcodes its own
38/// output list instead of deriving it from the program silently disagrees with the
39/// program the moment a buffer is added or removed. That is exactly how a
40/// three-output oracle came to be paired with a two-output DCE program: the oracle
41/// asserted a layout the program no longer had, and nothing compared the two.
42#[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
56/// One resident-buffer kernel launch in an ordered optimizer sequence.
57pub struct ResidentDispatchStep<'a> {
58    /// Program to launch.
59    pub program: &'a Program,
60    /// Resident handle ids in canonical buffer binding order.
61    pub handle_ids: &'a [u64],
62    /// Optional launch grid override.
63    pub grid_override: Option<[u32; 3]>,
64}
65
66/// One byte range to read from a resident buffer after an ordered sequence.
67#[derive(Clone, Copy, Debug, Eq, PartialEq)]
68pub struct ResidentReadRange {
69    /// Resident handle id.
70    pub handle_id: u64,
71    /// First byte to read from the device buffer.
72    pub byte_offset: usize,
73    /// Number of meaningful bytes to transfer.
74    pub byte_len: usize,
75}
76
77/// Resident handles for immutable payloads that may stay device-resident
78/// across optimizer calls.
79///
80/// `retained_by_dispatcher` means the dispatcher owns the handles after the
81/// caller is done with the current launch sequence. Call
82/// [`OptimizerDispatcher::release_resident_static_uploads`] instead of
83/// `free_resident` so CUDA can keep read-only graph/arena buffers hot while
84/// portable dispatchers free them immediately.
85#[derive(Debug)]
86pub struct ResidentStaticBufferSet {
87    /// Resident handle ids in the same order as the payload slice passed to
88    /// `acquire_resident_static_uploads`.
89    pub handles: Vec<u64>,
90    /// True when the handles were already resident and no host upload was paid.
91    pub cache_hit: bool,
92    /// True when the dispatcher retained ownership for future reuse.
93    pub retained_by_dispatcher: bool,
94}
95
96/// Errors a dispatcher may surface. Concrete backends compose their
97/// own error types into this; the orchestrator only needs the
98/// boundary message.
99#[derive(Debug)]
100pub enum DispatchError {
101    /// The dispatcher rejected the Program. The string carries the
102    /// backend's actionable message (must contain `Fix:`).
103    Rejected(String),
104    /// Input arity or shape did not match the Program's declared
105    /// buffer set. Hard error  -  not retryable.
106    BadInputs(String),
107    /// Backend raised an internal error. Same shape as `Rejected` but
108    /// the cause is in the backend, not the Program.
109    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
124/// Run a vyre Program with byte inputs, return byte outputs in the
125/// Program's declared output order.
126///
127/// This is the canonical dispatch boundary. Production impls go
128/// through `vyre-driver-wgpu` or `vyre-driver-cuda`; test impls use
129/// CPU oracles (gated to test-only builds).
130pub trait OptimizerDispatcher {
131    /// Dispatch `program` with the given byte inputs (one `Vec<u8>`
132    /// per declared input buffer in canonical buffer order). Returns
133    /// the declared outputs in the same canonical order.
134    ///
135    /// `grid_override` lets parallel kernels dispatch enough
136    /// workgroups to cover their input. `None` means "use the
137    /// backend's default grid" (typically `[1, 1, 1]`), which is what
138    /// sequential single-thread Programs want. Parallel passes
139    /// compute `Some([ceil(work/wg_x), 1, 1])` based on the input
140    /// size and their declared workgroup_size.
141    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    /// Whether this dispatcher supports the persistent-resident path.
149    /// Default: false. CUDA backend overrides to true. The orchestrator
150    /// uses this to decide whether to take the persistent fast-path
151    /// (encode arena once → upload once → dispatch many → readback once)
152    /// or use the non-resident per-call GPU dispatch path.
153    fn supports_persistent(&self) -> bool {
154        false
155    }
156
157    /// Device/lowering feature bits that affect reusable plan identity.
158    ///
159    /// Backends with feature-dependent lowering must override this so
160    /// self-substrate plan caches cannot replay a Program shape prepared for a
161    /// different hardware/lowering capability set. Test-only and reference
162    /// dispatchers keep the zero default because they do not specialize plans by
163    /// device.
164    fn device_feature_cache_key(&self) -> u64 {
165        0
166    }
167
168    /// Allocate a backend-resident buffer. Returns an opaque u64
169    /// handle. Callers must `free_resident` to release.
170    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    /// Allocate a logical group of resident buffers and roll back partial state
179    /// if any allocation fails.
180    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    /// Upload host bytes into a resident buffer.
211    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    /// Upload several resident buffers with one backend fence when supported.
218    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    /// Acquire resident handles for immutable payloads.
226    ///
227    /// Portable default behavior allocates and uploads exactly like
228    /// `alloc_resident` + `upload_resident_many`, then returns
229    /// `retained_by_dispatcher = false` so release frees the buffers. CUDA
230    /// overrides this to content-address immutable optimizer buffers and skip
231    /// H2D traffic on warmed identical programs.
232    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    /// Release a static resident buffer set acquired from
281    /// [`Self::acquire_resident_static_uploads`].
282    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    /// Download a resident buffer's current contents to host bytes.
296    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    /// Download several resident buffers with one backend fence when supported.
303    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    /// Download selected byte ranges from resident buffers.
311    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    /// Download selected byte ranges from resident buffers into caller-owned
321    /// byte slots.
322    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    /// Free a resident buffer previously returned by `alloc_resident`.
412    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    /// Dispatch a Program against resident-buffer handles. Each
419    /// handle is referenced from the Program's declared buffer in the
420    /// same canonical buffer order. RW buffers are not read back  -
421    /// caller invokes `read_resident` once at end of pipeline.
422    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    /// Dispatch an ordered sequence of resident-buffer Programs.
434    ///
435    /// Default implementation preserves correctness by fencing each step
436    /// through `dispatch_resident`. CUDA overrides this to enqueue the whole
437    /// dependent chain on one stream and synchronize once.
438    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    /// Dispatch an ordered resident sequence and read selected resident buffers.
449    ///
450    /// Default implementation keeps the portable contract: execute the ordered
451    /// sequence, then read buffers through `read_resident_many`. CUDA overrides
452    /// this to enqueue the D2H readbacks behind the kernels on the same stream
453    /// and pay one host fence.
454    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    /// Dispatch an ordered resident sequence and read selected byte ranges.
464    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    /// Upload resident buffers, dispatch an ordered resident sequence, then
474    /// read selected resident buffers.
475    ///
476    /// Default implementation fences at each portable boundary. CUDA overrides
477    /// this so H2D uploads, kernels, and D2H readbacks are ordered on one stream
478    /// with one host synchronization point.
479    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    /// Upload resident buffers, dispatch an ordered resident sequence, then
490    /// read selected byte ranges.
491    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    /// Same contract as [`Self::upload_resident_many_sequence_read_many`],
502    /// but writes readbacks into caller-owned byte slots.
503    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    /// Same contract as [`Self::upload_resident_many_sequence_read_many_into`],
525    /// but first clears full resident buffers to zero.
526    ///
527    /// Portable dispatchers emulate clears as zero-byte payload uploads and
528    /// still pay one upload/sequence/read boundary. CUDA overrides this to
529    /// enqueue device-side memset operations on the same stream before explicit
530    /// uploads and kernels, avoiding PCIe traffic for scratch initialization
531    /// without adding host fences.
532    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    /// Same contract as
568    /// [`Self::clear_upload_resident_many_sequence_read_many_into`], but fills
569    /// each resident buffer with an arbitrary byte value.
570    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    /// Same contract as [`Self::upload_resident_many_sequence_read_ranges_into`],
604    /// but fills resident buffers first. CUDA overrides this to use device
605    /// memset and compact D2H range copies on the same stream.
606    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    /// Same contract as [`Self::upload_resident_many_sequence_read_ranges`],
640    /// but writes compact readbacks into caller-owned byte slots.
641    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    //! CPU oracle dispatcher for tests and explicit CPU-parity builds. Maps a small allowlist of
708    //! self-hosted-optimizer Programs onto their `vyre_primitives`
709    //! `cpu_ref` reference implementations and reproduces the
710    //! dispatch byte contract.
711    //!
712    //! This module exists to prove the encoder/decoder are sound
713    //! against the same numerical contract the production GPU path
714    //! must honor. It is not compiled unless tests or `cpu-parity` are enabled.
715    //!
716    //! Adding a Program here means the oracle hand-writes the byte
717    //! marshalling that the WgpuBackend dispatcher infers from
718    //! `BufferDecl`s. That duplication is acceptable for tests; a
719    //! production dispatcher reflectively reads BufferDecls.
720    //!
721    //! For now we cover the Programs the orchestrator currently
722    //! invokes (DCE → `persistent_bfs`). When CSE / const-fold land
723    //! they each add a small case here.
724
725    use super::{DispatchError, OptimizerDispatcher, ParsedIfdsRules};
726    use vyre_foundation::ir::Program;
727
728    /// CPU oracle dispatcher. Recognizes only the optimizer's own
729    /// canonical Programs by matching the wrapping Region's generator
730    /// op-id and the declared buffer set.
731    pub struct CpuOracleDispatcher;
732
733    impl CpuOracleDispatcher {
734        /// Construct the oracle dispatcher. Cheap; does no backend
735        /// probing.
736        #[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            // Identify the optimizer Program by its top-level Region
756            // generator. Self-hosted Programs all wrap their bodies
757            // in a Region with a known op-id.
758            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        // Buffer order (per `persistent_bfs.rs::persistent_bfs`):
796        //   0 pg_nodes (RO)
797        //   1 pg_edge_offsets (RO)
798        //   2 pg_edge_targets (RO)
799        //   3 pg_edge_kind_mask (RO)
800        //   4 pg_node_tags (RO)
801        //   5 frontier_in (RO)
802        //   6 frontier_out (RW)
803        //   7 changed (RW)
804        //   8 converged (RW)
805        //   9 wg_scratch (workgroup)   -  not an input
806        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        // The Region carries the shape and max_iters in its body
820        // structure; rather than re-derive that from IR walks, the
821        // oracle re-computes via cpu_ref using the buffers' lengths.
822        let node_count = nodes.len() as u32;
823
824        // Iteration cap: if the caller declared `frontier_in` of length L
825        // (= bitset_words(node_count)) the oracle uses `node_count` as
826        // the saturation budget  -  same default the Program builder uses
827        // when callers want closure.
828        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        // Emit one buffer per DECLARED writable output, in declared order, rather
849        // than a hardcoded list. The oracle stands in for a device backend, so it
850        // must return exactly the layout the program declares; a pasted list goes
851        // stale the moment a buffer is added and then reports a shape the program
852        // does not have. The converged word matches the device readback
853        // bit-for-bit (proven in vyre-primitives converged_device_parity), so the
854        // oracle is a faithful stand-in for the real GPU signal.
855        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    /// Reconstruct the `changed` buffer the kernel would leave behind, sized to the
880    /// program's declared element count.
881    ///
882    /// The two BFS variants give slot 0 different meanings, and the declared count is
883    /// what distinguishes them, so deriving from it keeps the oracle honest for both:
884    ///
885    /// - count 1 (the DCE variant): slot 0 is the LAST iteration's progress, because
886    ///   the kernel zeroes it every iteration to drive its early exit. On a real
887    ///   fixpoint that leaves 0; on an exhausted budget the final iteration grew, so
888    ///   it leaves 1.
889    /// - count 2 (the sticky variant): slot 0 is that same per-iteration flag and
890    ///   slot 1 is the sticky OR across all iterations, which is what `cpu_ref`
891    ///   reports as `changed`.
892    ///
893    /// # Errors
894    /// Rejects any other declared count rather than guessing, since a wrong guess
895    /// here is a silently misread flag rather than a visible failure.
896    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        /// The DCE variant declares one `changed` word holding the LAST iteration's
931        /// progress, because the kernel zeroes it every iteration to drive the early
932        /// exit. Reaching a fixpoint means that final compare saw zero, even when
933        /// earlier iterations did grow the frontier.
934        #[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        /// The sticky variant declares two words: the same per-iteration flag, then the
949        /// OR across every iteration. Collapsing them would make a converged run look
950        /// like it never grew at all.
951        #[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        /// An unrecognized count is refused rather than guessed. Guessing would hand a
966        /// consumer a flag read from the wrong slot, which looks like valid data.
967        #[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    /// The dispatch ABI returns writable storage buffers in declared order. Workgroup
1398    /// scratch is not a dispatch output, and admitting it would shift every consumer's
1399    /// output index by one without any visible error.
1400    #[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    /// A program that declares no writable buffer has no dispatch outputs. Returning a
1427    /// non-empty list there would be the oracle inventing a layout, which is the defect
1428    /// this derivation exists to prevent.
1429    #[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}