cubecl_runtime/throughput/
base.rs

1use crate::throughput::{CmmaDims, ComputeCmmaConfig};
2use alloc::{format, string::String};
3use core::time::Duration;
4use cubecl_ir::{ElemType, FloatKind};
5
6/// Represents the mode of a throughput computation.
7#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
8#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
9pub enum ThroughputMode {
10    /// Compute direct calculation without special hardware acceleration.
11    ComputeDirect {
12        /// The data type of the computation.
13        dtype: ElemType,
14    },
15    /// Compute cmma calculation with CMMA hardware acceleration.
16    ComputeCmma {
17        /// The data type of the computation.
18        dtype: ElemType,
19        /// The configuration of the CMMA operation.
20        config: ComputeCmmaConfig,
21    },
22    /// Memory input reads and output writes.
23    Memory,
24    /// Launch overhead measurement.
25    Launch,
26}
27
28/// Represents a key/configuration used to identify the throughput of a computation.
29#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
30#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
31// Reject cached entries from an older key layout instead of silently ignoring their extra fields.
32#[cfg_attr(std_io, serde(deny_unknown_fields))]
33pub struct ThroughputKey {
34    /// The mode of the throughput computation.
35    pub mode: ThroughputMode,
36}
37
38impl ThroughputKey {
39    /// Returns the data type of the computation.
40    pub fn dtype(&self) -> ElemType {
41        match self.mode {
42            ThroughputMode::ComputeDirect { dtype } => dtype,
43            ThroughputMode::ComputeCmma { dtype, .. } => dtype,
44            // For memory and launch throughput, we use a default element type (F32).
45            ThroughputMode::Memory | ThroughputMode::Launch => ElemType::Float(FloatKind::F32),
46        }
47    }
48}
49
50/// Represents the throughput of a computation, including the number of operations and the duration.
51#[derive(Eq, PartialEq, Clone, Copy, Debug)]
52#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
53pub struct ThroughputValue {
54    /// The number of operations performed depending of the mode during the computation.
55    pub ops_count: usize,
56    /// The duration of the computation.
57    pub duration: Duration,
58}
59
60impl ThroughputValue {
61    /// A zero-initialized throughput value, representing no operations or duration.
62    pub const ZERO: Self = Self {
63        ops_count: 0,
64        duration: Duration::ZERO,
65    };
66
67    /// Returns the operations per second.
68    pub fn ops_per_s(&self) -> f64 {
69        if self.duration.is_zero() {
70            return f64::NAN;
71        }
72        self.ops_count as f64 / self.duration.as_secs_f64()
73    }
74
75    /// Returns the bytes per second.
76    pub fn bytes_per_s(&self, key: &ThroughputKey) -> f64 {
77        if self.duration.is_zero() {
78            return f64::NAN;
79        }
80        (self.ops_count * key.dtype().size()) as f64 / self.duration.as_secs_f64()
81    }
82
83    /// Returns the duration per operation.
84    pub fn duration_per_op(&self) -> Duration {
85        if self.ops_count == 0 {
86            Duration::ZERO
87        } else {
88            Duration::from_secs_f64(self.duration.as_secs_f64() / self.ops_count as f64)
89        }
90    }
91
92    /// Formats the throughput value as a clean human-readable string.
93    pub fn format(&self, key: &ThroughputKey) -> String {
94        let (mut val_per_s, unit) = match key.mode {
95            ThroughputMode::ComputeDirect { .. } | ThroughputMode::ComputeCmma { .. } => {
96                (self.ops_per_s(), "OPS")
97            }
98            ThroughputMode::Memory => (self.bytes_per_s(key), "bytes"),
99            ThroughputMode::Launch => {
100                let dur = self.duration_per_op();
101                if dur.is_zero() {
102                    return String::from("N/A");
103                }
104                return format!("{dur:?}/launch");
105            }
106        };
107
108        if val_per_s.is_nan() {
109            return String::from("N/A");
110        }
111
112        let suffixes = ["", "K", "M", "G", "T", "P", "E", "Z", "Y", "R", "Q"];
113        let mut suffix_idx = 0;
114
115        for _ in 0..suffixes.len() - 1 {
116            if val_per_s < 1000.0 {
117                break;
118            }
119            val_per_s /= 1000.0;
120            suffix_idx += 1;
121        }
122
123        format!("{val_per_s:.4} {}{unit}/s", suffixes[suffix_idx])
124    }
125}
126
127/// Constructs a compute [`ThroughputKey`] based on CMMA tile availability and types.
128pub fn compute_throughput_key(
129    cmma_tile: Option<(u32, u32, u32)>,
130    input_elem_type: ElemType,
131    acc_elem_type: ElemType,
132) -> ThroughputKey {
133    let mode = match cmma_tile {
134        Some((tile_m, tile_n, tile_k)) => ThroughputMode::ComputeCmma {
135            dtype: input_elem_type,
136            config: ComputeCmmaConfig {
137                accumulator_type: acc_elem_type,
138                cmma_dims: CmmaDims {
139                    m: tile_m as usize,
140                    n: tile_n as usize,
141                    k: tile_k as usize,
142                },
143            },
144        },
145        None => ThroughputMode::ComputeDirect {
146            dtype: acc_elem_type,
147        },
148    };
149
150    ThroughputKey { mode }
151}
cubecl_runtime/throughput/base.rs

cubecl_runtime/throughput/
base.rs
