Skip to main content

cubecl_common/
profile.rs

1use alloc::boxed::Box;
2use core::fmt::Display;
3
4pub use cubecl_environment::time::{Duration, Instant};
5
6use cubecl_environment::future::DynFut;
7
8/// How a benchmark's execution times are measured.
9#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
10#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
11pub enum TimingMethod {
12    /// Time measurements come from full timing of execution + sync
13    /// calls.
14    System,
15    /// Time measurements come from hardware reported timestamps
16    /// coming from a sync call.
17    Device,
18}
19
20impl Display for TimingMethod {
21    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
22        match self {
23            TimingMethod::System => f.write_str("system"),
24            TimingMethod::Device => f.write_str("device"),
25        }
26    }
27}
28
29/// Start and end point for a profile. Can be turned into a duration.
30#[derive(Debug)]
31pub struct ProfileTicks {
32    start: Instant,
33    end: Instant,
34}
35
36impl ProfileTicks {
37    /// Create a new `ProfileTicks` from a start and end time.
38    pub fn from_start_end(start: Instant, end: Instant) -> Self {
39        Self { start, end }
40    }
41
42    /// Get the duration contained in this `ProfileTicks`.
43    pub fn duration(&self) -> Duration {
44        self.end.duration_since(self.start)
45    }
46
47    /// Get the duration since the epoch start of this `ProfileTicks`.
48    pub fn start_duration_since(&self, epoch: Instant) -> Duration {
49        self.start.duration_since(epoch)
50    }
51
52    /// Get the duration since the epoch end of this `ProfileTicks`.
53    pub fn end_duration_since(&self, epoch: Instant) -> Duration {
54        self.end.duration_since(epoch)
55    }
56}
57
58/// Result from profiling between two measurements. This can either be a duration or a future that resolves to a duration.
59pub struct ProfileDuration {
60    // The future to read profiling data. For System profiling,
61    // this should be entirely synchronous.
62    future: DynFut<ProfileTicks>,
63    method: TimingMethod,
64}
65
66impl ProfileDuration {
67    /// The method used to measure the execution time.
68    pub fn timing_method(&self) -> TimingMethod {
69        self.method
70    }
71
72    /// Create a new `ProfileDuration` from a future that resolves to a duration.
73    pub fn new(future: DynFut<ProfileTicks>, method: TimingMethod) -> ProfileDuration {
74        Self { future, method }
75    }
76
77    /// Create a new `ProfileDuration` straight from a duration.
78    pub fn new_system_time(start: Instant, end: Instant) -> Self {
79        Self::new(
80            Box::pin(async move { ProfileTicks::from_start_end(start, end) }),
81            TimingMethod::System,
82        )
83    }
84
85    /// Create a new `ProfileDuration` from a future that resolves to a duration.
86    pub fn new_device_time(
87        future: impl Future<Output = ProfileTicks> + Send + 'static,
88    ) -> ProfileDuration {
89        Self::new(Box::pin(future), TimingMethod::Device)
90    }
91
92    /// Retrieve the future that resolves the profile.
93    pub fn into_future(self) -> DynFut<ProfileTicks> {
94        self.future
95    }
96
97    /// Resolve the actual duration of the profile, possibly by waiting for the future to complete.
98    pub async fn resolve(self) -> ProfileTicks {
99        self.future.await
100    }
101}