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 => 0x2,
220 }
221 }
222}
223
224#[derive(Debug, Clone)]
231pub struct EventPoolDesc {
232 pub capacity: u32,
234 pub scope: EventScope,
236 next_index: u32,
238}
239
240impl EventPoolDesc {
241 pub fn new(capacity: u32, scope: EventScope) -> LevelZeroResult<Self> {
243 if capacity == 0 {
244 return Err(LevelZeroError::InvalidArgument(
245 "event pool capacity must be > 0".into(),
246 ));
247 }
248 Ok(Self {
249 capacity,
250 scope,
251 next_index: 0,
252 })
253 }
254
255 pub fn alloc_event(&mut self) -> LevelZeroResult<u32> {
259 if self.next_index >= self.capacity {
260 return Err(LevelZeroError::OutOfMemory);
261 }
262 let idx = self.next_index;
263 self.next_index += 1;
264 Ok(idx)
265 }
266
267 #[must_use]
269 pub fn remaining(&self) -> u32 {
270 self.capacity - self.next_index
271 }
272
273 #[must_use]
275 pub fn is_host_visible(&self) -> bool {
276 matches!(self.scope, EventScope::HostVisible)
277 }
278}
279
280#[derive(Debug, Clone, PartialEq, Eq)]
286pub struct Barrier {
287 pub signal_event: Option<u32>,
289 pub wait_events: Vec<u32>,
291}
292
293impl Barrier {
294 #[must_use]
296 pub fn global() -> Self {
297 Self {
298 signal_event: None,
299 wait_events: Vec::new(),
300 }
301 }
302
303 #[must_use]
305 pub fn signalling(event: u32) -> Self {
306 Self {
307 signal_event: Some(event),
308 wait_events: Vec::new(),
309 }
310 }
311
312 #[must_use]
314 pub fn waiting(events: Vec<u32>) -> Self {
315 Self {
316 signal_event: None,
317 wait_events: events,
318 }
319 }
320}
321
322#[derive(Debug, Clone, PartialEq, Eq)]
326pub enum RecordedCommand {
327 MemoryCopy {
329 dst: u64,
331 src: u64,
333 bytes: u64,
335 },
336 MemoryFill {
338 dst: u64,
340 bytes: u64,
342 pattern_len: u32,
344 },
345 LaunchKernel {
347 kernel: String,
349 groups: GroupCount,
351 signal_event: Option<u32>,
353 },
354 Barrier(Barrier),
356}
357
358#[derive(Debug, Clone, Copy, PartialEq, Eq)]
360pub enum CommandListState {
361 Recording,
363 Closed,
365}
366
367#[derive(Debug, Clone)]
373pub struct CommandListRecorder {
374 ordinal: u32,
376 commands: Vec<RecordedCommand>,
377 state: CommandListState,
378}
379
380impl CommandListRecorder {
381 #[must_use]
383 pub fn new(ordinal: u32) -> Self {
384 Self {
385 ordinal,
386 commands: Vec::new(),
387 state: CommandListState::Recording,
388 }
389 }
390
391 #[must_use]
393 pub fn ordinal(&self) -> u32 {
394 self.ordinal
395 }
396
397 #[must_use]
399 pub fn state(&self) -> CommandListState {
400 self.state
401 }
402
403 #[must_use]
405 pub fn len(&self) -> usize {
406 self.commands.len()
407 }
408
409 #[must_use]
411 pub fn is_empty(&self) -> bool {
412 self.commands.is_empty()
413 }
414
415 #[must_use]
417 pub fn commands(&self) -> &[RecordedCommand] {
418 &self.commands
419 }
420
421 fn ensure_recording(&self) -> LevelZeroResult<()> {
422 if self.state != CommandListState::Recording {
423 return Err(LevelZeroError::CommandListError(
424 "command list is closed; reset before appending".into(),
425 ));
426 }
427 Ok(())
428 }
429
430 pub fn append_memory_copy(&mut self, dst: u64, src: u64, bytes: u64) -> LevelZeroResult<()> {
432 self.ensure_recording()?;
433 if bytes == 0 {
434 return Err(LevelZeroError::InvalidArgument(
435 "memory copy size must be > 0".into(),
436 ));
437 }
438 self.commands
439 .push(RecordedCommand::MemoryCopy { dst, src, bytes });
440 Ok(())
441 }
442
443 pub fn append_memory_fill(
445 &mut self,
446 dst: u64,
447 bytes: u64,
448 pattern_len: u32,
449 ) -> LevelZeroResult<()> {
450 self.ensure_recording()?;
451 if pattern_len == 0 {
452 return Err(LevelZeroError::InvalidArgument(
453 "fill pattern length must be > 0".into(),
454 ));
455 }
456 if bytes % u64::from(pattern_len) != 0 {
457 return Err(LevelZeroError::InvalidArgument(
458 "fill size must be a multiple of the pattern length".into(),
459 ));
460 }
461 self.commands.push(RecordedCommand::MemoryFill {
462 dst,
463 bytes,
464 pattern_len,
465 });
466 Ok(())
467 }
468
469 pub fn append_launch_kernel(
471 &mut self,
472 kernel: impl Into<String>,
473 groups: GroupCount,
474 signal_event: Option<u32>,
475 ) -> LevelZeroResult<()> {
476 self.ensure_recording()?;
477 if groups.total() == 0 {
478 return Err(LevelZeroError::InvalidArgument(
479 "kernel launch grid must contain at least one workgroup".into(),
480 ));
481 }
482 self.commands.push(RecordedCommand::LaunchKernel {
483 kernel: kernel.into(),
484 groups,
485 signal_event,
486 });
487 Ok(())
488 }
489
490 pub fn append_barrier(&mut self, barrier: Barrier) -> LevelZeroResult<()> {
492 self.ensure_recording()?;
493 self.commands.push(RecordedCommand::Barrier(barrier));
494 Ok(())
495 }
496
497 pub fn close(&mut self) -> LevelZeroResult<()> {
499 self.ensure_recording()?;
500 self.state = CommandListState::Closed;
501 Ok(())
502 }
503
504 pub fn reset(&mut self) {
506 self.commands.clear();
507 self.state = CommandListState::Recording;
508 }
509
510 #[must_use]
512 pub fn launch_count(&self) -> usize {
513 self.commands
514 .iter()
515 .filter(|c| matches!(c, RecordedCommand::LaunchKernel { .. }))
516 .count()
517 }
518}
519
520#[derive(Debug, Clone)]
529pub struct ReusableCommandList {
530 inner: CommandListRecorder,
531 reuse_count: u64,
532}
533
534impl ReusableCommandList {
535 #[must_use]
537 pub fn new(ordinal: u32) -> Self {
538 Self {
539 inner: CommandListRecorder::new(ordinal),
540 reuse_count: 0,
541 }
542 }
543
544 pub fn recorder_mut(&mut self) -> &mut CommandListRecorder {
546 &mut self.inner
547 }
548
549 #[must_use]
551 pub fn recorder(&self) -> &CommandListRecorder {
552 &self.inner
553 }
554
555 #[must_use]
557 pub fn reuse_count(&self) -> u64 {
558 self.reuse_count
559 }
560
561 pub fn recycle(&mut self) -> LevelZeroResult<()> {
566 if self.inner.state() != CommandListState::Closed {
567 return Err(LevelZeroError::CommandListError(
568 "close the command list before recycling it".into(),
569 ));
570 }
571 self.inner.reset();
572 self.reuse_count += 1;
573 Ok(())
574 }
575}
576
577#[cfg(test)]
580mod tests {
581 use super::*;
582
583 #[test]
584 fn group_count_total() {
585 let g = GroupCount { x: 4, y: 2, z: 3 };
586 assert_eq!(g.total(), 24);
587 }
588
589 #[test]
590 fn planner_1d_ceil_div() {
591 let p = GroupCountPlanner::new_1d(256).unwrap();
592 assert_eq!(p.local_size(), (256, 1, 1));
593 assert_eq!(p.plan_1d(1000).unwrap(), GroupCount { x: 4, y: 1, z: 1 });
595 assert_eq!(p.plan_1d(512).unwrap(), GroupCount { x: 2, y: 1, z: 1 });
597 assert_eq!(p.plan_1d(0).unwrap(), GroupCount { x: 0, y: 1, z: 1 });
599 }
600
601 #[test]
602 fn planner_3d_per_dimension() {
603 let p = GroupCountPlanner::new_3d(16, 16, 1).unwrap();
604 let g = p.plan_3d(100, 50, 8).unwrap();
605 assert_eq!(g.x, 7); assert_eq!(g.y, 4); assert_eq!(g.z, 8); }
609
610 #[test]
611 fn planner_rejects_zero_local_and_overflow() {
612 assert!(GroupCountPlanner::new_3d(0, 1, 1).is_err());
613 let p = GroupCountPlanner::new_1d(1).unwrap();
614 assert!(matches!(
616 p.plan_1d(1u64 << 33),
617 Err(LevelZeroError::InvalidArgument(_))
618 ));
619 }
620
621 #[test]
622 fn queue_group_selection() {
623 let table = CommandQueueGroupTable::new(vec![
624 CommandQueueGroupInfo {
625 ordinal: 0,
626 compute: true,
627 copy: true,
628 num_queues: 4,
629 },
630 CommandQueueGroupInfo {
631 ordinal: 1,
632 compute: false,
633 copy: true,
634 num_queues: 1,
635 },
636 ]);
637 assert_eq!(table.select_compute_ordinal().unwrap(), 0);
638 assert_eq!(table.select_copy_ordinal().unwrap(), 1);
640 assert!(table.has_dedicated_copy_engine());
641 assert_eq!(table.groups().len(), 2);
642 }
643
644 #[test]
645 fn queue_group_fallback_to_shared_copy() {
646 let table = CommandQueueGroupTable::new(vec![CommandQueueGroupInfo {
647 ordinal: 0,
648 compute: true,
649 copy: true,
650 num_queues: 1,
651 }]);
652 assert_eq!(table.select_copy_ordinal().unwrap(), 0);
654 assert!(!table.has_dedicated_copy_engine());
655 }
656
657 #[test]
658 fn queue_group_errors_when_missing_capability() {
659 let copy_only = CommandQueueGroupTable::new(vec![CommandQueueGroupInfo {
660 ordinal: 0,
661 compute: false,
662 copy: true,
663 num_queues: 1,
664 }]);
665 assert!(matches!(
666 copy_only.select_compute_ordinal(),
667 Err(LevelZeroError::Unsupported(_))
668 ));
669
670 let compute_only = CommandQueueGroupTable::new(vec![CommandQueueGroupInfo {
671 ordinal: 0,
672 compute: true,
673 copy: false,
674 num_queues: 1,
675 }]);
676 assert!(matches!(
677 compute_only.select_copy_ordinal(),
678 Err(LevelZeroError::Unsupported(_))
679 ));
680 }
681
682 #[test]
683 fn event_pool_allocation() {
684 let mut pool = EventPoolDesc::new(2, EventScope::HostVisible).unwrap();
685 assert!(pool.is_host_visible());
686 assert_eq!(pool.scope.ze_flags(), 0x1);
687 assert_eq!(pool.remaining(), 2);
688 assert_eq!(pool.alloc_event().unwrap(), 0);
689 assert_eq!(pool.alloc_event().unwrap(), 1);
690 assert_eq!(pool.remaining(), 0);
691 assert!(matches!(
692 pool.alloc_event(),
693 Err(LevelZeroError::OutOfMemory)
694 ));
695 }
696
697 #[test]
698 fn event_pool_rejects_zero_capacity() {
699 assert!(EventPoolDesc::new(0, EventScope::Device).is_err());
700 assert_eq!(EventScope::KernelTimestamp.ze_flags(), 0x2);
701 assert_eq!(EventScope::Device.ze_flags(), 0);
702 }
703
704 #[test]
705 fn barrier_constructors() {
706 let g = Barrier::global();
707 assert!(g.signal_event.is_none());
708 assert!(g.wait_events.is_empty());
709
710 let s = Barrier::signalling(3);
711 assert_eq!(s.signal_event, Some(3));
712
713 let w = Barrier::waiting(vec![1, 2]);
714 assert_eq!(w.wait_events, vec![1, 2]);
715 assert!(w.signal_event.is_none());
716 }
717
718 #[test]
719 fn recorder_records_in_order() {
720 let mut rec = CommandListRecorder::new(0);
721 assert_eq!(rec.ordinal(), 0);
722 assert!(rec.is_empty());
723 assert_eq!(rec.state(), CommandListState::Recording);
724
725 rec.append_memory_copy(2, 1, 256).unwrap();
726 rec.append_launch_kernel("gemm_f32", GroupCount { x: 4, y: 1, z: 1 }, Some(0))
727 .unwrap();
728 rec.append_barrier(Barrier::global()).unwrap();
729 assert_eq!(rec.len(), 3);
730 assert_eq!(rec.launch_count(), 1);
731
732 match &rec.commands()[0] {
733 RecordedCommand::MemoryCopy { dst, src, bytes } => {
734 assert_eq!((*dst, *src, *bytes), (2, 1, 256));
735 }
736 other => panic!("expected MemoryCopy, got {other:?}"),
737 }
738 }
739
740 #[test]
741 fn recorder_validates_args() {
742 let mut rec = CommandListRecorder::new(0);
743 assert!(rec.append_memory_copy(1, 0, 0).is_err());
744 assert!(rec.append_memory_fill(1, 10, 0).is_err());
745 assert!(
746 rec.append_memory_fill(1, 10, 4).is_err(),
747 "10 not a multiple of 4"
748 );
749 rec.append_memory_fill(1, 12, 4).unwrap();
750 assert!(
751 rec.append_launch_kernel("k", GroupCount { x: 0, y: 1, z: 1 }, None)
752 .is_err()
753 );
754 }
755
756 #[test]
757 fn recorder_close_blocks_append_until_reset() {
758 let mut rec = CommandListRecorder::new(0);
759 rec.append_memory_copy(2, 1, 128).unwrap();
760 rec.close().unwrap();
761 assert_eq!(rec.state(), CommandListState::Closed);
762 assert!(matches!(
763 rec.append_memory_copy(3, 1, 64),
764 Err(LevelZeroError::CommandListError(_))
765 ));
766 assert!(rec.close().is_err());
768
769 rec.reset();
770 assert_eq!(rec.state(), CommandListState::Recording);
771 assert!(rec.is_empty());
772 rec.append_memory_copy(3, 1, 64).unwrap();
773 }
774
775 #[test]
776 fn reusable_command_list_recycle_loop() {
777 let mut list = ReusableCommandList::new(0);
778 assert_eq!(list.reuse_count(), 0);
779
780 for expected in 0..3u64 {
781 assert_eq!(list.reuse_count(), expected);
782 let rec = list.recorder_mut();
783 rec.append_launch_kernel("k", GroupCount { x: 1, y: 1, z: 1 }, None)
784 .unwrap();
785 rec.close().unwrap();
786 list.recycle().unwrap();
787 }
788 assert_eq!(list.reuse_count(), 3);
789 assert!(list.recorder().is_empty());
791 assert_eq!(list.recorder().state(), CommandListState::Recording);
792 }
793
794 #[test]
795 fn reusable_rejects_recycle_while_open() {
796 let mut list = ReusableCommandList::new(0);
797 list.recorder_mut()
798 .append_launch_kernel("k", GroupCount { x: 1, y: 1, z: 1 }, None)
799 .unwrap();
800 assert!(matches!(
802 list.recycle(),
803 Err(LevelZeroError::CommandListError(_))
804 ));
805 assert_eq!(list.reuse_count(), 0);
806 }
807}