cubecl_runtime/tune/
local.rs

1use super::{AutotuneKey, AutotuneOperationSet, Tuner};
2use crate::{
3    channel::ComputeChannel, client::ComputeClient, server::ComputeServer, tune::TuneCacheResult,
4};
5use core::{fmt::Display, hash::Hash};
6use hashbrown::HashMap;
7
8#[cfg(not(feature = "std"))]
9use alloc::{boxed::Box, string::ToString};
10
11/// A local tuner allows to create a tuner for a specific key that can be different from the server
12/// key.
13pub struct LocalTuner<AK: AutotuneKey, ID> {
14    state: spin::RwLock<Option<HashMap<ID, Tuner<AK>>>>,
15    name: &'static str,
16}
17
18/// Create a local tuner with the provided name.
19#[macro_export]
20macro_rules! local_tuner {
21    ($name:expr) => {
22        LocalTuner::new(concat!(module_path!(), "-", $name));
23    };
24    () => {
25        LocalTuner::new(module_path!());
26    };
27}
28
29pub use local_tuner;
30
31impl<AK: AutotuneKey + 'static, ID: Hash + PartialEq + Eq + Clone + Display> LocalTuner<AK, ID> {
32    /// Create a new local tuner.
33    pub const fn new(name: &'static str) -> Self {
34        Self {
35            state: spin::RwLock::new(None),
36            name,
37        }
38    }
39
40    /// Clear the autotune state.
41    pub fn clear(&self) {
42        let mut state = self.state.write();
43        *state = None;
44    }
45
46    /// Execute the best operation in the provided [autotune operation set](AutotuneOperationSet)
47    pub fn execute<S, C, Out: Send + 'static>(
48        &self,
49        id: &ID,
50        client: &ComputeClient<S, C>,
51        autotune_operation_set: Box<dyn AutotuneOperationSet<AK, Out>>,
52    ) -> Out
53    where
54        S: ComputeServer + 'static,
55        C: ComputeChannel<S> + 'static,
56    {
57        let key = autotune_operation_set.key();
58
59        // If this is cached and ready, use the operation.
60        if let Some(map) = self.state.read().as_ref() {
61            if let Some(tuner) = map.get(id) {
62                if let TuneCacheResult::Hit { fastest_index } = tuner.fastest(&key) {
63                    let op = autotune_operation_set.fastest(fastest_index);
64                    return op.execute().expect("Should run when selected by autotune.");
65                }
66            }
67        }
68
69        // Create the tuner if needed, and update some state like
70        // checksums if need be.
71        let fastest = {
72            let mut state = self.state.write();
73            let map = state.get_or_insert_with(Default::default);
74            let tuner = map.entry(id.clone()).or_insert_with(move || {
75                let name = self.name.replace("::", "-");
76                Tuner::new(&name, &id.to_string())
77            });
78
79            #[allow(unused_mut)]
80            let mut fastest = tuner.fastest(&key);
81
82            // If the cache checksum hasn't been checked, do so now, and retry.
83            #[cfg(autotune_persistent_cache)]
84            if matches!(fastest, TuneCacheResult::Unchecked) {
85                let checksum = autotune_operation_set.compute_checksum();
86                tuner.validate_checksum(&key, &checksum);
87                fastest = tuner.fastest(&key);
88            }
89            fastest
90        };
91
92        match fastest {
93            TuneCacheResult::Hit { fastest_index } => {
94                return autotune_operation_set
95                    .fastest(fastest_index)
96                    .execute()
97                    .expect("Should run when selected by autotune.");
98            }
99            TuneCacheResult::Miss => {
100                // We don't know the results yet, start autotuning.
101                //
102                // Running benchmarks shound't lock the tuner, since an autotune operation can recursively use the
103                // same tuner.
104                //
105                // # Example
106                //
107                // ```
108                // - tune_1 start
109                //   - tune_2 start
110                //   - tune_2 save
111                // - tune_1 save
112                // ```
113                let state = self.state.read();
114                let state = state.as_ref().expect("Should be initialized");
115                let tuner = state.get(id).expect("Should be initialized");
116
117                tuner.execute_autotune(autotune_operation_set.as_ref(), client);
118            }
119            TuneCacheResult::Pending => {
120                // We're waiting for results to come in.
121            }
122            TuneCacheResult::Unchecked => {
123                panic!("Should have checked the cache already.")
124            }
125        };
126
127        let fastest = {
128            let mut state = self.state.write();
129            let state = state.as_mut().expect("Should be initialized");
130            let tuner = state.get_mut(id).expect("Should be initialized");
131            // Now read all results that have come in since.
132            tuner.resolve();
133
134            // Check again what the fastest option is, new results might have come in.
135            match tuner.fastest(&key) {
136                TuneCacheResult::Hit { fastest_index } => {
137                    // Theres a known good value - just run it.
138                    fastest_index
139                }
140                TuneCacheResult::Pending => {
141                    // If we still don't know, just execute a default index.
142                    0
143                }
144                TuneCacheResult::Miss => {
145                    panic!("Should have at least started autotuning");
146                }
147                TuneCacheResult::Unchecked => {
148                    panic!("Should have checked the cache.")
149                }
150            }
151        };
152
153        autotune_operation_set
154            .fastest(fastest)
155            .execute()
156            .expect("Should run when selected by autotune.")
157    }
158}
cubecl_runtime/tune/local.rs

cubecl_runtime/tune/
local.rs
