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, device::HardwareIdentity, tune::TuneCacheResult};
5#[cfg(feature = "autotune-checks")]
6use alloc::string::ToString;
7use alloc::sync::Arc;
8use alloc::vec::Vec;
9use core::{
10    any::{Any, TypeId},
11    hash::Hash,
12};
13use cubecl_environment::collections::HashMap;
14use cubecl_environment::sync::{Mutex, RwLock};
15
16/// The tunable sets a [`LocalTuner`] has built, keyed by device as well as by
17/// initializer: a set is built from the device it will run on — its client, its
18/// hardware properties — so one device's set cannot answer for another's. See
19/// [`LocalTuner::init`].
20type Sets<ID> = RwLock<Option<HashMap<(TypeId, ID), Arc<dyn Any + Send + Sync>>>>;
21
22/// A local tuner allows to create a tuner for a specific key that can be different from the server
23/// key.
24pub struct LocalTuner<AK: AutotuneKey, ID> {
25    state: Mutex<Option<Tuners<AK, ID>>>,
26    name: &'static str,
27    sets: Sets<ID>,
28}
29
30/// The tuners a [`LocalTuner`] opened. An id keeps a tuner per candidate list it executed, shared
31/// with the ids of its hardware, so ids of the same hardware running one list tune a key once for
32/// all of them.
33struct Tuners<AK: AutotuneKey, ID> {
34    /// An id's tuners, by the checksum of the list each runs.
35    by_id: HashMap<ID, HashMap<u64, Arc<Tuner<AK>>>>,
36    opened: HashMap<TunerIdentity, Arc<Tuner<AK>>>,
37}
38
39/// Which tuner a device's executes reach: one per hardware and candidate list.
40#[derive(Clone, Debug, PartialEq, Eq, Hash)]
41struct TunerIdentity {
42    hardware: HardwareIdentity,
43    /// The low half of the candidate list's [`compute_checksum`](TunableSet::compute_checksum): each
44    /// device builds its own list, and one of the same hardware may build another.
45    checksum: u64,
46}
47
48impl<AK: AutotuneKey, ID: Hash + Eq + Clone> Tuners<AK, ID> {
49    fn new() -> Self {
50        Self {
51            by_id: HashMap::new(),
52            opened: HashMap::new(),
53        }
54    }
55
56    /// The tuner `id` executes `operations` on, opened by the first id of its hardware to execute
57    /// the same candidate list.
58    fn get_or_open<I: TuneInputs, Out, Id>(
59        &mut self,
60        id: &ID,
61        name: &str,
62        client: &Client,
63        operations: &TunableSet<AK, I, Out, Id>,
64    ) -> Arc<Tuner<AK>> {
65        let checksum = operations.checksum();
66        if let Some(tuner) = self.by_id.get(id).and_then(|lists| lists.get(&checksum)) {
67            return tuner.clone();
68        }
69        let identity = TunerIdentity {
70            hardware: client.hardware_identity(),
71            checksum,
72        };
73        let tuner = self
74            .opened
75            .entry(identity)
76            .or_insert_with_key(|identity| {
77                Arc::new(Tuner::new(&name.replace("::", "-"), &identity.hardware))
78            })
79            .clone();
80        self.by_id
81            .entry(id.clone())
82            .or_default()
83            .insert(checksum, tuner.clone());
84        tuner
85    }
86}
87
88/// Create a local tuner with the provided name.
89#[macro_export]
90macro_rules! local_tuner {
91    ($name:expr) => {
92        LocalTuner::new(concat!(module_path!(), "-", $name))
93    };
94    () => {
95        LocalTuner::new(module_path!())
96    };
97}
98
99pub use local_tuner;
100
101impl<AK, ID> LocalTuner<AK, ID>
102where
103    AK: AutotuneKey + 'static,
104    ID: Hash + PartialEq + Eq + Clone,
105{
106    /// Create a new local tuner.
107    pub const fn new(name: &'static str) -> Self {
108        Self {
109            state: Mutex::new(None),
110            name,
111            sets: RwLock::new(None),
112        }
113    }
114
115    /// Get or initialize the [`TunableSet`] for `id`.
116    ///
117    /// Returns a cached `Arc<TunableSet>` keyed by the `TypeId` of `init_set`
118    /// *and* by `id`, so the initializer runs once per device rather than once
119    /// per process.
120    ///
121    /// The device is part of the key because a set is routinely built from the
122    /// device it will run on: a closure captures that device's
123    /// [`Client`] to ask what it supports, or reads its hardware
124    /// properties to decide which tunables are worth offering at all. Keyed by
125    /// the initializer alone, whichever device tuned first would answer those
126    /// questions for every device that followed — promoting kernels onto
127    /// hardware that cannot run them, or withholding kernels from hardware
128    /// that can.
129    pub fn init<I, Out, Id, F>(&self, id: &ID, init_set: F) -> Arc<TunableSet<AK, I, Out, Id>>
130    where
131        F: Fn() -> TunableSet<AK, I, Out, Id> + 'static + Send + Sync,
132        Id: Send + Sync + 'static,
133        I: TuneInputs,
134        Out: AutotuneOutput,
135    {
136        let key = (TypeId::of::<F>(), id.clone());
137        let sets = self.sets.read();
138
139        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.";
140
141        if let Some(sets) = sets.as_ref()
142            && let Some(set) = sets.get(&key)
143        {
144            return set.clone().downcast().expect(DOWNCAST_ERROR);
145        };
146
147        core::mem::drop(sets);
148
149        let mut sets = self.sets.write();
150
151        if let Some(sets) = sets.as_ref()
152            && let Some(set) = sets.get(&key)
153        {
154            return set.clone().downcast().expect(DOWNCAST_ERROR);
155        };
156
157        let content = Arc::new(init_set());
158
159        if let Some(sets) = sets.as_mut() {
160            sets.insert(key, content.clone());
161        } else {
162            let mut map = HashMap::<(TypeId, ID), Arc<dyn Any + Send + Sync>>::new();
163            map.insert(key, content.clone());
164            *sets = Some(map);
165        };
166
167        content
168    }
169
170    /// What the fastest tunable in `operations` for `key` on `id` is
171    /// [identified](super::Tunable::identified) by: `None` until a result for the key is
172    /// settled, tuned by a round in this process or read back from disk by an
173    /// [`execute`](Self::execute).
174    ///
175    /// Results are shared by every `id` of the same [hardware](crate::device::HardwareIdentity)
176    /// executing the same candidates, from the first [`execute`](Self::execute) of each: an `id`
177    /// that has not executed `operations` sees none.
178    ///
179    /// It reads results and never produces one: it never starts a round, never waits on one in
180    /// flight, never reads the persistent cache, and never resets the tuner's cache after an
181    /// environment switch. It reports nothing settled there until the next `execute`.
182    pub fn fastest_identity<I, Out, Id>(
183        &self,
184        id: &ID,
185        operations: &TunableSet<AK, I, Out, Id>,
186        key: &AK,
187    ) -> Option<Id>
188    where
189        I: TuneInputs,
190        Id: Clone,
191    {
192        let tuner = self
193            .state
194            .lock()
195            .as_ref()?
196            .by_id
197            .get(id)?
198            .get(&operations.checksum())?
199            .clone();
200        let fastest_index = tuner.settled(key)?;
201        Some(operations.identity(fastest_index).clone())
202    }
203
204    /// Clear the autotune state.
205    pub fn clear(&self) {
206        self.state.lock().take();
207    }
208
209    #[cfg(feature = "autotune-checks")]
210    fn checks<'a, I: TuneInputs, Out: AutotuneOutput, Id>(
211        &self,
212        operations: &TunableSet<AK, I, Out, Id>,
213        inputs: &<I as TuneInputs>::At<'a>,
214    ) -> alloc::vec::Vec<crate::tune::log::CheckResult>
215    where
216        <I as TuneInputs>::At<'a>: Clone + Send,
217    {
218        let mut checks_outputs = Vec::new();
219        for i in 0..operations.len() {
220            let op = operations.fastest(i);
221            let result = op.execute(inputs.clone());
222            checks_outputs.push((op.name.to_string(), result));
223        }
224        super::check_autotune_outputs(checks_outputs)
225    }
226
227    /// Execute the fastest operation in a [`TunableSet`], triggering a tuning pass on
228    /// the first call for a given key.
229    pub fn execute<'a, I: TuneInputs, Out, Id>(
230        &self,
231        id: &ID,
232        client: &Client,
233        operations: Arc<TunableSet<AK, I, Out, Id>>,
234        inputs: <I as TuneInputs>::At<'a>,
235    ) -> Out
236    where
237        <I as TuneInputs>::At<'a>: Clone + Send,
238        Out: AutotuneOutput,
239        Id: Send + Sync,
240    {
241        let key = operations.generate_key(&inputs);
242
243        let tuner = self
244            .state
245            .lock()
246            .get_or_insert_with(Tuners::new)
247            .get_or_open(id, self.name, client, &operations);
248
249        // Fast path: a cached hit skips straight to the fastest operation.
250        // `fastest` also resets the tuner cache if the environment switched, so
251        // a miss here falls through to `check_tune`, which reads the store.
252        if let TuneCacheResult::Hit { fastest_index } = tuner.fastest(&key) {
253            // `autotune-checks` cross-checks the candidates of every launch, a hit's included.
254            // Bound first, so the logger guard drops before the checks run every candidate.
255            #[cfg(feature = "autotune-checks")]
256            {
257                let mut log_context =
258                    crate::tune::AutotuneLogContext::new(&mut tuner.logger().lock());
259                log_context.set_checks(|| self.checks::<I, Out, Id>(&operations, &inputs));
260            }
261            return operations
262                .fastest(fastest_index)
263                .execute(inputs)
264                .expect("Should run when selected by autotune.");
265        }
266
267        // Built after the fast path: a hit logs nothing, so it takes no logger lock.
268        #[allow(unused_mut)]
269        let mut log_context = crate::tune::AutotuneLogContext::new(&mut tuner.logger().lock());
270
271        #[cfg(feature = "autotune-checks")]
272        log_context.set_checks(|| self.checks::<I, Out, Id>(&operations, &inputs));
273
274        let fastest =
275            tuner.check_tune::<I, Out, Id>(&key, &inputs, &operations, client, log_context);
276
277        // Run the execution depending on the cache state.
278        match fastest {
279            TuneCacheResult::Hit { fastest_index } => operations
280                .fastest(fastest_index)
281                .execute(inputs)
282                .expect("Should run when selected by autotune."),
283            TuneCacheResult::Miss => {
284                panic!(
285                    "Somehow we STILL didn't check a tuning checksum or start tuning, something has gone wrong."
286                )
287            }
288            // Still waiting (e.g. on wasm), stopped on a lost device, or its candidates' kernels
289            // queued under a `CompileOnly` override: run the first operation that serves the
290            // problem, which there only queues its kernels.
291            TuneCacheResult::Pending | TuneCacheResult::Compiled => {
292                let mut failures = Vec::new();
293                for i in 0..operations.len() {
294                    match operations.fastest(i).execute(inputs.clone()) {
295                        Ok(output) => return output,
296                        Err(err) => failures.push(err),
297                    }
298                }
299                panic!("All autotune operations failed, no viable operation found.\n{failures:?}");
300            }
301        }
302    }
303}
cubecl_runtime/tune/local.rs

cubecl_runtime/tune/
local.rs
