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.
26///
27/// For a fence the driver refused to create or record, its stream is synchronized
28/// on the spot. Only when that fails too does the fence hold the error.
29pub struct EventFence<A: EventApi> {
30    /// `None` when the fence was reached at creation.
31    event: Result<Option<Event<A>>, ServerError>,
32}
33
34impl<A: EventApi> EventFence<A> {
35    /// Record a fence at the current position of `stream`.
36    ///
37    /// When the driver refuses the event, `stream` is synchronized instead, which
38    /// keeps every guarantee a fence gives.
39    pub fn new(stream: A::Stream) -> Self {
40        let event = match Event::new().and_then(|event| event.record(stream).map(|()| event)) {
41            Ok(event) => Ok(Some(event)),
42            Err(err) => {
43                log::warn!(
44                    "a fence could not be recorded, synchronizing its stream instead: {err}"
45                );
46                A::stream_synchronize(stream)
47                    .map(|()| None)
48                    .map_err(Into::into)
49            }
50        };
51
52        Self { event }
53    }
54
55    /// Block until the device has reached this fence, so everything enqueued on
56    /// its stream beforehand is done.
57    ///
58    /// # Errors
59    ///
60    /// Returns a [`ServerError`] if the driver refused to create or record the
61    /// fence, or if the device was already faulty.
62    pub fn wait_sync(self) -> Result<(), ServerError> {
63        match self.event? {
64            Some(event) => Ok(event.wait()?),
65            None => Ok(()),
66        }
67    }
68
69    /// Make `stream` wait for this fence on the device, so work queued on it
70    /// afterwards runs behind the fenced stream's. Does not block the host,
71    /// unless the driver refuses the dependency: the host then waits for the
72    /// fence itself, which orders the work just the same.
73    ///
74    /// # Errors
75    ///
76    /// Returns a [`ServerError`] when the fenced stream's work failed.
77    pub fn wait_async(self, stream: A::Stream) -> Result<(), ServerError> {
78        let Some(event) = self.event? else {
79            return Ok(());
80        };
81        if let Err(err) = event.wait_async(stream) {
82            log::warn!("a stream could not be made to wait on a fence, waiting on the host: {err}");
83            event.wait()?;
84        }
85        Ok(())
86    }
87}
88
89impl<A: EventApi> core::fmt::Debug for EventFence<A> {
90    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
91        write!(f, "{}Fence", A::BACKEND)
92    }
93}
94
95/// The drop queue is where a fence pays off — a freed host buffer waits on one
96/// rather than on the server — and it is `multi_threading`-only, so the impl is
97/// too. Everything above it is the same fence either way.
98#[cfg(multi_threading)]
99impl<A: EventApi> drop_queue::Fence for EventFence<A> {
100    fn wait(self) -> Result<(), ServerError> {
101        self.wait_sync()
102    }
103}
cubecl_server/device_events/fence.rs

cubecl_server/device_events/
fence.rs
