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::{AutotuneBound, 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/// Pseudo generic bound type to avoid adding a generic parameter to [`TunableSet`].
41type B = AutotuneBound;
42
43/// A set of candidate tunable functions for autotune, sharing a key generator and an
44/// input generator. See [`TuneInputs`] for the `F` parameter.
45pub struct TunableSet<K: AutotuneKey, F: TuneInputs, Output: 'static> {
46    tunables: Vec<Tunable<K, F, Output>>,
47    key_gen: Arc<dyn KeyGenerator<K, F> + Send + Sync>,
48    input_gen: Arc<dyn InputGenerator<K, F> + Send + Sync>,
49    bounds_gen: Option<Arc<dyn BoundsGenerator<K, F, B> + Send + Sync>>,
50}
51
52impl<K: AutotuneKey, F: TuneInputs, Output: 'static> TunableSet<K, F, Output> {
53    /// The number of tunables in the set.
54    pub fn len(&self) -> usize {
55        self.tunables.len()
56    }
57
58    /// Whether this set contains no tunables.
59    pub fn is_empty(&self) -> bool {
60        self.tunables.is_empty()
61    }
62
63    /// Create a tunable set from a key generator and an input generator.
64    pub fn new(key_gen: impl KeyGenerator<K, F>, input_gen: impl InputGenerator<K, F>) -> Self {
65        Self {
66            tunables: Default::default(),
67            input_gen: Arc::new(input_gen),
68            key_gen: Arc::new(key_gen),
69            bounds_gen: None,
70        }
71    }
72
73    /// Shorthand for [`new`](Self::new) with a [`CloneInputGenerator`]: benchmarks run
74    /// on clones of the real call inputs.
75    pub fn new_cloning_inputs(key_gen: impl KeyGenerator<K, F>) -> Self {
76        Self::new(key_gen, super::CloneInputGenerator)
77    }
78
79    /// Register a tunable with this tunable set.
80    pub fn with(mut self, tunable: Tunable<K, F, Output>) -> Self {
81        self.tunables.push(tunable);
82        self
83    }
84
85    /// Sets the autotune bounds for this set.
86    pub fn with_bounds(mut self, bounds: Arc<dyn BoundsGenerator<K, F, B> + Send + Sync>) -> Self {
87        self.bounds_gen = Some(bounds);
88        self
89    }
90
91    /// All candidate operations in this set, in registration order.
92    pub fn autotunables(&self) -> impl Iterator<Item = &TuneFn<F, Output>> {
93        self.tunables.iter().map(|tunable| &tunable.function)
94    }
95
96    /// Returns the [autotune plan](TunePlan) for the given set.
97    pub(crate) fn plan(&self, key: &K) -> TunePlan {
98        TunePlan::new(key, &self.tunables)
99    }
100
101    /// Returns the operation for the given index, matching the order returned by
102    /// `autotunables`. Tunables are tried in order, so index 0 should be a good default.
103    pub fn fastest(&self, fastest_index: usize) -> &TuneFn<F, Output> {
104        &self.tunables[fastest_index].function
105    }
106
107    /// Compute a checksum that invalidates outdated cached auto-tune results when the
108    /// set of tunable names changes.
109    pub fn compute_checksum(&self) -> String {
110        let mut checksum = String::new();
111        for tune in &self.tunables {
112            checksum += &tune.function.name;
113        }
114        format!("{:x}", md5::compute(checksum))
115    }
116
117    /// Generate a key from a set of inputs
118    pub fn generate_key<'a>(&self, inputs: &F::At<'a>) -> K {
119        self.key_gen.generate(inputs)
120    }
121
122    /// Generate a set of test inputs from a key and reference inputs.
123    pub fn generate_inputs<'a>(&self, key: &K, inputs: &F::At<'a>) -> F::At<'a> {
124        self.input_gen.generate(key, inputs)
125    }
126
127    /// The throughput bounds registered on this set, if any.
128    pub fn bounds(&self, key: &K, inputs: &F::At<'_>) -> Option<Bounds<B>> {
129        self.bounds_gen.as_ref().map(|f| f.generate(key, inputs))
130    }
131}
132
133#[cfg(std_io)]
134/// Trait alias with support for persistent caching
135pub trait AutotuneKey:
136    Clone
137    + Debug
138    + PartialEq
139    + Eq
140    + Hash
141    + Display
142    + serde::Serialize
143    + serde::de::DeserializeOwned
144    + Send
145    + Sync
146    + 'static
147{
148}
149#[cfg(not(std_io))]
150/// Trait alias
151pub trait AutotuneKey:
152    Clone + Debug + PartialEq + Eq + Hash + Display + Send + Sync + 'static
153{
154}
155
156impl AutotuneKey for String {}
cubecl_runtime/tune/operation.rs

cubecl_runtime/tune/
operation.rs
