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(crate) fn settled(&self, key: &K) -> Option<usize> {
228 self.cache.lock().settled(key)
229 }
230
231 pub fn logger(&self) -> Arc<Mutex<Logger>> {
233 self.logger.clone()
234 }
235
236 pub fn check_tune<'a, F: TuneInputs, Out: AutotuneOutput, Id>(
239 &self,
240 key: &K,
241 inputs: &F::At<'a>,
242 tunables: &TunableSet<K, F, Out, Id>,
243 #[cfg_attr(not(persistence), allow(unused))] checksum: impl FnOnce() -> String + Send + Sync,
244 client: &Client,
245 mut log_context: Option<crate::tune::AutotuneLogContext>,
246 ) -> TuneCacheResult
247 where
248 <F as TuneInputs>::At<'a>: Clone + Send,
249 {
250 {
251 let mut cache = self.cache.lock();
252 #[cfg(persistence)]
253 cache.reset_if_environment_switched();
254 let cur = cache.fastest(key);
255
256 #[cfg(persistence)]
260 let cur = if matches!(cur, TuneCacheResult::Miss) {
261 cache.sync_persistent();
262 cache.fastest(key)
263 } else {
264 cur
265 };
266
267 #[cfg(persistence)]
268 let cur = if matches!(cur, TuneCacheResult::Unchecked) {
269 let mut log = self.logger.lock();
270 let checksum = checksum();
271 if let AutotuneLogLevel::Full = log.log_level_autotune() {
272 log.log_autotune(&format!("validate checksum key={key}, checksum={checksum}"));
273 }
274 cache.validate_checksum(key, &checksum)
275 } else {
276 cur
277 };
278
279 match cur {
280 TuneCacheResult::Hit { .. } | TuneCacheResult::Pending => return cur,
281 TuneCacheResult::Miss | TuneCacheResult::Unchecked => {
282 cache.mark_pending(key.clone())
283 }
284 }
285 }
288
289 log::info!("Tuning {key}");
290
291 let autotunables = tunables.autotunables().collect::<Vec<_>>();
292 let results: Vec<AutotuneResult> = autotunables
293 .iter()
294 .map(|a| {
295 AutotuneResult::error(AutotuneError::Skip {
296 name: a.name.to_string(),
297 })
298 })
299 .collect();
300
301 #[cfg(persistence)]
302 let checksum = tunables.compute_checksum();
303
304 if results.len() == 1 {
306 self.cache.lock().cache_insert(key.clone(), 0);
307 return TuneCacheResult::Hit { fastest_index: 0 };
308 }
309
310 #[cfg(persistence)]
313 let mut recording =
314 crate::tune::record::TuneRecording::new(&self.cache.lock(), key, &checksum);
315 #[cfg(persistence)]
318 if recording.is_open() {
319 log_context.get_or_insert_with(Default::default);
320 }
321
322 let test_inputs = tunables.generate_inputs(key, inputs);
323 let plan = tunables.plan(key);
324 #[cfg(persistence)]
325 if recording.is_open() {
326 let names: Vec<&str> = autotunables.iter().map(|tune| tune.name.as_str()).collect();
327 recording.plan(&plan, &names);
328 }
329 let bounds = tunables.bounds(key, inputs);
330 let limit = bounds.as_ref().and_then(|bounds| bounds.time_limit());
331
332 log_context.set_bounds(bounds.clone());
333 log_context.set_limit(limit);
334
335 #[cfg(not(target_family = "wasm"))]
338 let short_circuit = limit.is_some()
339 && tunables.is_short_circuit_enabled()
340 && !crate::config::CubeClRuntimeConfig::get()
341 .autotune
342 .disable_short_circuit;
343
344 let job = TuneJob {
345 key: key.clone(),
346 autotunables,
347 test_inputs,
348 evictor: tunables.evictor(key, inputs),
349 plan,
350 results,
351 #[cfg(any(not(target_family = "wasm"), persistence))]
352 limit,
353 #[cfg(persistence)]
354 bounds,
355 #[cfg(not(target_family = "wasm"))]
356 short_circuit,
357 #[cfg(persistence)]
358 checksum,
359 log_context,
360 #[cfg(persistence)]
361 recording,
362 };
363
364 #[cfg(not(target_family = "wasm"))]
365 if crate::config::CubeClRuntimeConfig::get()
366 .autotune
367 .bench
368 .adaptive
369 {
370 return self.tune_adaptive(job, client);
371 }
372
373 self.tune_fixed_samples(job, client)
374 }
375
376 #[cfg(not(target_family = "wasm"))]
379 fn tune_adaptive<'i, F: TuneInputs, Out: AutotuneOutput>(
380 &self,
381 mut job: TuneJob<'_, 'i, K, F, Out>,
382 client: &Client,
383 ) -> TuneCacheResult
384 where
385 <F as TuneInputs>::At<'i>: Clone + Send,
386 {
387 let mut schedule = crate::tune::schedule::Schedule {
388 config: crate::config::CubeClRuntimeConfig::get()
389 .autotune
390 .bench
391 .clone(),
392 limit: job.limit,
393 short_circuit: job.short_circuit,
394 track_steps: job.log_context.is_some(),
395 evictor: job.evictor.take(),
396 };
397
398 let outcome = schedule.run_plan(
399 &job.key,
400 &mut job.plan,
401 &job.autotunables,
402 &job.test_inputs,
403 client,
404 &mut job.results,
405 );
406
407 for (name, duration) in outcome.steps {
408 job.log_context.push_tuning_step(name, duration);
409 }
410 if let Some(name) = outcome.short_circuit {
411 job.log_context.push_short_circuit(name);
412 }
413
414 let request = job.into_request(Vec::new(), outcome.decided);
415
416 cubecl_environment::future::block_on(process_request(request, &self.cache, &self.logger))
417 }
418
419 fn tune_fixed_samples<'i, F: TuneInputs, Out: AutotuneOutput>(
422 &self,
423 mut job: TuneJob<'_, 'i, K, F, Out>,
424 client: &Client,
425 ) -> TuneCacheResult
426 where
427 <F as TuneInputs>::At<'i>: Clone + Send,
428 {
429 #[cfg(not(target_family = "wasm"))]
433 let mut batch_success = false;
434 #[cfg(target_family = "wasm")]
435 let batch_success = false;
436
437 let mut pending = Vec::<PendingBench>::new();
442 loop {
443 let batch = job.plan.next();
444
445 if batch.is_empty() {
446 let key = &job.key;
447 panic!(
448 "Can't execute the autotune plan for key: {key:?}\n - plan: {:?}\n - results: {:?}",
449 job.plan, job.results
450 );
451 }
452
453 for index in batch.indices() {
455 let op = job.autotunables[index];
456
457 let start_time = job
458 .log_context
459 .is_some()
460 .then(cubecl_common::profile::Instant::now);
461
462 match tune_benchmark(
463 op,
464 job.test_inputs.clone(),
465 client.clone(),
466 job.evictor.as_deref_mut(),
467 ) {
468 Ok(profiles) => {
469 let bench = PendingBench {
470 index,
471 name: op.name.clone(),
472 profiles,
473 launch: start_time.map(|start| start.elapsed()),
474 };
475
476 #[cfg(not(target_family = "wasm"))]
477 if job.short_circuit {
478 let result = cubecl_environment::future::block_on(resolve_bench(bench));
479
480 let close_enough = result
482 .outcome
483 .as_ref()
484 .is_ok_and(|out| out.computation.median <= job.limit.unwrap());
485
486 batch_success |= result.outcome.is_ok();
487 job.results[index] = result;
488
489 if let Some(start) = start_time {
490 job.log_context
491 .push_tuning_step(op.name.to_string(), start.elapsed());
492 }
493
494 if close_enough {
495 job.log_context.push_short_circuit(op.name.to_string());
496 break;
497 }
498
499 continue;
500 }
501
502 pending.push(bench);
505 }
506 Err(err) => {
507 job.results[index] = AutotuneResult::error(err);
508 if let Some(start) = start_time {
509 job.log_context
510 .push_tuning_step(op.name.to_string(), start.elapsed());
511 }
512 }
513 }
514 }
515
516 #[cfg(not(target_family = "wasm"))]
517 if !pending.is_empty() || batch_success {
518 break;
519 }
520 #[cfg(target_family = "wasm")]
521 if !pending.is_empty() {
522 break;
523 }
524 }
525
526 let request = job.into_request(pending, None);
529
530 #[cfg(target_family = "wasm")]
533 {
534 let cache = self.cache.clone();
535 let logger = self.logger.clone();
536 wasm_bindgen_futures::spawn_local(async move {
537 process_request(request, &cache, &logger).await;
538 });
539
540 return TuneCacheResult::Pending;
541 }
542
543 #[cfg(not(target_family = "wasm"))]
544 cubecl_environment::future::block_on(process_request(request, &self.cache, &self.logger))
545 }
546}
547
548async fn resolve_bench(bench: PendingBench) -> AutotuneResult {
554 let PendingBench {
555 index,
556 name,
557 profiles,
558 launch: _,
559 } = bench;
560
561 let Some(first) = profiles.first() else {
562 return AutotuneResult::error(AutotuneError::Unknown {
563 name: name.to_string(),
564 err: "No profiling available".to_string(),
565 });
566 };
567 let timing_method = first.timing_method();
568
569 let Some(durations) =
573 futures_util::future::join_all(profiles.into_iter().map(ProfileDuration::resolve))
574 .await
575 .into_iter()
576 .map(|ticks| ticks.map(|ticks| ticks.duration()))
577 .collect::<Option<Vec<Duration>>>()
578 else {
579 return AutotuneResult::error(AutotuneError::NotMeasured {
580 name: name.to_string(),
581 });
582 };
583
584 AutotuneResult::success(AutotuneOutcome::new(
585 name,
586 index,
587 BenchmarkComputations::new(&BenchmarkDurations::from_durations(
588 timing_method,
589 durations,
590 )),
591 ))
592}
593
594async fn process_request<K: AutotuneKey>(
596 request: TuneRequest<K>,
597 cache: &Mutex<TuneCache<K>>,
598 logger: &Mutex<Logger>,
599) -> TuneCacheResult {
600 let TuneRequest {
601 key,
602 mut results,
603 #[cfg(persistence)]
604 checksum,
605 mut log_context,
606 pending,
607 decided,
608 #[cfg(persistence)]
609 limit,
610 #[cfg(persistence)]
611 bounds,
612 #[cfg(persistence)]
613 recording,
614 } = request;
615
616 let resolved = futures_util::future::join_all(pending.into_iter().map(|bench| {
621 let index = bench.index;
622 let name = bench.name.clone();
623 let launch = bench.launch;
624
625 async move {
626 let started = cubecl_common::profile::Instant::now();
627 let result = resolve_bench(bench).await;
628 let step = launch.map(|launch| (name, launch + started.elapsed()));
629
630 (index, step, result)
631 }
632 }))
633 .await;
634
635 for (index, step, result) in resolved {
636 if let Some((name, duration)) = step {
637 log_context.push_tuning_step(name, duration);
638 }
639
640 results[index] = result;
641 }
642
643 #[cfg(persistence)]
647 let unmeasured = decided.is_some_and(|index| results[index].outcome.is_err());
648
649 results.sort_by(|a, b| {
650 let a = a
651 .outcome
652 .as_ref()
653 .map(|r| r.computation.score())
654 .unwrap_or(u64::MAX);
655 let b = b
656 .outcome
657 .as_ref()
658 .map(|r| r.computation.score())
659 .unwrap_or(u64::MAX);
660 a.cmp(&b)
661 });
662
663 let fastest_index = match decided {
667 Some(index) => index,
668 None => {
669 results
670 .first()
671 .expect("At least one kernel needed.")
672 .outcome
673 .as_ref()
674 .expect("At least one kernel has to succeed.")
675 .index
676 }
677 };
678
679 {
680 log_context.log_result(&mut logger.lock(), &key, &results);
681 cache.lock().cache_insert(key.clone(), fastest_index);
684
685 #[cfg(persistence)]
692 let stored = !unmeasured
693 && cache.lock().persistent_cache_insert(
694 key,
695 checksum,
696 crate::tune::PersistentCacheValue {
697 fastest_index,
698 results,
699 bounds,
700 limit,
701 },
702 );
703
704 #[cfg(persistence)]
705 recording.finish(
706 crate::tune::record::Answer {
707 winner: fastest_index,
708 stored,
709 },
710 log_context.as_ref(),
711 );
712 }
713
714 TuneCacheResult::Hit { fastest_index }
715}
716
717#[cfg(feature = "autotune-checks")]
718pub(crate) fn check_autotune_outputs<O: AutotuneOutput>(
719 mut checks_outputs: Vec<(String, Result<O, AutotuneError>)>,
720) -> Vec<crate::tune::log::CheckResult> {
721 if checks_outputs.is_empty() {
722 return Vec::new();
723 }
724
725 let reference_idx = checks_outputs
726 .iter()
727 .position(|(_, res)| res.is_ok())
728 .unwrap_or(checks_outputs.len() - 1);
729 let reference = checks_outputs.remove(reference_idx);
730 let reference_result = reference.1;
731 #[cfg(std_io)]
732 let reference_name = reference.0;
733
734 let decisions_enabled = is_decisions_enabled();
735
736 #[cfg(std_io)]
737 {
738 let reference_passed = reference_result.is_ok();
739 let mut check_results = execute_checks(checks_outputs, reference_result, decisions_enabled);
740 check_results.push(crate::tune::log::CheckResult {
741 name: reference_name,
742 passed: reference_passed,
743 });
744
745 check_results
746 }
747
748 #[cfg(not(std_io))]
749 {
750 execute_checks(checks_outputs, reference_result, decisions_enabled)
751 }
752}
753
754#[cfg(feature = "autotune-checks")]
758fn is_decisions_enabled() -> bool {
759 crate::config::CubeClRuntimeConfig::get()
760 .autotune
761 .decisions_enabled()
762}
763
764#[cfg(feature = "autotune-checks")]
765fn execute_checks<O: AutotuneOutput>(
766 checks_outputs: Vec<(String, Result<O, AutotuneError>)>,
767 reference_result: Result<O, AutotuneError>,
768 decisions_enabled: bool,
769) -> Vec<crate::tune::log::CheckResult> {
770 let mut check_results = Vec::new();
771
772 let Ok(reference) = reference_result else {
773 for (name, _) in checks_outputs.into_iter() {
774 check_results.push(crate::tune::log::CheckResult {
775 name,
776 passed: false,
777 });
778 }
779 return check_results;
780 };
781
782 for (name, other_result) in checks_outputs.into_iter() {
783 if let Ok(other) = other_result {
784 let passed = check_equivalence(&reference, other, decisions_enabled);
785 check_results.push(crate::tune::log::CheckResult { name, passed });
786 } else {
787 check_results.push(crate::tune::log::CheckResult {
788 name,
789 passed: false,
790 });
791 }
792 }
793
794 check_results
795}
796
797#[cfg(feature = "autotune-checks")]
798fn check_equivalence<O: AutotuneOutput>(reference: &O, other: O, decisions_enabled: bool) -> bool {
799 if decisions_enabled {
802 #[cfg(std_io)]
803 {
804 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
805 reference.check_equivalence(other);
806 }))
807 .is_ok()
808 }
809 #[cfg(not(std_io))]
810 {
811 reference.check_equivalence(other);
812 true
813 }
814 } else {
815 reference.check_equivalence(other);
816 true
817 }
818}