cubecl_runtime/tune/
local.rs

1use super::{AutotuneKey, AutotuneOutput, TunableSet, TuneInputs, Tuner};
2#[cfg(feature = "autotune-checks")]
3use crate::tune::AutotuneLoggerExt;
4use crate::{client::Client, tune::TuneCacheResult};
5use alloc::string::ToString;
6use alloc::sync::Arc;
7use core::{
8    any::{Any, TypeId},
9    fmt::Display,
10    hash::Hash,
11};
12use cubecl_environment::collections::HashMap;
13use cubecl_environment::sync::{Mutex, RwLock};
14
15/// The tunable sets a [`LocalTuner`] has built, keyed by device as well as by
16/// initializer: a set is built from the device it will run on — its client, its
17/// hardware properties — so one device's set cannot answer for another's. See
18/// [`LocalTuner::init`].
19type Sets<ID> = RwLock<Option<HashMap<(TypeId, ID), Arc<dyn Any + Send + Sync>>>>;
20
21/// A local tuner allows to create a tuner for a specific key that can be different from the server
22/// key.
23pub struct LocalTuner<AK: AutotuneKey, ID> {
24    state: Mutex<Option<HashMap<ID, Arc<Tuner<AK>>>>>,
25    name: &'static str,
26    sets: Sets<ID>,
27}
28
29/// Create a local tuner with the provided name.
30#[macro_export]
31macro_rules! local_tuner {
32    ($name:expr) => {
33        LocalTuner::new(concat!(module_path!(), "-", $name))
34    };
35    () => {
36        LocalTuner::new(module_path!());
37    };
38}
39
40pub use local_tuner;
41
42impl<AK, ID> LocalTuner<AK, ID>
43where
44    AK: AutotuneKey + 'static,
45    ID: Hash + PartialEq + Eq + Clone + Display,
46{
47    /// Create a new local tuner.
48    pub const fn new(name: &'static str) -> Self {
49        Self {
50            state: Mutex::new(None),
51            name,
52            sets: RwLock::new(None),
53        }
54    }
55
56    /// Get or initialize the [`TunableSet`] for `id`.
57    ///
58    /// Returns a cached `Arc<TunableSet>` keyed by the `TypeId` of `init_set`
59    /// *and* by `id`, so the initializer runs once per device rather than once
60    /// per process.
61    ///
62    /// The device is part of the key because a set is routinely built from the
63    /// device it will run on: a closure captures that device's
64    /// [`Client`] to ask what it supports, or reads its hardware
65    /// properties to decide which tunables are worth offering at all. Keyed by
66    /// the initializer alone, whichever device tuned first would answer those
67    /// questions for every device that followed — promoting kernels onto
68    /// hardware that cannot run them, or withholding kernels from hardware
69    /// that can.
70    pub fn init<I, Out, Id, F>(&self, id: &ID, init_set: F) -> Arc<TunableSet<AK, I, Out, Id>>
71    where
72        F: Fn() -> TunableSet<AK, I, Out, Id> + 'static + Send + Sync,
73        Id: Send + Sync + 'static,
74        I: TuneInputs,
75        Out: AutotuneOutput,
76    {
77        let key = (TypeId::of::<F>(), id.clone());
78        let sets = self.sets.read();
79
80        static DOWNCAST_ERROR: &str = "Local tuner only support one set of tunable that must work on the same input and output declared with the init function.";
81
82        if let Some(sets) = sets.as_ref()
83            && let Some(set) = sets.get(&key)
84        {
85            return set.clone().downcast().expect(DOWNCAST_ERROR);
86        };
87
88        core::mem::drop(sets);
89
90        let mut sets = self.sets.write();
91
92        if let Some(sets) = sets.as_ref()
93            && let Some(set) = sets.get(&key)
94        {
95            return set.clone().downcast().expect(DOWNCAST_ERROR);
96        };
97
98        let content = Arc::new(init_set());
99
100        if let Some(sets) = sets.as_mut() {
101            sets.insert(key, content.clone());
102        } else {
103            let mut map = HashMap::<(TypeId, ID), Arc<dyn Any + Send + Sync>>::new();
104            map.insert(key, content.clone());
105            *sets = Some(map);
106        };
107
108        content
109    }
110
111    /// What the fastest tunable in `operations` for `key` on `id` is
112    /// [identified](super::Tunable::identified) by: `None` until a result for the key is
113    /// settled — tuned by a round in this process, or read back from disk and validated by an
114    /// [`execute`](Self::execute).
115    ///
116    /// `operations` must be the set `id` executes, as [`init`](Self::init) returns it for the
117    /// initializer `execute` is handed: a result is an index into that set, and results are kept
118    /// per `id` alone, so a set another initializer built under the same `id` would map it onto
119    /// another tunable.
120    ///
121    /// It reads results and never produces one: it never starts a round, never waits on one in
122    /// flight, never validates a persisted result, and never resets the tuner's cache after an
123    /// environment switch — it reports nothing settled there until the next `execute`.
124    pub fn fastest_identity<I, Out, Id>(
125        &self,
126        id: &ID,
127        operations: &TunableSet<AK, I, Out, Id>,
128        key: &AK,
129    ) -> Option<Id>
130    where
131        I: TuneInputs,
132        Id: Clone,
133    {
134        let tuner = self.state.lock().as_ref()?.get(id)?.clone();
135        let fastest_index = tuner.settled(key)?;
136        Some(operations.identity(fastest_index).clone())
137    }
138
139    /// Clear the autotune state.
140    pub fn clear(&self) {
141        if let Some(s) = self.state.lock().as_mut() {
142            s.clear()
143        }
144    }
145
146    #[cfg(feature = "autotune-checks")]
147    fn checks<'a, I: TuneInputs, Out: AutotuneOutput, Id>(
148        &self,
149        operations: &TunableSet<AK, I, Out, Id>,
150        inputs: &<I as TuneInputs>::At<'a>,
151    ) -> alloc::vec::Vec<crate::tune::log::CheckResult>
152    where
153        <I as TuneInputs>::At<'a>: Clone + Send,
154    {
155        use alloc::vec::Vec;
156
157        let mut checks_outputs = Vec::new();
158        for i in 0..operations.len() {
159            let op = operations.fastest(i);
160            let result = op.execute(inputs.clone());
161            checks_outputs.push((op.name.to_string(), result));
162        }
163        super::check_autotune_outputs(checks_outputs)
164    }
165
166    /// Execute the fastest operation in a [`TunableSet`], triggering a tuning pass on
167    /// the first call for a given key.
168    pub fn execute<'a, I: TuneInputs, Out, Id>(
169        &self,
170        id: &ID,
171        client: &Client,
172        operations: Arc<TunableSet<AK, I, Out, Id>>,
173        inputs: <I as TuneInputs>::At<'a>,
174    ) -> Out
175    where
176        <I as TuneInputs>::At<'a>: Clone + Send,
177        Out: AutotuneOutput,
178        Id: Send + Sync,
179    {
180        let key = operations.generate_key(&inputs);
181
182        let tuner = {
183            let mut state_lock = self.state.lock();
184            let state_map = state_lock.get_or_insert_with(|| HashMap::new());
185            state_map
186                .entry(id.clone())
187                .or_insert_with(move || {
188                    let name = self.name.replace("::", "-");
189                    Arc::new(Tuner::new(&name, &id.to_string()))
190                })
191                .clone()
192        };
193
194        #[allow(unused_mut)]
195        let mut log_context = crate::tune::AutotuneLogContext::new(&mut tuner.logger().lock());
196
197        #[cfg(feature = "autotune-checks")]
198        log_context.set_checks(|| self.checks::<I, Out, Id>(&operations, &inputs));
199
200        // Fast path: a cached hit skips straight to the fastest operation.
201        // `fastest` also resets the tuner cache if the environment switched, so
202        // a miss here falls through to `check_tune`, which re-hydrates.
203        if let TuneCacheResult::Hit { fastest_index } = tuner.fastest(&key) {
204            return operations
205                .fastest(fastest_index)
206                .execute(inputs)
207                .expect("Should run when selected by autotune.");
208        }
209
210        let fastest = tuner.check_tune::<I, Out, Id>(
211            &key,
212            &inputs,
213            &operations,
214            || operations.compute_checksum(),
215            client,
216            log_context,
217        );
218
219        // Run the execution depending on the cache state.
220        match fastest {
221            TuneCacheResult::Hit { fastest_index } => operations
222                .fastest(fastest_index)
223                .execute(inputs)
224                .expect("Should run when selected by autotune."),
225            TuneCacheResult::Unchecked | TuneCacheResult::Miss => {
226                panic!(
227                    "Somehow we STILL didn't check a tuning checksum or start tuning, something has gone wrong."
228                )
229            }
230            TuneCacheResult::Pending => {
231                // Still waiting (e.g. on wasm). Try all operations as a fallback.
232                for i in 0..operations.len() {
233                    if let Ok(output) = operations.fastest(i).execute(inputs.clone()) {
234                        return output;
235                    }
236                }
237                panic!("All autotune operations failed, no viable operation found.");
238            }
239        }
240    }
241}
cubecl_runtime/tune/local.rs

cubecl_runtime/tune/
local.rs
