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
21/// The buffers an [`OptimizerDispatcher`] returns, in declared order.
22///
23/// The trait contract is "the declared outputs in the same canonical order", which
24/// means every writable storage buffer. Workgroup scratch is never a dispatch
25/// output, so it is excluded.
26///
27/// This is the single owner of that rule. A dispatcher that hardcodes its own
28/// output list instead of deriving it from the program silently disagrees with the
29/// program the moment a buffer is added or removed. That is exactly how a
30/// three-output oracle came to be paired with a two-output DCE program: the oracle
31/// asserted a layout the program no longer had, and nothing compared the two.
32#[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
46/// One resident-buffer kernel launch in an ordered optimizer sequence.
47pub struct ResidentDispatchStep<'a> {
48    /// Program to launch.
49    pub program: &'a Program,
50    /// Resident handle ids in canonical buffer binding order.
51    pub handle_ids: &'a [u64],
52    /// Optional launch grid override.
53    pub grid_override: Option<[u32; 3]>,
54}
55
56/// One byte range to read from a resident buffer after an ordered sequence.
57#[derive(Clone, Copy, Debug, Eq, PartialEq)]
58pub struct ResidentReadRange {
59    /// Resident handle id.
60    pub handle_id: u64,
61    /// First byte to read from the device buffer.
62    pub byte_offset: usize,
63    /// Number of meaningful bytes to transfer.
64    pub byte_len: usize,
65}
66
67/// Resident handles for immutable payloads that may stay device-resident
68/// across optimizer calls.
69///
70/// `retained_by_dispatcher` means the dispatcher owns the handles after the
71/// caller is done with the current launch sequence. Call
72/// [`OptimizerDispatcher::release_resident_static_uploads`] instead of
73/// `free_resident` so CUDA can keep read-only graph/arena buffers hot while
74/// portable dispatchers free them immediately.
75#[derive(Debug)]
76pub struct ResidentStaticBufferSet {
77    /// Resident handle ids in the same order as the payload slice passed to
78    /// `acquire_resident_static_uploads`.
79    pub handles: Vec<u64>,
80    /// True when the handles were already resident and no host upload was paid.
81    pub cache_hit: bool,
82    /// True when the dispatcher retained ownership for future reuse.
83    pub retained_by_dispatcher: bool,
84}
85
86/// Errors a dispatcher may surface. Concrete backends compose their
87/// own error types into this; the orchestrator only needs the
88/// boundary message.
89#[derive(Debug)]
90pub enum DispatchError {
91    /// The dispatcher rejected the Program. The string carries the
92    /// backend's actionable message (must contain `Fix:`).
93    Rejected(String),
94    /// Input arity or shape did not match the Program's declared
95    /// buffer set. Hard error  -  not retryable.
96    BadInputs(String),
97    /// Backend raised an internal error. Same shape as `Rejected` but
98    /// the cause is in the backend, not the Program.
99    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
114/// Run a vyre Program with byte inputs, return byte outputs in the
115/// Program's declared output order.
116///
117/// This is the canonical dispatch boundary. Production impls go
118/// through `vyre-driver-wgpu` or `vyre-driver-cuda`; test impls use
119/// CPU oracles (gated to test-only builds).
120pub trait OptimizerDispatcher {
121    /// Dispatch `program` with the given byte inputs (one `Vec<u8>`
122    /// per declared input buffer in canonical buffer order). Returns
123    /// the declared outputs in the same canonical order.
124    ///
125    /// `grid_override` lets parallel kernels dispatch enough
126    /// workgroups to cover their input. `None` means "use the
127    /// backend's default grid" (typically `[1, 1, 1]`), which is what
128    /// sequential single-thread Programs want. Parallel passes
129    /// compute `Some([ceil(work/wg_x), 1, 1])` based on the input
130    /// size and their declared workgroup_size.
131    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    /// Whether this dispatcher supports the persistent-resident path.
139    /// Default: false. CUDA backend overrides to true. The orchestrator
140    /// uses this to decide whether to take the persistent fast-path
141    /// (encode arena once → upload once → dispatch many → readback once)
142    /// or use the non-resident per-call GPU dispatch path.
143    fn supports_persistent(&self) -> bool {
144        false
145    }
146
147    /// Device/lowering feature bits that affect reusable plan identity.
148    ///
149    /// Backends with feature-dependent lowering must override this so
150    /// self-substrate plan caches cannot replay a Program shape prepared for a
151    /// different hardware/lowering capability set. Test-only and reference
152    /// dispatchers keep the zero default because they do not specialize plans by
153    /// device.
154    fn device_feature_cache_key(&self) -> u64 {
155        0
156    }
157
158    /// Allocate a backend-resident buffer. Returns an opaque u64
159    /// handle. Callers must `free_resident` to release.
160    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    /// Allocate a logical group of resident buffers and roll back partial state
169    /// if any allocation fails.
170    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    /// Upload host bytes into a resident buffer.
201    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    /// Upload several resident buffers with one backend fence when supported.
208    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    /// Acquire resident handles for immutable payloads.
216    ///
217    /// Portable default behavior allocates and uploads exactly like
218    /// `alloc_resident` + `upload_resident_many`, then returns
219    /// `retained_by_dispatcher = false` so release frees the buffers. CUDA
220    /// overrides this to content-address immutable optimizer buffers and skip
221    /// H2D traffic on warmed identical programs.
222    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    /// Release a static resident buffer set acquired from
271    /// [`Self::acquire_resident_static_uploads`].
272    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    /// Download a resident buffer's current contents to host bytes.
286    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    /// Download several resident buffers with one backend fence when supported.
293    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    /// Download selected byte ranges from resident buffers.
301    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    /// Download selected byte ranges from resident buffers into caller-owned
311    /// byte slots.
312    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    /// Free a resident buffer previously returned by `alloc_resident`.
402    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    /// Dispatch a Program against resident-buffer handles. Each
409    /// handle is referenced from the Program's declared buffer in the
410    /// same canonical buffer order. RW buffers are not read back  -
411    /// caller invokes `read_resident` once at end of pipeline.
412    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    /// Dispatch an ordered sequence of resident-buffer Programs.
424    ///
425    /// Default implementation preserves correctness by fencing each step
426    /// through `dispatch_resident`. CUDA overrides this to enqueue the whole
427    /// dependent chain on one stream and synchronize once.
428    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    /// Dispatch an ordered resident sequence and read selected resident buffers.
439    ///
440    /// Default implementation keeps the portable contract: execute the ordered
441    /// sequence, then read buffers through `read_resident_many`. CUDA overrides
442    /// this to enqueue the D2H readbacks behind the kernels on the same stream
443    /// and pay one host fence.
444    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    /// Dispatch an ordered resident sequence and read selected byte ranges.
454    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    /// Upload resident buffers, dispatch an ordered resident sequence, then
464    /// read selected resident buffers.
465    ///
466    /// Default implementation fences at each portable boundary. CUDA overrides
467    /// this so H2D uploads, kernels, and D2H readbacks are ordered on one stream
468    /// with one host synchronization point.
469    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    /// Upload resident buffers, dispatch an ordered resident sequence, then
480    /// read selected byte ranges.
481    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    /// Same contract as [`Self::upload_resident_many_sequence_read_many`],
492    /// but writes readbacks into caller-owned byte slots.
493    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    /// Same contract as [`Self::upload_resident_many_sequence_read_many_into`],
515    /// but first clears full resident buffers to zero.
516    ///
517    /// Portable dispatchers emulate clears as zero-byte payload uploads and
518    /// still pay one upload/sequence/read boundary. CUDA overrides this to
519    /// enqueue device-side memset operations on the same stream before explicit
520    /// uploads and kernels, avoiding PCIe traffic for scratch initialization
521    /// without adding host fences.
522    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    /// Same contract as
558    /// [`Self::clear_upload_resident_many_sequence_read_many_into`], but fills
559    /// each resident buffer with an arbitrary byte value.
560    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    /// Same contract as [`Self::upload_resident_many_sequence_read_ranges_into`],
594    /// but fills resident buffers first. CUDA overrides this to use device
595    /// memset and compact D2H range copies on the same stream.
596    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    /// Same contract as [`Self::upload_resident_many_sequence_read_ranges`],
630    /// but writes compact readbacks into caller-owned byte slots.
631    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    //! CPU oracle dispatcher for tests and explicit CPU-parity builds. Maps a small allowlist of
698    //! self-hosted-optimizer Programs onto their `vyre_primitives`
699    //! `cpu_ref` reference implementations and reproduces the
700    //! dispatch byte contract.
701    //!
702    //! This module exists to prove the encoder/decoder are sound
703    //! against the same numerical contract the production GPU path
704    //! must honor. It is not compiled unless tests or `cpu-parity` are enabled.
705    //!
706    //! Adding a Program here means the oracle hand-writes the byte
707    //! marshalling that the WgpuBackend dispatcher infers from
708    //! `BufferDecl`s. That duplication is acceptable for tests; a
709    //! production dispatcher reflectively reads BufferDecls.
710    //!
711    //! For now we cover the Programs the orchestrator currently
712    //! invokes (DCE → `persistent_bfs`). When CSE / const-fold land
713    //! they each add a small case here.
714
715    use super::{DispatchError, OptimizerDispatcher};
716    use vyre_foundation::ir::Program;
717
718    /// CPU oracle dispatcher. Recognizes only the optimizer's own
719    /// canonical Programs by matching the wrapping Region's generator
720    /// op-id and the declared buffer set.
721    pub struct CpuOracleDispatcher;
722
723    impl CpuOracleDispatcher {
724        /// Construct the oracle dispatcher. Cheap; does no backend
725        /// probing.
726        #[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            // Identify the optimizer Program by its top-level Region
746            // generator. Self-hosted Programs all wrap their bodies
747            // in a Region with a known op-id.
748            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        // Buffer order (per `persistent_bfs.rs::persistent_bfs`):
786        //   0 pg_nodes (RO)
787        //   1 pg_edge_offsets (RO)
788        //   2 pg_edge_targets (RO)
789        //   3 pg_edge_kind_mask (RO)
790        //   4 pg_node_tags (RO)
791        //   5 frontier_in (RO)
792        //   6 frontier_out (RW)
793        //   7 changed (RW)
794        //   8 converged (RW)
795        //   9 wg_scratch (workgroup)   -  not an input
796        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        // The Region carries the shape and max_iters in its body
810        // structure; rather than re-derive that from IR walks, the
811        // oracle re-computes via cpu_ref using the buffers' lengths.
812        let node_count = nodes.len() as u32;
813
814        // Iteration cap: if the caller declared `frontier_in` of length L
815        // (= bitset_words(node_count)) the oracle uses `node_count` as
816        // the saturation budget  -  same default the Program builder uses
817        // when callers want closure.
818        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        // Emit one buffer per DECLARED writable output, in declared order, rather
839        // than a hardcoded list. The oracle stands in for a device backend, so it
840        // must return exactly the layout the program declares; a pasted list goes
841        // stale the moment a buffer is added and then reports a shape the program
842        // does not have. The converged word matches the device readback
843        // bit-for-bit (proven in vyre-primitives converged_device_parity), so the
844        // oracle is a faithful stand-in for the real GPU signal.
845        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    /// Reconstruct the `changed` buffer the kernel would leave behind, sized to the
870    /// program's declared element count.
871    ///
872    /// The two BFS variants give slot 0 different meanings, and the declared count is
873    /// what distinguishes them, so deriving from it keeps the oracle honest for both:
874    ///
875    /// - count 1 (the DCE variant): slot 0 is the LAST iteration's progress, because
876    ///   the kernel zeroes it every iteration to drive its early exit. On a real
877    ///   fixpoint that leaves 0; on an exhausted budget the final iteration grew, so
878    ///   it leaves 1.
879    /// - count 2 (the sticky variant): slot 0 is that same per-iteration flag and
880    ///   slot 1 is the sticky OR across all iterations, which is what `cpu_ref`
881    ///   reports as `changed`.
882    ///
883    /// # Errors
884    /// Rejects any other declared count rather than guessing, since a wrong guess
885    /// here is a silently misread flag rather than a visible failure.
886    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        /// The DCE variant declares one `changed` word holding the LAST iteration's
920        /// progress, because the kernel zeroes it every iteration to drive the early
921        /// exit. Reaching a fixpoint means that final compare saw zero, even when
922        /// earlier iterations did grow the frontier.
923        #[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        /// The sticky variant declares two words: the same per-iteration flag, then the
938        /// OR across every iteration. Collapsing them would make a converged run look
939        /// like it never grew at all.
940        #[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        /// An unrecognized count is refused rather than guessed. Guessing would hand a
955        /// consumer a flag read from the wrong slot, which looks like valid data.
956        #[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    /// The dispatch ABI returns writable storage buffers in declared order. Workgroup
1457    /// scratch is not a dispatch output, and admitting it would shift every consumer's
1458    /// output index by one without any visible error.
1459    #[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    /// A program that declares no writable buffer has no dispatch outputs. Returning a
1486    /// non-empty list there would be the oracle inventing a layout, which is the defect
1487    /// this derivation exists to prevent.
1488    #[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}