cubecl_server/device_events/
event.rs

1use alloc::vec::Vec;
2use core::ops::Deref;
3
4use cubecl_environment::sync::{Arc, Mutex};
5
6use crate::device_events::EventApi;
7use crate::driver::DriverError;
8
9/// A device event that owns itself: created with timing enabled, released when
10/// dropped.
11pub struct Event<A: EventApi> {
12    sys: A::Event,
13}
14
15impl<A: EventApi> Event<A> {
16    /// Create an event, with timing enabled.
17    ///
18    /// # Errors
19    ///
20    /// [`DriverError`] when the driver refuses it.
21    pub fn new() -> Result<Self, DriverError> {
22        Ok(Self {
23            sys: A::event_create()?,
24        })
25    }
26
27    /// Enqueue the timestamp on `stream`, returning once it is queued.
28    ///
29    /// # Errors
30    ///
31    /// [`DriverError`] when the stream will not take the event.
32    pub fn record(&self, stream: A::Stream) -> Result<(), DriverError> {
33        A::event_record(&self.sys, stream)
34    }
35
36    /// Block until the device has reached this event.
37    ///
38    /// # Errors
39    ///
40    /// [`DriverError`], the fault the wait revealed on the stream.
41    pub fn wait(&self) -> Result<(), DriverError> {
42        A::event_wait(&self.sys)
43    }
44
45    /// Make `stream` wait for this event on the device, without blocking here.
46    ///
47    /// # Errors
48    ///
49    /// [`DriverError`] when the driver refuses the dependency.
50    pub fn wait_async(&self, stream: A::Stream) -> Result<(), DriverError> {
51        A::stream_wait_event(stream, &self.sys)
52    }
53
54    /// Device time from `self` to `other`. Both must have been reached.
55    ///
56    /// # Errors
57    ///
58    /// [`DriverError`], including "not ready" when either event has not been
59    /// reached.
60    pub fn elapsed(&self, other: &Self) -> Result<cubecl_common::profile::Duration, DriverError> {
61        A::event_elapsed(&self.sys, &other.sys)
62    }
63}
64
65impl<A: EventApi> Drop for Event<A> {
66    fn drop(&mut self) {
67        if let Err(err) = A::event_destroy(&mut self.sys) {
68            log::warn!("Failed to release a {} event: {err}", A::BACKEND);
69        }
70    }
71}
72
73impl<A: EventApi> core::fmt::Debug for Event<A> {
74    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
75        write!(f, "{}Event", A::BACKEND)
76    }
77}
78
79/// Recycled events, shared between whoever holds the pool and every event it
80/// has handed out.
81///
82/// A profiling window's events outlive the profiler's own map: they are read by
83/// the future that closed the window, which may be awaited much later and on
84/// another thread. So the pool is shared rather than owned, and an event
85/// returns to it when the [`Pooled`] holding it is dropped.
86pub(crate) struct EventPool<A: EventApi> {
87    free: Arc<Mutex<Vec<Event<A>>>>,
88}
89
90impl<A: EventApi> EventPool<A> {
91    /// Take an event from the pool, creating one if it is empty.
92    ///
93    /// The event returns to the pool when the [`Pooled`] is dropped, on every
94    /// path — including the ones that never record it.
95    ///
96    /// # Errors
97    ///
98    /// [`DriverError`] when the pool was empty and the driver refused a new
99    /// event.
100    pub(crate) fn acquire(&self) -> Result<Pooled<A>, DriverError> {
101        let event = match self.free.lock().pop() {
102            Some(event) => event,
103            None => Event::new()?,
104        };
105
106        Ok(Pooled {
107            event: Some(event),
108            pool: self.clone(),
109        })
110    }
111
112    fn release(&self, event: Event<A>) {
113        self.free.lock().push(event);
114    }
115}
116
117impl<A: EventApi> Clone for EventPool<A> {
118    fn clone(&self) -> Self {
119        Self {
120            free: self.free.clone(),
121        }
122    }
123}
124
125impl<A: EventApi> Default for EventPool<A> {
126    fn default() -> Self {
127        Self {
128            free: Arc::new(Mutex::new(Vec::new())),
129        }
130    }
131}
132
133impl<A: EventApi> core::fmt::Debug for EventPool<A> {
134    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
135        write!(f, "EventPool<{}>({})", A::BACKEND, self.free.lock().len())
136    }
137}
138
139/// An [`Event`] borrowed from an [`EventPool`], returned to it when dropped.
140///
141/// Acquire and release have to happen in pairs, and a window is abandoned,
142/// errored out or dropped mid-setup on paths that have no obvious release
143/// site. Tying the release to the drop is what removes the step there is to
144/// forget.
145pub(crate) struct Pooled<A: EventApi> {
146    /// `Some` for the whole life of the value; emptied only by `Drop`, which
147    /// is how the event moves back into the pool without cloning it.
148    event: Option<Event<A>>,
149    pool: EventPool<A>,
150}
151
152impl<A: EventApi> Deref for Pooled<A> {
153    type Target = Event<A>;
154
155    fn deref(&self) -> &Self::Target {
156        self.event.as_ref().expect("emptied only by Drop")
157    }
158}
159
160impl<A: EventApi> Drop for Pooled<A> {
161    fn drop(&mut self) {
162        if let Some(event) = self.event.take() {
163            self.pool.release(event);
164        }
165    }
166}
167
168impl<A: EventApi> core::fmt::Debug for Pooled<A> {
169    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
170        write!(f, "Pooled<{}>", A::BACKEND)
171    }
172}
cubecl_server/device_events/event.rs

cubecl_server/device_events/
event.rs
