Skip to main content

nuvec_gpu/metal/
command.rs

1//! Command encoding with minimal barriers, and submission.
2
3use 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
19/// Metal's limit on buffer argument slots per kernel.
20const MAX_ARGS: usize = 31;
21/// Metal's limit for inline bytes (`setBytes`).
22const MAX_INLINE_BYTES: usize = 4096;
23
24/// A kernel argument, bound to the slot of its position in the argument list.
25#[derive(Debug)]
26pub enum Arg<'a> {
27    /// A buffer the kernel only reads, from byte `offset` on.
28    Read(&'a Buffer, usize),
29    /// A buffer the kernel writes (and may read), from byte `offset` on.
30    Write(&'a mut Buffer, usize),
31    /// Small constant data copied into the command stream (at most 4 KiB).
32    Bytes(&'a [u8]),
33    /// The buffer of the argument at this index bound again, from byte `offset` on: an
34    /// in-place kernel reading and writing one buffer through two slots. That argument must
35    /// be an [`Arg::Write`] (whose hazards cover both slots); anything else is an
36    /// [`GpuError::InvalidArgument`].
37    Same(usize, usize),
38}
39
40impl<'a> Arg<'a> {
41    /// The whole buffer, read-only.
42    pub fn read(buffer: &'a Buffer) -> Self {
43        Self::Read(buffer, 0)
44    }
45
46    /// The whole buffer, written.
47    pub fn write(buffer: &'a mut Buffer) -> Self {
48        Self::Write(buffer, 0)
49    }
50
51    /// Inline constant bytes.
52    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
67/// Orders command buffers that are committed without waiting: each command buffer created
68/// on a timeline starts its work only after the previous one's work has finished. Buffers
69/// are untracked, so consecutive command buffers would otherwise overlap on the GPU (as
70/// [`Device::ordered_command_buffer`]s do when they are independent).
71///
72/// Implemented with one fence (as in MLX): a command buffer's compute pass waits for the
73/// fence and updates it at the end. Command buffers without work leave it untouched, so
74/// empty or abandoned ones need no special handling.
75pub 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
86/// Work being recorded for the GPU: one concurrent compute encoder, with memory barriers
87/// inserted only where a dispatch depends on an earlier one.
88pub struct CommandBuffer {
89    device: Device,
90    raw: Retained<ProtocolObject<dyn MTLCommandBuffer>>,
91    /// The timeline fence, and whether the compute pass must wait for it.
92    fence: Option<(Fence, bool)>,
93    /// For ordered command buffers: this command buffer's fence (updated at the end of its
94    /// work and recorded on the buffers it writes), and the fences already waited for.
95    own: Option<Fence>,
96    waited: Vec<usize>,
97    encoder: Option<Retained<ProtocolObject<dyn MTLComputeCommandEncoder>>>,
98    /// Buffers read and written since the last barrier.
99    reads: Vec<usize>,
100    writes: Vec<usize>,
101    barriers: usize,
102    /// In-flight guards of the buffers used, and the buffers themselves (kept alive).
103    guards: Vec<(Arc<Guard>, bool)>,
104    /// Per buffer (by id): its uses in this command buffer, all and the writing ones.
105    uses: HashMap<usize, (usize, usize)>,
106    retained: Vec<Retained<ProtocolObject<dyn MTLBuffer>>>,
107    /// The kernels dispatched, while profiling.
108    kernels: Vec<Arc<str>>,
109}
110
111impl Device {
112    /// A new command buffer on this device's queue. It is *unordered*: its work does not
113    /// wait for other command buffers, so binding reused memory whose previous owner's GPU
114    /// work is still running ([`Device::new_gpu_buffer`]) is refused with
115    /// [`GpuError::BufferInUse`]; use an
116    /// [`ordered_command_buffer`](Self::ordered_command_buffer) for that.
117    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    /// A new, empty timeline.
140    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    /// A new command buffer whose work reads the outputs of earlier *ordered* command
150    /// buffers (committed on this device from any thread) safely while they are in flight:
151    /// before a dispatch reads or writes a buffer with a pending GPU write, it waits for the
152    /// writing command buffer's fence. Independent command buffers overlap on the GPU, as
153    /// in MLX.
154    ///
155    /// Reads wait for the pending writer; writes wait for every pending user (writing a
156    /// buffer that earlier work still reads). Other command buffers do not wait.
157    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    /// A new command buffer whose work runs after the work of every command buffer created
166    /// earlier on `timeline` (commit them in creation order). Like
167    /// [`command_buffer`](Self::command_buffer), it refuses reused memory whose previous
168    /// owner's GPU work is still running.
169    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
177/// Reports a finished command buffer's GPU error, if any.
178fn 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    /// Encodes `pipeline` over `threads` threads (a 1-D grid) with `args` bound to slots
191    /// `0..args.len()`. The threadgroup size is derived from the pipeline's limits.
192    ///
193    /// Arguments are validated (slot count, inline size, offsets within their buffers,
194    /// [`Arg::Same`] targets, ordering with other command buffers), but what the kernel does
195    /// with them is not: nuvec cannot see which elements a kernel touches.
196    ///
197    /// # Safety
198    ///
199    /// The kernel, run over this grid, must access each bound buffer only within its length
200    /// from the bound offset (and write only the buffers bound with [`Arg::Write`] or
201    /// [`Arg::Same`]). An out-of-bounds access reads or writes memory that other buffers,
202    /// possibly referenced by safe code, own.
203    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    /// Encodes `pipeline` over a grid of `groups` threadgroups of `threads` threads each, for
224    /// kernels that cooperate within a threadgroup (reductions, tiled matrix products).
225    ///
226    /// # Safety
227    ///
228    /// As for [`dispatch`](Self::dispatch): the kernel must access each bound buffer only
229    /// within its length from the bound offset.
230    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    /// Encodes `pipeline` over exactly `grid` threads in up to three dimensions, in
255    /// threadgroups of `group` threads (the edge threadgroups are partial).
256    ///
257    /// # Safety
258    ///
259    /// As for [`dispatch`](Self::dispatch): the kernel must access each bound buffer only
260    /// within its length from the bound offset.
261    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    /// Validates and binds `args`, inserts a barrier if needed, then lets `dispatch` encode
285    /// over `extent` threads (for profiling). Runs in an autorelease pool: Metal
286    /// autoreleases objects while encoding (the compute encoder's context among them), and
287    /// a thread other than the main one drains none of them otherwise.
288    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                    // SAFETY: `validate` checked `index < 31` and `bytes.len() <= 4096`;
340                    // Metal copies the bytes during this call, so the pointer need only be
341                    // valid for `bytes.len()` bytes now, which the slice guarantees.
342                    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                    // SAFETY: `validate` checked `index < 31` and `offset <= buffer.len()`, so
347                    // the binding stays within the buffer; the buffer is retained by this
348                    // command buffer until the GPU has finished.
349                    unsafe { encoder.setBuffer_offset_atIndex(buffer.raw(), offset, index) };
350                }
351                Arg::Same(earlier, offset) => {
352                    let (buffer, _, _) = args[*earlier].buffer().expect("validated");
353                    // SAFETY: as above: `validate` checked that `earlier` names another
354                    // buffer argument and `offset <= buffer.len()`.
355                    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    /// Validates `args`, then refuses uses that nothing would order after earlier GPU work:
379    /// - reused memory whose previous owner's work is pending, in a command buffer without
380    ///   per-buffer fences (only ordered command buffers wait for that work);
381    /// - a buffer that another command buffer, not submitted yet, writes (or, for a write
382    ///   here, uses at all): its fence orders later work only once it is submitted.
383    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    /// Memory barriers inserted so far.
399    pub fn barriers(&self) -> usize {
400        self.barriers
401    }
402
403    /// Ends encoding, updating the timeline fence after this command buffer's work.
404    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    /// Submits the work to the GPU. Dropping the returned [`Submitted`] waits for it and
425    /// discards its result; [`Submitted::wait`] reports GPU errors.
426    #[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    /// Submits the work and returns at once. `on_complete` runs exactly once with the
435    /// command buffer's result: on a Metal thread when the GPU finishes, or earlier on the
436    /// thread that calls [`Pending::wait`].
437    ///
438    /// `on_complete` must not block, in particular not on other GPU work: Metal runs a
439    /// queue's completion handlers one at a time, so waiting there for another command
440    /// buffer's completion can deadlock.
441    /// A panic in `on_complete` on Metal's thread is caught and reported on standard error,
442    /// so it never aborts the process; the result is still published to [`Pending::wait`].
443    #[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            // A panic must not unwind into Metal's Objective-C caller, which aborts the
463            // process: report it instead (the result is published either way).
464            let completed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
465                // SAFETY: Metal passes the completed command buffer, valid for this call.
466                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            // SAFETY: the block is valid, and Metal copies it before this call returns. It
474            // is called at most once, on any thread: its captures are `Send` (the buffers
475            // are only retained or released, and the guards are atomics).
476            unsafe { self.raw.addCompletedHandler(RcBlock::as_ptr(&handler)) };
477            self.raw.commit();
478        });
479        for (guard, write) in guards {
480            guard.submitted(write);
481            // With a fence, these uses now order work committed later; without one, they
482            // count as committed once complete.
483            if self.own.is_some() {
484                guard.committed();
485            }
486        }
487        Pending {
488            raw: self.raw.clone(),
489            finish,
490        }
491    }
492
493    /// Whether a dispatch with `args` depends on dispatches since the last barrier:
494    /// read-after-write, write-after-write or write-after-read.
495    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    /// Abandoned work is never submitted; release the buffers' guards. (Metal objects are
519    /// released in an autorelease pool: their deallocation may autorelease others.)
520    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); // abandoned before its fences were recorded
529            }
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
583/// What remains to do when a command buffer committed with
584/// [`CommandBuffer::commit_then`] completes.
585struct Completion {
586    on_complete: Box<dyn FnOnce(Result<()>) + Send>,
587    /// The command buffer's fence, recorded on the buffers it uses.
588    own: Option<Fence>,
589    /// The kernels dispatched, while profiling.
590    kernels: Vec<Arc<str>>,
591    guards: Vec<(Arc<Guard>, bool)>,
592    /// The command buffer had no timeline fence: its uses count as committed only now.
593    unfenced: bool,
594    _retained: Vec<Retained<ProtocolObject<dyn MTLBuffer>>>,
595    device: Device,
596}
597
598// SAFETY: the retained Metal buffers are never accessed, only released (Metal objects may be
599// released from any thread); the guards are atomics, the callback is `Send`, and `Device` is
600// shared across threads by design (`Device::shared`).
601unsafe impl Send for Completion {}
602
603/// The completion of a command buffer committed with [`CommandBuffer::commit_then`] (run
604/// once, by whoever gets there first: a waiter or Metal's handler), and its result.
605struct Finish {
606    completion: Mutex<Option<Completion>>,
607    /// The command buffer's result, set once the completion has run.
608    result: Mutex<Option<Result<()>>>,
609    done: Condvar,
610}
611
612impl Finish {
613    /// Blocks until the completion has run (on any thread), and returns the result.
614    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
628/// Publishes the result when dropped, so that waiters wake up even if `on_complete` panics.
629struct 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
639/// Runs the completion once, whoever gets there first (a waiter or Metal's handler):
640/// releases the buffers' uses, ends the flight, runs `on_complete`, then publishes the
641/// result. This is the only place where submitted uses are released.
642fn 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/// A command buffer committed with [`CommandBuffer::commit_then`]. Dropping it does not
666/// wait; [`wait`](Self::wait) blocks until the GPU has finished and the completion has run
667/// (on this thread, if Metal's handler has not started it yet).
668#[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    /// Blocks until the GPU has finished; the completion has run when this returns (its
676    /// buffers are released and [`Device::in_flight`] no longer counts it).
677    pub fn wait(&self) {
678        let _ = self.result();
679    }
680
681    /// Waits as [`wait`](Self::wait) does and returns the command buffer's result.
682    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/// Submitted GPU work. [`wait`](Self::wait) for it to finish; dropping it also waits (and
698/// discards a GPU error).
699#[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    /// Blocks until the GPU has finished and the buffers are released, then reports any
707    /// GPU error.
708    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}