cubecl_runtime/
timestamp_profiler.rs

1use cubecl_common::profile::{Instant, ProfileDuration};
2use cubecl_environment::backtrace::BackTrace;
3use cubecl_environment::collections::HashMap;
4
5use crate::server::{ProfileError, ProfilingToken};
6
7#[derive(Default, Debug)]
8/// A simple struct to keep track of timestamps for kernel execution.
9/// This should be used for servers that do not have native device profiling.
10pub struct TimestampProfiler {
11    state: HashMap<ProfilingToken, State>,
12    counter: u64,
13}
14
15#[derive(Debug)]
16enum State {
17    Start(Instant),
18    Error(ProfileError),
19}
20
21impl TimestampProfiler {
22    /// If there is some profiling registered.
23    pub fn is_empty(&self) -> bool {
24        self.state.is_empty()
25    }
26    /// Start measuring
27    pub fn start(&mut self) -> ProfilingToken {
28        let token = ProfilingToken { id: self.counter };
29        self.counter += 1;
30        self.state.insert(token, State::Start(Instant::now()));
31        token
32    }
33
34    /// Stop measuring
35    pub fn stop(&mut self, token: ProfilingToken) -> Result<ProfileDuration, ProfileError> {
36        let state = self.state.remove(&token);
37        let start = match state {
38            Some(val) => match val {
39                State::Start(instant) => instant,
40                State::Error(profile_error) => return Err(profile_error),
41            },
42            None => {
43                return Err(ProfileError::NotRegistered {
44                    backtrace: BackTrace::capture(),
45                });
46            }
47        };
48        Ok(ProfileDuration::new_system_time(start, Instant::now()))
49    }
50
51    /// Register an error during profiling.
52    pub fn error(&mut self, error: ProfileError) {
53        self.state
54            .iter_mut()
55            .for_each(|(_, state)| *state = State::Error(error.clone()));
56    }
57}
cubecl_runtime/timestamp_profiler.rs

cubecl_runtime/
timestamp_profiler.rs
