cubecl_runtime/throughput/
base.rs

1use crate::throughput::{CmmaDims, ComputeCmmaConfig};
2use alloc::{format, string::String};
3use core::time::Duration;
4use cubecl_ir::{ElemType, FloatKind};
5use thiserror::Error;
6
7/// What the probes measure, as opposed to which release ran them. Bump it when
8/// a probe changes what it reports.
9pub const PROBE_VERSION: u32 = 2;
10
11/// Bytes one buffer of a [`ThroughputMode::Memory`] probe moves per pass at its
12/// default working set. The probe's buffer is a multiple of this, and both are
13/// clamped to the device's maximum allocation when the probe runs.
14pub const DEFAULT_WORKING_SET_BYTES: u64 = 512 * 1024 * 1024;
15
16/// Which directions of traffic a memory probe issues.
17#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
18#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
19pub enum MemoryAccess {
20    /// Reads every byte and writes it back out. The ceiling for a kernel that
21    /// both loads and stores.
22    Copy,
23    /// Reads only, storing nothing. The ceiling for a kernel that streams data
24    /// it does not write back — a weight stream, a reduction, a gather. Such a
25    /// kernel legitimately exceeds [`Copy`](Self::Copy), because half of the
26    /// copy's traffic is a direction it never uses.
27    Read,
28    /// Writes only, reading nothing at the software level. The ceiling for a
29    /// kernel that streams stores it never reads back: an RNG fill, a memset,
30    /// a broadcast. Ordinary stores still carry read-for-ownership traffic on
31    /// cache-coherent hardware, and this probe's stores do too, which is what
32    /// makes it the honest ceiling for a kernel that uses ordinary stores
33    /// rather than a non-temporal one.
34    Write,
35}
36
37impl MemoryAccess {
38    /// How many buffers of equal size one pass touches: two for a copy (one in,
39    /// one out), one for a read or a write.
40    pub const fn buffers(&self) -> u64 {
41        match self {
42            Self::Copy => 2,
43            Self::Read | Self::Write => 1,
44        }
45    }
46
47    /// The working set of the single-size probe for this access, in bytes moved
48    /// per pass: [`DEFAULT_WORKING_SET_BYTES`] per buffer touched.
49    pub const fn default_working_set(&self) -> u64 {
50        DEFAULT_WORKING_SET_BYTES * self.buffers()
51    }
52}
53
54/// Represents the mode of a throughput computation.
55#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
56#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
57pub enum ThroughputMode {
58    /// Compute direct calculation without special hardware acceleration.
59    ///
60    /// The ceiling is for operands of `dtype`, not arithmetic performed in it:
61    /// where converting to f32 retires more, that is what the probe reports.
62    ComputeDirect {
63        /// The data type of the computation.
64        dtype: ElemType,
65    },
66    /// Compute cmma calculation with CMMA hardware acceleration.
67    ComputeCmma {
68        /// The data type of the computation.
69        dtype: ElemType,
70        /// The configuration of the CMMA operation.
71        config: ComputeCmmaConfig,
72    },
73    /// Traffic across the memory interface, as described by its [`MemorySpec`].
74    ///
75    /// `bytes` is the total one pass moves, so a
76    /// [`Copy`](MemoryAccess::Copy) splits it across two buffers where a
77    /// [`Read`](MemoryAccess::Read) takes it all from one.
78    Memory(MemorySpec),
79    /// Launch overhead measurement.
80    Launch,
81}
82
83/// What a memory mode asks of a probe.
84#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
85#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
86#[cfg_attr(serializable, serde(deny_unknown_fields))]
87pub struct MemorySpec {
88    /// Which directions of traffic to issue.
89    pub access: MemoryAccess,
90    /// The bytes one pass moves across the interface.
91    pub bytes: u64,
92}
93
94impl ThroughputMode {
95    /// `access` over as much as the interface will take at once.
96    pub const fn memory(access: MemoryAccess) -> Self {
97        Self::Memory(MemorySpec::new(access, access.default_working_set()))
98    }
99
100    /// What this mode asks of a memory probe, or `None` for the modes that do
101    /// not measure memory.
102    pub const fn memory_probe(&self) -> Option<MemorySpec> {
103        match self {
104            Self::Memory(spec) => Some(*spec),
105            Self::ComputeDirect { .. } | Self::ComputeCmma { .. } | Self::Launch => None,
106        }
107    }
108}
109
110impl MemorySpec {
111    /// A probe moving `bytes` per pass in the directions `access` names.
112    pub const fn new(access: MemoryAccess, bytes: u64) -> Self {
113        Self { access, bytes }
114    }
115}
116
117/// Represents a key/configuration used to identify the throughput of a computation.
118#[derive(Eq, PartialEq, Clone, Hash, Debug, Copy)]
119#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
120// Reject cached entries from an older key layout instead of silently ignoring their extra fields.
121#[cfg_attr(serializable, serde(deny_unknown_fields))]
122pub struct ThroughputKey {
123    /// The mode of the throughput computation.
124    pub mode: ThroughputMode,
125}
126
127impl ThroughputKey {
128    /// Returns the data type of the computation.
129    pub fn dtype(&self) -> ElemType {
130        match self.mode {
131            ThroughputMode::ComputeDirect { dtype } => dtype,
132            ThroughputMode::ComputeCmma { dtype, .. } => dtype,
133            // For memory and launch throughput, we use a default element type (F32).
134            ThroughputMode::Memory(_) | ThroughputMode::Launch => ElemType::Float(FloatKind::F32),
135        }
136    }
137}
138
139/// Why a device has no peak to report for a probe.
140#[derive(Error, Eq, PartialEq, Clone, Copy, Debug)]
141pub enum ThroughputError {
142    /// The device implements no such operation.
143    #[error("unsupported")]
144    Unsupported,
145    /// The device's timer reported no elapsed time for any shape of the probe.
146    #[error("no timing")]
147    NoTiming,
148    /// The device could not allocate the buffers the probe runs over.
149    #[error("allocation failed")]
150    Allocation,
151    /// The probe's kernel did not run: it failed to compile or to launch, or
152    /// the device faulted under it.
153    #[error("launch failed")]
154    Launch,
155}
156
157/// Represents the throughput of a computation, including the number of operations and the duration.
158#[derive(Eq, PartialEq, Clone, Copy, Debug)]
159#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
160pub struct ThroughputValue {
161    /// The number of operations performed depending of the mode during the computation.
162    pub ops_count: usize,
163    /// The duration of the computation.
164    pub duration: Duration,
165}
166
167impl ThroughputValue {
168    /// A zero-initialized throughput value, representing no operations or duration.
169    pub const ZERO: Self = Self {
170        ops_count: 0,
171        duration: Duration::ZERO,
172    };
173
174    /// Returns the operations per second.
175    pub fn ops_per_s(&self) -> f64 {
176        if self.duration.is_zero() {
177            return f64::NAN;
178        }
179        self.ops_count as f64 / self.duration.as_secs_f64()
180    }
181
182    /// Returns the bytes per second.
183    pub fn bytes_per_s(&self, key: &ThroughputKey) -> f64 {
184        if self.duration.is_zero() {
185            return f64::NAN;
186        }
187        (self.ops_count * key.dtype().size()) as f64 / self.duration.as_secs_f64()
188    }
189
190    /// Returns the duration per operation.
191    pub fn duration_per_op(&self) -> Duration {
192        if self.ops_count == 0 {
193            Duration::ZERO
194        } else {
195            Duration::from_secs_f64(self.duration.as_secs_f64() / self.ops_count as f64)
196        }
197    }
198
199    /// Formats the throughput value as a clean human-readable string.
200    pub fn format(&self, key: &ThroughputKey) -> String {
201        let (mut val_per_s, unit) = match key.mode {
202            ThroughputMode::ComputeDirect { .. } | ThroughputMode::ComputeCmma { .. } => {
203                (self.ops_per_s(), "OPS")
204            }
205            ThroughputMode::Memory(_) => (self.bytes_per_s(key), "bytes"),
206            ThroughputMode::Launch => {
207                let dur = self.duration_per_op();
208                if dur.is_zero() {
209                    return String::from("N/A");
210                }
211                return format!("{dur:?}/launch");
212            }
213        };
214
215        if val_per_s.is_nan() {
216            return String::from("N/A");
217        }
218
219        let suffixes = ["", "K", "M", "G", "T", "P", "E", "Z", "Y", "R", "Q"];
220        let mut suffix_idx = 0;
221
222        for _ in 0..suffixes.len() - 1 {
223            if val_per_s < 1000.0 {
224                break;
225            }
226            val_per_s /= 1000.0;
227            suffix_idx += 1;
228        }
229
230        format!("{val_per_s:.4} {}{unit}/s", suffixes[suffix_idx])
231    }
232}
233
234/// Constructs a compute [`ThroughputKey`] based on CMMA tile availability and types.
235pub fn compute_throughput_key(
236    cmma_tile: Option<(u32, u32, u32)>,
237    input_elem_type: ElemType,
238    acc_elem_type: ElemType,
239) -> ThroughputKey {
240    let mode = match cmma_tile {
241        Some((tile_m, tile_n, tile_k)) => ThroughputMode::ComputeCmma {
242            dtype: input_elem_type,
243            config: ComputeCmmaConfig {
244                accumulator_type: acc_elem_type,
245                cmma_dims: CmmaDims {
246                    m: tile_m as usize,
247                    n: tile_n as usize,
248                    k: tile_k as usize,
249                },
250            },
251        },
252        None => ThroughputMode::ComputeDirect {
253            dtype: acc_elem_type,
254        },
255    };
256
257    ThroughputKey { mode }
258}
259
260#[cfg(all(test, std_io))]
261mod tests {
262    use super::*;
263
264    /// The throughput cache keys on the serialized key, so a layout change
265    /// drops every measurement users have already paid for. That is what a
266    /// version bump is for; a change that is not one must leave these alone.
267    #[test]
268    fn a_memory_key_keeps_its_serialized_form() {
269        let encode = |mode| serde_json::to_string(&ThroughputKey { mode }).unwrap();
270
271        assert_eq!(
272            encode(ThroughputMode::memory(MemoryAccess::Copy)),
273            r#"{"mode":{"Memory":{"access":"Copy","bytes":1073741824}}}"#
274        );
275        assert_eq!(
276            encode(ThroughputMode::Memory(MemorySpec::new(
277                MemoryAccess::Read,
278                8192
279            ))),
280            r#"{"mode":{"Memory":{"access":"Read","bytes":8192}}}"#
281        );
282    }
283
284    /// A working set is part of the key, so two sizes of the same access are
285    /// separate measurements rather than one overwriting the other.
286    #[test]
287    fn a_working_set_is_part_of_the_key() {
288        let small = ThroughputMode::Memory(MemorySpec::new(MemoryAccess::Read, 8192));
289        let large = ThroughputMode::Memory(MemorySpec::new(MemoryAccess::Read, 16384));
290
291        assert_ne!(small, large);
292    }
293}
cubecl_runtime/throughput/base.rs

cubecl_runtime/throughput/
base.rs
