cubecl_runtime/throughput/
base.rs

1use alloc::{format, string::String};
2
3use cubecl_ir::ElemType;
4
5use crate::throughput::ComputeCmmaConfig;
6
7/// Represents the mode of a throughput computation.
8#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
9#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
10pub enum ThroughputMode {
11    /// Compute direct calculation without special hardware acceleration.
12    ComputeDirect,
13    /// Compute cmma calculation with CMMA hardware acceleration.
14    ComputeCmma(ComputeCmmaConfig),
15    /// Memory input reads and output writes.
16    Memory,
17}
18
19/// Represents a key/configuration used to identify the throughput of a computation.
20#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
21#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
22pub struct ThroughputKey {
23    /// The mode of the throughput computation.
24    pub mode: ThroughputMode,
25    /// The data type of the computation.
26    pub dtype: ElemType,
27}
28
29/// Represents the throughput of a computation, including the number of operations and the duration.
30#[derive(Eq, PartialEq, Clone, Copy, Debug)]
31#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
32pub struct ThroughputValue {
33    /// The number of operations performed depending of the mode during the computation.
34    pub ops_count: usize,
35    /// The duration of the computation.
36    pub duration: core::time::Duration,
37}
38
39impl ThroughputValue {
40    /// A zero-initialized throughput value, representing no operations or duration.
41    pub const ZERO: Self = Self {
42        ops_count: 0,
43        duration: core::time::Duration::ZERO,
44    };
45
46    /// Returns the operations per second.
47    pub fn ops_per_s(&self) -> f64 {
48        if self.duration.is_zero() {
49            return f64::NAN;
50        }
51        self.ops_count as f64 / self.duration.as_secs_f64()
52    }
53
54    /// Returns the bytes per second.
55    pub fn bytes_per_s(&self, key: &ThroughputKey) -> f64 {
56        if self.duration.is_zero() {
57            return f64::NAN;
58        }
59        (self.ops_count * key.dtype.size()) as f64 / self.duration.as_secs_f64()
60    }
61
62    /// Formats the throughput value as a clean human-readable string.
63    pub fn format(&self, key: &ThroughputKey) -> String {
64        let unit = match key.mode {
65            ThroughputMode::ComputeDirect | ThroughputMode::ComputeCmma(_) => "OPS",
66            ThroughputMode::Memory => "bytes",
67        };
68
69        let mut val_per_s = match key.mode {
70            ThroughputMode::ComputeDirect | ThroughputMode::ComputeCmma(_) => self.ops_per_s(),
71            ThroughputMode::Memory => self.bytes_per_s(key),
72        };
73
74        if val_per_s.is_nan() {
75            return String::from("N/A");
76        }
77
78        let suffixes = ["", "K", "M", "G", "T", "P", "E", "Z", "Y", "R", "Q"];
79        let mut suffix_idx = 0;
80
81        for _ in 0..suffixes.len() - 1 {
82            if val_per_s < 1000.0 {
83                break;
84            }
85            val_per_s /= 1000.0;
86            suffix_idx += 1;
87        }
88
89        format!("{val_per_s:.4} {}{unit}/s", suffixes[suffix_idx])
90    }
91}
cubecl_runtime/throughput/base.rs

cubecl_runtime/throughput/
base.rs
