cubecl_server/device_events/
fence.rs

1use crate::device_events::{Event, EventApi};
2#[cfg(multi_threading)]
3use crate::memory_management::drop_queue;
4use crate::server::ServerError;
5
6/// An event recorded on a stream and handed out, so a caller can wait for that
7/// stream's work from outside the server.
8///
9/// The server sits behind a mutex or a channel, so a synchronize that blocked
10/// while holding it would stall every other logical stream too. Recording an
11/// event costs the host nothing: the server records one and returns, and
12/// whoever holds the fence waits on its own time.
13///
14/// Named for the trait it is not: `memory_management::drop_queue::Fence` is the
15/// contract, and this is the implementation of it that a device event gives
16/// you. Backends alias it back to `Fence` for their own call sites. Named in
17/// prose rather than linked because that trait is `multi_threading`-only while
18/// this type is not — the fence itself is two driver calls and wants no
19/// threads.
20///
21/// Its event is created and destroyed outright rather than recycled through the
22/// pool an [`EventProfiler`](super::EventProfiler) keeps: a fence is raised from
23/// `StreamBackend::flush`, which is handed a stream and nothing else, so there
24/// is no pool in reach without threading one through that trait. Pool these too
25/// once something else needs the same argument.
26pub struct EventFence<A: EventApi> {
27    event: Event<A>,
28}
29
30impl<A: EventApi> EventFence<A> {
31    /// Record a fence at the current position of `stream`.
32    ///
33    /// # Panics
34    ///
35    /// A fence that never recorded cannot be waited on, and every caller takes
36    /// one by value expecting to be able to. There is no useful weaker answer
37    /// than failing here.
38    pub fn new(stream: A::Stream) -> Self {
39        let event = Event::new().expect("the fence needs an event");
40        event
41            .record(stream)
42            .expect("the fence needs its event recorded");
43
44        Self { event }
45    }
46
47    /// Block until the device has reached this fence, so everything enqueued on
48    /// its stream beforehand is done.
49    ///
50    /// # Errors
51    ///
52    /// The fault the wait reveals, when the stream itself failed.
53    pub fn wait_sync(self) -> Result<(), ServerError> {
54        Ok(self.event.wait()?)
55    }
56
57    /// Make `stream` wait for this fence on the device, so work queued on it
58    /// afterwards runs behind the fenced stream's. Does not block the host.
59    ///
60    /// # Panics
61    ///
62    /// A refused dependency would let `stream` run ahead of work it must
63    /// follow, which is a wrong answer rather than a slow one.
64    pub fn wait_async(self, stream: A::Stream) {
65        self.event
66            .wait_async(stream)
67            .expect("the stream has to wait on the fence");
68    }
69}
70
71impl<A: EventApi> core::fmt::Debug for EventFence<A> {
72    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
73        write!(f, "{}Fence", A::BACKEND)
74    }
75}
76
77/// The drop queue is where a fence pays off — a freed host buffer waits on one
78/// rather than on the server — and it is `multi_threading`-only, so the impl is
79/// too. Everything above it is the same fence either way.
80#[cfg(multi_threading)]
81impl<A: EventApi> drop_queue::Fence for EventFence<A> {
82    fn wait(self) -> Result<(), ServerError> {
83        self.wait_sync()
84    }
85}
cubecl_server/device_events/fence.rs

cubecl_server/device_events/
fence.rs
