1use alloc::boxed::Box;
2use core::fmt::Display;
3
4pub use cubecl_environment::time::{Duration, Instant};
5
6use cubecl_environment::future::DynFut;
7
8#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
10#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
11pub enum TimingMethod {
12 System,
15 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#[derive(Debug)]
31pub struct ProfileTicks {
32 start: Instant,
33 end: Instant,
34}
35
36impl ProfileTicks {
37 pub fn from_start_end(start: Instant, end: Instant) -> Self {
39 Self { start, end }
40 }
41
42 pub fn duration(&self) -> Duration {
44 self.end.duration_since(self.start)
45 }
46
47 pub fn start_duration_since(&self, epoch: Instant) -> Duration {
49 self.start.duration_since(epoch)
50 }
51
52 pub fn end_duration_since(&self, epoch: Instant) -> Duration {
54 self.end.duration_since(epoch)
55 }
56}
57
58pub struct ProfileDuration {
60 future: DynFut<ProfileTicks>,
63 method: TimingMethod,
64}
65
66impl ProfileDuration {
67 pub fn timing_method(&self) -> TimingMethod {
69 self.method
70 }
71
72 pub fn new(future: DynFut<ProfileTicks>, method: TimingMethod) -> ProfileDuration {
74 Self { future, method }
75 }
76
77 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 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 pub fn into_future(self) -> DynFut<ProfileTicks> {
94 self.future
95 }
96
97 pub async fn resolve(self) -> ProfileTicks {
99 self.future.await
100 }
101}