1use alloc::boxed::Box;
2#[cfg(persistence)]
3use alloc::format;
4use alloc::sync::Arc;
5use alloc::vec::Vec;
6use cubecl_common::profile::ProfileDuration;
7use derive_more::Display;
8
9use core::time::Duration;
10
11use cubecl_environment::sync::Mutex;
12
13use alloc::string::{String, ToString};
14use cubecl_common::benchmark::{BenchmarkComputations, BenchmarkDurations};
15
16use crate::client::Client;
17use crate::config::Logger;
18#[cfg(persistence)]
19use crate::config::autotune::AutotuneLogLevel;
20use crate::server::LaunchError;
21use crate::tune::{AutotuneLoggerExt, AutotuneResult, TimeBound, TuneCache, tune_benchmark};
22use cubecl_environment::config::RuntimeConfig;
23
24use super::{
25 AutotuneKey, AutotuneOutput, TunableSet, TuneCacheResult, TuneFn, TuneInputs, TunePlan,
26};
27
28#[derive(Debug)]
29pub struct Tuner<K: AutotuneKey> {
35 cache: Arc<Mutex<TuneCache<K>>>,
36 logger: Arc<Mutex<Logger>>,
37}
38
39#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
41#[derive(new, Debug, Clone, PartialEq, Eq)]
42pub struct AutotuneOutcome {
43 pub name: String,
45 pub index: usize,
47 pub computation: BenchmarkComputations,
49}
50
51impl core::fmt::Display for AutotuneOutcome {
52 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
53 write!(
54 f,
55 "Autotune[{}] name {} => {:?}",
56 self.index, self.name, self.computation
57 )
58 }
59}
60
61#[derive(Clone, Display)]
63#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
64pub enum AutotuneError {
65 #[display("{name}: An unknown error happened.\n{err}")]
67 Unknown {
68 name: String,
70 err: String,
72 },
73 #[display("{name}: All samples are invalid.")]
75 InvalidSamples {
76 name: String,
78 },
79 #[display("No autotune was flagged as valid for the problem.\n{context}")]
85 NoValidKernelFound {
86 context: String,
88 },
89 #[display("{name}: A profiled sample carried no measurement.")]
95 NotMeasured {
96 name: String,
98 },
99 #[display("{name}: The autotune is skipped manually.")]
101 Skip {
102 name: String,
104 },
105
106 Launch(LaunchError),
108}
109
110impl core::fmt::Debug for AutotuneError {
111 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
112 write!(f, "{self}")
113 }
114}
115
116impl From<LaunchError> for AutotuneError {
117 fn from(value: LaunchError) -> Self {
118 Self::Launch(value)
119 }
120}
121
122struct PendingBench {
124 index: usize,
125 name: String,
126 profiles: Vec<ProfileDuration>,
127 launch: Option<Duration>,
130}
131
132struct TuneJob<'t, 'i, K: AutotuneKey, F: TuneInputs, Out> {
135 key: K,
136 autotunables: Vec<&'t TuneFn<F, Out>>,
137 test_inputs: <F as TuneInputs>::At<'i>,
138 evictor: Option<Box<crate::tune::Evictor<'i>>>,
140 plan: TunePlan,
141 results: Vec<AutotuneResult>,
142 #[cfg(any(not(target_family = "wasm"), persistence))]
143 limit: Option<Duration>,
144 #[cfg(persistence)]
145 bounds: Option<crate::tune::Bounds>,
146 #[cfg(not(target_family = "wasm"))]
147 short_circuit: bool,
148 #[cfg(persistence)]
149 checksum: String,
150 log_context: Option<crate::tune::AutotuneLogContext>,
151 #[cfg(persistence)]
152 recording: crate::tune::record::TuneRecording<K>,
153}
154
155impl<K: AutotuneKey, F: TuneInputs, Out> TuneJob<'_, '_, K, F, Out> {
156 fn into_request(self, pending: Vec<PendingBench>, decided: Option<usize>) -> TuneRequest<K> {
157 TuneRequest {
158 key: self.key,
159 results: self.results,
160 #[cfg(persistence)]
161 checksum: self.checksum,
162 log_context: self.log_context,
163 pending,
164 decided,
165 #[cfg(persistence)]
166 limit: self.limit,
167 #[cfg(persistence)]
168 bounds: self.bounds,
169 #[cfg(persistence)]
170 recording: self.recording,
171 }
172 }
173}
174
175struct TuneRequest<K: AutotuneKey> {
178 key: K,
179 results: Vec<AutotuneResult>,
180 #[cfg(persistence)]
181 checksum: String,
182 log_context: Option<crate::tune::AutotuneLogContext>,
183 pending: Vec<PendingBench>,
184 decided: Option<usize>,
187 #[cfg(persistence)]
188 limit: Option<Duration>,
189 #[cfg(persistence)]
190 bounds: Option<crate::tune::Bounds>,
191 #[cfg(persistence)]
192 recording: crate::tune::record::TuneRecording<K>,
193}
194
195#[allow(clippy::new_without_default)]
196impl<K: AutotuneKey> Tuner<K> {
197 pub fn new(name: &str, device_id: &str) -> Self {
201 Self {
202 cache: Arc::new(Mutex::new(TuneCache::new(name, device_id))),
203 logger: Arc::new(Mutex::new(Logger::new())),
204 }
205 }
206
207 pub fn fastest(&self, key: &K) -> TuneCacheResult {
216 #[cfg_attr(not(persistence), allow(unused_mut))]
217 let mut cache = self.cache.lock();
218 #[cfg(persistence)]
219 cache.reset_if_environment_switched();
220
221 cache.fastest(key)
222 }
223
224 pub fn logger(&self) -> Arc<Mutex<Logger>> {
226 self.logger.clone()
227 }
228
229 pub fn check_tune<'a, F: TuneInputs, Out: AutotuneOutput>(
232 &self,
233 key: &K,
234 inputs: &F::At<'a>,
235 tunables: &TunableSet<K, F, Out>,
236 #[cfg_attr(not(persistence), allow(unused))] checksum: impl FnOnce() -> String + Send + Sync,
237 client: &Client,
238 mut log_context: Option<crate::tune::AutotuneLogContext>,
239 ) -> TuneCacheResult
240 where
241 <F as TuneInputs>::At<'a>: Clone + Send,
242 {
243 {
244 let mut cache = self.cache.lock();
245 #[cfg(persistence)]
246 cache.reset_if_environment_switched();
247 let cur = cache.fastest(key);
248
249 #[cfg(persistence)]
253 let cur = if matches!(cur, TuneCacheResult::Miss) {
254 cache.sync_persistent();
255 cache.fastest(key)
256 } else {
257 cur
258 };
259
260 #[cfg(persistence)]
261 let cur = if matches!(cur, TuneCacheResult::Unchecked) {
262 let mut log = self.logger.lock();
263 let checksum = checksum();
264 if let AutotuneLogLevel::Full = log.log_level_autotune() {
265 log.log_autotune(&format!("validate checksum key={key}, checksum={checksum}"));
266 }
267 cache.validate_checksum(key, &checksum)
268 } else {
269 cur
270 };
271
272 match cur {
273 TuneCacheResult::Hit { .. } | TuneCacheResult::Pending => return cur,
274 TuneCacheResult::Miss | TuneCacheResult::Unchecked => {
275 cache.mark_pending(key.clone())
276 }
277 }
278 }
281
282 log::info!("Tuning {key}");
283
284 let autotunables = tunables.autotunables().collect::<Vec<_>>();
285 let results: Vec<AutotuneResult> = autotunables
286 .iter()
287 .map(|a| {
288 AutotuneResult::error(AutotuneError::Skip {
289 name: a.name.to_string(),
290 })
291 })
292 .collect();
293
294 #[cfg(persistence)]
295 let checksum = tunables.compute_checksum();
296
297 if results.len() == 1 {
299 self.cache.lock().cache_insert(key.clone(), 0);
300 return TuneCacheResult::Hit { fastest_index: 0 };
301 }
302
303 #[cfg(persistence)]
306 let recording = crate::tune::record::TuneRecording::new(&self.cache.lock(), key, &checksum);
307 #[cfg(persistence)]
310 if recording.is_open() {
311 log_context.get_or_insert_with(Default::default);
312 }
313
314 let test_inputs = tunables.generate_inputs(key, inputs);
315 let plan = tunables.plan(key);
316 let bounds = tunables.bounds(key, inputs);
317 let limit = bounds.as_ref().and_then(|bounds| bounds.time_limit());
318
319 log_context.set_bounds(bounds.clone());
320 log_context.set_limit(limit);
321
322 #[cfg(not(target_family = "wasm"))]
325 let short_circuit = limit.is_some()
326 && tunables.is_short_circuit_enabled()
327 && !crate::config::CubeClRuntimeConfig::get()
328 .autotune
329 .disable_short_circuit;
330
331 let job = TuneJob {
332 key: key.clone(),
333 autotunables,
334 test_inputs,
335 evictor: tunables.evictor(key, inputs),
336 plan,
337 results,
338 #[cfg(any(not(target_family = "wasm"), persistence))]
339 limit,
340 #[cfg(persistence)]
341 bounds,
342 #[cfg(not(target_family = "wasm"))]
343 short_circuit,
344 #[cfg(persistence)]
345 checksum,
346 log_context,
347 #[cfg(persistence)]
348 recording,
349 };
350
351 #[cfg(not(target_family = "wasm"))]
352 if crate::config::CubeClRuntimeConfig::get()
353 .autotune
354 .bench
355 .adaptive
356 {
357 return self.tune_adaptive(job, client);
358 }
359
360 self.tune_fixed_samples(job, client)
361 }
362
363 #[cfg(not(target_family = "wasm"))]
366 fn tune_adaptive<'i, F: TuneInputs, Out: AutotuneOutput>(
367 &self,
368 mut job: TuneJob<'_, 'i, K, F, Out>,
369 client: &Client,
370 ) -> TuneCacheResult
371 where
372 <F as TuneInputs>::At<'i>: Clone + Send,
373 {
374 let mut schedule = crate::tune::schedule::Schedule {
375 config: crate::config::CubeClRuntimeConfig::get()
376 .autotune
377 .bench
378 .clone(),
379 limit: job.limit,
380 short_circuit: job.short_circuit,
381 track_steps: job.log_context.is_some(),
382 evictor: job.evictor.take(),
383 };
384
385 let outcome = schedule.run_plan(
386 &job.key,
387 &mut job.plan,
388 &job.autotunables,
389 &job.test_inputs,
390 client,
391 &mut job.results,
392 );
393
394 for (name, duration) in outcome.steps {
395 job.log_context.push_tuning_step(name, duration);
396 }
397 if let Some(name) = outcome.short_circuit {
398 job.log_context.push_short_circuit(name);
399 }
400
401 let request = job.into_request(Vec::new(), outcome.decided);
402
403 cubecl_environment::future::block_on(process_request(request, &self.cache, &self.logger))
404 }
405
406 fn tune_fixed_samples<'i, F: TuneInputs, Out: AutotuneOutput>(
409 &self,
410 mut job: TuneJob<'_, 'i, K, F, Out>,
411 client: &Client,
412 ) -> TuneCacheResult
413 where
414 <F as TuneInputs>::At<'i>: Clone + Send,
415 {
416 #[cfg(not(target_family = "wasm"))]
420 let mut batch_success = false;
421 #[cfg(target_family = "wasm")]
422 let batch_success = false;
423
424 let mut pending = Vec::<PendingBench>::new();
429 loop {
430 let tunable_indices = job.plan.next();
431
432 if tunable_indices.is_empty() {
433 let key = &job.key;
434 panic!(
435 "Can't execute the autotune plan for key: {key:?}\n - plan: {:?}\n - results: {:?}",
436 job.plan, job.results
437 );
438 }
439
440 for index in tunable_indices {
441 let op = job.autotunables[index];
442
443 let start_time = job
444 .log_context
445 .is_some()
446 .then(cubecl_common::profile::Instant::now);
447
448 match tune_benchmark(
449 op,
450 job.test_inputs.clone(),
451 client.clone(),
452 job.evictor.as_deref_mut(),
453 ) {
454 Ok(profiles) => {
455 let bench = PendingBench {
456 index,
457 name: op.name.clone(),
458 profiles,
459 launch: start_time.map(|start| start.elapsed()),
460 };
461
462 #[cfg(not(target_family = "wasm"))]
463 if job.short_circuit {
464 let result = cubecl_environment::future::block_on(resolve_bench(bench));
465
466 let close_enough = result
468 .outcome
469 .as_ref()
470 .is_ok_and(|out| out.computation.median <= job.limit.unwrap());
471
472 batch_success |= result.outcome.is_ok();
473 job.results[index] = result;
474
475 if let Some(start) = start_time {
476 job.log_context
477 .push_tuning_step(op.name.to_string(), start.elapsed());
478 }
479
480 if close_enough {
481 job.log_context.push_short_circuit(op.name.to_string());
482 break;
483 }
484
485 continue;
486 }
487
488 pending.push(bench);
491 }
492 Err(err) => {
493 job.results[index] = AutotuneResult::error(err);
494 if let Some(start) = start_time {
495 job.log_context
496 .push_tuning_step(op.name.to_string(), start.elapsed());
497 }
498 }
499 }
500 }
501
502 #[cfg(not(target_family = "wasm"))]
503 if !pending.is_empty() || batch_success {
504 break;
505 }
506 #[cfg(target_family = "wasm")]
507 if !pending.is_empty() {
508 break;
509 }
510 }
511
512 let request = job.into_request(pending, None);
515
516 #[cfg(target_family = "wasm")]
519 {
520 let cache = self.cache.clone();
521 let logger = self.logger.clone();
522 wasm_bindgen_futures::spawn_local(async move {
523 process_request(request, &cache, &logger).await;
524 });
525
526 return TuneCacheResult::Pending;
527 }
528
529 #[cfg(not(target_family = "wasm"))]
530 cubecl_environment::future::block_on(process_request(request, &self.cache, &self.logger))
531 }
532}
533
534async fn resolve_bench(bench: PendingBench) -> AutotuneResult {
540 let PendingBench {
541 index,
542 name,
543 profiles,
544 launch: _,
545 } = bench;
546
547 let Some(first) = profiles.first() else {
548 return AutotuneResult::error(AutotuneError::Unknown {
549 name: name.to_string(),
550 err: "No profiling available".to_string(),
551 });
552 };
553 let timing_method = first.timing_method();
554
555 let Some(durations) =
559 futures_util::future::join_all(profiles.into_iter().map(ProfileDuration::resolve))
560 .await
561 .into_iter()
562 .map(|ticks| ticks.map(|ticks| ticks.duration()))
563 .collect::<Option<Vec<Duration>>>()
564 else {
565 return AutotuneResult::error(AutotuneError::NotMeasured {
566 name: name.to_string(),
567 });
568 };
569
570 AutotuneResult::success(AutotuneOutcome::new(
571 name,
572 index,
573 BenchmarkComputations::new(&BenchmarkDurations::from_durations(
574 timing_method,
575 durations,
576 )),
577 ))
578}
579
580async fn process_request<K: AutotuneKey>(
582 request: TuneRequest<K>,
583 cache: &Mutex<TuneCache<K>>,
584 logger: &Mutex<Logger>,
585) -> TuneCacheResult {
586 let TuneRequest {
587 key,
588 mut results,
589 #[cfg(persistence)]
590 checksum,
591 mut log_context,
592 pending,
593 decided,
594 #[cfg(persistence)]
595 limit,
596 #[cfg(persistence)]
597 bounds,
598 #[cfg(persistence)]
599 recording,
600 } = request;
601
602 let resolved = futures_util::future::join_all(pending.into_iter().map(|bench| {
607 let index = bench.index;
608 let name = bench.name.clone();
609 let launch = bench.launch;
610
611 async move {
612 let started = cubecl_common::profile::Instant::now();
613 let result = resolve_bench(bench).await;
614 let step = launch.map(|launch| (name, launch + started.elapsed()));
615
616 (index, step, result)
617 }
618 }))
619 .await;
620
621 for (index, step, result) in resolved {
622 if let Some((name, duration)) = step {
623 log_context.push_tuning_step(name, duration);
624 }
625
626 results[index] = result;
627 }
628
629 #[cfg(persistence)]
633 let unmeasured = decided.is_some_and(|index| results[index].outcome.is_err());
634
635 results.sort_by(|a, b| {
636 let a = a
637 .outcome
638 .as_ref()
639 .map(|r| r.computation.score())
640 .unwrap_or(u64::MAX);
641 let b = b
642 .outcome
643 .as_ref()
644 .map(|r| r.computation.score())
645 .unwrap_or(u64::MAX);
646 a.cmp(&b)
647 });
648
649 let fastest_index = match decided {
653 Some(index) => index,
654 None => {
655 results
656 .first()
657 .expect("At least one kernel needed.")
658 .outcome
659 .as_ref()
660 .expect("At least one kernel has to succeed.")
661 .index
662 }
663 };
664
665 {
666 log_context.log_result(&mut logger.lock(), &key, &results);
667 cache.lock().cache_insert(key.clone(), fastest_index);
670
671 #[cfg(persistence)]
678 let stored = !unmeasured
679 && cache.lock().persistent_cache_insert(
680 key,
681 checksum,
682 crate::tune::PersistentCacheValue {
683 fastest_index,
684 results,
685 bounds,
686 limit,
687 },
688 );
689
690 #[cfg(persistence)]
691 recording.finish(
692 crate::tune::record::Answer {
693 winner: fastest_index,
694 stored,
695 },
696 log_context.as_ref(),
697 );
698 }
699
700 TuneCacheResult::Hit { fastest_index }
701}
702
703#[cfg(feature = "autotune-checks")]
704pub(crate) fn check_autotune_outputs<O: AutotuneOutput>(
705 mut checks_outputs: Vec<(String, Result<O, AutotuneError>)>,
706) -> Vec<crate::tune::log::CheckResult> {
707 if checks_outputs.is_empty() {
708 return Vec::new();
709 }
710
711 let reference_idx = checks_outputs
712 .iter()
713 .position(|(_, res)| res.is_ok())
714 .unwrap_or(checks_outputs.len() - 1);
715 let reference = checks_outputs.remove(reference_idx);
716 let reference_result = reference.1;
717 #[cfg(std_io)]
718 let reference_name = reference.0;
719
720 let decisions_enabled = is_decisions_enabled();
721
722 #[cfg(std_io)]
723 {
724 let reference_passed = reference_result.is_ok();
725 let mut check_results = execute_checks(checks_outputs, reference_result, decisions_enabled);
726 check_results.push(crate::tune::log::CheckResult {
727 name: reference_name,
728 passed: reference_passed,
729 });
730
731 check_results
732 }
733
734 #[cfg(not(std_io))]
735 {
736 execute_checks(checks_outputs, reference_result, decisions_enabled)
737 }
738}
739
740#[cfg(feature = "autotune-checks")]
744fn is_decisions_enabled() -> bool {
745 crate::config::CubeClRuntimeConfig::get()
746 .autotune
747 .decisions_enabled()
748}
749
750#[cfg(feature = "autotune-checks")]
751fn execute_checks<O: AutotuneOutput>(
752 checks_outputs: Vec<(String, Result<O, AutotuneError>)>,
753 reference_result: Result<O, AutotuneError>,
754 decisions_enabled: bool,
755) -> Vec<crate::tune::log::CheckResult> {
756 let mut check_results = Vec::new();
757
758 let Ok(reference) = reference_result else {
759 for (name, _) in checks_outputs.into_iter() {
760 check_results.push(crate::tune::log::CheckResult {
761 name,
762 passed: false,
763 });
764 }
765 return check_results;
766 };
767
768 for (name, other_result) in checks_outputs.into_iter() {
769 if let Ok(other) = other_result {
770 let passed = check_equivalence(&reference, other, decisions_enabled);
771 check_results.push(crate::tune::log::CheckResult { name, passed });
772 } else {
773 check_results.push(crate::tune::log::CheckResult {
774 name,
775 passed: false,
776 });
777 }
778 }
779
780 check_results
781}
782
783#[cfg(feature = "autotune-checks")]
784fn check_equivalence<O: AutotuneOutput>(reference: &O, other: O, decisions_enabled: bool) -> bool {
785 if decisions_enabled {
788 #[cfg(std_io)]
789 {
790 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
791 reference.check_equivalence(other);
792 }))
793 .is_ok()
794 }
795 #[cfg(not(std_io))]
796 {
797 reference.check_equivalence(other);
798 true
799 }
800 } else {
801 reference.check_equivalence(other);
802 true
803 }
804}