cubecl_runtime/benchmark/
multi_device.rs

1//! A [`Benchmark`] timed on several devices at once, each member on its own stream.
2//!
3//! [`Benchmark::run`] times one client. Work that spans devices, such as a collective, has one
4//! duration per member and a slowest one that bounds the step, and must be queued on every member
5//! before any is waited on, or a member blocks in the collective waiting for a peer that was never
6//! launched.
7
8use super::{Benchmark, BenchmarkComputations, BenchmarkDurations, TimingMethod};
9use crate::client::Client;
10use alloc::{format, string::String, vec::Vec};
11use core::time::Duration;
12use cubecl_common::profile::{Instant, ProfileDuration};
13use cubecl_environment::future::block_on;
14
15/// Launches the warmup makes however long they take, as [`Benchmark::run`] does.
16const MIN_WARMUP_RUNS: usize = 5;
17
18/// `benchmark` timed on every one of `clients`, member `i` being `clients[i]`.
19///
20/// The benchmark's [`execute`](Benchmark::execute) queues its work on every member and waits on
21/// none of them. Its [`sync`](Benchmark::sync) is never called: the run syncs every client itself.
22pub struct MultiDeviceRun<'a, B: Benchmark> {
23    benchmark: &'a B,
24    clients: &'a [Client],
25}
26
27impl<'a, B: Benchmark> MultiDeviceRun<'a, B> {
28    /// A run of `benchmark` on `clients`.
29    pub fn new(benchmark: &'a B, clients: &'a [Client]) -> Self {
30        Self { benchmark, clients }
31    }
32
33    /// Warm the devices up for the benchmark's warmup budget, then time every member of each of
34    /// its samples.
35    pub fn run(&self) -> Result<MultiDeviceDurations, String> {
36        let input = self.benchmark.prepare();
37        let budget = self.benchmark.warmup_budget();
38        let warmup = Instant::now();
39        let mut warmups = 0;
40        while warmups < MIN_WARMUP_RUNS || warmup.elapsed() < budget {
41            self.sample(&input)?;
42            warmups += 1;
43        }
44        let samples = (0..self.benchmark.num_samples())
45            .map(|_| self.sample(&input))
46            .collect::<Result<Vec<_>, _>>()?;
47        Ok(MultiDeviceDurations::new(samples))
48    }
49
50    /// One execution: every member synced, a window opened on each, the work queued on all of
51    /// them, every window closed, then every member synced again.
52    fn sample(&self, input: &B::Input) -> Result<MultiDeviceSample, String> {
53        let clients = self.clients;
54        sync_all(clients)?;
55        let started = Instant::now();
56        let windows = clients
57            .iter()
58            .map(|client| client.profile_start())
59            .collect::<Result<Vec<_>, _>>()
60            .map_err(|error| format!("{error:?}"))?;
61        let output = self.benchmark.execute(input.clone());
62        let profiles = clients
63            .iter()
64            .zip(windows)
65            .map(|(client, window)| client.profile_end(window))
66            .collect::<Result<Vec<_>, _>>()
67            .map_err(|error| format!("{error:?}"));
68        sync_all(clients)?;
69        let wall = started.elapsed();
70        // Held to here so the work it names is not dropped before it runs.
71        drop(output?);
72        let profiles = profiles?;
73        let timing_method = profiles
74            .first()
75            .map_or(TimingMethod::System, ProfileDuration::timing_method);
76        let members = profiles
77            .into_iter()
78            .map(|profile| {
79                block_on(profile.resolve())
80                    .map(|ticks| ticks.duration())
81                    .ok_or_else(|| String::from("a member's window carried no measurement"))
82            })
83            .collect::<Result<Vec<_>, _>>()?;
84        Ok(MultiDeviceSample {
85            members,
86            wall,
87            timing_method,
88        })
89    }
90}
91
92/// One execution's durations: each member's window on its own stream, and the wall around all.
93#[derive(Debug, Clone)]
94pub struct MultiDeviceSample {
95    /// Member `i`'s window, from before its part was queued to after.
96    pub members: Vec<Duration>,
97    /// The host's wall from before the first window opened to every member synced.
98    pub wall: Duration,
99    /// How the members' windows were timed.
100    pub timing_method: TimingMethod,
101}
102
103/// Every sample of a multi-device benchmark, per member, slowest member and wall.
104#[derive(Debug, Clone)]
105pub struct MultiDeviceDurations {
106    /// Member `i`'s duration in each sample.
107    pub members: Vec<BenchmarkDurations>,
108    /// The slowest member's duration in each sample: what bounds a step all members take part in.
109    pub slowest: BenchmarkDurations,
110    /// The host's wall around each sample, syncs included.
111    pub wall: BenchmarkDurations,
112}
113
114impl MultiDeviceDurations {
115    /// The durations of `samples`, which all have one member count.
116    ///
117    /// # Panics
118    ///
119    /// When `samples` is empty or their member counts differ.
120    pub fn new(samples: Vec<MultiDeviceSample>) -> Self {
121        let first = samples
122            .first()
123            .expect("a multi-device benchmark takes samples");
124        let timing_method = first.timing_method;
125        let count = first.members.len();
126        assert!(
127            samples.iter().all(|sample| sample.members.len() == count),
128            "every sample of a multi-device benchmark has the same members"
129        );
130        let column = |pick: &dyn Fn(&MultiDeviceSample) -> Duration| {
131            BenchmarkDurations::from_durations(timing_method, samples.iter().map(pick).collect())
132        };
133        Self {
134            members: (0..count)
135                .map(|member| column(&|sample| sample.members[member]))
136                .collect(),
137            slowest: column(&|sample| sample.members.iter().copied().max().unwrap_or_default()),
138            wall: BenchmarkDurations::from_durations(
139                TimingMethod::System,
140                samples.iter().map(|sample| sample.wall).collect(),
141            ),
142        }
143    }
144
145    /// The median of the slowest member, the figure a step's cost is read from.
146    pub fn median(&self) -> Duration {
147        BenchmarkComputations::new(&self.slowest).median
148    }
149
150    /// Slowest sample of the slowest member over its fastest: a lower bound on the run's noise.
151    pub fn spread(&self) -> f64 {
152        let computed = BenchmarkComputations::new(&self.slowest);
153        computed.max.as_secs_f64() / computed.min.as_secs_f64().max(f64::MIN_POSITIVE)
154    }
155}
156
157fn sync_all(clients: &[Client]) -> Result<(), String> {
158    clients
159        .iter()
160        .try_for_each(|client| block_on(client.sync()).map_err(|error| format!("{error:?}")))
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166    use alloc::vec;
167
168    fn sample(members: &[u64], wall: u64) -> MultiDeviceSample {
169        MultiDeviceSample {
170            members: members
171                .iter()
172                .map(|&us| Duration::from_micros(us))
173                .collect(),
174            wall: Duration::from_micros(wall),
175            timing_method: TimingMethod::Device,
176        }
177    }
178
179    #[test]
180    fn the_slowest_member_of_each_sample_bounds_it() {
181        let durations = MultiDeviceDurations::new(vec![
182            sample(&[10, 30], 40),
183            sample(&[50, 20], 60),
184            sample(&[15, 25], 30),
185        ]);
186
187        let slowest: Vec<u64> = durations
188            .slowest
189            .durations
190            .iter()
191            .map(|d| d.as_micros() as u64)
192            .collect();
193        assert_eq!(slowest, [30, 50, 25]);
194        assert_eq!(durations.members.len(), 2);
195        assert_eq!(durations.median(), Duration::from_micros(30));
196        assert_eq!(durations.spread(), 2.0);
197    }
198}
cubecl_runtime/benchmark/multi_device.rs

cubecl_runtime/benchmark/
multi_device.rs
