Skip to main content

hrx/
runtime.rs

1//! GPU conveniences implemented over the owned native fabric.
2use crate::{Error, Result, Target, fabric};
3use std::{
4    cell::RefCell,
5    collections::BTreeMap,
6    path::Path,
7    sync::{Arc, Mutex},
8};
9
10/// Native entry's checked argument and workgroup layout.
11#[derive(Clone, Debug, Default)]
12pub struct ExportInfo {
13    /// Scalar bytes packed in declaration order, excluding alignment padding.
14    pub constant_byte_length: u32,
15    /// Number of explicit buffer arguments.
16    pub binding_count: u32,
17    /// Number of explicit native arguments.
18    pub parameter_count: u32,
19    /// Required workgroup dimensions; zero means unconstrained.
20    pub workgroup_size: [u32; 3],
21}
22struct Inner {
23    device: fabric::Device,
24    queue: fabric::Queue,
25    last: Mutex<Option<fabric::Completion>>,
26    free: Mutex<Vec<fabric::Buffer>>,
27    dispatch_cache: Mutex<Vec<CachedDispatch>>,
28    transfer_cache: Mutex<Vec<CachedTransfer>>,
29}
30impl Inner {
31    fn submit(&self, command: &fabric::PreparedGpu) -> Result<()> {
32        let done = match unsafe { command.dispatch() } {
33            Ok(done) => done,
34            Err(Error::Busy(_)) => {
35                self.wait()?;
36                unsafe { command.dispatch() }?
37            }
38            Err(error) => return Err(error),
39        };
40        *self
41            .last
42            .lock()
43            .map_err(|_| Error::DeviceLost("stream timeline poisoned".into()))? = Some(done);
44        Ok(())
45    }
46    fn drain_for_drop(&self) -> bool {
47        self.last.lock().ok().is_some_and(|last| {
48            last.as_ref().is_none_or(|done| {
49                done.wait_timeout(std::time::Duration::from_secs(10))
50                    .unwrap_or(false)
51            })
52        })
53    }
54    fn wait(&self) -> Result<()> {
55        if let Some(done) = self
56            .last
57            .lock()
58            .map_err(|_| Error::DeviceLost("stream timeline poisoned".into()))?
59            .as_ref()
60        {
61            done.wait()?;
62        }
63        Ok(())
64    }
65}
66/// An owned GPU address and execution domain.
67#[derive(Clone)]
68pub struct Device {
69    native: fabric::Device,
70}
71impl Device {
72    /// Activate an exact GPU ordinal in the native provider.
73    pub fn open(index: i32) -> Result<Self> {
74        let index = usize::try_from(index)
75            .map_err(|_| Error::Message("GPU index must be nonnegative".into()))?;
76        Ok(Self {
77            native: fabric::Device::open(fabric::Engine::Gpu, index)?,
78        })
79    }
80    /// Require the workload's exact target before creating execution resources.
81    pub fn open_for(index: i32, target: &str) -> Result<Self> {
82        let device = Self::open(index)?;
83        if device.target().as_str() != target {
84            return Err(Error::Unsupported(format!(
85                "workload requires {target}, found {}",
86                device.target().as_str()
87            )));
88        }
89        Ok(device)
90    }
91    /// Exact compiler deployment key.
92    pub fn target(&self) -> &Target {
93        self.native.target()
94    }
95    /// Create an independent ordered native command queue.
96    pub fn stream(&self) -> Result<Stream> {
97        Ok(Stream {
98            inner: Arc::new(Inner {
99                device: self.native.clone(),
100                queue: self.native.queue()?,
101                last: Mutex::new(None),
102                free: Mutex::new(Vec::with_capacity(16)),
103                dispatch_cache: Mutex::new(Vec::with_capacity(64)),
104                transfer_cache: Mutex::new(Vec::with_capacity(128)),
105            }),
106            staging: Vec::new(),
107            staging_pool: Vec::new(),
108            scratch: BTreeMap::new(),
109            scratch_bytes: 0,
110            scratch_limit: 256 * 1024 * 1024,
111            budget: None,
112            budget_uses: RefCell::new(BudgetUses::default()),
113        })
114    }
115    /// Access the owned native domain for explicit cross-engine allocations.
116    pub fn native(&self) -> &fabric::Device {
117        &self.native
118    }
119}
120/// Device backing with allocation-budget ownership.
121pub struct Buffer {
122    pub(crate) native: fabric::Buffer,
123    bytes: usize,
124    owner: Arc<Inner>,
125    reservation: Option<Arc<crate::residency::MemoryReservation>>,
126    poolable: bool,
127}
128impl Drop for Buffer {
129    fn drop(&mut self) {
130        // Cached commands are only useful while the application owns their
131        // bindings. Eviction also keeps pooled or budgeted backing reclaimable.
132        if let Ok(mut cache) = self.owner.dispatch_cache.lock() {
133            cache.retain(|entry| {
134                !entry
135                    .bindings
136                    .iter()
137                    .any(|(buffer, _, _)| buffer.same_backing(&self.native))
138            });
139        }
140        if let Ok(mut cache) = self.owner.transfer_cache.lock() {
141            cache.retain(|entry| {
142                !entry.destination.same_backing(&self.native)
143                    && !entry
144                        .source
145                        .as_ref()
146                        .is_some_and(|(buffer, _)| buffer.same_backing(&self.native))
147            });
148        }
149        if self.poolable
150            && self.reservation.is_none()
151            && self.native.exclusively_owned()
152            && self.bytes <= 64 * 1024 * 1024
153            && let Ok(mut pool) = self.owner.free.lock()
154        {
155            let used: usize = pool.iter().map(fabric::Buffer::len).sum();
156            if pool.len() < 16 && self.bytes <= (64 * 1024 * 1024usize).saturating_sub(used) {
157                pool.push(self.native.clone());
158            }
159        }
160    }
161}
162/// Checked borrowed logical range in an owned buffer.
163#[derive(Clone, Copy, Debug)]
164pub struct View<'a> {
165    owner: &'a Buffer,
166    offset: usize,
167    length: usize,
168}
169impl<'a> View<'a> {
170    /// Select a checked range relative to this view.
171    pub fn slice(self, offset: usize, length: usize) -> Result<Self> {
172        checked_span(offset, length, self.length)?;
173        Ok(Self {
174            owner: self.owner,
175            offset: self.offset + offset,
176            length,
177        })
178    }
179    /// Owning allocation.
180    pub fn owner(self) -> &'a Buffer {
181        self.owner
182    }
183    /// Absolute offset in the allocation.
184    pub fn offset(self) -> usize {
185        self.offset
186    }
187    /// Logical byte length.
188    pub fn len(self) -> usize {
189        self.length
190    }
191    /// Whether the range is empty.
192    pub fn is_empty(self) -> bool {
193        self.length == 0
194    }
195}
196impl Buffer {
197    // Tracked storage accounts for its own lifetime; its backing must not move
198    // into an uncharged stream pool when that owner releases its reservation.
199    pub(crate) fn into_unpooled(mut self) -> Self {
200        self.poolable = false;
201        self
202    }
203    #[cfg(feature = "npu")]
204    pub(crate) fn share_with(mut self, device: &fabric::Device) -> Result<Self> {
205        self.poolable = false;
206        if self.native.device_address(device).is_err() {
207            self.native = self
208                .owner
209                .device
210                .fabric()
211                .share_owned(&self.native, &[self.owner.device.clone(), device.clone()])?;
212        }
213        Ok(self)
214    }
215    pub(crate) fn charged_to(&self, budget: &crate::residency::MemoryBudget) -> bool {
216        self.reservation
217            .as_ref()
218            .is_some_and(|reservation| budget.contains(reservation))
219    }
220    pub(crate) fn device_id(&self) -> usize {
221        self.owner.device.id()
222    }
223    /// Allocation's exposed byte length.
224    pub fn bytes(&self) -> usize {
225        self.bytes
226    }
227    /// Whole allocation as a logical binding.
228    pub fn binding(&self) -> View<'_> {
229        View {
230            owner: self,
231            offset: 0,
232            length: self.bytes,
233        }
234    }
235    /// Checked subrange.
236    pub fn try_slice(&self, offset: usize, length: usize) -> Result<View<'_>> {
237        self.binding().slice(offset, length)
238    }
239    /// Checked subrange, panicking on invalid bounds.
240    pub fn slice(&self, offset: usize, length: usize) -> View<'_> {
241        self.try_slice(offset, length)
242            .expect("slice exceeds allocation")
243    }
244    /// Mapped host address. Dereferencing requires external synchronization.
245    /// Ordinary allocations additionally require `cache_control` before host
246    /// reads (including partial cache-line updates) and after host writes;
247    /// `allocate_shared` omits that requirement.
248    pub fn device_ptr(&self) -> Result<*mut std::ffi::c_void> {
249        Ok(self.native.host_pointer().cast())
250    }
251    /// Publish host writes (`flush = true`) or acquire device writes (`false`).
252    /// Stream upload/read operations perform these transitions automatically.
253    ///
254    /// # Safety
255    /// The caller must exclude conflicting host/device access for the entire
256    /// transition and retain the allocation until all accesses have completed.
257    pub unsafe fn cache_control(&self, flush: bool, offset: usize, length: usize) -> Result<()> {
258        unsafe { self.native.cache_control(flush, offset, length) }
259    }
260    pub(crate) fn allocation_address(&self) -> Result<u64> {
261        self.native.device_address(&self.owner.device)
262    }
263}
264/// Trusted native GPU entry and its declaration-order argument layout.
265#[derive(Clone)]
266pub struct Kernel {
267    native: fabric::Kernel,
268    info: ExportInfo,
269    layout: Arc<[(u32, usize)]>,
270    symbol: Arc<str>,
271}
272impl Kernel {
273    pub(crate) fn device_id(&self) -> usize {
274        self.native.device().id()
275    }
276    /// Checked argument and launch metadata.
277    pub fn info(&self) -> &ExportInfo {
278        &self.info
279    }
280    /// Selected native entry name.
281    pub fn symbol(&self) -> &str {
282        &self.symbol
283    }
284    fn from_native(native: fabric::Kernel, symbol: &str) -> Result<Self> {
285        let layout = native.argument_layout()?;
286        let mut info = ExportInfo {
287            workgroup_size: native.workgroup_size(),
288            parameter_count: layout.len() as u32,
289            ..Default::default()
290        };
291        for &(kind, size) in &layout {
292            match kind {
293                1 => {
294                    info.constant_byte_length =
295                        info.constant_byte_length
296                            .checked_add(size as u32)
297                            .ok_or_else(|| Error::Message("constant layout overflow".into()))?
298                }
299                2 => info.binding_count += 1,
300                _ => {
301                    return Err(Error::Unsupported(
302                        "native argument kind is unsupported".into(),
303                    ));
304                }
305            }
306        }
307        Ok(Self {
308            native,
309            info,
310            layout: layout.into(),
311            symbol: symbol.into(),
312        })
313    }
314}
315#[derive(Default)]
316struct BudgetUses(std::collections::HashMap<usize, Arc<crate::residency::MemoryReservation>>);
317impl BudgetUses {
318    fn retain(&mut self, views: &[View<'_>]) {
319        for view in views {
320            if let Some(reservation) = &view.owner.reservation {
321                self.0
322                    .entry(Arc::as_ptr(reservation) as usize)
323                    .or_insert_with(|| reservation.clone());
324            }
325        }
326    }
327    fn clear(&mut self) {
328        self.0.clear();
329    }
330}
331const DISPATCH_CACHE_ENTRIES: usize = 1024;
332const DISPATCH_CACHE_BYTES: usize = 16 << 20;
333
334struct CachedDispatch {
335    kernel: Kernel,
336    grid: [u32; 3],
337    block: [u32; 3],
338    constants: Constants,
339    bindings: Vec<(fabric::Buffer, usize, usize)>,
340    command: fabric::PreparedGpu,
341}
342impl CachedDispatch {
343    fn matches(
344        &self,
345        kernel: &Kernel,
346        grid: [u32; 3],
347        block: [u32; 3],
348        constants: &Constants,
349        bindings: &[View<'_>],
350    ) -> bool {
351        self.kernel.native.same_entry(&kernel.native)
352            && self.grid == grid
353            && self.block == block
354            && self.constants.as_bytes() == constants.as_bytes()
355            && self.bindings.len() == bindings.len()
356            && self
357                .bindings
358                .iter()
359                .zip(bindings)
360                .all(|((buffer, offset, length), view)| {
361                    buffer.same_backing(&view.owner.native)
362                        && *offset == view.offset
363                        && *length == view.length
364                })
365    }
366}
367struct CachedTransfer {
368    destination: fabric::Buffer,
369    destination_offset: usize,
370    length: usize,
371    source: Option<(fabric::Buffer, usize)>,
372    value: u8,
373    command: fabric::PreparedGpu,
374}
375/// An ordered GPU stream with explicit prepared commands and bounded pools.
376pub struct Stream {
377    inner: Arc<Inner>,
378    staging: Vec<Buffer>,
379    staging_pool: Vec<Buffer>,
380    scratch: BTreeMap<usize, Vec<Buffer>>,
381    scratch_bytes: usize,
382    scratch_limit: usize,
383    budget: Option<crate::residency::MemoryBudget>,
384    budget_uses: RefCell<BudgetUses>,
385}
386impl Stream {
387    /// Open the first qualified GPU.
388    pub fn open() -> Result<Self> {
389        Device::open(0)?.stream()
390    }
391    /// Stable native device identity.
392    pub fn device_id(&self) -> usize {
393        self.inner.device.id()
394    }
395    /// Stream identity, unique while it remains live.
396    pub fn id(&self) -> usize {
397        Arc::as_ptr(&self.inner) as usize
398    }
399    /// Compiler deployment key.
400    pub fn target(&self) -> &Target {
401        self.inner.device.target()
402    }
403    /// Attach a shared allocation ceiling.
404    pub fn with_memory_budget(mut self, budget: crate::residency::MemoryBudget) -> Self {
405        self.budget = Some(budget);
406        self
407    }
408    /// Selected allocation ceiling.
409    pub fn memory_budget(&self) -> Option<&crate::residency::MemoryBudget> {
410        self.budget.as_ref()
411    }
412    fn reserve(&self, bytes: usize) -> Result<Option<Arc<crate::residency::MemoryReservation>>> {
413        self.budget
414            .as_ref()
415            .map(|budget| budget.reserve(bytes.max(1)).map(Arc::new))
416            .transpose()
417    }
418    fn owns(&self, buffer: &Buffer) -> Result<()> {
419        owns(&self.inner, buffer)
420    }
421    /// Allocate initialized native GPU-visible storage with unspecified contents.
422    /// Empty requests round to one.
423    pub fn allocate(&self, bytes: usize) -> Result<Buffer> {
424        let bytes = bytes.max(1);
425        let reservation = self.reserve(bytes)?;
426        let reused = if reservation.is_none() {
427            let mut pool = self
428                .inner
429                .free
430                .lock()
431                .map_err(|_| Error::DeviceLost("allocation pool poisoned".into()))?;
432            pool.iter()
433                .position(|buffer| buffer.len() == bytes)
434                .map(|index| pool.swap_remove(index))
435        } else {
436            None
437        };
438        let native = match reused {
439            Some(native) => native,
440            None => self
441                .inner
442                .device
443                .fabric()
444                .allocate(bytes, std::slice::from_ref(&self.inner.device))?,
445        };
446        Ok(Buffer {
447            native,
448            bytes,
449            owner: self.inner.clone(),
450            reservation,
451            poolable: true,
452        })
453    }
454    /// Allocate initialized storage containing zero bytes.
455    pub fn allocate_zeroed(&self, bytes: usize) -> Result<Buffer> {
456        let buffer = self.allocate(bytes)?;
457        buffer.native.zero()?;
458        Ok(buffer)
459    }
460    /// Allocate GPU-coherent host storage for direct, synchronized host access.
461    pub fn allocate_shared(&self, bytes: usize) -> Result<Buffer> {
462        let bytes = bytes.max(1);
463        let reservation = self.reserve(bytes)?;
464        let native = self
465            .inner
466            .device
467            .fabric()
468            .allocate_shared(bytes, std::slice::from_ref(&self.inner.device))?;
469        Ok(Buffer {
470            native,
471            bytes,
472            owner: self.inner.clone(),
473            reservation,
474            poolable: false,
475        })
476    }
477    #[cfg(feature = "npu")]
478    pub(crate) fn allocate_for(&self, bytes: usize, devices: &[fabric::Device]) -> Result<Buffer> {
479        let reservation = self.reserve(bytes)?;
480        let native = self.inner.device.fabric().allocate(bytes, devices)?;
481        native.device_address(&self.inner.device)?;
482        Ok(Buffer {
483            native,
484            bytes,
485            owner: self.inner.clone(),
486            reservation,
487            poolable: false,
488        })
489    }
490    /// Wait for all preceding native work and reclaim transfer staging.
491    pub fn synchronize(&mut self) -> Result<()> {
492        self.inner.wait()?;
493        self.reclaim_staging();
494        self.budget_uses.get_mut().clear();
495        Ok(())
496    }
497    /// Snapshot preceding work without a host wait.
498    pub fn record_event(&mut self) -> Result<Event> {
499        Ok(Event {
500            device: self.device_id(),
501            done: self
502                .inner
503                .last
504                .lock()
505                .map_err(|_| Error::DeviceLost("stream timeline poisoned".into()))?
506                .clone(),
507        })
508    }
509    /// Order subsequent stream work after an immutable event.
510    pub fn wait_event(&mut self, event: &Event) -> Result<()> {
511        if event.device != self.device_id() {
512            return Err(Error::Message("event belongs to another device".into()));
513        }
514        if let Some(done) = &event.done {
515            self.inner.submit(&self.inner.queue.prepare_wait(done)?)?;
516        }
517        Ok(())
518    }
519    /// Observe the current native submission prefix.
520    pub fn submit(&mut self) -> Result<Submission<'_>> {
521        Ok(Submission { stream: self })
522    }
523    /// Write a checked range after draining this stream.
524    pub fn upload_blocking(&mut self, dst: View<'_>, bytes: &[u8]) -> Result<()> {
525        self.owns(dst.owner)?;
526        checked_span(0, bytes.len(), dst.len())?;
527        self.synchronize()?;
528        dst.owner.native.write(dst.offset, bytes)
529    }
530    /// Write at an allocation-relative byte offset.
531    pub fn upload_blocking_at(&mut self, dst: &Buffer, offset: usize, bytes: &[u8]) -> Result<()> {
532        self.upload_blocking(dst.try_slice(offset, bytes.len())?, bytes)
533    }
534    /// Read a checked range after draining this stream.
535    pub fn read_blocking(&mut self, src: View<'_>, bytes: &mut [u8]) -> Result<()> {
536        self.owns(src.owner)?;
537        checked_span(0, bytes.len(), src.len())?;
538        self.synchronize()?;
539        src.owner.native.read(src.offset, bytes)
540    }
541    /// Read at an allocation-relative byte offset.
542    pub fn read_blocking_at(
543        &mut self,
544        src: &Buffer,
545        offset: usize,
546        bytes: &mut [u8],
547    ) -> Result<()> {
548        self.read_blocking(src.try_slice(offset, bytes.len())?, bytes)
549    }
550    fn prepare_fill(&self, dst: View<'_>, value: u8) -> Result<fabric::PreparedGpu> {
551        self.owns(dst.owner)?;
552        let mut cache = self
553            .inner
554            .transfer_cache
555            .lock()
556            .map_err(|_| Error::DeviceLost("transfer cache poisoned".into()))?;
557        if let Some(entry) = cache.iter().find(|entry| {
558            entry.destination.same_backing(&dst.owner.native)
559                && entry.destination_offset == dst.offset
560                && entry.length == dst.length
561                && entry.source.is_none()
562                && entry.value == value
563        }) {
564            return Ok(entry.command.clone());
565        }
566        let command =
567            self.inner
568                .queue
569                .prepare_fill(&dst.owner.native, dst.offset, dst.length, value)?;
570        // Other streams may use this device-scoped allocation. Only cache when
571        // its owning stream can evict the entry when the public buffer drops.
572        if Arc::ptr_eq(&dst.owner.owner, &self.inner) {
573            if cache.len() == 128 {
574                cache.remove(0);
575            }
576            cache.push(CachedTransfer {
577                destination: dst.owner.native.clone(),
578                destination_offset: dst.offset,
579                length: dst.length,
580                source: None,
581                value,
582                command: command.clone(),
583            });
584        }
585        Ok(command)
586    }
587    fn prepare_copy(&self, dst: View<'_>, src: View<'_>) -> Result<fabric::PreparedGpu> {
588        self.owns(dst.owner)?;
589        self.owns(src.owner)?;
590        if dst.len() != src.len() {
591            return Err(Error::Message("copy requires equal spans".into()));
592        }
593        let mut cache = self
594            .inner
595            .transfer_cache
596            .lock()
597            .map_err(|_| Error::DeviceLost("transfer cache poisoned".into()))?;
598        if let Some(entry) = cache.iter().find(|entry| {
599            entry.destination.same_backing(&dst.owner.native)
600                && entry.destination_offset == dst.offset
601                && entry.length == dst.length
602                && entry.source.as_ref().is_some_and(|(buffer, offset)| {
603                    buffer.same_backing(&src.owner.native) && *offset == src.offset
604                })
605        }) {
606            return Ok(entry.command.clone());
607        }
608        let command = self.inner.queue.prepare_copy(
609            &dst.owner.native,
610            dst.offset,
611            &src.owner.native,
612            src.offset,
613            src.length,
614        )?;
615        if Arc::ptr_eq(&dst.owner.owner, &self.inner) && Arc::ptr_eq(&src.owner.owner, &self.inner)
616        {
617            if cache.len() == 128 {
618                cache.remove(0);
619            }
620            cache.push(CachedTransfer {
621                destination: dst.owner.native.clone(),
622                destination_offset: dst.offset,
623                length: dst.length,
624                source: Some((src.owner.native.clone(), src.offset)),
625                value: 0,
626                command: command.clone(),
627            });
628        }
629        Ok(command)
630    }
631    /// Enqueue a native byte-pattern fill.
632    pub fn fill(&self, dst: View<'_>, value: u8) -> Result<()> {
633        self.budget_uses.borrow_mut().retain(&[dst]);
634        self.inner.submit(&self.prepare_fill(dst, value)?)
635    }
636    /// Enqueue a native copy between equal, non-overlapping ranges.
637    pub fn copy(&self, dst: View<'_>, src: View<'_>) -> Result<()> {
638        self.budget_uses.borrow_mut().retain(&[dst, src]);
639        self.inner.submit(&self.prepare_copy(dst, src)?)
640    }
641    /// Retain an owned staging copy of host bytes and enqueue a GPU transfer.
642    pub fn upload(&mut self, dst: View<'_>, bytes: &[u8]) -> Result<()> {
643        self.owns(dst.owner)?;
644        checked_span(0, bytes.len(), dst.len())?;
645        if bytes.is_empty() {
646            return Ok(());
647        }
648        if self.staging.len() >= 8
649            || self
650                .staging
651                .iter()
652                .map(|b| b.bytes)
653                .sum::<usize>()
654                .saturating_add(bytes.len())
655                > STAGING_LIMIT
656        {
657            self.synchronize()?;
658        }
659        let staging = if let Some(index) = self
660            .staging_pool
661            .iter()
662            .enumerate()
663            .filter(|(_, b)| b.bytes >= bytes.len())
664            .min_by_key(|(_, b)| b.bytes)
665            .map(|(index, _)| index)
666        {
667            self.staging_pool.swap_remove(index)
668        } else {
669            self.allocate(bytes.len())?
670        };
671        staging.native.write(0, bytes)?;
672        self.copy(
673            dst.slice(0, bytes.len())?,
674            staging.try_slice(0, bytes.len())?,
675        )?;
676        self.staging.push(staging);
677        Ok(())
678    }
679    /// Enqueue a checked upload at an allocation-relative offset.
680    pub fn upload_at(&mut self, dst: &Buffer, offset: usize, bytes: &[u8]) -> Result<()> {
681        self.upload(dst.try_slice(offset, bytes.len())?, bytes)
682    }
683    fn reclaim_staging(&mut self) {
684        let mut cached = self.staging_pool.iter().map(|b| b.bytes).sum::<usize>();
685        for buffer in self.staging.drain(..) {
686            if buffer.bytes > STAGING_LIMIT {
687                continue;
688            }
689            while self.staging_pool.len() >= 8
690                || buffer.bytes > STAGING_LIMIT.saturating_sub(cached)
691            {
692                let index = self
693                    .staging_pool
694                    .iter()
695                    .enumerate()
696                    .min_by_key(|(_, b)| b.bytes)
697                    .map(|(i, _)| i)
698                    .unwrap();
699                cached -= self.staging_pool.swap_remove(index).bytes;
700            }
701            cached += buffer.bytes;
702            self.staging_pool.push(buffer);
703        }
704    }
705    /// Load a trusted native code object from disk.
706    ///
707    /// # Safety
708    /// Native code must obey its declared memory and argument contract.
709    pub unsafe fn load(&self, path: &Path, symbol: &str) -> Result<Kernel> {
710        Kernel::from_native(
711            unsafe { self.inner.device.load_bytes(&std::fs::read(path)?, symbol) }?,
712            symbol,
713        )
714    }
715    /// Load a compiler artifact for this exact device target.
716    ///
717    /// # Safety
718    /// Native code must obey its declared memory and argument contract.
719    pub unsafe fn load_artifact(&self, artifact: &crate::loom::Artifact) -> Result<Kernel> {
720        Kernel::from_native(
721            unsafe { self.inner.device.load(artifact) }?,
722            artifact.symbol(),
723        )
724    }
725    unsafe fn prepare_dispatch(
726        &self,
727        kernel: &Kernel,
728        grid: [u32; 3],
729        block: [u32; 3],
730        constants: &Constants,
731        bindings: &[View<'_>],
732    ) -> Result<fabric::PreparedGpu> {
733        if kernel.device_id() != self.device_id() {
734            return Err(Error::Message("kernel belongs to another device".into()));
735        }
736        validate_export_launch(&kernel.info, grid, block)?;
737        if constants.len != kernel.info.constant_byte_length as usize
738            || bindings.len() != kernel.info.binding_count as usize
739        {
740            return Err(Error::Message(
741                "kernel binding or constant byte count mismatch".into(),
742            ));
743        }
744        for view in bindings {
745            self.owns(view.owner)?;
746        }
747        let mut args = Vec::with_capacity(kernel.layout.len());
748        let mut scalar = 0;
749        let mut binding = 0;
750        for &(kind, size) in kernel.layout.iter() {
751            if kind == 1 {
752                args.push(fabric::Argument::Value(
753                    &constants.bytes[scalar..scalar + size],
754                ));
755                scalar += size;
756            } else {
757                let view = bindings[binding];
758                args.push(fabric::Argument::Buffer(&view.owner.native, view.offset));
759                binding += 1;
760            }
761        }
762        unsafe {
763            self.inner
764                .queue
765                .prepare(&kernel.native, grid, block.map(|v| v as u16), &args)
766        }
767    }
768    /// Enqueue a kernel invocation with declaration-order constants and buffers.
769    ///
770    /// # Safety
771    /// Dimensions, constants, and actual kernel accesses must obey each binding's range.
772    pub unsafe fn dispatch(
773        &self,
774        kernel: &Kernel,
775        grid: [u32; 3],
776        block: [u32; 3],
777        constants: &Constants,
778        bindings: &[View<'_>],
779    ) -> Result<()> {
780        self.budget_uses.borrow_mut().retain(bindings);
781        // Owned buffers evict their cached commands before releasing backing or
782        // budget charges. Queued uses keep their reservations until the fence.
783        if bindings
784            .iter()
785            .all(|view| Arc::ptr_eq(&view.owner.owner, &self.inner))
786        {
787            let mut cache = self
788                .inner
789                .dispatch_cache
790                .lock()
791                .map_err(|_| Error::DeviceLost("dispatch cache poisoned".into()))?;
792            if let Some(entry) = cache
793                .iter()
794                .find(|entry| entry.matches(kernel, grid, block, constants, bindings))
795            {
796                return self.inner.submit(&entry.command);
797            }
798            let command =
799                unsafe { self.prepare_dispatch(kernel, grid, block, constants, bindings) }?;
800            self.inner.submit(&command)?;
801            let bytes = command.storage_bytes();
802            if bytes > DISPATCH_CACHE_BYTES {
803                return Ok(());
804            }
805            let mut retained: usize = cache
806                .iter()
807                .map(|entry| entry.command.storage_bytes())
808                .sum();
809            while cache.len() >= DISPATCH_CACHE_ENTRIES || retained + bytes > DISPATCH_CACHE_BYTES {
810                retained -= cache.remove(0).command.storage_bytes();
811            }
812            cache.push(CachedDispatch {
813                kernel: kernel.clone(),
814                grid,
815                block,
816                constants: constants.clone(),
817                bindings: bindings
818                    .iter()
819                    .map(|view| (view.owner.native.clone(), view.offset, view.length))
820                    .collect(),
821                command,
822            });
823            Ok(())
824        } else {
825            self.inner.submit(&unsafe {
826                self.prepare_dispatch(kernel, grid, block, constants, bindings)
827            }?)
828        }
829    }
830    /// Reuse scratch storage under this queue's execution order.
831    pub fn scratch(&mut self, bytes: usize) -> Result<Buffer> {
832        let choice = self
833            .scratch
834            .range(bytes.max(1)..=bytes.saturating_mul(2).max(1))
835            .next()
836            .map(|(&size, _)| size);
837        if let Some(size) = choice {
838            let list = self.scratch.get_mut(&size).unwrap();
839            let buffer = list.pop().unwrap();
840            if list.is_empty() {
841                self.scratch.remove(&size);
842            }
843            self.scratch_bytes -= buffer.bytes;
844            Ok(buffer)
845        } else {
846            self.allocate(bytes)
847        }
848    }
849    /// Return this stream's allocation after recording its last use.
850    pub fn recycle(&mut self, buffer: Buffer) -> Result<()> {
851        if !Arc::ptr_eq(&buffer.owner, &self.inner) {
852            return Err(Error::Message("scratch belongs to another stream".into()));
853        }
854        if buffer.bytes <= self.scratch_limit.saturating_sub(self.scratch_bytes) {
855            self.scratch_bytes += buffer.bytes;
856            self.scratch.entry(buffer.bytes).or_default().push(buffer);
857        }
858        Ok(())
859    }
860    /// Bound retained scratch storage, evicting large entries first.
861    pub fn set_scratch_limit(&mut self, bytes: usize) {
862        self.scratch_limit = bytes;
863        while self.scratch_bytes > bytes {
864            if let Some((_, list)) = self.scratch.pop_last() {
865                self.scratch_bytes -= list.iter().map(|b| b.bytes).sum::<usize>();
866            } else {
867                break;
868            }
869        }
870    }
871    /// Queue an owned readback.
872    pub fn read(&mut self, source: View<'_>) -> Result<Readback> {
873        let buffer = self.allocate(source.len())?;
874        if !source.is_empty() {
875            self.copy(buffer.try_slice(0, source.len())?, source)?;
876        }
877        Ok(Readback {
878            buffer,
879            length: source.len(),
880        })
881    }
882}
883impl Drop for Stream {
884    fn drop(&mut self) {
885        if !self.inner.drain_for_drop() {
886            std::mem::forget(std::mem::take(&mut self.staging));
887            std::mem::forget(std::mem::take(self.budget_uses.get_mut()));
888        }
889    }
890}
891/// Immutable native completion point.
892#[derive(Clone)]
893pub struct Event {
894    device: usize,
895    done: Option<fabric::Completion>,
896}
897impl Event {
898    /// Poll terminal completion.
899    pub fn is_complete(&self) -> Result<bool> {
900        self.done
901            .as_ref()
902            .map_or(Ok(true), fabric::Completion::is_complete)
903    }
904    /// Wait on the host for preceding work.
905    pub fn synchronize(&self) -> Result<()> {
906        self.done.as_ref().map_or(Ok(()), fabric::Completion::wait)
907    }
908}
909/// Borrowed completion observer that also reclaims completed staging.
910#[must_use]
911pub struct Submission<'a> {
912    stream: &'a mut Stream,
913}
914impl Submission<'_> {
915    /// Poll submitted work and reclaim its completed staging.
916    pub fn is_complete(&mut self) -> Result<bool> {
917        let done = self
918            .stream
919            .inner
920            .last
921            .lock()
922            .map_err(|_| Error::DeviceLost("stream timeline poisoned".into()))?
923            .as_ref()
924            .map_or(Ok(true), fabric::Completion::is_complete)?;
925        if done {
926            self.stream.reclaim_staging();
927            self.stream.budget_uses.get_mut().clear();
928        }
929        Ok(done)
930    }
931    /// Wait and reclaim staging.
932    pub fn wait(self) -> Result<()> {
933        self.stream.synchronize()
934    }
935}
936/// Owned destination retained through native readback completion.
937#[must_use]
938pub struct Readback {
939    buffer: Buffer,
940    length: usize,
941}
942impl Readback {
943    /// Wait on the originating stream and return initialized bytes.
944    pub fn wait(self, stream: &mut Stream) -> Result<Vec<u8>> {
945        if !Arc::ptr_eq(&self.buffer.owner, &stream.inner) {
946            return Err(Error::Message("readback belongs to another stream".into()));
947        }
948        stream.synchronize()?;
949        let mut bytes = vec![0; self.length];
950        self.buffer.native.read(0, &mut bytes)?;
951        Ok(bytes)
952    }
953}
954fn owns(inner: &Arc<Inner>, buffer: &Buffer) -> Result<()> {
955    if inner.device.id() == buffer.device_id() {
956        Ok(())
957    } else {
958        Err(Error::Message("buffer belongs to another device".into()))
959    }
960}
961pub(crate) fn checked_span(offset: usize, length: usize, capacity: usize) -> Result<()> {
962    if offset.checked_add(length).is_none_or(|end| end > capacity) {
963        Err(Error::Message(format!(
964            "span {offset}+{length} exceeds {capacity} bytes"
965        )))
966    } else {
967        Ok(())
968    }
969}
970pub(crate) fn validate_launch(grid: [u32; 3], block: [u32; 3]) -> Result<()> {
971    if grid.contains(&0)
972        || block.contains(&0)
973        || block
974            .iter()
975            .try_fold(1u32, |n, b| n.checked_mul(*b))
976            .is_none_or(|n| n > 1024)
977    {
978        Err(Error::Message("invalid GPU launch dimensions".into()))
979    } else {
980        Ok(())
981    }
982}
983pub(crate) fn validate_export_launch(
984    info: &ExportInfo,
985    grid: [u32; 3],
986    block: [u32; 3],
987) -> Result<()> {
988    validate_launch(grid, block)?;
989    if info
990        .workgroup_size
991        .iter()
992        .zip(block)
993        .any(|(&expected, actual)| expected != 0 && expected != actual)
994    {
995        Err(Error::Message(format!(
996            "workgroup size {block:?} disagrees with compiled size {:?}",
997            info.workgroup_size
998        )))
999    } else {
1000        Ok(())
1001    }
1002}
1003const STAGING_LIMIT: usize = 64 * 1024 * 1024;
1004
1005/// Explicit scalar widths, packed in declaration order for HRX binding dispatch.
1006/// Unlike direct kernargs this format has no implicit alignment or pointer slots.
1007/// Export metadata gives only the total byte count. Mixed scalar widths must
1008/// come from the kernel's argument contract.
1009#[derive(Clone, Debug)]
1010pub struct Constants {
1011    bytes: [u8; 256],
1012    len: usize,
1013}
1014impl Default for Constants {
1015    fn default() -> Self {
1016        Self {
1017            bytes: [0; 256],
1018            len: 0,
1019        }
1020    }
1021}
1022impl Constants {
1023    /// Pack an entry point whose scalar arguments are all unsigned Loom indices.
1024    ///
1025    /// Loom lowers a homogeneous index list to 32- or 64-bit slots depending on
1026    /// its range analysis. This adapter is only for that model contract; mixed
1027    /// scalar types must use `push` with their explicit native widths.
1028    ///
1029    /// Neither [`ExportInfo`] nor the compiler's resource report carries per-slot
1030    /// scalar types — only `constant_byte_length` and `parameter_count` — so no
1031    /// API here can build a mixed-width constant block by construction. A width
1032    /// inferred by dividing the byte count is wrong for a mixed `u32`/`u64`
1033    /// signature, and wrong quietly. Take each width from the source that
1034    /// declared it.
1035    pub fn indices(kernel: &Kernel, values: &[u32]) -> Result<Self> {
1036        let size = kernel.info.constant_byte_length as usize;
1037        let mut constants = Self::new();
1038        if values.is_empty() {
1039            if size != 0 {
1040                return Err(Error::Message("missing index constants".into()));
1041            }
1042            return Ok(constants);
1043        }
1044        if size == values.len() * 4 {
1045            for &value in values {
1046                constants.push(value)?;
1047            }
1048        } else if size == values.len() * 8 {
1049            for &value in values {
1050                constants.push(u64::from(value))?;
1051            }
1052        } else {
1053            return Err(Error::Message(
1054                "export is not a homogeneous 32/64-bit index list".into(),
1055            ));
1056        }
1057        Ok(constants)
1058    }
1059    /// Create an empty constant block.
1060    pub fn new() -> Self {
1061        Self::default()
1062    }
1063    #[must_use = "constant packing can fail when capacity is exceeded"]
1064    /// Append a scalar at its explicit width; fails if the 256-byte capacity is exceeded.
1065    pub fn push<T: Scalar>(&mut self, value: T) -> Result<()> {
1066        let bytes = value.bytes();
1067        checked_span(self.len, bytes.as_ref().len(), self.bytes.len())?;
1068        self.bytes[self.len..self.len + bytes.as_ref().len()].copy_from_slice(bytes.as_ref());
1069        self.len += bytes.as_ref().len();
1070        Ok(())
1071    }
1072    /// Construct an explicitly contracted scalar block from packed bytes.
1073    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
1074        let mut result = Self::new();
1075        checked_span(0, bytes.len(), result.bytes.len())?;
1076        result.bytes[..bytes.len()].copy_from_slice(bytes);
1077        result.len = bytes.len();
1078        Ok(result)
1079    }
1080    /// The packed constant bytes in declaration order.
1081    pub fn as_bytes(&self) -> &[u8] {
1082        &self.bytes[..self.len]
1083    }
1084}
1085mod sealed {
1086    pub trait Sealed {}
1087}
1088/// A supported fixed-width scalar for little-endian binding constants.
1089pub trait Scalar: sealed::Sealed {
1090    /// The scalar's fixed-size byte representation.
1091    type Bytes: AsRef<[u8]>;
1092    /// Encode this value in little-endian order.
1093    fn bytes(self) -> Self::Bytes;
1094}
1095macro_rules! scalars {
1096    ($($ty:ty => $n:literal),*) => {
1097        $(
1098            impl sealed::Sealed for $ty {}
1099            impl Scalar for $ty {
1100                type Bytes = [u8; $n];
1101                fn bytes(self) -> Self::Bytes { self.to_le_bytes() }
1102            }
1103        )*
1104    };
1105}
1106scalars!(u32 => 4, i32 => 4, f32 => 4, u64 => 8, i64 => 8, f64 => 8);
1107
1108/// A dependency node branded with its recording's identity.
1109#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
1110pub struct Node {
1111    graph: u64,
1112    index: u32,
1113}
1114static NEXT_GRAPH_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1115/// Owned native commands being recorded under checked dependency edges.
1116pub struct Graph<'a> {
1117    stream: &'a Stream,
1118    id: u64,
1119    nodes: Vec<(Recorded<'a>, Option<u32>)>,
1120    budget_uses: BudgetUses,
1121}
1122enum Recorded<'a> {
1123    Fill(View<'a>, u8),
1124    Copy(View<'a>, View<'a>),
1125    Prepared(fabric::PreparedGpu),
1126    Join,
1127}
1128/// Reusable prepared GPU graph and its last immutable completion point.
1129pub struct GraphExec {
1130    inner: Arc<Inner>,
1131    commands: Option<fabric::PreparedGpu>,
1132    profile: Option<fabric::profile::ProfileCapture>,
1133    last: Option<fabric::Completion>,
1134    failed: bool,
1135    budget_uses: BudgetUses,
1136}
1137impl Stream {
1138    /// Begin a dependency-checked recording borrowing this stream.
1139    pub fn graph(&self) -> Result<Graph<'_>> {
1140        Ok(Graph {
1141            stream: self,
1142            id: NEXT_GRAPH_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
1143            nodes: Vec::new(),
1144            budget_uses: BudgetUses::default(),
1145        })
1146    }
1147    /// Replay an instrumented graph, wait for retirement, and read GPU times.
1148    /// Ordinary graphs return an error; construct with `Graph::finish_profiled`.
1149    pub fn launch_profiled(&mut self, graph: &mut GraphExec) -> Result<fabric::DeviceProfile> {
1150        if graph.profile.is_none() {
1151            return Err(Error::Message(
1152                "graph was not prepared for profiling".into(),
1153            ));
1154        }
1155        self.launch(graph)?;
1156        self.synchronize()?;
1157        graph.profile.as_mut().unwrap().read()
1158    }
1159    /// Replay fixed native commands on the originating stream.
1160    pub fn launch(&mut self, graph: &mut GraphExec) -> Result<()> {
1161        if !Arc::ptr_eq(&graph.inner, &self.inner) {
1162            return Err(Error::Message("graph belongs to another stream".into()));
1163        }
1164        if graph.failed {
1165            return Err(Error::DeviceLost(
1166                "graph has an unretired failed launch".into(),
1167            ));
1168        }
1169        graph.failed = true;
1170        if let Some(command) = &graph.commands {
1171            self.inner.submit(command)?;
1172        }
1173        graph.last = self
1174            .inner
1175            .last
1176            .lock()
1177            .map_err(|_| Error::DeviceLost("stream timeline poisoned".into()))?
1178            .clone();
1179        graph.failed = false;
1180        Ok(())
1181    }
1182}
1183impl<'a> Graph<'a> {
1184    fn validate_dependencies(&self, after: &[Node]) -> Result<()> {
1185        for node in after {
1186            if node.graph != self.id || node.index as usize >= self.nodes.len() {
1187                return Err(Error::Message(
1188                    "dependency node belongs to another graph".into(),
1189                ));
1190            }
1191        }
1192        if self.nodes.len() >= u32::MAX as usize {
1193            return Err(Error::Message("graph node capacity exceeded".into()));
1194        }
1195        Ok(())
1196    }
1197    fn record(&mut self, command: Recorded<'a>, after: &[Node]) -> Node {
1198        let index = self.nodes.len() as u32;
1199        self.nodes
1200            .push((command, after.iter().map(|node| node.index).max()));
1201        Node {
1202            graph: self.id,
1203            index,
1204        }
1205    }
1206    /// Record a byte fill after checked predecessor nodes.
1207    pub fn fill(&mut self, after: &[Node], dst: View<'a>, pattern: u8) -> Result<Node> {
1208        self.validate_dependencies(after)?;
1209        self.stream.owns(dst.owner)?;
1210        if dst.is_empty() {
1211            return Err(Error::Message("empty graph fill".into()));
1212        }
1213        self.budget_uses.retain(&[dst]);
1214        Ok(self.record(Recorded::Fill(dst, pattern), after))
1215    }
1216    /// Record a non-overlapping byte copy after checked predecessor nodes.
1217    pub fn copy(&mut self, after: &[Node], dst: View<'a>, src: View<'a>) -> Result<Node> {
1218        self.validate_dependencies(after)?;
1219        self.stream.owns(dst.owner)?;
1220        self.stream.owns(src.owner)?;
1221        if dst.len() != src.len() || dst.is_empty() {
1222            return Err(Error::Message(
1223                "graph copy requires equal nonempty spans".into(),
1224            ));
1225        }
1226        self.budget_uses.retain(&[dst, src]);
1227        Ok(self.record(Recorded::Copy(dst, src), after))
1228    }
1229    /// Record a fixed native kernel invocation.
1230    ///
1231    /// # Safety
1232    /// Kernel accesses must obey bindings and constants; dependency edges must
1233    /// order conflicting accesses for every replay.
1234    pub unsafe fn dispatch(
1235        &mut self,
1236        after: &[Node],
1237        kernel: &'a Kernel,
1238        grid: [u32; 3],
1239        block: [u32; 3],
1240        constants: &Constants,
1241        bindings: &[View<'a>],
1242    ) -> Result<Node> {
1243        self.validate_dependencies(after)?;
1244        let command = unsafe {
1245            self.stream
1246                .prepare_dispatch(kernel, grid, block, constants, bindings)
1247        }?;
1248        self.budget_uses.retain(bindings);
1249        Ok(self.record(Recorded::Prepared(command), after))
1250    }
1251    /// Record a dependency join with no native payload.
1252    pub fn join(&mut self, after: &[Node]) -> Result<Node> {
1253        self.validate_dependencies(after)?;
1254        Ok(self.record(Recorded::Join, after))
1255    }
1256    /// Transfer recorded commands and allocation charges into owned replay state.
1257    pub fn finish(self) -> Result<GraphExec> {
1258        self.finish_inner(None)
1259    }
1260    /// Prepare a diagnostic graph with device-clock markers around each payload.
1261    /// Supply one label per non-join node. Instrumentation adds completion
1262    /// barriers and must not be used as end-to-end performance evidence.
1263    pub fn finish_profiled(self, labels: &[String]) -> Result<GraphExec> {
1264        self.finish_inner(Some(labels))
1265    }
1266    fn finish_inner(self, labels: Option<&[String]>) -> Result<GraphExec> {
1267        let mut ordered = Vec::new();
1268        let mut completed = 0usize;
1269        for (index, (node, dependency)) in self.nodes.into_iter().enumerate() {
1270            let command = match node {
1271                Recorded::Fill(dst, value) => self.stream.prepare_fill(dst, value)?,
1272                Recorded::Copy(dst, src) => self.stream.prepare_copy(dst, src)?,
1273                Recorded::Prepared(command) => command,
1274                Recorded::Join => continue,
1275            };
1276            let barrier = dependency.is_some_and(|node| node as usize >= completed);
1277            if barrier {
1278                completed = index;
1279            }
1280            ordered.push((command, barrier));
1281        }
1282        let (commands, profile) = if let Some(labels) = labels {
1283            // The same validated edges and retained bindings as normal replay.
1284            let profiled = unsafe {
1285                self.stream
1286                    .inner
1287                    .queue
1288                    .prepare_profiled_batch(&ordered, labels)
1289            }?;
1290            (Some(profiled.commands), Some(profiled.capture))
1291        } else if ordered.is_empty() {
1292            (None, None)
1293        } else {
1294            // Graph edges order conflicting accesses; all binding backing is retained.
1295            (
1296                Some(unsafe { self.stream.inner.queue.prepare_batch(&ordered) }?),
1297                None,
1298            )
1299        };
1300        Ok(GraphExec {
1301            inner: self.stream.inner.clone(),
1302            commands,
1303            profile,
1304            last: None,
1305            failed: false,
1306            budget_uses: self.budget_uses,
1307        })
1308    }
1309}
1310impl Drop for GraphExec {
1311    fn drop(&mut self) {
1312        if self.failed
1313            || self.last.as_ref().is_some_and(|done| {
1314                !done
1315                    .wait_timeout(std::time::Duration::from_secs(10))
1316                    .unwrap_or(false)
1317            })
1318        {
1319            std::mem::forget(std::mem::take(&mut self.budget_uses));
1320        }
1321    }
1322}
1323impl std::fmt::Debug for Device {
1324    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1325        f.debug_struct("Device")
1326            .field("target", self.target())
1327            .finish_non_exhaustive()
1328    }
1329}
1330impl std::fmt::Debug for Buffer {
1331    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1332        f.debug_struct("Buffer")
1333            .field("bytes", &self.bytes)
1334            .finish_non_exhaustive()
1335    }
1336}
1337impl std::fmt::Debug for Kernel {
1338    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1339        f.debug_struct("Kernel")
1340            .field("symbol", &self.symbol)
1341            .finish_non_exhaustive()
1342    }
1343}
1344impl std::fmt::Debug for Stream {
1345    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1346        f.debug_struct("Stream")
1347            .field("scratch_bytes", &self.scratch_bytes)
1348            .finish_non_exhaustive()
1349    }
1350}
1351impl std::fmt::Debug for Graph<'_> {
1352    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1353        f.debug_struct("Graph")
1354            .field("nodes", &self.nodes.len())
1355            .finish_non_exhaustive()
1356    }
1357}
1358impl std::fmt::Debug for GraphExec {
1359    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1360        f.debug_struct("GraphExec").finish_non_exhaustive()
1361    }
1362}
1363impl std::fmt::Debug for Readback {
1364    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1365        f.debug_struct("Readback").finish_non_exhaustive()
1366    }
1367}
1368impl std::fmt::Debug for Submission<'_> {
1369    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1370        f.debug_struct("Submission").finish_non_exhaustive()
1371    }
1372}
1373
1374#[cfg(test)]
1375mod tests {
1376    #[test]
1377    fn compiled_workgroup_dimensions_are_enforced() {
1378        let info = super::ExportInfo {
1379            workgroup_size: [64, 2, 1],
1380            ..Default::default()
1381        };
1382        assert!(super::validate_export_launch(&info, [1; 3], [64, 2, 1]).is_ok());
1383        assert!(super::validate_export_launch(&info, [1; 3], [128, 1, 1]).is_err());
1384        assert!(super::validate_export_launch(&info, [0, 1, 1], [64, 2, 1]).is_err());
1385        let dynamic = super::ExportInfo {
1386            workgroup_size: [0, 2, 1],
1387            ..Default::default()
1388        };
1389        assert!(super::validate_export_launch(&dynamic, [1; 3], [32, 2, 1]).is_ok());
1390        assert!(super::validate_export_launch(&dynamic, [1; 3], [32, 1, 1]).is_err());
1391    }
1392    /// The bounds arithmetic, without a device: `offset + length` must not wrap into a pass.
1393    #[test]
1394    fn a_slice_past_the_allocation_is_refused_even_when_the_sum_wraps() {
1395        let checked = |offset, length, bytes| super::checked_span(offset, length, bytes).is_ok();
1396        assert!(checked(0, 8, 8));
1397        assert!(checked(4, 4, 8));
1398        assert!(!checked(4, 5, 8));
1399        // the wrapping case: 8 bytes must not accept an offset near the top of the address space
1400        assert!(!checked(usize::MAX, 2, 8));
1401        assert!(!checked(usize::MAX - 1, 4, 8));
1402    }
1403}
1404
1405#[cfg(test)]
1406mod staging_tests {
1407    use super::*;
1408
1409    #[test]
1410    #[ignore = "requires gfx1151"]
1411    fn budgeted_commands_replay_and_evict_without_releasing_pending_charges() -> Result<()> {
1412        let residency = crate::residency::ResidencyManager::new(4096)?;
1413        let device = Device::open(0)?;
1414        let mut stream = device.stream()?.with_memory_budget(residency.budget());
1415        let mut foreign = device.stream()?;
1416        let compiler = crate::loom::Compiler::for_stream(None, &stream)?;
1417        let mut spec = crate::loom::Specialization::new("add_f32");
1418        spec.set_config("add.grid", "1");
1419        let artifact = compiler
1420            .module(include_str!("../tests/kernels/add_f32.loom"))
1421            .compile(&spec)?;
1422        // SAFETY: the trusted fixture adds a scalar to exactly 64 f32 elements.
1423        let kernel = unsafe { stream.load_artifact(&artifact)? };
1424        let buffer = stream.allocate(256)?;
1425        stream.fill(buffer.binding(), 0)?;
1426        let constants = |delta: f32| -> Result<Constants> {
1427            let mut constants = Constants::new();
1428            match kernel.layout[0].1 {
1429                4 => constants.push(64u32)?,
1430                8 => constants.push(64u64)?,
1431                _ => unreachable!("index width"),
1432            }
1433            constants.push(delta)?;
1434            Ok(constants)
1435        };
1436        // More entries than the old cache could hold, with distinct constants.
1437        for _ in 0..2 {
1438            for delta in 1..=100 {
1439                unsafe {
1440                    stream.dispatch(
1441                        &kernel,
1442                        [1; 3],
1443                        [256, 1, 1],
1444                        &constants(delta as f32)?,
1445                        &[buffer.binding()],
1446                    )?;
1447                }
1448            }
1449            stream.synchronize()?;
1450            assert_eq!(stream.inner.dispatch_cache.lock().unwrap().len(), 100);
1451        }
1452        let mut actual = [0u8; 256];
1453        stream.read_blocking(buffer.binding(), &mut actual)?;
1454        assert!(
1455            actual
1456                .as_chunks::<4>()
1457                .0
1458                .iter()
1459                .all(|bytes| f32::from_le_bytes(*bytes) == 10100.)
1460        );
1461        assert_eq!(residency.statistics().reserved_bytes, 256);
1462        for delta in 101..=1124 {
1463            unsafe {
1464                stream.dispatch(
1465                    &kernel,
1466                    [1; 3],
1467                    [256, 1, 1],
1468                    &constants(delta as f32)?,
1469                    &[buffer.binding()],
1470                )?;
1471            }
1472        }
1473        stream.synchronize()?;
1474        {
1475            let cache = stream.inner.dispatch_cache.lock().unwrap();
1476            assert_eq!(cache.len(), DISPATCH_CACHE_ENTRIES);
1477            assert!(
1478                cache
1479                    .iter()
1480                    .map(|entry| entry.command.storage_bytes())
1481                    .sum::<usize>()
1482                    <= DISPATCH_CACHE_BYTES
1483            );
1484        }
1485        // A foreign stream cannot retain a command whose owner cannot evict it.
1486        unsafe {
1487            foreign.dispatch(
1488                &kernel,
1489                [1; 3],
1490                [256, 1, 1],
1491                &constants(1.)?,
1492                &[buffer.binding()],
1493            )?;
1494        }
1495        foreign.synchronize()?;
1496        assert!(foreign.inner.dispatch_cache.lock().unwrap().is_empty());
1497        // Removing public backing purges both caches; its queued reservation
1498        // survives until the stream explicitly observes completion.
1499        unsafe {
1500            stream.dispatch(
1501                &kernel,
1502                [1; 3],
1503                [256, 1, 1],
1504                &constants(1.)?,
1505                &[buffer.binding()],
1506            )?;
1507        }
1508        drop(buffer);
1509        assert!(stream.inner.dispatch_cache.lock().unwrap().is_empty());
1510        assert!(stream.inner.transfer_cache.lock().unwrap().is_empty());
1511        assert_eq!(residency.statistics().reserved_bytes, 256);
1512        stream.synchronize()?;
1513        assert_eq!(residency.statistics().reserved_bytes, 0);
1514        Ok(())
1515    }
1516
1517    #[test]
1518    #[ignore = "requires gfx1151"]
1519    fn a_scratch_pool_only_accepts_the_stream_that_allocated_the_buffer() -> Result<()> {
1520        let device = Device::open(0)?;
1521        let mut owner = device.stream()?;
1522        let mut other = device.stream()?;
1523        let private = other.scratch(4096)?;
1524        let private_raw = private.allocation_address()?;
1525        other.recycle(private)?;
1526
1527        // Buffers are device-scoped, so another stream may read and write this one.
1528        let buffer = owner.allocate(4096)?;
1529        owner.fill(buffer.binding(), 0x5a)?;
1530        let ready = owner.record_event()?;
1531        other.wait_event(&ready)?;
1532        let mut seen = [0u8; 16];
1533        other.read_blocking(buffer.binding().slice(0, 16)?, &mut seen)?;
1534        assert_eq!(seen, [0x5a; 16], "a foreign stream can use the allocation");
1535
1536        // A scratch pool is still one stream's private free list, not shared state.
1537        assert!(
1538            other
1539                .recycle(buffer)
1540                .unwrap_err()
1541                .to_string()
1542                .contains("another stream")
1543        );
1544        assert!(owner.scratch.is_empty());
1545        assert_eq!(other.scratch_bytes, 4096);
1546        assert_eq!(other.scratch(4096)?.allocation_address()?, private_raw);
1547        Ok(())
1548    }
1549
1550    #[test]
1551    #[ignore = "requires gfx1151"]
1552    fn completed_staging_is_reclaimed_and_reused_on_every_completion_path() -> Result<()> {
1553        let mut stream = Stream::open()?;
1554        let buffer = stream.allocate(1024)?;
1555        stream.upload(buffer.binding(), &[7; 1024])?;
1556        let original = stream.staging[0].allocation_address()?;
1557        let mut output = [0; 1024];
1558        stream.read_blocking(buffer.binding(), &mut output)?;
1559        assert_eq!(output, [7; 1024]);
1560        assert!(stream.staging.is_empty());
1561        stream.upload(buffer.binding(), &[9; 1024])?;
1562        assert_eq!(stream.staging[0].allocation_address()?, original);
1563        stream.upload_blocking(buffer.binding(), &[11; 1024])?;
1564        assert!(stream.staging.is_empty());
1565        stream.upload(buffer.binding(), &[12; 512])?;
1566        assert_eq!(stream.staging[0].allocation_address()?, original);
1567        stream.synchronize()?;
1568        stream.upload(buffer.binding(), &[13; 1024])?;
1569        stream.inner.wait()?; // Establish completion without the cleanup being tested.
1570        assert!(stream.submit()?.is_complete()?);
1571        assert!(stream.staging.is_empty());
1572        for _ in 0..32 {
1573            stream.upload(buffer.binding(), &[1; 1024])?;
1574            assert!(stream.staging.len() <= 8);
1575        }
1576        stream.synchronize()?;
1577        assert!(stream.staging_pool.len() <= 8);
1578        assert!(stream.staging_pool.iter().map(|b| b.bytes).sum::<usize>() <= STAGING_LIMIT);
1579        Ok(())
1580    }
1581
1582    #[test]
1583    #[ignore = "requires gfx1151"]
1584    fn uploads_batch_until_pressure_and_reuse_the_smallest_capacity() -> Result<()> {
1585        let mut stream = Stream::open()?;
1586        let buffer = stream.allocate(4096)?;
1587        for size in [1024, 2048, 4096] {
1588            stream.upload(buffer.binding(), &vec![7; size])?;
1589        }
1590        assert_eq!(stream.staging.len(), 3);
1591        let medium = stream.staging[1].allocation_address()?;
1592        stream.synchronize()?;
1593        stream.upload(buffer.binding(), &[9; 1500])?;
1594        assert_eq!(stream.staging[0].allocation_address()?, medium);
1595        for _ in 0..7 {
1596            stream.upload(buffer.binding(), &[11; 1024])?;
1597        }
1598        assert_eq!(stream.staging.len(), 8);
1599        stream.upload(buffer.binding(), &[13; 1024])?;
1600        assert_eq!(stream.staging.len(), 1);
1601        stream.synchronize()?;
1602        Ok(())
1603    }
1604
1605    #[test]
1606    #[ignore = "requires gfx1151"]
1607    fn larger_uploads_replace_small_completed_staging() -> Result<()> {
1608        let mut stream = Stream::open()?;
1609        let buffer = stream.allocate(16 << 20)?;
1610        for _ in 0..8 {
1611            stream.upload(buffer.binding(), &[7; 1024])?;
1612        }
1613        stream.synchronize()?;
1614        assert_eq!(stream.staging_pool.len(), 8);
1615
1616        let source = vec![0x35; 1 << 20];
1617        stream.upload(buffer.binding(), &source)?;
1618        let large = stream.staging[0].allocation_address()?;
1619        stream.synchronize()?;
1620        let replacement = vec![0xa9; source.len()];
1621        stream.upload(buffer.binding(), &replacement)?;
1622        assert_eq!(stream.staging[0].allocation_address()?, large);
1623        let mut output = vec![0; source.len()];
1624        stream.read_blocking(buffer.try_slice(0, output.len())?, &mut output)?;
1625        assert_eq!(output, replacement);
1626        assert!(stream.staging_pool.len() <= 8);
1627        assert!(stream.staging_pool.iter().map(|b| b.bytes).sum::<usize>() <= STAGING_LIMIT);
1628
1629        // Larger packets also exercise eviction at the byte ceiling, before
1630        // the cache reaches its eight-entry ceiling.
1631        let source = vec![0x63; 16 << 20];
1632        for _ in 0..9 {
1633            stream.upload(buffer.binding(), &source)?;
1634        }
1635        let mut output = vec![0; source.len()];
1636        stream.read_blocking(buffer.binding(), &mut output)?;
1637        assert_eq!(output, source);
1638        assert!(stream.staging_pool.len() <= 8);
1639        assert!(stream.staging_pool.iter().map(|b| b.bytes).sum::<usize>() <= STAGING_LIMIT);
1640        Ok(())
1641    }
1642}
1643
1644#[cfg(test)]
1645mod dag_probe {
1646    use super::*;
1647
1648    /// Declared edges cost scheduling time, so the graph API must be able to omit
1649    /// the ones a workload does not need. Latency-bound by construction: 64 tiny
1650    /// fills, so bandwidth cannot confound the comparison.
1651    #[test]
1652    #[ignore = "requires gfx1151"]
1653    fn independent_nodes_beat_a_serial_chain() -> Result<()> {
1654        const NODES: usize = 64;
1655        let mut stream = Stream::open()?;
1656        let buffers: Vec<Buffer> = (0..NODES)
1657            .map(|_| stream.allocate(4096))
1658            .collect::<Result<_>>()?;
1659        let mut timings = Vec::new();
1660        for chained in [true, false] {
1661            let mut graph = stream.graph()?;
1662            let mut previous: Option<Node> = None;
1663            for buffer in &buffers {
1664                let after: &[Node] = match (chained, &previous) {
1665                    (true, Some(node)) => std::slice::from_ref(node),
1666                    _ => &[],
1667                };
1668                previous = Some(graph.fill(after, buffer.binding(), 0x3c)?);
1669            }
1670            let mut exec = graph.finish()?;
1671            for _ in 0..3 {
1672                stream.launch(&mut exec)?;
1673                stream.synchronize()?;
1674            }
1675            let mut samples = Vec::new();
1676            for _ in 0..9 {
1677                let start = std::time::Instant::now();
1678                stream.launch(&mut exec)?;
1679                stream.synchronize()?;
1680                samples.push(start.elapsed().as_secs_f64() * 1e6);
1681            }
1682            samples.sort_by(f64::total_cmp);
1683            println!(
1684                "{:<12} {NODES} nodes: {:8.1} us total, {:6.2} us/node",
1685                if chained { "serial" } else { "independent" },
1686                samples[4],
1687                samples[4] / NODES as f64
1688            );
1689            timings.push(samples[4]);
1690            let mut seen = [0u8; 4096];
1691            stream.read_blocking(buffers[NODES - 1].binding(), &mut seen)?;
1692            assert_eq!(seen, [0x3c; 4096], "every node ran");
1693        }
1694        // A margin rather than a strict inequality: the property under test is
1695        // that omitting edges still reaches the runtime, not that a shared CI
1696        // runner is quiet. The measured gap is ~1.7x, so 1.2x is a wide floor.
1697        assert!(
1698            timings[1] * 1.2 < timings[0],
1699            "omitting edges should be materially cheaper: {timings:?}"
1700        );
1701        Ok(())
1702    }
1703}