1use std::collections::HashMap;
4use std::ptr::NonNull;
5use std::sync::{Arc, Condvar, Mutex, PoisonError};
6
7use block2::RcBlock;
8use objc2::rc::{Retained, autoreleasepool};
9use objc2::runtime::ProtocolObject;
10use objc2_metal::{
11 MTLBarrierScope, MTLBuffer, MTLCommandBuffer, MTLCommandBufferStatus, MTLCommandEncoder,
12 MTLCommandQueue, MTLComputeCommandEncoder, MTLDevice, MTLDispatchType, MTLFence, MTLSize,
13};
14
15use super::buffer::Guard;
16use super::{Buffer, Device, Pipeline};
17use crate::{GpuError, Result};
18
19const MAX_ARGS: usize = 31;
21const MAX_INLINE_BYTES: usize = 4096;
23
24#[derive(Debug)]
26pub enum Arg<'a> {
27 Read(&'a Buffer, usize),
29 Write(&'a mut Buffer, usize),
31 Bytes(&'a [u8]),
33 Same(usize, usize),
38}
39
40impl<'a> Arg<'a> {
41 pub fn read(buffer: &'a Buffer) -> Self {
43 Self::Read(buffer, 0)
44 }
45
46 pub fn write(buffer: &'a mut Buffer) -> Self {
48 Self::Write(buffer, 0)
49 }
50
51 pub fn bytes(bytes: &'a [u8]) -> Self {
53 Self::Bytes(bytes)
54 }
55
56 fn buffer(&self) -> Option<(&Buffer, usize, bool)> {
57 match self {
58 Self::Read(b, offset) => Some((b, *offset, false)),
59 Self::Write(b, offset) => Some((b, *offset, true)),
60 Self::Bytes(_) | Self::Same(..) => None,
61 }
62 }
63}
64
65pub(crate) type Fence = Retained<ProtocolObject<dyn MTLFence>>;
66
67pub struct Timeline {
76 fence: Fence,
77 started: bool,
78}
79
80impl std::fmt::Debug for Timeline {
81 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82 f.debug_struct("Timeline").finish_non_exhaustive()
83 }
84}
85
86pub struct CommandBuffer {
89 device: Device,
90 raw: Retained<ProtocolObject<dyn MTLCommandBuffer>>,
91 fence: Option<(Fence, bool)>,
93 own: Option<Fence>,
96 waited: Vec<Fence>,
98 encoder: Option<Retained<ProtocolObject<dyn MTLComputeCommandEncoder>>>,
99 reads: Vec<usize>,
101 writes: Vec<usize>,
102 barriers: usize,
103 guards: Vec<(Arc<Guard>, bool)>,
105 uses: HashMap<usize, (usize, usize)>,
107 retained: Vec<Retained<ProtocolObject<dyn MTLBuffer>>>,
108 kernels: Vec<Arc<str>>,
110}
111
112impl Device {
113 pub fn command_buffer(&self) -> Result<CommandBuffer> {
119 let raw = autoreleasepool(|_| self.queue().commandBuffer())
120 .ok_or_else(|| GpuError::CommandBuffer("cannot create a command buffer".into()))?;
121 Ok(CommandBuffer {
122 device: self.clone(),
123 raw,
124 fence: None,
125 own: None,
126 waited: Vec::new(),
127 encoder: None,
128 reads: Vec::new(),
129 writes: Vec::new(),
130 barriers: 0,
131 guards: Vec::new(),
132 uses: HashMap::new(),
133 retained: Vec::new(),
134 kernels: Vec::new(),
135 })
136 }
137}
138
139impl Device {
140 pub fn timeline(&self) -> Result<Timeline> {
142 let fence = autoreleasepool(|_| self.raw().newFence())
143 .ok_or_else(|| GpuError::CommandBuffer("cannot create a fence".into()))?;
144 Ok(Timeline {
145 fence,
146 started: false,
147 })
148 }
149
150 pub fn ordered_command_buffer(&self) -> Result<CommandBuffer> {
159 let fence = autoreleasepool(|_| self.raw().newFence())
160 .ok_or_else(|| GpuError::CommandBuffer("cannot create a fence".into()))?;
161 let mut cmd = self.command_buffer()?;
162 cmd.own = Some(fence);
163 Ok(cmd)
164 }
165
166 pub fn command_buffer_on(&self, timeline: &mut Timeline) -> Result<CommandBuffer> {
171 let mut cmd = self.command_buffer()?;
172 cmd.fence = Some((timeline.fence.clone(), timeline.started));
173 timeline.started = true;
174 Ok(cmd)
175 }
176}
177
178fn status(raw: &ProtocolObject<dyn MTLCommandBuffer>) -> Result<()> {
180 if raw.status() == MTLCommandBufferStatus::Error {
181 let message = raw.error().map_or_else(
182 || "unknown error".to_owned(),
183 |e| e.localizedDescription().to_string(),
184 );
185 return Err(GpuError::CommandBuffer(message));
186 }
187 Ok(())
188}
189
190impl CommandBuffer {
191 pub unsafe fn dispatch(
205 &mut self,
206 pipeline: &Pipeline,
207 args: &[Arg<'_>],
208 threads: usize,
209 ) -> Result<()> {
210 if threads == 0 {
211 return self.check(args);
212 }
213 let width = pipeline.execution_width().max(1);
214 let mut group = pipeline.max_threads_per_threadgroup().min(threads).max(1);
215 if group > width {
216 group -= group % width;
217 }
218 self.encode(pipeline, args, [threads, 1, 1], |encoder| {
219 encoder
220 .dispatchThreads_threadsPerThreadgroup(size([threads, 1, 1]), size([group, 1, 1]));
221 })
222 }
223
224 pub unsafe fn dispatch_groups(
232 &mut self,
233 pipeline: &Pipeline,
234 args: &[Arg<'_>],
235 groups: [usize; 3],
236 threads: [usize; 3],
237 ) -> Result<()> {
238 let per_group: usize = threads.iter().product();
239 if per_group > pipeline.max_threads_per_threadgroup() {
240 return Err(GpuError::InvalidArgument(format!(
241 "{per_group} threads per threadgroup; `{}` allows {}",
242 pipeline.name(),
243 pipeline.max_threads_per_threadgroup()
244 )));
245 }
246 if groups.contains(&0) || per_group == 0 {
247 return self.check(args);
248 }
249 let extent = [0, 1, 2].map(|d| groups[d] * threads[d]);
250 self.encode(pipeline, args, extent, |encoder| {
251 encoder.dispatchThreadgroups_threadsPerThreadgroup(size(groups), size(threads));
252 })
253 }
254
255 pub unsafe fn dispatch_threads(
263 &mut self,
264 pipeline: &Pipeline,
265 args: &[Arg<'_>],
266 grid: [usize; 3],
267 group: [usize; 3],
268 ) -> Result<()> {
269 let per_group: usize = group.iter().product();
270 if per_group > pipeline.max_threads_per_threadgroup() {
271 return Err(GpuError::InvalidArgument(format!(
272 "{per_group} threads per threadgroup; `{}` allows {}",
273 pipeline.name(),
274 pipeline.max_threads_per_threadgroup()
275 )));
276 }
277 if grid.contains(&0) || per_group == 0 {
278 return self.check(args);
279 }
280 self.encode(pipeline, args, grid, |encoder| {
281 encoder.dispatchThreads_threadsPerThreadgroup(size(grid), size(group));
282 })
283 }
284
285 fn encode(
290 &mut self,
291 pipeline: &Pipeline,
292 args: &[Arg<'_>],
293 extent: [usize; 3],
294 dispatch: impl FnOnce(&ProtocolObject<dyn MTLComputeCommandEncoder>),
295 ) -> Result<()> {
296 autoreleasepool(|_| self.encode_now(pipeline, args, extent, dispatch))
297 }
298
299 fn encode_now(
300 &mut self,
301 pipeline: &Pipeline,
302 args: &[Arg<'_>],
303 extent: [usize; 3],
304 dispatch: impl FnOnce(&ProtocolObject<dyn MTLComputeCommandEncoder>),
305 ) -> Result<()> {
306 self.check(args)?;
307 if super::profile::enabled() {
308 self.kernels.push(match super::profile::shapes() {
309 true => format!("{} {extent:?}", pipeline.name()).into(),
310 false => pipeline.name.clone(),
311 });
312 }
313 if self.conflicts(args) {
314 self.encoder()?
315 .memoryBarrierWithScope(MTLBarrierScope::Buffers);
316 self.barriers += 1;
317 self.reads.clear();
318 self.writes.clear();
319 }
320 let encoder = self.encoder()?.clone();
321 if let Some(own) = &self.own {
322 let own = Retained::as_ptr(own) as usize;
323 for (buffer, _, write) in args.iter().filter_map(Arg::buffer) {
324 for fence in buffer.pending_fences(write) {
325 let key = Retained::as_ptr(&fence) as usize;
326 let seen = self
327 .waited
328 .iter()
329 .any(|f| Retained::as_ptr(f) as usize == key);
330 if key != own && !seen {
331 encoder.waitForFence(&fence);
332 self.waited.push(fence);
333 }
334 }
335 }
336 }
337 encoder.setComputePipelineState(pipeline.raw());
338 for (index, arg) in args.iter().enumerate() {
339 match arg {
340 Arg::Bytes(bytes) => {
341 let data = NonNull::new(bytes.as_ptr().cast_mut())
342 .unwrap_or(NonNull::dangling())
343 .cast();
344 unsafe { encoder.setBytes_length_atIndex(data, bytes.len(), index) };
348 }
349 Arg::Read(..) | Arg::Write(..) => {
350 let (buffer, offset, _) = arg.buffer().expect("buffer argument");
351 unsafe { encoder.setBuffer_offset_atIndex(buffer.raw(), offset, index) };
355 }
356 Arg::Same(earlier, offset) => {
357 let (buffer, _, _) = args[*earlier].buffer().expect("validated");
358 unsafe { encoder.setBuffer_offset_atIndex(buffer.raw(), *offset, index) };
361 }
362 }
363 }
364 for (buffer, _, write) in args.iter().filter_map(Arg::buffer) {
365 (if write {
366 &mut self.writes
367 } else {
368 &mut self.reads
369 })
370 .push(buffer.id());
371 let uses = self.uses.entry(buffer.id()).or_default();
372 uses.0 += 1;
373 uses.1 += usize::from(write);
374 let guard = buffer.guard();
375 guard.acquire(write);
376 self.guards.push((guard, write));
377 self.retained.extend(buffer.retained());
378 }
379 dispatch(&encoder);
380 Ok(())
381 }
382
383 fn check(&self, args: &[Arg<'_>]) -> Result<()> {
389 validate(args)?;
390 for (buffer, _, write) in args.iter().filter_map(Arg::buffer) {
391 if self.own.is_none() && buffer.previous_busy() {
392 return Err(GpuError::BufferInUse);
393 }
394 let (uses, writes) = buffer.unsubmitted();
395 let (own_uses, own_writes) = self.uses.get(&buffer.id()).copied().unwrap_or_default();
396 if writes > own_writes || (write && uses > own_uses) {
397 return Err(GpuError::BufferInUse);
398 }
399 }
400 Ok(())
401 }
402
403 pub fn barriers(&self) -> usize {
405 self.barriers
406 }
407
408 fn close(&mut self) {
410 autoreleasepool(|_| self.end_encoding());
411 self.device.flights().begin();
412 }
413
414 fn end_encoding(&mut self) {
415 if let Some(encoder) = self.encoder.take() {
416 if let Some((fence, _)) = &self.fence {
417 encoder.updateFence(fence);
418 }
419 if let Some(own) = &self.own {
420 encoder.updateFence(own);
421 for (guard, write) in &self.guards {
422 guard.record(own, *write);
423 }
424 }
425 encoder.endEncoding();
426 }
427 }
428
429 #[must_use = "dropping `Submitted` waits for the GPU and discards its errors; call `wait`"]
432 pub fn commit(self) -> Submitted {
433 Submitted {
434 pending: self.commit_then(|_| {}),
435 waited: false,
436 }
437 }
438
439 #[must_use = "dropping `Pending` does not wait; keep it to `wait`, or drop it explicitly"]
449 pub fn commit_then(mut self, on_complete: impl FnOnce(Result<()>) + Send + 'static) -> Pending {
450 self.close();
451 let guards = self.guards.clone();
452 let finish = Arc::new(Finish {
453 completion: Mutex::new(Some(Completion {
454 on_complete: Box::new(on_complete),
455 own: self.own.clone(),
456 kernels: std::mem::take(&mut self.kernels),
457 unfenced: self.own.is_none(),
458 guards: std::mem::take(&mut self.guards),
459 _retained: std::mem::take(&mut self.retained),
460 device: self.device.clone(),
461 })),
462 result: Mutex::new(None),
463 done: Condvar::new(),
464 });
465 let shared = finish.clone();
466 let handler = RcBlock::new(move |raw: NonNull<ProtocolObject<dyn MTLCommandBuffer>>| {
467 let completed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
470 complete(&shared, unsafe { raw.as_ref() });
472 }));
473 if completed.is_err() {
474 eprintln!("nuvec-gpu: a command buffer's completion handler panicked");
475 }
476 });
477 autoreleasepool(|_| {
478 unsafe { self.raw.addCompletedHandler(RcBlock::as_ptr(&handler)) };
482 self.raw.commit();
483 });
484 let sequence = super::buffer::commit_sequence();
485 for (guard, write) in guards {
486 guard.committed_as(sequence);
487 guard.submitted(write);
488 if self.own.is_some() {
491 guard.committed();
492 }
493 }
494 Pending {
495 raw: self.raw.clone(),
496 finish,
497 }
498 }
499
500 fn conflicts(&self, args: &[Arg<'_>]) -> bool {
503 args.iter().filter_map(Arg::buffer).any(|(b, _, write)| {
504 let id = b.id();
505 self.writes.contains(&id) || (write && self.reads.contains(&id))
506 })
507 }
508
509 fn encoder(&mut self) -> Result<&Retained<ProtocolObject<dyn MTLComputeCommandEncoder>>> {
510 if self.encoder.is_none() {
511 let encoder = self
512 .raw
513 .computeCommandEncoderWithDispatchType(MTLDispatchType::Concurrent)
514 .ok_or_else(|| GpuError::CommandBuffer("cannot create a compute encoder".into()))?;
515 if let Some((fence, true)) = &self.fence {
516 encoder.waitForFence(fence);
517 }
518 self.encoder = Some(encoder);
519 }
520 Ok(self.encoder.as_ref().expect("just created"))
521 }
522}
523
524impl Drop for CommandBuffer {
525 fn drop(&mut self) {
528 autoreleasepool(|_| {
529 if let Some(encoder) = self.encoder.take() {
530 encoder.endEncoding();
531 }
532 for (guard, write) in self.guards.drain(..) {
533 guard.committed();
534 guard.submitted(write);
535 guard.release(write, None); }
537 self.retained.clear();
538 });
539 }
540}
541
542fn size([width, height, depth]: [usize; 3]) -> MTLSize {
543 MTLSize {
544 width,
545 height,
546 depth,
547 }
548}
549
550fn validate(args: &[Arg<'_>]) -> Result<()> {
551 if args.len() > MAX_ARGS {
552 return Err(GpuError::InvalidArgument(format!(
553 "{} arguments; Metal allows {MAX_ARGS}",
554 args.len()
555 )));
556 }
557 for (index, arg) in args.iter().enumerate() {
558 match arg {
559 Arg::Bytes(bytes) if bytes.len() > MAX_INLINE_BYTES => {
560 return Err(GpuError::InvalidArgument(format!(
561 "argument {index}: {} inline bytes; Metal allows {MAX_INLINE_BYTES}",
562 bytes.len()
563 )));
564 }
565 Arg::Write(b, _) if b.is_mapped() => return Err(GpuError::ReadOnly),
566 Arg::Same(earlier, offset) => match args.get(*earlier) {
567 Some(Arg::Write(b, _)) if *offset <= b.len() => {}
568 _ => {
569 return Err(GpuError::InvalidArgument(format!(
570 "argument {index}: `Same({earlier}, {offset})` needs a written \
571 (`Arg::Write`) argument and an offset within it"
572 )));
573 }
574 },
575 Arg::Read(..) | Arg::Write(..)
576 if arg.buffer().is_some_and(|(b, offset, _)| offset > b.len()) =>
577 {
578 let (b, offset, _) = arg.buffer().expect("buffer argument");
579 return Err(GpuError::InvalidArgument(format!(
580 "argument {index}: offset {offset} is past the end of a {}-byte buffer",
581 b.len()
582 )));
583 }
584 _ => {}
585 }
586 }
587 Ok(())
588}
589
590struct Completion {
593 on_complete: Box<dyn FnOnce(Result<()>) + Send>,
594 own: Option<Fence>,
596 kernels: Vec<Arc<str>>,
598 guards: Vec<(Arc<Guard>, bool)>,
599 unfenced: bool,
601 _retained: Vec<Retained<ProtocolObject<dyn MTLBuffer>>>,
602 device: Device,
603}
604
605unsafe impl Send for Completion {}
609
610struct Finish {
613 completion: Mutex<Option<Completion>>,
614 result: Mutex<Option<Result<()>>>,
616 done: Condvar,
617}
618
619impl Finish {
620 fn wait(&self) -> Result<()> {
622 let mut result = self.result.lock().unwrap_or_else(PoisonError::into_inner);
623 loop {
624 if let Some(r) = &*result {
625 return r.clone();
626 }
627 result = self
628 .done
629 .wait(result)
630 .unwrap_or_else(PoisonError::into_inner);
631 }
632 }
633}
634
635struct Publish<'a>(&'a Finish, Option<Result<()>>);
637
638impl Drop for Publish<'_> {
639 fn drop(&mut self) {
640 let mut result = self.0.result.lock().unwrap_or_else(PoisonError::into_inner);
641 *result = self.1.take();
642 self.0.done.notify_all();
643 }
644}
645
646fn complete(finish: &Finish, raw: &ProtocolObject<dyn MTLCommandBuffer>) {
650 let Some(c) = finish
651 .completion
652 .lock()
653 .unwrap_or_else(PoisonError::into_inner)
654 .take()
655 else {
656 return;
657 };
658 let result = status(raw);
659 let _publish = Publish(finish, Some(result.clone()));
660 super::profile::record(&c.kernels, raw.GPUEndTime() - raw.GPUStartTime());
661 for (guard, write) in c.guards {
662 if c.unfenced {
663 guard.committed();
664 }
665 guard.release(write, c.own.as_ref());
666 }
667 drop(c._retained);
668 c.device.flights().end();
669 (c.on_complete)(result);
670}
671
672#[must_use = "dropping `Pending` does not wait; keep it to `wait`, or drop it explicitly"]
676pub struct Pending {
677 raw: Retained<ProtocolObject<dyn MTLCommandBuffer>>,
678 finish: Arc<Finish>,
679}
680
681impl Pending {
682 pub fn wait(&self) {
685 let _ = self.result();
686 }
687
688 fn result(&self) -> Result<()> {
690 autoreleasepool(|_| {
691 self.raw.waitUntilCompleted();
692 complete(&self.finish, &self.raw);
693 });
694 self.finish.wait()
695 }
696}
697
698impl std::fmt::Debug for Pending {
699 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
700 f.debug_struct("Pending").finish_non_exhaustive()
701 }
702}
703
704#[must_use = "dropping `Submitted` waits for the GPU and discards its errors; call `wait`"]
707pub struct Submitted {
708 pending: Pending,
709 waited: bool,
710}
711
712impl Submitted {
713 pub fn wait(mut self) -> Result<()> {
716 self.waited = true;
717 self.pending.result()
718 }
719}
720
721impl std::fmt::Debug for Submitted {
722 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
723 f.debug_struct("Submitted").finish_non_exhaustive()
724 }
725}
726
727impl Drop for Submitted {
728 fn drop(&mut self) {
729 if !self.waited {
730 self.pending.wait();
731 }
732 }
733}