1use crate::{Error, Result, Target, fabric};
3use std::{
4 cell::RefCell,
5 collections::BTreeMap,
6 path::Path,
7 sync::{Arc, Mutex},
8};
9
10#[derive(Clone, Debug, Default)]
12pub struct ExportInfo {
13 pub constant_byte_length: u32,
15 pub binding_count: u32,
17 pub parameter_count: u32,
19 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#[derive(Clone)]
68pub struct Device {
69 native: fabric::Device,
70}
71impl Device {
72 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 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 pub fn target(&self) -> &Target {
93 self.native.target()
94 }
95 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 pub fn native(&self) -> &fabric::Device {
117 &self.native
118 }
119}
120pub 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 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#[derive(Clone, Copy, Debug)]
164pub struct View<'a> {
165 owner: &'a Buffer,
166 offset: usize,
167 length: usize,
168}
169impl<'a> View<'a> {
170 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 pub fn owner(self) -> &'a Buffer {
181 self.owner
182 }
183 pub fn offset(self) -> usize {
185 self.offset
186 }
187 pub fn len(self) -> usize {
189 self.length
190 }
191 pub fn is_empty(self) -> bool {
193 self.length == 0
194 }
195}
196impl Buffer {
197 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 pub fn bytes(&self) -> usize {
225 self.bytes
226 }
227 pub fn binding(&self) -> View<'_> {
229 View {
230 owner: self,
231 offset: 0,
232 length: self.bytes,
233 }
234 }
235 pub fn try_slice(&self, offset: usize, length: usize) -> Result<View<'_>> {
237 self.binding().slice(offset, length)
238 }
239 pub fn slice(&self, offset: usize, length: usize) -> View<'_> {
241 self.try_slice(offset, length)
242 .expect("slice exceeds allocation")
243 }
244 pub fn device_ptr(&self) -> Result<*mut std::ffi::c_void> {
249 Ok(self.native.host_pointer().cast())
250 }
251 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#[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 pub fn info(&self) -> &ExportInfo {
278 &self.info
279 }
280 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}
375pub 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 pub fn open() -> Result<Self> {
389 Device::open(0)?.stream()
390 }
391 pub fn device_id(&self) -> usize {
393 self.inner.device.id()
394 }
395 pub fn id(&self) -> usize {
397 Arc::as_ptr(&self.inner) as usize
398 }
399 pub fn target(&self) -> &Target {
401 self.inner.device.target()
402 }
403 pub fn with_memory_budget(mut self, budget: crate::residency::MemoryBudget) -> Self {
405 self.budget = Some(budget);
406 self
407 }
408 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 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 pub fn allocate_zeroed(&self, bytes: usize) -> Result<Buffer> {
456 let buffer = self.allocate(bytes)?;
457 buffer.native.zero()?;
458 Ok(buffer)
459 }
460 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 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 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 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 pub fn submit(&mut self) -> Result<Submission<'_>> {
521 Ok(Submission { stream: self })
522 }
523 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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#[derive(Clone)]
893pub struct Event {
894 device: usize,
895 done: Option<fabric::Completion>,
896}
897impl Event {
898 pub fn is_complete(&self) -> Result<bool> {
900 self.done
901 .as_ref()
902 .map_or(Ok(true), fabric::Completion::is_complete)
903 }
904 pub fn synchronize(&self) -> Result<()> {
906 self.done.as_ref().map_or(Ok(()), fabric::Completion::wait)
907 }
908}
909#[must_use]
911pub struct Submission<'a> {
912 stream: &'a mut Stream,
913}
914impl Submission<'_> {
915 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 pub fn wait(self) -> Result<()> {
933 self.stream.synchronize()
934 }
935}
936#[must_use]
938pub struct Readback {
939 buffer: Buffer,
940 length: usize,
941}
942impl Readback {
943 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#[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 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 pub fn new() -> Self {
1061 Self::default()
1062 }
1063 #[must_use = "constant packing can fail when capacity is exceeded"]
1064 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 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 pub fn as_bytes(&self) -> &[u8] {
1082 &self.bytes[..self.len]
1083 }
1084}
1085mod sealed {
1086 pub trait Sealed {}
1087}
1088pub trait Scalar: sealed::Sealed {
1090 type Bytes: AsRef<[u8]>;
1092 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#[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);
1115pub 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}
1128pub 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 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 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 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 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 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 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 pub fn join(&mut self, after: &[Node]) -> Result<Node> {
1253 self.validate_dependencies(after)?;
1254 Ok(self.record(Recorded::Join, after))
1255 }
1256 pub fn finish(self) -> Result<GraphExec> {
1258 self.finish_inner(None)
1259 }
1260 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 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 (
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 #[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 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 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 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 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 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 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 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()?; 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 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 #[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 assert!(
1698 timings[1] * 1.2 < timings[0],
1699 "omitting edges should be materially cheaper: {timings:?}"
1700 );
1701 Ok(())
1702 }
1703}