nuvec_gpu/metal/
event.rs

1//! Shared events: GPU work that waits for a value the CPU signals (MLX's `metal/event.cpp`).
2
3use objc2::rc::Retained;
4use objc2::runtime::ProtocolObject;
5use objc2_metal::{MTLDevice, MTLSharedEvent};
6
7use super::Device;
8use crate::{GpuError, Result};
9
10/// How long [`Event::wait`] spins before it blocks.
11const SPIN: std::time::Duration = std::time::Duration::from_micros(50);
12
13/// A counter command buffers wait on
14/// ([`CommandBuffer::wait_event`](super::CommandBuffer::wait_event)) and the CPU raises
15/// ([`signal`](Self::signal)). Waiters for a value run once the counter reaches it.
16pub struct Event {
17    raw: Retained<ProtocolObject<dyn MTLSharedEvent>>,
18}
19
20// SAFETY: MTLSharedEvent is thread-safe (its value is read and set from any thread, and
21// command buffers on any queue wait on it); the handle is only retained and released.
22unsafe impl Send for Event {}
23// SAFETY: as above; every method takes `&self` and is safe to call concurrently.
24unsafe impl Sync for Event {}
25
26impl Event {
27    /// A new event on `device`, at value 0.
28    pub fn new(device: &Device) -> Result<Self> {
29        let raw = device
30            .raw()
31            .newSharedEvent()
32            .ok_or_else(|| GpuError::CommandBuffer("cannot create a shared event".into()))?;
33        Ok(Self { raw })
34    }
35
36    /// Raises the counter to `value`, releasing GPU work that waits for it.
37    pub fn signal(&self, value: u64) {
38        self.raw.setSignaledValue(value);
39    }
40
41    /// Waits on the CPU until the counter reaches `value`, for at most `timeout`; whether it
42    /// did. Spins briefly first: a collective's input is usually ready within microseconds.
43    pub fn wait(&self, value: u64, timeout: std::time::Duration) -> bool {
44        let spin = std::time::Instant::now();
45        while spin.elapsed() < SPIN.min(timeout) {
46            if self.value() >= value {
47                return true;
48            }
49            std::hint::spin_loop();
50        }
51        let ms = timeout.saturating_sub(SPIN).as_millis().max(1) as u64;
52        self.raw.waitUntilSignaledValue_timeoutMS(value, ms)
53    }
54
55    /// The counter.
56    pub fn value(&self) -> u64 {
57        self.raw.signaledValue()
58    }
59
60    pub(crate) fn raw(&self) -> &ProtocolObject<dyn MTLSharedEvent> {
61        &self.raw
62    }
63}
64
65impl std::fmt::Debug for Event {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        f.debug_struct("Event")
68            .field("value", &self.value())
69            .finish()
70    }
71}
nuvec_gpu/metal/event.rs

nuvec_gpu/metal/
event.rs
