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    /// Held (not their addresses): a dropped fence's address could be reused by a new one.
97    waited: Vec<Fence>,
98    encoder: Option<Retained<ProtocolObject<dyn MTLComputeCommandEncoder>>>,
99    /// Buffers read and written since the last barrier.
100    reads: Vec<usize>,
101    writes: Vec<usize>,
102    barriers: usize,
103    /// In-flight guards of the buffers used, and the buffers themselves (kept alive).
104    guards: Vec<(Arc<Guard>, bool)>,
105    /// Per buffer (by id): its uses in this command buffer, all and the writing ones.
106    uses: HashMap<usize, (usize, usize)>,
107    retained: Vec<Retained<ProtocolObject<dyn MTLBuffer>>>,
108    /// The kernels dispatched, while profiling.
109    kernels: Vec<Arc<str>>,
110}
111
112impl Device {
113    /// A new command buffer on this device's queue. It is *unordered*: its work does not
114    /// wait for other command buffers, so binding reused memory whose previous owner's GPU
115    /// work is still running ([`Device::new_gpu_buffer`]) is refused with
116    /// [`GpuError::BufferInUse`]; use an
117    /// [`ordered_command_buffer`](Self::ordered_command_buffer) for that.
118    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    /// A new, empty timeline.
141    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    /// A new command buffer whose work reads the outputs of earlier *ordered* command
151    /// buffers (committed on this device from any thread) safely while they are in flight:
152    /// before a dispatch reads or writes a buffer with a pending GPU write, it waits for the
153    /// writing command buffer's fence. Independent command buffers overlap on the GPU, as
154    /// in MLX.
155    ///
156    /// Reads wait for the pending writer; writes wait for every pending user (writing a
157    /// buffer that earlier work still reads). Other command buffers do not wait.
158    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    /// A new command buffer whose work runs after the work of every command buffer created
167    /// earlier on `timeline` (commit them in creation order). Like
168    /// [`command_buffer`](Self::command_buffer), it refuses reused memory whose previous
169    /// owner's GPU work is still running.
170    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
178/// Reports a finished command buffer's GPU error, if any.
179fn 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    /// Encodes `pipeline` over `threads` threads (a 1-D grid) with `args` bound to slots
192    /// `0..args.len()`. The threadgroup size is derived from the pipeline's limits.
193    ///
194    /// Arguments are validated (slot count, inline size, offsets within their buffers,
195    /// [`Arg::Same`] targets, ordering with other command buffers), but what the kernel does
196    /// with them is not: nuvec cannot see which elements a kernel touches.
197    ///
198    /// # Safety
199    ///
200    /// The kernel, run over this grid, must access each bound buffer only within its length
201    /// from the bound offset (and write only the buffers bound with [`Arg::Write`] or
202    /// [`Arg::Same`]). An out-of-bounds access reads or writes memory that other buffers,
203    /// possibly referenced by safe code, own.
204    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    /// Encodes `pipeline` over a grid of `groups` threadgroups of `threads` threads each, for
225    /// kernels that cooperate within a threadgroup (reductions, tiled matrix products).
226    ///
227    /// # Safety
228    ///
229    /// As for [`dispatch`](Self::dispatch): the kernel must access each bound buffer only
230    /// within its length from the bound offset.
231    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    /// Encodes `pipeline` over exactly `grid` threads in up to three dimensions, in
256    /// threadgroups of `group` threads (the edge threadgroups are partial).
257    ///
258    /// # Safety
259    ///
260    /// As for [`dispatch`](Self::dispatch): the kernel must access each bound buffer only
261    /// within its length from the bound offset.
262    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    /// Validates and binds `args`, inserts a barrier if needed, then lets `dispatch` encode
286    /// over `extent` threads (for profiling). Runs in an autorelease pool: Metal
287    /// autoreleases objects while encoding (the compute encoder's context among them), and
288    /// a thread other than the main one drains none of them otherwise.
289    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                    // SAFETY: `validate` checked `index < 31` and `bytes.len() <= 4096`;
345                    // Metal copies the bytes during this call, so the pointer need only be
346                    // valid for `bytes.len()` bytes now, which the slice guarantees.
347                    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                    // SAFETY: `validate` checked `index < 31` and `offset <= buffer.len()`, so
352                    // the binding stays within the buffer; the buffer is retained by this
353                    // command buffer until the GPU has finished.
354                    unsafe { encoder.setBuffer_offset_atIndex(buffer.raw(), offset, index) };
355                }
356                Arg::Same(earlier, offset) => {
357                    let (buffer, _, _) = args[*earlier].buffer().expect("validated");
358                    // SAFETY: as above: `validate` checked that `earlier` names another
359                    // buffer argument and `offset <= buffer.len()`.
360                    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    /// Validates `args`, then refuses uses that nothing would order after earlier GPU work:
384    /// - reused memory whose previous owner's work is pending, in a command buffer without
385    ///   per-buffer fences (only ordered command buffers wait for that work);
386    /// - a buffer that another command buffer, not submitted yet, writes (or, for a write
387    ///   here, uses at all): its fence orders later work only once it is submitted.
388    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    /// Memory barriers inserted so far.
404    pub fn barriers(&self) -> usize {
405        self.barriers
406    }
407
408    /// Ends encoding, updating the timeline fence after this command buffer's work.
409    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    /// Submits the work to the GPU. Dropping the returned [`Submitted`] waits for it and
430    /// discards its result; [`Submitted::wait`] reports GPU errors.
431    #[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    /// Submits the work and returns at once. `on_complete` runs exactly once with the
440    /// command buffer's result: on a Metal thread when the GPU finishes, or earlier on the
441    /// thread that calls [`Pending::wait`].
442    ///
443    /// `on_complete` must not block, in particular not on other GPU work: Metal runs a
444    /// queue's completion handlers one at a time, so waiting there for another command
445    /// buffer's completion can deadlock.
446    /// A panic in `on_complete` on Metal's thread is caught and reported on standard error,
447    /// so it never aborts the process; the result is still published to [`Pending::wait`].
448    #[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            // A panic must not unwind into Metal's Objective-C caller, which aborts the
468            // process: report it instead (the result is published either way).
469            let completed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
470                // SAFETY: Metal passes the completed command buffer, valid for this call.
471                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            // SAFETY: the block is valid, and Metal copies it before this call returns. It
479            // is called at most once, on any thread: its captures are `Send` (the buffers
480            // are only retained or released, and the guards are atomics).
481            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            // With a fence, these uses now order work committed later; without one, they
489            // count as committed once complete.
490            if self.own.is_some() {
491                guard.committed();
492            }
493        }
494        Pending {
495            raw: self.raw.clone(),
496            finish,
497        }
498    }
499
500    /// Whether a dispatch with `args` depends on dispatches since the last barrier:
501    /// read-after-write, write-after-write or write-after-read.
502    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    /// Abandoned work is never submitted; release the buffers' guards. (Metal objects are
526    /// released in an autorelease pool: their deallocation may autorelease others.)
527    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); // abandoned before its fences were recorded
536            }
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
590/// What remains to do when a command buffer committed with
591/// [`CommandBuffer::commit_then`] completes.
592struct Completion {
593    on_complete: Box<dyn FnOnce(Result<()>) + Send>,
594    /// The command buffer's fence, recorded on the buffers it uses.
595    own: Option<Fence>,
596    /// The kernels dispatched, while profiling.
597    kernels: Vec<Arc<str>>,
598    guards: Vec<(Arc<Guard>, bool)>,
599    /// The command buffer had no timeline fence: its uses count as committed only now.
600    unfenced: bool,
601    _retained: Vec<Retained<ProtocolObject<dyn MTLBuffer>>>,
602    device: Device,
603}
604
605// SAFETY: the retained Metal buffers are never accessed, only released (Metal objects may be
606// released from any thread); the guards are atomics, the callback is `Send`, and `Device` is
607// shared across threads by design (`Device::shared`).
608unsafe impl Send for Completion {}
609
610/// The completion of a command buffer committed with [`CommandBuffer::commit_then`] (run
611/// once, by whoever gets there first: a waiter or Metal's handler), and its result.
612struct Finish {
613    completion: Mutex<Option<Completion>>,
614    /// The command buffer's result, set once the completion has run.
615    result: Mutex<Option<Result<()>>>,
616    done: Condvar,
617}
618
619impl Finish {
620    /// Blocks until the completion has run (on any thread), and returns the result.
621    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
635/// Publishes the result when dropped, so that waiters wake up even if `on_complete` panics.
636struct 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
646/// Runs the completion once, whoever gets there first (a waiter or Metal's handler):
647/// releases the buffers' uses, ends the flight, runs `on_complete`, then publishes the
648/// result. This is the only place where submitted uses are released.
649fn 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/// A command buffer committed with [`CommandBuffer::commit_then`]. Dropping it does not
673/// wait; [`wait`](Self::wait) blocks until the GPU has finished and the completion has run
674/// (on this thread, if Metal's handler has not started it yet).
675#[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    /// Blocks until the GPU has finished; the completion has run when this returns (its
683    /// buffers are released and [`Device::in_flight`] no longer counts it).
684    pub fn wait(&self) {
685        let _ = self.result();
686    }
687
688    /// Waits as [`wait`](Self::wait) does and returns the command buffer's result.
689    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/// Submitted GPU work. [`wait`](Self::wait) for it to finish; dropping it also waits (and
705/// discards a GPU error).
706#[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    /// Blocks until the GPU has finished and the buffers are released, then reports any
714    /// GPU error.
715    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}