1use crate::error::{LevelZeroError, LevelZeroResult};
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub struct GroupCount {
29 pub x: u32,
31 pub y: u32,
33 pub z: u32,
35}
36
37impl GroupCount {
38 #[must_use]
40 pub fn total(&self) -> u64 {
41 u64::from(self.x) * u64::from(self.y) * u64::from(self.z)
42 }
43}
44
45#[derive(Debug, Clone, Copy)]
52pub struct GroupCountPlanner {
53 local_x: u32,
54 local_y: u32,
55 local_z: u32,
56}
57
58impl GroupCountPlanner {
59 pub fn new_1d(local_x: u32) -> LevelZeroResult<Self> {
61 Self::new_3d(local_x, 1, 1)
62 }
63
64 pub fn new_3d(local_x: u32, local_y: u32, local_z: u32) -> LevelZeroResult<Self> {
66 if local_x == 0 || local_y == 0 || local_z == 0 {
67 return Err(LevelZeroError::InvalidArgument(
68 "workgroup size must be non-zero in every dimension".into(),
69 ));
70 }
71 Ok(Self {
72 local_x,
73 local_y,
74 local_z,
75 })
76 }
77
78 #[must_use]
80 pub fn local_size(&self) -> (u32, u32, u32) {
81 (self.local_x, self.local_y, self.local_z)
82 }
83
84 pub fn plan_1d(&self, global_x: u64) -> LevelZeroResult<GroupCount> {
86 self.plan_3d(global_x, 1, 1)
87 }
88
89 pub fn plan_3d(
95 &self,
96 global_x: u64,
97 global_y: u64,
98 global_z: u64,
99 ) -> LevelZeroResult<GroupCount> {
100 let x = ceil_div_u32(global_x, self.local_x)?;
101 let y = ceil_div_u32(global_y, self.local_y)?;
102 let z = ceil_div_u32(global_z, self.local_z)?;
103 Ok(GroupCount { x, y, z })
104 }
105}
106
107fn ceil_div_u32(num: u64, den: u32) -> LevelZeroResult<u32> {
109 if num == 0 {
110 return Ok(0);
111 }
112 let groups = num.div_ceil(u64::from(den));
113 u32::try_from(groups).map_err(|_| {
114 LevelZeroError::InvalidArgument(format!("group count {groups} exceeds u32 dispatch limit"))
115 })
116}
117
118#[derive(Debug, Clone, PartialEq, Eq)]
122pub struct CommandQueueGroupInfo {
123 pub ordinal: u32,
125 pub compute: bool,
127 pub copy: bool,
129 pub num_queues: u32,
131}
132
133#[derive(Debug, Clone, Default)]
139pub struct CommandQueueGroupTable {
140 groups: Vec<CommandQueueGroupInfo>,
141}
142
143impl CommandQueueGroupTable {
144 #[must_use]
146 pub fn new(groups: Vec<CommandQueueGroupInfo>) -> Self {
147 Self { groups }
148 }
149
150 #[must_use]
152 pub fn groups(&self) -> &[CommandQueueGroupInfo] {
153 &self.groups
154 }
155
156 pub fn select_compute_ordinal(&self) -> LevelZeroResult<u32> {
158 self.groups
159 .iter()
160 .filter(|g| g.compute)
161 .min_by_key(|g| g.ordinal)
162 .map(|g| g.ordinal)
163 .ok_or_else(|| {
164 LevelZeroError::Unsupported("device exposes no compute queue group".into())
165 })
166 }
167
168 pub fn select_copy_ordinal(&self) -> LevelZeroResult<u32> {
174 if let Some(g) = self
176 .groups
177 .iter()
178 .filter(|g| g.copy && !g.compute)
179 .min_by_key(|g| g.ordinal)
180 {
181 return Ok(g.ordinal);
182 }
183 self.groups
185 .iter()
186 .filter(|g| g.copy)
187 .min_by_key(|g| g.ordinal)
188 .map(|g| g.ordinal)
189 .ok_or_else(|| LevelZeroError::Unsupported("device exposes no copy queue group".into()))
190 }
191
192 #[must_use]
194 pub fn has_dedicated_copy_engine(&self) -> bool {
195 self.groups.iter().any(|g| g.copy && !g.compute)
196 }
197}
198
199#[derive(Debug, Clone, Copy, PartialEq, Eq)]
203pub enum EventScope {
204 Device,
206 HostVisible,
208 KernelTimestamp,
210}
211
212impl EventScope {
213 #[must_use]
215 pub fn ze_flags(self) -> u32 {
216 match self {
217 EventScope::Device => 0,
218 EventScope::HostVisible => 0x1,
219 EventScope::KernelTimestamp => 0x4,
222 }
223 }
224}
225
226#[derive(Debug, Clone)]
233pub struct EventPoolDesc {
234 pub capacity: u32,
236 pub scope: EventScope,
238 next_index: u32,
240}
241
242impl EventPoolDesc {
243 pub fn new(capacity: u32, scope: EventScope) -> LevelZeroResult<Self> {
245 if capacity == 0 {
246 return Err(LevelZeroError::InvalidArgument(
247 "event pool capacity must be > 0".into(),
248 ));
249 }
250 Ok(Self {
251 capacity,
252 scope,
253 next_index: 0,
254 })
255 }
256
257 pub fn alloc_event(&mut self) -> LevelZeroResult<u32> {
261 if self.next_index >= self.capacity {
262 return Err(LevelZeroError::OutOfMemory);
263 }
264 let idx = self.next_index;
265 self.next_index += 1;
266 Ok(idx)
267 }
268
269 #[must_use]
271 pub fn remaining(&self) -> u32 {
272 self.capacity - self.next_index
273 }
274
275 #[must_use]
277 pub fn is_host_visible(&self) -> bool {
278 matches!(self.scope, EventScope::HostVisible)
279 }
280}
281
282#[derive(Debug, Clone, PartialEq, Eq)]
288pub struct Barrier {
289 pub signal_event: Option<u32>,
291 pub wait_events: Vec<u32>,
293}
294
295impl Barrier {
296 #[must_use]
298 pub fn global() -> Self {
299 Self {
300 signal_event: None,
301 wait_events: Vec::new(),
302 }
303 }
304
305 #[must_use]
307 pub fn signalling(event: u32) -> Self {
308 Self {
309 signal_event: Some(event),
310 wait_events: Vec::new(),
311 }
312 }
313
314 #[must_use]
316 pub fn waiting(events: Vec<u32>) -> Self {
317 Self {
318 signal_event: None,
319 wait_events: events,
320 }
321 }
322}
323
324#[derive(Debug, Clone, PartialEq, Eq)]
328pub enum RecordedCommand {
329 MemoryCopy {
331 dst: u64,
333 src: u64,
335 bytes: u64,
337 },
338 MemoryFill {
340 dst: u64,
342 bytes: u64,
344 pattern_len: u32,
346 },
347 LaunchKernel {
349 kernel: String,
351 groups: GroupCount,
353 signal_event: Option<u32>,
355 },
356 Barrier(Barrier),
358}
359
360#[derive(Debug, Clone, Copy, PartialEq, Eq)]
362pub enum CommandListState {
363 Recording,
365 Closed,
367}
368
369#[derive(Debug, Clone)]
375pub struct CommandListRecorder {
376 ordinal: u32,
378 commands: Vec<RecordedCommand>,
379 state: CommandListState,
380}
381
382impl CommandListRecorder {
383 #[must_use]
385 pub fn new(ordinal: u32) -> Self {
386 Self {
387 ordinal,
388 commands: Vec::new(),
389 state: CommandListState::Recording,
390 }
391 }
392
393 #[must_use]
395 pub fn ordinal(&self) -> u32 {
396 self.ordinal
397 }
398
399 #[must_use]
401 pub fn state(&self) -> CommandListState {
402 self.state
403 }
404
405 #[must_use]
407 pub fn len(&self) -> usize {
408 self.commands.len()
409 }
410
411 #[must_use]
413 pub fn is_empty(&self) -> bool {
414 self.commands.is_empty()
415 }
416
417 #[must_use]
419 pub fn commands(&self) -> &[RecordedCommand] {
420 &self.commands
421 }
422
423 fn ensure_recording(&self) -> LevelZeroResult<()> {
424 if self.state != CommandListState::Recording {
425 return Err(LevelZeroError::CommandListError(
426 "command list is closed; reset before appending".into(),
427 ));
428 }
429 Ok(())
430 }
431
432 pub fn append_memory_copy(&mut self, dst: u64, src: u64, bytes: u64) -> LevelZeroResult<()> {
434 self.ensure_recording()?;
435 if bytes == 0 {
436 return Err(LevelZeroError::InvalidArgument(
437 "memory copy size must be > 0".into(),
438 ));
439 }
440 self.commands
441 .push(RecordedCommand::MemoryCopy { dst, src, bytes });
442 Ok(())
443 }
444
445 pub fn append_memory_fill(
447 &mut self,
448 dst: u64,
449 bytes: u64,
450 pattern_len: u32,
451 ) -> LevelZeroResult<()> {
452 self.ensure_recording()?;
453 if pattern_len == 0 {
454 return Err(LevelZeroError::InvalidArgument(
455 "fill pattern length must be > 0".into(),
456 ));
457 }
458 if bytes % u64::from(pattern_len) != 0 {
459 return Err(LevelZeroError::InvalidArgument(
460 "fill size must be a multiple of the pattern length".into(),
461 ));
462 }
463 self.commands.push(RecordedCommand::MemoryFill {
464 dst,
465 bytes,
466 pattern_len,
467 });
468 Ok(())
469 }
470
471 pub fn append_launch_kernel(
473 &mut self,
474 kernel: impl Into<String>,
475 groups: GroupCount,
476 signal_event: Option<u32>,
477 ) -> LevelZeroResult<()> {
478 self.ensure_recording()?;
479 if groups.total() == 0 {
480 return Err(LevelZeroError::InvalidArgument(
481 "kernel launch grid must contain at least one workgroup".into(),
482 ));
483 }
484 self.commands.push(RecordedCommand::LaunchKernel {
485 kernel: kernel.into(),
486 groups,
487 signal_event,
488 });
489 Ok(())
490 }
491
492 pub fn append_barrier(&mut self, barrier: Barrier) -> LevelZeroResult<()> {
494 self.ensure_recording()?;
495 self.commands.push(RecordedCommand::Barrier(barrier));
496 Ok(())
497 }
498
499 pub fn close(&mut self) -> LevelZeroResult<()> {
501 self.ensure_recording()?;
502 self.state = CommandListState::Closed;
503 Ok(())
504 }
505
506 pub fn reset(&mut self) {
508 self.commands.clear();
509 self.state = CommandListState::Recording;
510 }
511
512 #[must_use]
514 pub fn launch_count(&self) -> usize {
515 self.commands
516 .iter()
517 .filter(|c| matches!(c, RecordedCommand::LaunchKernel { .. }))
518 .count()
519 }
520}
521
522#[derive(Debug, Clone)]
531pub struct ReusableCommandList {
532 inner: CommandListRecorder,
533 reuse_count: u64,
534}
535
536impl ReusableCommandList {
537 #[must_use]
539 pub fn new(ordinal: u32) -> Self {
540 Self {
541 inner: CommandListRecorder::new(ordinal),
542 reuse_count: 0,
543 }
544 }
545
546 pub fn recorder_mut(&mut self) -> &mut CommandListRecorder {
548 &mut self.inner
549 }
550
551 #[must_use]
553 pub fn recorder(&self) -> &CommandListRecorder {
554 &self.inner
555 }
556
557 #[must_use]
559 pub fn reuse_count(&self) -> u64 {
560 self.reuse_count
561 }
562
563 pub fn recycle(&mut self) -> LevelZeroResult<()> {
568 if self.inner.state() != CommandListState::Closed {
569 return Err(LevelZeroError::CommandListError(
570 "close the command list before recycling it".into(),
571 ));
572 }
573 self.inner.reset();
574 self.reuse_count += 1;
575 Ok(())
576 }
577}
578
579#[cfg(test)]
582mod tests {
583 use super::*;
584
585 #[test]
586 fn group_count_total() {
587 let g = GroupCount { x: 4, y: 2, z: 3 };
588 assert_eq!(g.total(), 24);
589 }
590
591 #[test]
592 fn planner_1d_ceil_div() {
593 let p = GroupCountPlanner::new_1d(256).unwrap();
594 assert_eq!(p.local_size(), (256, 1, 1));
595 assert_eq!(p.plan_1d(1000).unwrap(), GroupCount { x: 4, y: 1, z: 1 });
597 assert_eq!(p.plan_1d(512).unwrap(), GroupCount { x: 2, y: 1, z: 1 });
599 assert_eq!(p.plan_1d(0).unwrap(), GroupCount { x: 0, y: 1, z: 1 });
601 }
602
603 #[test]
604 fn planner_3d_per_dimension() {
605 let p = GroupCountPlanner::new_3d(16, 16, 1).unwrap();
606 let g = p.plan_3d(100, 50, 8).unwrap();
607 assert_eq!(g.x, 7); assert_eq!(g.y, 4); assert_eq!(g.z, 8); }
611
612 #[test]
613 fn planner_rejects_zero_local_and_overflow() {
614 assert!(GroupCountPlanner::new_3d(0, 1, 1).is_err());
615 let p = GroupCountPlanner::new_1d(1).unwrap();
616 assert!(matches!(
618 p.plan_1d(1u64 << 33),
619 Err(LevelZeroError::InvalidArgument(_))
620 ));
621 }
622
623 #[test]
624 fn queue_group_selection() {
625 let table = CommandQueueGroupTable::new(vec![
626 CommandQueueGroupInfo {
627 ordinal: 0,
628 compute: true,
629 copy: true,
630 num_queues: 4,
631 },
632 CommandQueueGroupInfo {
633 ordinal: 1,
634 compute: false,
635 copy: true,
636 num_queues: 1,
637 },
638 ]);
639 assert_eq!(table.select_compute_ordinal().unwrap(), 0);
640 assert_eq!(table.select_copy_ordinal().unwrap(), 1);
642 assert!(table.has_dedicated_copy_engine());
643 assert_eq!(table.groups().len(), 2);
644 }
645
646 #[test]
647 fn queue_group_fallback_to_shared_copy() {
648 let table = CommandQueueGroupTable::new(vec![CommandQueueGroupInfo {
649 ordinal: 0,
650 compute: true,
651 copy: true,
652 num_queues: 1,
653 }]);
654 assert_eq!(table.select_copy_ordinal().unwrap(), 0);
656 assert!(!table.has_dedicated_copy_engine());
657 }
658
659 #[test]
660 fn queue_group_errors_when_missing_capability() {
661 let copy_only = CommandQueueGroupTable::new(vec![CommandQueueGroupInfo {
662 ordinal: 0,
663 compute: false,
664 copy: true,
665 num_queues: 1,
666 }]);
667 assert!(matches!(
668 copy_only.select_compute_ordinal(),
669 Err(LevelZeroError::Unsupported(_))
670 ));
671
672 let compute_only = CommandQueueGroupTable::new(vec![CommandQueueGroupInfo {
673 ordinal: 0,
674 compute: true,
675 copy: false,
676 num_queues: 1,
677 }]);
678 assert!(matches!(
679 compute_only.select_copy_ordinal(),
680 Err(LevelZeroError::Unsupported(_))
681 ));
682 }
683
684 #[test]
685 fn event_pool_allocation() {
686 let mut pool = EventPoolDesc::new(2, EventScope::HostVisible).unwrap();
687 assert!(pool.is_host_visible());
688 assert_eq!(pool.scope.ze_flags(), 0x1);
689 assert_eq!(pool.remaining(), 2);
690 assert_eq!(pool.alloc_event().unwrap(), 0);
691 assert_eq!(pool.alloc_event().unwrap(), 1);
692 assert_eq!(pool.remaining(), 0);
693 assert!(matches!(
694 pool.alloc_event(),
695 Err(LevelZeroError::OutOfMemory)
696 ));
697 }
698
699 #[test]
700 fn event_pool_rejects_zero_capacity() {
701 assert!(EventPoolDesc::new(0, EventScope::Device).is_err());
702 assert_eq!(EventScope::KernelTimestamp.ze_flags(), 0x4);
703 assert_eq!(EventScope::HostVisible.ze_flags(), 0x1);
704 assert_eq!(EventScope::Device.ze_flags(), 0);
705 }
706
707 #[test]
708 fn barrier_constructors() {
709 let g = Barrier::global();
710 assert!(g.signal_event.is_none());
711 assert!(g.wait_events.is_empty());
712
713 let s = Barrier::signalling(3);
714 assert_eq!(s.signal_event, Some(3));
715
716 let w = Barrier::waiting(vec![1, 2]);
717 assert_eq!(w.wait_events, vec![1, 2]);
718 assert!(w.signal_event.is_none());
719 }
720
721 #[test]
722 fn recorder_records_in_order() {
723 let mut rec = CommandListRecorder::new(0);
724 assert_eq!(rec.ordinal(), 0);
725 assert!(rec.is_empty());
726 assert_eq!(rec.state(), CommandListState::Recording);
727
728 rec.append_memory_copy(2, 1, 256).unwrap();
729 rec.append_launch_kernel("gemm_f32", GroupCount { x: 4, y: 1, z: 1 }, Some(0))
730 .unwrap();
731 rec.append_barrier(Barrier::global()).unwrap();
732 assert_eq!(rec.len(), 3);
733 assert_eq!(rec.launch_count(), 1);
734
735 match &rec.commands()[0] {
736 RecordedCommand::MemoryCopy { dst, src, bytes } => {
737 assert_eq!((*dst, *src, *bytes), (2, 1, 256));
738 }
739 other => panic!("expected MemoryCopy, got {other:?}"),
740 }
741 }
742
743 #[test]
744 fn recorder_validates_args() {
745 let mut rec = CommandListRecorder::new(0);
746 assert!(rec.append_memory_copy(1, 0, 0).is_err());
747 assert!(rec.append_memory_fill(1, 10, 0).is_err());
748 assert!(
749 rec.append_memory_fill(1, 10, 4).is_err(),
750 "10 not a multiple of 4"
751 );
752 rec.append_memory_fill(1, 12, 4).unwrap();
753 assert!(
754 rec.append_launch_kernel("k", GroupCount { x: 0, y: 1, z: 1 }, None)
755 .is_err()
756 );
757 }
758
759 #[test]
760 fn recorder_close_blocks_append_until_reset() {
761 let mut rec = CommandListRecorder::new(0);
762 rec.append_memory_copy(2, 1, 128).unwrap();
763 rec.close().unwrap();
764 assert_eq!(rec.state(), CommandListState::Closed);
765 assert!(matches!(
766 rec.append_memory_copy(3, 1, 64),
767 Err(LevelZeroError::CommandListError(_))
768 ));
769 assert!(rec.close().is_err());
771
772 rec.reset();
773 assert_eq!(rec.state(), CommandListState::Recording);
774 assert!(rec.is_empty());
775 rec.append_memory_copy(3, 1, 64).unwrap();
776 }
777
778 #[test]
779 fn reusable_command_list_recycle_loop() {
780 let mut list = ReusableCommandList::new(0);
781 assert_eq!(list.reuse_count(), 0);
782
783 for expected in 0..3u64 {
784 assert_eq!(list.reuse_count(), expected);
785 let rec = list.recorder_mut();
786 rec.append_launch_kernel("k", GroupCount { x: 1, y: 1, z: 1 }, None)
787 .unwrap();
788 rec.close().unwrap();
789 list.recycle().unwrap();
790 }
791 assert_eq!(list.reuse_count(), 3);
792 assert!(list.recorder().is_empty());
794 assert_eq!(list.recorder().state(), CommandListState::Recording);
795 }
796
797 #[test]
798 fn reusable_rejects_recycle_while_open() {
799 let mut list = ReusableCommandList::new(0);
800 list.recorder_mut()
801 .append_launch_kernel("k", GroupCount { x: 1, y: 1, z: 1 }, None)
802 .unwrap();
803 assert!(matches!(
805 list.recycle(),
806 Err(LevelZeroError::CommandListError(_))
807 ));
808 assert_eq!(list.reuse_count(), 0);
809 }
810}