cubecl_runtime/tune/
operation.rs

1use alloc::boxed::Box;
2use alloc::string::String;
3use alloc::sync::Arc;
4use alloc::vec::Vec;
5use core::fmt::{Debug, Display};
6use core::hash::Hash;
7use cubecl_common::hash::StableHasher;
8
9use alloc::format;
10
11use crate::tune::{Bounds, BoundsGenerator, Eviction, Evictor};
12
13use super::{
14    AutotuneError, input_generator::InputGenerator, key_generator::KeyGenerator,
15    tune_inputs::TuneInputs,
16};
17use super::{Tunable, TunePlan};
18
19/// A type-erased delegate for a tunable function.
20///
21/// The lifetime `'inp` is the lifetime of the input data, the function must be defined such that
22/// it can be called for any lifetime `inp` and produce a `Result<Out, AutotuneError>`.
23type TuneDelegate<I, Out> =
24    dyn for<'inp> Fn(<I as TuneInputs>::At<'inp>) -> Result<Out, AutotuneError> + Send + Sync;
25
26/// A named, type-erased tunable function stored in a [`TunableSet`]. Constructed via
27/// [`Tunable::new`](super::Tunable::new); callers don't name this type directly.
28#[derive(new)]
29pub struct TuneFn<I: TuneInputs, Out> {
30    pub(crate) name: String,
31    func: Box<TuneDelegate<I, Out>>,
32}
33
34impl<I: TuneInputs, Out: 'static> TuneFn<I, Out> {
35    /// Run the wrapped function on the given inputs.
36    pub fn execute<'a>(&self, inputs: <I as TuneInputs>::At<'a>) -> Result<Out, AutotuneError> {
37        (self.func)(inputs)
38    }
39}
40
41/// A set of candidate tunable functions for autotune, sharing a key generator and an
42/// input generator. See [`TuneInputs`] for the `F` parameter.
43///
44/// `Id` is what its tunables are [identified](Tunable::identified) by, inferred from the
45/// tunables it is given: `()` for tunables built by [`Tunable::new`], named and nothing more.
46pub struct TunableSet<K: AutotuneKey, F: TuneInputs, Output: 'static, Id = ()> {
47    tunables: Vec<Tunable<K, F, Output, Id>>,
48    key_gen: Arc<dyn KeyGenerator<K, F> + Send + Sync>,
49    input_gen: Arc<dyn InputGenerator<K, F> + Send + Sync>,
50    bounds_gen: Option<Arc<dyn BoundsGenerator<K, F> + Send + Sync>>,
51    eviction: Option<Arc<dyn Eviction<K, F> + Send + Sync>>,
52    short_circuit: bool,
53}
54
55impl<K: AutotuneKey, F: TuneInputs, Output: 'static, Id> TunableSet<K, F, Output, Id> {
56    /// Create a tunable set from a key generator and an input generator.
57    pub fn new(key_gen: impl KeyGenerator<K, F>, input_gen: impl InputGenerator<K, F>) -> Self {
58        Self {
59            tunables: Default::default(),
60            input_gen: Arc::new(input_gen),
61            key_gen: Arc::new(key_gen),
62            bounds_gen: None,
63            eviction: None,
64            short_circuit: true,
65        }
66    }
67
68    /// Shorthand for [`new`](Self::new) with a [`CloneInputGenerator`](super::CloneInputGenerator): benchmarks run
69    /// on clones of the real call inputs.
70    pub fn new_cloning_inputs(key_gen: impl KeyGenerator<K, F>) -> Self {
71        Self::new(key_gen, super::CloneInputGenerator)
72    }
73
74    /// The number of tunables in the set.
75    pub fn len(&self) -> usize {
76        self.tunables.len()
77    }
78
79    /// Whether this set contains no tunables.
80    pub fn is_empty(&self) -> bool {
81        self.tunables.is_empty()
82    }
83
84    /// Register a tunable with this tunable set.
85    pub fn with(mut self, tunable: Tunable<K, F, Output, Id>) -> Self {
86        self.tunables.push(tunable);
87        self
88    }
89
90    /// Sets the autotune bounds for this set.
91    pub fn with_bounds(mut self, bounds: Arc<dyn BoundsGenerator<K, F> + Send + Sync>) -> Self {
92        self.bounds_gen = Some(bounds);
93        self
94    }
95
96    /// Sets what runs before every measured sample, so the candidates are timed reading
97    /// memory rather than the cache the previous sample left warm. See [`Eviction`].
98    pub fn with_eviction(mut self, eviction: Arc<dyn Eviction<K, F> + Send + Sync>) -> Self {
99        self.eviction = Some(eviction);
100        self
101    }
102
103    /// Set whether bounds-based short-circuiting is enabled.
104    pub fn with_short_circuit(mut self, enabled: bool) -> Self {
105        self.short_circuit = enabled;
106        self
107    }
108
109    /// Whether short-circuiting is enabled for this set.
110    pub fn is_short_circuit_enabled(&self) -> bool {
111        self.short_circuit
112    }
113
114    /// All candidate operations in this set, in registration order.
115    pub fn autotunables(&self) -> impl Iterator<Item = &TuneFn<F, Output>> {
116        self.tunables.iter().map(|tunable| &tunable.function)
117    }
118
119    /// Returns the [autotune plan](TunePlan) for the given set.
120    pub(crate) fn plan(&self, key: &K) -> TunePlan {
121        TunePlan::new(key, &self.tunables)
122    }
123
124    /// Returns the operation for the given index, matching the order returned by
125    /// `autotunables`. Tunables are tried in order, so index 0 should be a good default.
126    pub fn fastest(&self, fastest_index: usize) -> &TuneFn<F, Output> {
127        &self.tunables[fastest_index].function
128    }
129
130    /// What the tunable at `index` is [identified](Tunable::identified) by. The index is one
131    /// the tuner stored for this set, so it panics out of range as [`fastest`](Self::fastest)
132    /// does.
133    pub(crate) fn identity(&self, index: usize) -> &Id {
134        self.tunables[index].identity()
135    }
136
137    /// Compute a checksum that invalidates outdated cached auto-tune results when the
138    /// set of tunable names changes.
139    pub fn compute_checksum(&self) -> String {
140        let mut checksum = String::new();
141        for tune in &self.tunables {
142            checksum += &tune.function.name;
143        }
144        format!("{:x}", StableHasher::hash_one(&checksum))
145    }
146
147    /// Generate a key from a set of inputs
148    pub fn generate_key<'a>(&self, inputs: &F::At<'a>) -> K {
149        self.key_gen.generate(inputs)
150    }
151
152    /// Generate a set of test inputs from a key and reference inputs.
153    pub fn generate_inputs<'a>(&self, key: &K, inputs: &F::At<'a>) -> F::At<'a> {
154        self.input_gen.generate(key, inputs)
155    }
156
157    /// The eviction registered on this set, bound to `key` and the reference `inputs`, if
158    /// any: what a benchmark loop runs before each of its samples.
159    pub(crate) fn evictor<'i>(&self, key: &K, inputs: &F::At<'i>) -> Option<Box<Evictor<'i>>> {
160        let eviction = self.eviction.clone()?;
161        let key = key.clone();
162        let inputs = inputs.clone();
163        Some(Box::new(move || eviction.evict(&key, &inputs)))
164    }
165
166    /// The throughput bounds registered on this set, if any.
167    pub fn bounds(&self, key: &K, inputs: &F::At<'_>) -> Option<Bounds> {
168        self.bounds_gen.as_ref().map(|f| f.generate(key, inputs))
169    }
170}
171
172#[cfg(serializable)]
173/// Trait alias, serializable for the persistent cache and the autotune log
174pub trait AutotuneKey:
175    Clone
176    + Debug
177    + PartialEq
178    + Eq
179    + Hash
180    + Display
181    + serde::Serialize
182    + serde::de::DeserializeOwned
183    + Send
184    + Sync
185    + 'static
186{
187}
188#[cfg(not(serializable))]
189/// Trait alias
190pub trait AutotuneKey:
191    Clone + Debug + PartialEq + Eq + Hash + Display + Send + Sync + 'static
192{
193}
194
195impl AutotuneKey for String {}
cubecl_runtime/tune/operation.rs

cubecl_runtime/tune/
operation.rs
