cubecl_runtime/tune/
eviction.rs

1use crate::tune::{AutotuneError, TuneInputs};
2use alloc::string::{String, ToString};
3
4/// Runs before every measured sample of a set's candidates, so a candidate whose operands fit
5/// the device's last-level cache is timed reading memory the way a real call does.
6///
7/// A tuner samples one candidate several times over the same inputs, and a warm-up runs
8/// first: from the second launch on, an operand small enough for the cache is served from it.
9/// A memory-bound kernel then reads as several times faster than the bus, the round's
10/// throughput bound is met by the first candidate tried, and the round ends there — with a
11/// winner chosen on a read production never gets. The set says how to evict, since it knows
12/// what its operands weigh; the tuner says when.
13///
14/// What the eviction runs is the set's business — typically a launch that writes past the
15/// cache's capacity through the inputs' client. It receives the reference inputs the tune was
16/// called with rather than the generated ones the candidates are measured on, so it can scrub
17/// through buffers the caller already holds instead of allocating its own: an output the
18/// winner overwrites afterwards is scratch until then. Whether that reaches past the cache is
19/// the set's call — an operand-sized write evicts an operand-sized cache, no more.
20///
21/// The eviction is issued outside the profiled region, so it costs the round time rather than
22/// the sample, as long as it goes through the stream the sample is profiled on: the tuner
23/// orders nothing between the two, and a launch on another stream can overlap the sample or
24/// land after it.
25///
26/// A failed eviction is logged and the sample is measured anyway, warm: an eviction is a
27/// measurement aid, and losing one is not worth failing the tune for.
28#[diagnostic::on_unimplemented(
29    message = "`{Self}` is not a valid eviction",
30    label = "invalid eviction"
31)]
32pub trait Eviction<K, I: TuneInputs>: Send + Sync + 'static {
33    /// Evict what the last measured launch left in cache, for a given key and the reference
34    /// inputs the tune was called with.
35    fn evict<'a>(&self, key: &K, inputs: &I::At<'a>) -> Result<(), AutotuneError>;
36}
37
38/// `Fn(&K, &A) -> Result<(), E>` acts as an [`Eviction`] when `A` is an owned type. For
39/// multi-input kernels, `A` is a tuple that the closure destructures internally.
40impl<K, Func, A, Err> Eviction<K, A> for Func
41where
42    A: Clone + Send + Sync + 'static,
43    K: 'static,
44    Err: Into<String> + 'static,
45    Func: Send + Sync + 'static + Fn(&K, &A) -> Result<(), Err>,
46{
47    #[inline]
48    fn evict<'a>(&self, key: &K, inputs: &<A as TuneInputs>::At<'a>) -> Result<(), AutotuneError> {
49        (self)(key, inputs).map_err(|err| AutotuneError::Unknown {
50            name: "eviction".to_string(),
51            err: err.into(),
52        })
53    }
54}
55
56/// An [`Eviction`] bound to one key and the reference inputs of one tune: what a benchmark
57/// loop calls before each sample.
58///
59/// `FnMut` rather than `Fn` because the loop lends it out exclusively: the closure holds the
60/// inputs, which are only `Send`, and a shared borrow would need them `Sync` to reach the
61/// device thread the samples are launched from.
62pub type Evictor<'i> = dyn FnMut() -> Result<(), AutotuneError> + Send + 'i;
cubecl_runtime/tune/eviction.rs

cubecl_runtime/tune/
eviction.rs
