burn_backend/backend/
profile.rs

1pub use burn_std::profile::{ProfileDuration, ProfileTicks, TimingMethod};
2use burn_std::{ExecutionError, backtrace::BackTrace, profile::Instant};
3
4/// How a [profiled window](crate::Backend::profile) treats the work a backend
5/// still holds in a queue when the window closes.
6#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
7pub struct ProfileOptions {
8    flush: bool,
9}
10
11impl ProfileOptions {
12    /// Execute everything queued before closing the window, so the
13    /// measurement holds all the work the closure registered.
14    ///
15    /// A backend that batches operations (fusion) otherwise leaves the
16    /// closure's last operations queued past the window, and cuts nothing
17    /// short — its batching is what production runs. Flushing forces the batch
18    /// out at the window's end, which measures the closure's whole work at the
19    /// cost of running it as production would not.
20    pub fn flush(mut self) -> Self {
21        self.flush = true;
22        self
23    }
24
25    /// Whether the window flushes queued work before closing.
26    pub fn flushes(&self) -> bool {
27        self.flush
28    }
29}
30
31/// An open profiling window on a backend, handed out by
32/// [`profile_start`](crate::Backend::profile_start) and passed back to
33/// [`profile_end`](crate::Backend::profile_end) to close it.
34///
35/// Windows nest and overlap: a backend keys each on its token.
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
37pub struct ProfileToken {
38    /// The backend's own identifier for the window.
39    pub id: u64,
40    /// The backend's own identifier for **where** the window was opened — a
41    /// stream, a queue — carried so it can be closed there.
42    ///
43    /// Opaque to this crate, and `0` for a backend that has only one such
44    /// place. A window belongs to where it was opened, not to whoever closes
45    /// it: the split pair exists for work launched from somewhere else, so
46    /// the two calls can land on different threads, and a close that guessed
47    /// from the *calling* thread would pair a start recorded on one stream
48    /// with an end recorded on another — which reads as a plausible duration
49    /// and means nothing.
50    #[serde(default)]
51    pub opened_on: u64,
52}
53
54/// Measure `func` in wall-clock time between two syncs of `device`.
55///
56/// What a backend with no device clock reports: the default of
57/// [`Backend::profile`](crate::Backend::profile), and what a backend that
58/// forwards its windows falls back to when the backend behind it opens none.
59/// The syncs are what make the number mean something, so unlike a device
60/// window this waits, and an inner window's syncs are charged to the outer.
61pub fn profile_system_time<B: crate::Backend, O: Send + 'static>(
62    device: &B::Device,
63    func: impl FnOnce() -> O + Send,
64) -> Result<(O, ProfileDuration), crate::ExecutionError> {
65    B::sync(device)?;
66    let start = Instant::now();
67    let out = func();
68    B::sync(device)?;
69    Ok((out, ProfileDuration::new_system_time(start, Instant::now())))
70}
71
72/// A closure-bracketed window over a backend's split
73/// [`profile_start`](crate::Backend::profile_start) / [`profile_end`](crate::Backend::profile_end),
74/// falling back to [`profile_system_time`] when the backend opens none.
75///
76/// What a forwarding backend measures with, when what it forwards to may or
77/// may not have a device clock, and what a backend with a device clock
78/// measures with when the closure must not run under a hold of the device.
79pub fn profile_with_tokens<B: crate::Backend, O: Send + 'static>(
80    device: &B::Device,
81    options: ProfileOptions,
82    func: impl FnOnce() -> O + Send,
83) -> Result<(O, ProfileDuration), ExecutionError> {
84    let opened = match B::profile_start(device) {
85        Ok(Some(token)) => token,
86        Ok(None) => return profile_system_time::<B, O>(device, func),
87        // The window could not be opened — a remote device reached from a
88        // browser thread, which cannot wait on the server. `func` still runs:
89        // a caller asked for their work to be measured, and handing back an
90        // error having quietly skipped the work is the one outcome they
91        // cannot recover from. Measuring is what failed, so only the
92        // measurement is lost.
93        Err(err) => {
94            func();
95            return Err(err);
96        }
97    };
98
99    // Held so an unwinding `func` abandons the window instead of leaving it
100    // open on the server for the rest of the process.
101    let mut window = OpenWindow::<B> {
102        device,
103        token: Some(opened),
104    };
105    let out = func();
106    let token = window
107        .token
108        .take()
109        .expect("the window is held from opening until here");
110    let duration = B::profile_end(device, token, options)?;
111    Ok((out, duration))
112}
113
114/// An open profiling window, abandoned on drop unless it was taken to be
115/// closed. See [`crate::Backend::profile_abandon`].
116struct OpenWindow<'a, B: crate::Backend> {
117    device: &'a B::Device,
118    token: Option<ProfileToken>,
119}
120
121impl<B: crate::Backend> Drop for OpenWindow<'_, B> {
122    fn drop(&mut self) {
123        if let Some(token) = self.token.take() {
124            B::profile_abandon(self.device, token);
125        }
126    }
127}
128
129pub(crate) fn profile_unsupported() -> ExecutionError {
130    ExecutionError::Generic {
131        reason: alloc::string::String::from(
132            "profiling windows are not supported by this backend; use `profile`",
133        ),
134        backtrace: BackTrace::capture(),
135    }
136}
burn_backend/backend/profile.rs

burn_backend/backend/
profile.rs
