cubecl_runtime/tune/
local.rs

1use super::{AutotuneError, 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 execute_settled(&operations, &key, fastest_index, inputs);
262        }
263
264        // Built after the fast path: a hit logs nothing, so it takes no logger lock.
265        #[allow(unused_mut)]
266        let mut log_context = crate::tune::AutotuneLogContext::new(&mut tuner.logger().lock());
267
268        #[cfg(feature = "autotune-checks")]
269        log_context.set_checks(|| self.checks::<I, Out, Id>(&operations, &inputs));
270
271        let fastest =
272            tuner.check_tune::<I, Out, Id>(&key, &inputs, &operations, client, log_context);
273
274        // Run the execution depending on the cache state.
275        match fastest {
276            TuneCacheResult::Hit { fastest_index } => {
277                execute_settled(&operations, &key, fastest_index, inputs)
278            }
279            TuneCacheResult::Miss => {
280                panic!(
281                    "Somehow we STILL didn't check a tuning checksum or start tuning, something has gone wrong."
282                )
283            }
284            // Still waiting (e.g. on wasm), stopped on a lost device, or its candidates' kernels
285            // queued under a `CompileOnly` override, where an operation only queues its kernels:
286            // the candidates run in the plan's order, the one the measurement follows, so the
287            // best ranked candidate that launches runs rather than the first declared.
288            TuneCacheResult::Pending | TuneCacheResult::Compiled => {
289                match execute_in_plan_order(&operations, &key, inputs) {
290                    Ok(output) => output,
291                    Err(failures) => panic!(
292                        "All autotune operations failed, no viable operation found.\n{failures:?}"
293                    ),
294                }
295            }
296        }
297    }
298}
299
300/// Runs the candidates in the plan's [order](super::TunePlan::order) and hands back the
301/// first output, or every candidate's error when none launches. A candidate that fails to
302/// launch is skipped, as the measurement skips it.
303/// The candidate `key` settled on, run on `inputs`; where it refuses them, the candidates in the
304/// plan's order, the first that launches.
305///
306/// A key stands for every call of its bucket, and its candidates are prepared and timed on one
307/// representative of it: a winner can still refuse a call the representative did not show, as a
308/// launch whose grid grows with a batch the key does not tell apart passes the device's limit.
309/// Every candidate computes the same output, so the call runs the next one that serves it, as a
310/// call does while its key is still tuning.
311///
312/// # Panics
313///
314/// When no candidate launches on `inputs`.
315fn execute_settled<'a, AK: AutotuneKey, I: TuneInputs, Out: 'static, Id>(
316    operations: &TunableSet<AK, I, Out, Id>,
317    key: &AK,
318    fastest_index: usize,
319    inputs: <I as TuneInputs>::At<'a>,
320) -> Out
321where
322    <I as TuneInputs>::At<'a>: Clone,
323{
324    let refused = match operations.fastest(fastest_index).execute(inputs.clone()) {
325        Ok(output) => return output,
326        Err(refused) => refused,
327    };
328    match execute_in_plan_order(operations, key, inputs) {
329        Ok(output) => output,
330        Err(failures) => panic!(
331            "The operation autotune settled on refused this call ({refused:?}), and so did every \
332             other candidate.\n{failures:?}"
333        ),
334    }
335}
336
337fn execute_in_plan_order<'a, AK: AutotuneKey, I: TuneInputs, Out: 'static, Id>(
338    operations: &TunableSet<AK, I, Out, Id>,
339    key: &AK,
340    inputs: <I as TuneInputs>::At<'a>,
341) -> Result<Out, Vec<AutotuneError>>
342where
343    <I as TuneInputs>::At<'a>: Clone,
344{
345    let mut failures = Vec::new();
346    for (index, _) in operations.plan(key).order(operations.len()) {
347        match operations.fastest(index).execute(inputs.clone()) {
348            Ok(output) => return Ok(output),
349            Err(err) => failures.push(err),
350        }
351    }
352    Err(failures)
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358    use crate::tune::{Tunable, TuneGroup};
359    use alloc::string::String;
360    use core::fmt::Display;
361    use serde::{Deserialize, Serialize};
362
363    #[derive(Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize, Debug)]
364    struct FakeAutotuneKey;
365
366    impl Display for FakeAutotuneKey {
367        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
368            f.write_str("FakeAutotuneKey")
369        }
370    }
371
372    impl AutotuneKey for FakeAutotuneKey {}
373
374    fn set() -> TunableSet<FakeAutotuneKey, (), usize> {
375        TunableSet::new(|_: &()| FakeAutotuneKey, |_: &FakeAutotuneKey, _: &()| ())
376    }
377
378    fn answers(index: usize) -> impl Fn(()) -> Result<usize, String> {
379        move |_| Ok(index)
380    }
381
382    fn fails(_: ()) -> Result<usize, String> {
383        Err("does not launch here".into())
384    }
385
386    #[test]
387    fn the_pending_fallback_runs_the_plan_s_first_candidate_not_the_first_declared() {
388        let group = TuneGroup::<FakeAutotuneKey>::new("group", |_| 1);
389        let operations = set()
390            .with(Tunable::<FakeAutotuneKey, (), usize>::new(
391                "declared_first",
392                answers(0),
393            ))
394            .with(
395                Tunable::<FakeAutotuneKey, (), usize>::new("planned_first", answers(1))
396                    .group(&group, |_| 1),
397            );
398
399        assert_eq!(
400            execute_in_plan_order(&operations, &FakeAutotuneKey, ()).ok(),
401            Some(1)
402        );
403    }
404
405    #[test]
406    fn the_pending_fallback_moves_down_the_plan_when_a_candidate_does_not_launch() {
407        let group = TuneGroup::<FakeAutotuneKey>::new("group", |_| 1);
408        let operations = set()
409            .with(
410                Tunable::<FakeAutotuneKey, (), usize>::new("planned_last", answers(0))
411                    .group(&group, |_| 0),
412            )
413            .with(
414                Tunable::<FakeAutotuneKey, (), usize>::new("planned_first", fails)
415                    .group(&group, |_| 2),
416            )
417            .with(
418                Tunable::<FakeAutotuneKey, (), usize>::new("planned_second", answers(2))
419                    .group(&group, |_| 1),
420            );
421
422        assert_eq!(
423            execute_in_plan_order(&operations, &FakeAutotuneKey, ()).ok(),
424            Some(2)
425        );
426    }
427
428    #[test]
429    fn a_settled_winner_that_refuses_a_call_falls_back_down_the_plan() {
430        let group = TuneGroup::<FakeAutotuneKey>::new("group", |_| 1);
431        let operations = set()
432            .with(
433                Tunable::<FakeAutotuneKey, (), usize>::new("settled_winner", fails)
434                    .group(&group, |_| 2),
435            )
436            .with(
437                Tunable::<FakeAutotuneKey, (), usize>::new("planned_second", answers(1))
438                    .group(&group, |_| 1),
439            );
440
441        assert_eq!(execute_settled(&operations, &FakeAutotuneKey, 0, ()), 1);
442    }
443
444    #[test]
445    fn a_settled_winner_that_serves_the_call_runs() {
446        let operations = set()
447            .with(Tunable::<FakeAutotuneKey, (), usize>::new(
448                "first",
449                answers(0),
450            ))
451            .with(Tunable::<FakeAutotuneKey, (), usize>::new(
452                "settled_winner",
453                answers(1),
454            ));
455
456        assert_eq!(execute_settled(&operations, &FakeAutotuneKey, 1, ()), 1);
457    }
458
459    #[test]
460    fn the_pending_fallback_tries_what_the_plan_left_out_last() {
461        let never = TuneGroup::<FakeAutotuneKey>::new("never", |_| -1);
462        let operations = set()
463            .with(Tunable::<FakeAutotuneKey, (), usize>::new(
464                "declared_first",
465                fails,
466            ))
467            .with(
468                Tunable::<FakeAutotuneKey, (), usize>::new("left_out", answers(1))
469                    .group(&never, |_| 1),
470            );
471
472        assert_eq!(
473            execute_in_plan_order(&operations, &FakeAutotuneKey, ()).ok(),
474            Some(1)
475        );
476        let nothing = set().with(Tunable::<FakeAutotuneKey, (), usize>::new("fails", fails));
477        assert!(execute_in_plan_order(&nothing, &FakeAutotuneKey, ()).is_err());
478    }
479}
cubecl_runtime/tune/local.rs

cubecl_runtime/tune/
local.rs
