cubecl_runtime/logging/
profiling.rs

1use alloc::{
2    format,
3    string::{String, ToString},
4    vec::Vec,
5};
6use core::fmt::Display;
7use hashbrown::HashMap;
8
9#[derive(Debug, Default)]
10pub(crate) struct Profiled {
11    durations: HashMap<String, ProfileItem>,
12}
13
14#[derive(Debug, Default, Clone)]
15pub(crate) struct ProfileItem {
16    total_duration: core::time::Duration,
17    num_computed: usize,
18}
19
20#[derive(Debug, Copy, Clone)]
21/// Control the amount of info being display when profiling.
22pub enum ProfileLevel {
23    /// Provide only the summary information about kernels being run.
24    Basic,
25    /// Provide the summary information about kernels being run with their trace.
26    Medium,
27    /// Provide more information about kernels being run.
28    Full,
29    /// Only the execution are logged.
30    ExecutionOnly,
31}
32
33impl Profiled {
34    /// If some computation was profiled.
35    pub fn is_empty(&self) -> bool {
36        self.durations.is_empty()
37    }
38    pub fn update(&mut self, name: &String, duration: core::time::Duration) {
39        let name = if name.contains("\n") {
40            name.split("\n").next().unwrap()
41        } else {
42            name
43        };
44        if let Some(item) = self.durations.get_mut(name) {
45            item.update(duration);
46        } else {
47            self.durations.insert(
48                name.to_string(),
49                ProfileItem {
50                    total_duration: duration,
51                    num_computed: 1,
52                },
53            );
54        }
55    }
56}
57
58impl Display for Profiled {
59    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
60        let header_name = "Name";
61        let header_num_computed = "Num Computed";
62        let header_duration = "Duration";
63        let header_ratio = "Ratio";
64
65        let mut ratio_len = header_ratio.len();
66        let mut name_len = header_name.len();
67        let mut num_computed_len = header_num_computed.len();
68        let mut duration_len = header_duration.len();
69
70        let mut total_duration = core::time::Duration::from_secs(0);
71        let mut total_computed = 0;
72
73        let mut items: Vec<(String, String, String, core::time::Duration)> = self
74            .durations
75            .iter()
76            .map(|(key, item)| {
77                let name = key.clone();
78                let num_computed = format!("{}", item.num_computed);
79                let duration = format!("{:?}", item.total_duration);
80
81                name_len = usize::max(name_len, name.len());
82                num_computed_len = usize::max(num_computed_len, num_computed.len());
83                duration_len = usize::max(duration_len, duration.len());
84
85                total_duration += item.total_duration;
86                total_computed += item.num_computed;
87
88                (name, num_computed, duration, item.total_duration)
89            })
90            .collect();
91
92        let total_duration_fmt = format!("{total_duration:?}");
93        let total_compute_fmt = format!("{total_computed:?}");
94        let total_ratio_fmt = "100 %";
95
96        duration_len = usize::max(duration_len, total_duration_fmt.len());
97        num_computed_len = usize::max(num_computed_len, total_compute_fmt.len());
98        ratio_len = usize::max(ratio_len, total_ratio_fmt.len());
99
100        let line_length = name_len + duration_len + num_computed_len + ratio_len + 11;
101
102        let write_line = |char: &str, f: &mut core::fmt::Formatter<'_>| {
103            writeln!(f, "|{}| ", char.repeat(line_length))
104        };
105        items.sort_by(|(_, _, _, a), (_, _, _, b)| b.cmp(a));
106
107        write_line("⎺", f)?;
108
109        writeln!(
110            f,
111            "| {header_name:<name_len$} | {header_duration:<duration_len$} | {header_num_computed:<num_computed_len$} | {header_ratio:<ratio_len$} |",
112        )?;
113
114        write_line("⎼", f)?;
115
116        for (name, num_computed, duration, num) in items {
117            let ratio = (100 * num.as_micros()) / total_duration.as_micros();
118
119            writeln!(
120                f,
121                "| {:<width_name$} | {:<width_duration$} | {:<width_num_computed$} | {:<width_ratio$} |",
122                name,
123                duration,
124                num_computed,
125                format!("{} %", ratio),
126                width_name = name_len,
127                width_duration = duration_len,
128                width_num_computed = num_computed_len,
129                width_ratio = ratio_len,
130            )?;
131        }
132
133        write_line("⎼", f)?;
134
135        writeln!(
136            f,
137            "| {:<width_name$} | {:<width_duration$} | {:<width_num_computed$} | {:<width_ratio$} |",
138            "Total",
139            total_duration_fmt,
140            total_compute_fmt,
141            total_ratio_fmt,
142            width_name = name_len,
143            width_duration = duration_len,
144            width_num_computed = num_computed_len,
145            width_ratio = ratio_len,
146        )?;
147
148        write_line("⎯", f)?;
149
150        Ok(())
151    }
152}
153
154impl ProfileItem {
155    pub fn update(&mut self, duration: core::time::Duration) {
156        self.total_duration += duration;
157        self.num_computed += 1;
158    }
159}
cubecl_runtime/logging/profiling.rs

cubecl_runtime/logging/
profiling.rs
