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;
7
8use alloc::format;
9
10use crate::tune::{Bounds, BoundsGenerator};
11
12use super::{
13    AutotuneError, input_generator::InputGenerator, key_generator::KeyGenerator,
14    tune_inputs::TuneInputs,
15};
16use super::{Tunable, TunePlan};
17
18/// A type-erased delegate for a tunable function.
19///
20/// The lifetime `'inp` is the lifetime of the input data, the function must be defined such that
21/// it can be called for any lifetime `inp` and produce a `Result<Out, AutotuneError>`.
22type TuneDelegate<I, Out> =
23    dyn for<'inp> Fn(<I as TuneInputs>::At<'inp>) -> Result<Out, AutotuneError> + Send + Sync;
24
25/// A named, type-erased tunable function stored in a [`TunableSet`]. Constructed via
26/// [`Tunable::new`](super::Tunable::new); callers don't name this type directly.
27#[derive(new)]
28pub struct TuneFn<I: TuneInputs, Out> {
29    pub(crate) name: String,
30    func: Box<TuneDelegate<I, Out>>,
31}
32
33impl<I: TuneInputs, Out: 'static> TuneFn<I, Out> {
34    /// Run the wrapped function on the given inputs.
35    pub fn execute<'a>(&self, inputs: <I as TuneInputs>::At<'a>) -> Result<Out, AutotuneError> {
36        (self.func)(inputs)
37    }
38}
39
40/// A set of candidate tunable functions for autotune, sharing a key generator and an
41/// input generator. See [`TuneInputs`] for the `F` parameter.
42pub struct TunableSet<K: AutotuneKey, F: TuneInputs, Output: 'static> {
43    tunables: Vec<Tunable<K, F, Output>>,
44    key_gen: Arc<dyn KeyGenerator<K, F> + Send + Sync>,
45    input_gen: Arc<dyn InputGenerator<K, F> + Send + Sync>,
46    bounds_gen: Option<Arc<dyn BoundsGenerator<K, F> + Send + Sync>>,
47    short_circuit: bool,
48}
49
50impl<K: AutotuneKey, F: TuneInputs, Output: 'static> TunableSet<K, F, Output> {
51    /// The number of tunables in the set.
52    pub fn len(&self) -> usize {
53        self.tunables.len()
54    }
55
56    /// Whether this set contains no tunables.
57    pub fn is_empty(&self) -> bool {
58        self.tunables.is_empty()
59    }
60
61    /// Create a tunable set from a key generator and an input generator.
62    pub fn new(key_gen: impl KeyGenerator<K, F>, input_gen: impl InputGenerator<K, F>) -> Self {
63        Self {
64            tunables: Default::default(),
65            input_gen: Arc::new(input_gen),
66            key_gen: Arc::new(key_gen),
67            bounds_gen: None,
68            short_circuit: true,
69        }
70    }
71
72    /// Shorthand for [`new`](Self::new) with a [`CloneInputGenerator`]: benchmarks run
73    /// on clones of the real call inputs.
74    pub fn new_cloning_inputs(key_gen: impl KeyGenerator<K, F>) -> Self {
75        Self::new(key_gen, super::CloneInputGenerator)
76    }
77
78    /// Register a tunable with this tunable set.
79    pub fn with(mut self, tunable: Tunable<K, F, Output>) -> Self {
80        self.tunables.push(tunable);
81        self
82    }
83
84    /// Sets the autotune bounds for this set.
85    pub fn with_bounds(mut self, bounds: Arc<dyn BoundsGenerator<K, F> + Send + Sync>) -> Self {
86        self.bounds_gen = Some(bounds);
87        self
88    }
89
90    /// Set whether bounds-based short-circuiting is enabled.
91    pub fn with_short_circuit(mut self, enabled: bool) -> Self {
92        self.short_circuit = enabled;
93        self
94    }
95
96    /// Whether short-circuiting is enabled for this set.
97    pub fn is_short_circuit_enabled(&self) -> bool {
98        self.short_circuit
99    }
100
101    /// All candidate operations in this set, in registration order.
102    pub fn autotunables(&self) -> impl Iterator<Item = &TuneFn<F, Output>> {
103        self.tunables.iter().map(|tunable| &tunable.function)
104    }
105
106    /// Returns the [autotune plan](TunePlan) for the given set.
107    pub(crate) fn plan(&self, key: &K) -> TunePlan {
108        TunePlan::new(key, &self.tunables)
109    }
110
111    /// Returns the operation for the given index, matching the order returned by
112    /// `autotunables`. Tunables are tried in order, so index 0 should be a good default.
113    pub fn fastest(&self, fastest_index: usize) -> &TuneFn<F, Output> {
114        &self.tunables[fastest_index].function
115    }
116
117    /// Compute a checksum that invalidates outdated cached auto-tune results when the
118    /// set of tunable names changes.
119    pub fn compute_checksum(&self) -> String {
120        let mut checksum = String::new();
121        for tune in &self.tunables {
122            checksum += &tune.function.name;
123        }
124        format!("{:x}", md5::compute(checksum))
125    }
126
127    /// Generate a key from a set of inputs
128    pub fn generate_key<'a>(&self, inputs: &F::At<'a>) -> K {
129        self.key_gen.generate(inputs)
130    }
131
132    /// Generate a set of test inputs from a key and reference inputs.
133    pub fn generate_inputs<'a>(&self, key: &K, inputs: &F::At<'a>) -> F::At<'a> {
134        self.input_gen.generate(key, inputs)
135    }
136
137    /// The throughput bounds registered on this set, if any.
138    pub fn bounds(&self, key: &K, inputs: &F::At<'_>) -> Option<Bounds> {
139        self.bounds_gen.as_ref().map(|f| f.generate(key, inputs))
140    }
141}
142
143#[cfg(autotune_persistence)]
144/// Trait alias with support for persistent caching
145pub trait AutotuneKey:
146    Clone
147    + Debug
148    + PartialEq
149    + Eq
150    + Hash
151    + Display
152    + serde::Serialize
153    + serde::de::DeserializeOwned
154    + Send
155    + Sync
156    + 'static
157{
158}
159#[cfg(not(autotune_persistence))]
160/// Trait alias
161pub trait AutotuneKey:
162    Clone + Debug + PartialEq + Eq + Hash + Display + Send + Sync + 'static
163{
164}
165
166impl AutotuneKey for String {}
cubecl_runtime/tune/operation.rs

cubecl_runtime/tune/
operation.rs
