cubecl-std 0.11.0-pre.4

CubeCL Standard Library.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use cubecl_runtime::throughput::{
    KernelConfig, ThroughputBenchmarker, ThroughputError, ThroughputValue,
};

/// The shapes a probe can be launched in, ranked against each other so that
/// only the winner is measured in full.
pub(super) struct ShapeSweep<S> {
    shapes: alloc::vec::Vec<S>,
}

impl<S: Copy> ShapeSweep<S> {
    pub(super) fn new(shapes: alloc::vec::Vec<S>) -> Self {
        Self { shapes }
    }

    /// What the device answers fastest, and the shape that answered it.
    ///
    /// Shapes are built one at a time: a memory probe's pool is a large fraction
    /// of what the device will allocate. A rate that is not finite did not run.
    ///
    /// # Errors
    ///
    /// Whatever `build` reports, and [`ThroughputError::NoTiming`] where no
    /// shape ran.
    pub(super) fn fastest(
        &self,
        build: impl Fn(S) -> Result<KernelConfig, ThroughputError>,
    ) -> Result<(ThroughputValue, S), ThroughputError> {
        let (fastest, warmed) = match self.shapes.len() {
            0 => return Err(ThroughputError::NoTiming),
            1 => (self.shapes[0], None),
            _ => {
                let (fastest, iterations) = self.ranked(&build)?;

                (fastest, Some(iterations))
            }
        };

        let config = build(fastest)?;
        let value = match warmed {
            Some(iterations) => ThroughputBenchmarker::sample_at(&config, iterations),
            None => ThroughputBenchmarker::sample(config),
        };

        value
            .ops_per_s()
            .is_finite()
            .then_some((value, fastest))
            .ok_or(ThroughputError::NoTiming)
    }

    /// The shape that answers fastest over a ranking pass, and the count the
    /// device was warmed at. Every shape is timed at that same count, so they
    /// are ordered on what they do rather than on which was warmed.
    fn ranked(
        &self,
        build: impl Fn(S) -> Result<KernelConfig, ThroughputError>,
    ) -> Result<(S, usize), ThroughputError> {
        let mut warmed = None;
        let mut fastest: Option<(f64, S)> = None;

        for shape in &self.shapes {
            let config = build(*shape)?;
            let start = *warmed.get_or_insert_with(|| ThroughputBenchmarker::warm(&config));
            let ranked = ThroughputBenchmarker::rank(&config, start);
            let rate = ranked.value.ops_per_s();

            // The next shape starts from what this one settled on, so a sweep
            // whose shapes drift in cost keeps its settling launch near target.
            warmed = Some(ranked.iterations);

            if rate.is_finite() && fastest.is_none_or(|(best, _)| rate > best) {
                fastest = Some((rate, *shape));
            }
        }

        fastest
            .zip(warmed)
            .map(|((_, shape), iterations)| (shape, iterations))
            .ok_or(ThroughputError::NoTiming)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The reported value is a full measurement of the winner, not the
    /// ranking pass that found it.
    #[test]
    fn the_shape_that_ranks_fastest_is_the_one_measured() {
        // A shape here is its rate: the faster it is, the less time a pass takes.
        let build = |rate: u64| {
            Ok(KernelConfig {
                sample: alloc::boxed::Box::new(move |iterations| {
                    core::time::Duration::from_nanos(iterations as u64 * 1000 / rate)
                }),
                ops_count: 1,
                min_iterations: 1,
            })
        };

        let (value, fastest) = ShapeSweep::new(alloc::vec![1, 4, 2])
            .fastest(build)
            .expect("a shape ran");

        assert_eq!(fastest, 4);
        assert!(value.ops_per_s().is_finite());
    }

    /// A shape that cannot be built stops the sweep with its own error, rather
    /// than being outranked by the shapes that could.
    #[test]
    fn a_shape_that_cannot_be_built_fails_the_sweep() {
        let build = |_: u64| Err(ThroughputError::Allocation);

        for shapes in [alloc::vec![1], alloc::vec![1, 2]] {
            let failed = ShapeSweep::new(shapes).fastest(build);

            assert_eq!(failed.err(), Some(ThroughputError::Allocation));
        }
    }
}