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, F>(&self, id: &ID, init_set: F) -> Arc<TunableSet<AK, I, Out>>
71    where
72        F: Fn() -> TunableSet<AK, I, Out> + 'static + Send + Sync,
73        I: TuneInputs,
74        Out: AutotuneOutput,
75    {
76        let key = (TypeId::of::<F>(), id.clone());
77        let sets = self.sets.read();
78
79        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.";
80
81        if let Some(sets) = sets.as_ref()
82            && let Some(set) = sets.get(&key)
83        {
84            return set.clone().downcast().expect(DOWNCAST_ERROR);
85        };
86
87        core::mem::drop(sets);
88
89        let mut sets = self.sets.write();
90
91        if let Some(sets) = sets.as_ref()
92            && let Some(set) = sets.get(&key)
93        {
94            return set.clone().downcast().expect(DOWNCAST_ERROR);
95        };
96
97        let content = Arc::new(init_set());
98
99        if let Some(sets) = sets.as_mut() {
100            sets.insert(key, content.clone());
101        } else {
102            let mut map = HashMap::<(TypeId, ID), Arc<dyn Any + Send + Sync>>::new();
103            map.insert(key, content.clone());
104            *sets = Some(map);
105        };
106
107        content
108    }
109
110    /// Clear the autotune state.
111    pub fn clear(&self) {
112        if let Some(s) = self.state.lock().as_mut() {
113            s.clear()
114        }
115    }
116
117    #[cfg(feature = "autotune-checks")]
118    fn checks<'a, I: TuneInputs, Out: AutotuneOutput>(
119        &self,
120        operations: &TunableSet<AK, I, Out>,
121        inputs: &<I as TuneInputs>::At<'a>,
122    ) -> alloc::vec::Vec<crate::tune::log::CheckResult>
123    where
124        <I as TuneInputs>::At<'a>: Clone + Send,
125    {
126        use alloc::vec::Vec;
127
128        let mut checks_outputs = Vec::new();
129        for i in 0..operations.len() {
130            let op = operations.fastest(i);
131            let result = op.execute(inputs.clone());
132            checks_outputs.push((op.name.to_string(), result));
133        }
134        super::check_autotune_outputs(checks_outputs)
135    }
136
137    /// Execute the fastest operation in a [`TunableSet`], triggering a tuning pass on
138    /// the first call for a given key.
139    pub fn execute<'a, I: TuneInputs, Out>(
140        &self,
141        id: &ID,
142        client: &Client,
143        operations: Arc<TunableSet<AK, I, Out>>,
144        inputs: <I as TuneInputs>::At<'a>,
145    ) -> Out
146    where
147        <I as TuneInputs>::At<'a>: Clone + Send,
148        Out: AutotuneOutput,
149    {
150        let key = operations.generate_key(&inputs);
151
152        let tuner = {
153            let mut state_lock = self.state.lock();
154            let state_map = state_lock.get_or_insert_with(|| HashMap::new());
155            state_map
156                .entry(id.clone())
157                .or_insert_with(move || {
158                    let name = self.name.replace("::", "-");
159                    Arc::new(Tuner::new(&name, &id.to_string()))
160                })
161                .clone()
162        };
163
164        #[allow(unused_mut)]
165        let mut log_context = crate::tune::AutotuneLogContext::new(&mut tuner.logger().lock());
166
167        #[cfg(feature = "autotune-checks")]
168        log_context.set_checks(|| self.checks::<I, Out>(&operations, &inputs));
169
170        // Fast path: a cached hit skips straight to the fastest operation.
171        // `fastest` also resets the tuner cache if the environment switched, so
172        // a miss here falls through to `check_tune`, which re-hydrates.
173        if let TuneCacheResult::Hit { fastest_index } = tuner.fastest(&key) {
174            return operations
175                .fastest(fastest_index)
176                .execute(inputs)
177                .expect("Should run when selected by autotune.");
178        }
179
180        let fastest = tuner.check_tune::<I, Out>(
181            &key,
182            &inputs,
183            &operations,
184            || operations.compute_checksum(),
185            client,
186            log_context,
187        );
188
189        // Run the execution depending on the cache state.
190        match fastest {
191            TuneCacheResult::Hit { fastest_index } => operations
192                .fastest(fastest_index)
193                .execute(inputs)
194                .expect("Should run when selected by autotune."),
195            TuneCacheResult::Unchecked | TuneCacheResult::Miss => {
196                panic!(
197                    "Somehow we STILL didn't check a tuning checksum or start tuning, something has gone wrong."
198                )
199            }
200            TuneCacheResult::Pending => {
201                // Still waiting (e.g. on wasm). Try all operations as a fallback.
202                for i in 0..operations.len() {
203                    if let Ok(output) = operations.fastest(i).execute(inputs.clone()) {
204                        return output;
205                    }
206                }
207                panic!("All autotune operations failed, no viable operation found.");
208            }
209        }
210    }
211}
cubecl_runtime/tune/local.rs

cubecl_runtime/tune/
local.rs
