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<usize>,
97 encoder: Option<Retained<ProtocolObject<dyn MTLComputeCommandEncoder>>>,
98 reads: Vec<usize>,
100 writes: Vec<usize>,
101 barriers: usize,
102 guards: Vec<(Arc<Guard>, bool)>,
104 uses: HashMap<usize, (usize, usize)>,
106 retained: Vec<Retained<ProtocolObject<dyn MTLBuffer>>>,
107 kernels: Vec<Arc<str>>,
109}
110
111impl Device {
112 pub fn command_buffer(&self) -> Result<CommandBuffer> {
118 let raw = autoreleasepool(|_| self.queue().commandBuffer())
119 .ok_or_else(|| GpuError::CommandBuffer("cannot create a command buffer".into()))?;
120 Ok(CommandBuffer {
121 device: self.clone(),
122 raw,
123 fence: None,
124 own: None,
125 waited: Vec::new(),
126 encoder: None,
127 reads: Vec::new(),
128 writes: Vec::new(),
129 barriers: 0,
130 guards: Vec::new(),
131 uses: HashMap::new(),
132 retained: Vec::new(),
133 kernels: Vec::new(),
134 })
135 }
136}
137
138impl Device {
139 pub fn timeline(&self) -> Result<Timeline> {
141 let fence = autoreleasepool(|_| self.raw().newFence())
142 .ok_or_else(|| GpuError::CommandBuffer("cannot create a fence".into()))?;
143 Ok(Timeline {
144 fence,
145 started: false,
146 })
147 }
148
149 pub fn ordered_command_buffer(&self) -> Result<CommandBuffer> {
158 let fence = autoreleasepool(|_| self.raw().newFence())
159 .ok_or_else(|| GpuError::CommandBuffer("cannot create a fence".into()))?;
160 let mut cmd = self.command_buffer()?;
161 cmd.own = Some(fence);
162 Ok(cmd)
163 }
164
165 pub fn command_buffer_on(&self, timeline: &mut Timeline) -> Result<CommandBuffer> {
170 let mut cmd = self.command_buffer()?;
171 cmd.fence = Some((timeline.fence.clone(), timeline.started));
172 timeline.started = true;
173 Ok(cmd)
174 }
175}
176
177fn status(raw: &ProtocolObject<dyn MTLCommandBuffer>) -> Result<()> {
179 if raw.status() == MTLCommandBufferStatus::Error {
180 let message = raw.error().map_or_else(
181 || "unknown error".to_owned(),
182 |e| e.localizedDescription().to_string(),
183 );
184 return Err(GpuError::CommandBuffer(message));
185 }
186 Ok(())
187}
188
189impl CommandBuffer {
190 pub unsafe fn dispatch(
204 &mut self,
205 pipeline: &Pipeline,
206 args: &[Arg<'_>],
207 threads: usize,
208 ) -> Result<()> {
209 if threads == 0 {
210 return self.check(args);
211 }
212 let width = pipeline.execution_width().max(1);
213 let mut group = pipeline.max_threads_per_threadgroup().min(threads).max(1);
214 if group > width {
215 group -= group % width;
216 }
217 self.encode(pipeline, args, [threads, 1, 1], |encoder| {
218 encoder
219 .dispatchThreads_threadsPerThreadgroup(size([threads, 1, 1]), size([group, 1, 1]));
220 })
221 }
222
223 pub unsafe fn dispatch_groups(
231 &mut self,
232 pipeline: &Pipeline,
233 args: &[Arg<'_>],
234 groups: [usize; 3],
235 threads: [usize; 3],
236 ) -> Result<()> {
237 let per_group: usize = threads.iter().product();
238 if per_group > pipeline.max_threads_per_threadgroup() {
239 return Err(GpuError::InvalidArgument(format!(
240 "{per_group} threads per threadgroup; `{}` allows {}",
241 pipeline.name(),
242 pipeline.max_threads_per_threadgroup()
243 )));
244 }
245 if groups.contains(&0) || per_group == 0 {
246 return self.check(args);
247 }
248 let extent = [0, 1, 2].map(|d| groups[d] * threads[d]);
249 self.encode(pipeline, args, extent, |encoder| {
250 encoder.dispatchThreadgroups_threadsPerThreadgroup(size(groups), size(threads));
251 })
252 }
253
254 pub unsafe fn dispatch_threads(
262 &mut self,
263 pipeline: &Pipeline,
264 args: &[Arg<'_>],
265 grid: [usize; 3],
266 group: [usize; 3],
267 ) -> Result<()> {
268 let per_group: usize = group.iter().product();
269 if per_group > pipeline.max_threads_per_threadgroup() {
270 return Err(GpuError::InvalidArgument(format!(
271 "{per_group} threads per threadgroup; `{}` allows {}",
272 pipeline.name(),
273 pipeline.max_threads_per_threadgroup()
274 )));
275 }
276 if grid.contains(&0) || per_group == 0 {
277 return self.check(args);
278 }
279 self.encode(pipeline, args, grid, |encoder| {
280 encoder.dispatchThreads_threadsPerThreadgroup(size(grid), size(group));
281 })
282 }
283
284 fn encode(
289 &mut self,
290 pipeline: &Pipeline,
291 args: &[Arg<'_>],
292 extent: [usize; 3],
293 dispatch: impl FnOnce(&ProtocolObject<dyn MTLComputeCommandEncoder>),
294 ) -> Result<()> {
295 autoreleasepool(|_| self.encode_now(pipeline, args, extent, dispatch))
296 }
297
298 fn encode_now(
299 &mut self,
300 pipeline: &Pipeline,
301 args: &[Arg<'_>],
302 extent: [usize; 3],
303 dispatch: impl FnOnce(&ProtocolObject<dyn MTLComputeCommandEncoder>),
304 ) -> Result<()> {
305 self.check(args)?;
306 if super::profile::enabled() {
307 self.kernels.push(match super::profile::shapes() {
308 true => format!("{} {extent:?}", pipeline.name()).into(),
309 false => pipeline.name.clone(),
310 });
311 }
312 if self.conflicts(args) {
313 self.encoder()?
314 .memoryBarrierWithScope(MTLBarrierScope::Buffers);
315 self.barriers += 1;
316 self.reads.clear();
317 self.writes.clear();
318 }
319 let encoder = self.encoder()?.clone();
320 if let Some(own) = &self.own {
321 let own = Retained::as_ptr(own) as usize;
322 for (buffer, _, write) in args.iter().filter_map(Arg::buffer) {
323 for fence in buffer.pending_fences(write) {
324 let key = Retained::as_ptr(&fence) as usize;
325 if key != own && !self.waited.contains(&key) {
326 encoder.waitForFence(&fence);
327 self.waited.push(key);
328 }
329 }
330 }
331 }
332 encoder.setComputePipelineState(pipeline.raw());
333 for (index, arg) in args.iter().enumerate() {
334 match arg {
335 Arg::Bytes(bytes) => {
336 let data = NonNull::new(bytes.as_ptr().cast_mut())
337 .unwrap_or(NonNull::dangling())
338 .cast();
339 unsafe { encoder.setBytes_length_atIndex(data, bytes.len(), index) };
343 }
344 Arg::Read(..) | Arg::Write(..) => {
345 let (buffer, offset, _) = arg.buffer().expect("buffer argument");
346 unsafe { encoder.setBuffer_offset_atIndex(buffer.raw(), offset, index) };
350 }
351 Arg::Same(earlier, offset) => {
352 let (buffer, _, _) = args[*earlier].buffer().expect("validated");
353 unsafe { encoder.setBuffer_offset_atIndex(buffer.raw(), *offset, index) };
356 }
357 }
358 }
359 for (buffer, _, write) in args.iter().filter_map(Arg::buffer) {
360 (if write {
361 &mut self.writes
362 } else {
363 &mut self.reads
364 })
365 .push(buffer.id());
366 let uses = self.uses.entry(buffer.id()).or_default();
367 uses.0 += 1;
368 uses.1 += usize::from(write);
369 let guard = buffer.guard();
370 guard.acquire(write);
371 self.guards.push((guard, write));
372 self.retained.extend(buffer.retained());
373 }
374 dispatch(&encoder);
375 Ok(())
376 }
377
378 fn check(&self, args: &[Arg<'_>]) -> Result<()> {
384 validate(args)?;
385 for (buffer, _, write) in args.iter().filter_map(Arg::buffer) {
386 if self.own.is_none() && buffer.previous_busy() {
387 return Err(GpuError::BufferInUse);
388 }
389 let (uses, writes) = buffer.unsubmitted();
390 let (own_uses, own_writes) = self.uses.get(&buffer.id()).copied().unwrap_or_default();
391 if writes > own_writes || (write && uses > own_uses) {
392 return Err(GpuError::BufferInUse);
393 }
394 }
395 Ok(())
396 }
397
398 pub fn barriers(&self) -> usize {
400 self.barriers
401 }
402
403 fn close(&mut self) {
405 autoreleasepool(|_| self.end_encoding());
406 self.device.flights().begin();
407 }
408
409 fn end_encoding(&mut self) {
410 if let Some(encoder) = self.encoder.take() {
411 if let Some((fence, _)) = &self.fence {
412 encoder.updateFence(fence);
413 }
414 if let Some(own) = &self.own {
415 encoder.updateFence(own);
416 for (guard, write) in &self.guards {
417 guard.record(own, *write);
418 }
419 }
420 encoder.endEncoding();
421 }
422 }
423
424 #[must_use = "dropping `Submitted` waits for the GPU and discards its errors; call `wait`"]
427 pub fn commit(self) -> Submitted {
428 Submitted {
429 pending: self.commit_then(|_| {}),
430 waited: false,
431 }
432 }
433
434 #[must_use = "dropping `Pending` does not wait; keep it to `wait`, or drop it explicitly"]
444 pub fn commit_then(mut self, on_complete: impl FnOnce(Result<()>) + Send + 'static) -> Pending {
445 self.close();
446 let guards = self.guards.clone();
447 let finish = Arc::new(Finish {
448 completion: Mutex::new(Some(Completion {
449 on_complete: Box::new(on_complete),
450 own: self.own.clone(),
451 kernels: std::mem::take(&mut self.kernels),
452 unfenced: self.own.is_none(),
453 guards: std::mem::take(&mut self.guards),
454 _retained: std::mem::take(&mut self.retained),
455 device: self.device.clone(),
456 })),
457 result: Mutex::new(None),
458 done: Condvar::new(),
459 });
460 let shared = finish.clone();
461 let handler = RcBlock::new(move |raw: NonNull<ProtocolObject<dyn MTLCommandBuffer>>| {
462 let completed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
465 complete(&shared, unsafe { raw.as_ref() });
467 }));
468 if completed.is_err() {
469 eprintln!("nuvec-gpu: a command buffer's completion handler panicked");
470 }
471 });
472 autoreleasepool(|_| {
473 unsafe { self.raw.addCompletedHandler(RcBlock::as_ptr(&handler)) };
477 self.raw.commit();
478 });
479 for (guard, write) in guards {
480 guard.submitted(write);
481 if self.own.is_some() {
484 guard.committed();
485 }
486 }
487 Pending {
488 raw: self.raw.clone(),
489 finish,
490 }
491 }
492
493 fn conflicts(&self, args: &[Arg<'_>]) -> bool {
496 args.iter().filter_map(Arg::buffer).any(|(b, _, write)| {
497 let id = b.id();
498 self.writes.contains(&id) || (write && self.reads.contains(&id))
499 })
500 }
501
502 fn encoder(&mut self) -> Result<&Retained<ProtocolObject<dyn MTLComputeCommandEncoder>>> {
503 if self.encoder.is_none() {
504 let encoder = self
505 .raw
506 .computeCommandEncoderWithDispatchType(MTLDispatchType::Concurrent)
507 .ok_or_else(|| GpuError::CommandBuffer("cannot create a compute encoder".into()))?;
508 if let Some((fence, true)) = &self.fence {
509 encoder.waitForFence(fence);
510 }
511 self.encoder = Some(encoder);
512 }
513 Ok(self.encoder.as_ref().expect("just created"))
514 }
515}
516
517impl Drop for CommandBuffer {
518 fn drop(&mut self) {
521 autoreleasepool(|_| {
522 if let Some(encoder) = self.encoder.take() {
523 encoder.endEncoding();
524 }
525 for (guard, write) in self.guards.drain(..) {
526 guard.committed();
527 guard.submitted(write);
528 guard.release(write, None); }
530 self.retained.clear();
531 });
532 }
533}
534
535fn size([width, height, depth]: [usize; 3]) -> MTLSize {
536 MTLSize {
537 width,
538 height,
539 depth,
540 }
541}
542
543fn validate(args: &[Arg<'_>]) -> Result<()> {
544 if args.len() > MAX_ARGS {
545 return Err(GpuError::InvalidArgument(format!(
546 "{} arguments; Metal allows {MAX_ARGS}",
547 args.len()
548 )));
549 }
550 for (index, arg) in args.iter().enumerate() {
551 match arg {
552 Arg::Bytes(bytes) if bytes.len() > MAX_INLINE_BYTES => {
553 return Err(GpuError::InvalidArgument(format!(
554 "argument {index}: {} inline bytes; Metal allows {MAX_INLINE_BYTES}",
555 bytes.len()
556 )));
557 }
558 Arg::Write(b, _) if b.is_mapped() => return Err(GpuError::ReadOnly),
559 Arg::Same(earlier, offset) => match args.get(*earlier) {
560 Some(Arg::Write(b, _)) if *offset <= b.len() => {}
561 _ => {
562 return Err(GpuError::InvalidArgument(format!(
563 "argument {index}: `Same({earlier}, {offset})` needs a written \
564 (`Arg::Write`) argument and an offset within it"
565 )));
566 }
567 },
568 Arg::Read(..) | Arg::Write(..)
569 if arg.buffer().is_some_and(|(b, offset, _)| offset > b.len()) =>
570 {
571 let (b, offset, _) = arg.buffer().expect("buffer argument");
572 return Err(GpuError::InvalidArgument(format!(
573 "argument {index}: offset {offset} is past the end of a {}-byte buffer",
574 b.len()
575 )));
576 }
577 _ => {}
578 }
579 }
580 Ok(())
581}
582
583struct Completion {
586 on_complete: Box<dyn FnOnce(Result<()>) + Send>,
587 own: Option<Fence>,
589 kernels: Vec<Arc<str>>,
591 guards: Vec<(Arc<Guard>, bool)>,
592 unfenced: bool,
594 _retained: Vec<Retained<ProtocolObject<dyn MTLBuffer>>>,
595 device: Device,
596}
597
598unsafe impl Send for Completion {}
602
603struct Finish {
606 completion: Mutex<Option<Completion>>,
607 result: Mutex<Option<Result<()>>>,
609 done: Condvar,
610}
611
612impl Finish {
613 fn wait(&self) -> Result<()> {
615 let mut result = self.result.lock().unwrap_or_else(PoisonError::into_inner);
616 loop {
617 if let Some(r) = &*result {
618 return r.clone();
619 }
620 result = self
621 .done
622 .wait(result)
623 .unwrap_or_else(PoisonError::into_inner);
624 }
625 }
626}
627
628struct Publish<'a>(&'a Finish, Option<Result<()>>);
630
631impl Drop for Publish<'_> {
632 fn drop(&mut self) {
633 let mut result = self.0.result.lock().unwrap_or_else(PoisonError::into_inner);
634 *result = self.1.take();
635 self.0.done.notify_all();
636 }
637}
638
639fn complete(finish: &Finish, raw: &ProtocolObject<dyn MTLCommandBuffer>) {
643 let Some(c) = finish
644 .completion
645 .lock()
646 .unwrap_or_else(PoisonError::into_inner)
647 .take()
648 else {
649 return;
650 };
651 let result = status(raw);
652 let _publish = Publish(finish, Some(result.clone()));
653 super::profile::record(&c.kernels, raw.GPUEndTime() - raw.GPUStartTime());
654 for (guard, write) in c.guards {
655 if c.unfenced {
656 guard.committed();
657 }
658 guard.release(write, c.own.as_ref());
659 }
660 drop(c._retained);
661 c.device.flights().end();
662 (c.on_complete)(result);
663}
664
665#[must_use = "dropping `Pending` does not wait; keep it to `wait`, or drop it explicitly"]
669pub struct Pending {
670 raw: Retained<ProtocolObject<dyn MTLCommandBuffer>>,
671 finish: Arc<Finish>,
672}
673
674impl Pending {
675 pub fn wait(&self) {
678 let _ = self.result();
679 }
680
681 fn result(&self) -> Result<()> {
683 autoreleasepool(|_| {
684 self.raw.waitUntilCompleted();
685 complete(&self.finish, &self.raw);
686 });
687 self.finish.wait()
688 }
689}
690
691impl std::fmt::Debug for Pending {
692 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
693 f.debug_struct("Pending").finish_non_exhaustive()
694 }
695}
696
697#[must_use = "dropping `Submitted` waits for the GPU and discards its errors; call `wait`"]
700pub struct Submitted {
701 pending: Pending,
702 waited: bool,
703}
704
705impl Submitted {
706 pub fn wait(mut self) -> Result<()> {
709 self.waited = true;
710 self.pending.result()
711 }
712}
713
714impl std::fmt::Debug for Submitted {
715 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
716 f.debug_struct("Submitted").finish_non_exhaustive()
717 }
718}
719
720impl Drop for Submitted {
721 fn drop(&mut self) {
722 if !self.waited {
723 self.pending.wait();
724 }
725 }
726}