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