cubecl_runtime/tune/
util.rs

1use core::sync::atomic::{AtomicI32, Ordering};
2
3use crate::config::GlobalConfig;
4
5/// Autotune levels:
6///
7/// '0' => Minimal autotune: scaled anchor of '1.25'.
8/// '1' => Medium autotune: normal anchor.
9/// '2' => More autotune: scaled anchor of '0.75'.
10/// '3' => Autotune everything without anchor.
11static AUTOTUNE_LEVEL: AtomicI32 = AtomicI32::new(-1);
12
13/// Anchor a number to a power of the provided base.
14///
15/// Useful when creating autotune keys.
16pub fn anchor(x: usize, max: Option<usize>, min: Option<usize>, base: Option<usize>) -> usize {
17    let autotune_level = load_autotune_level();
18    let factor = match autotune_level {
19        3 => return x, // Autotune everything, there is no anchor.
20        2 => 0.75,
21        1 => 1.0,
22        0 => 1.25,
23        _ => panic!("Invalid autotune level {autotune_level:?}"),
24    };
25
26    let base = base.unwrap_or(2) as f64 * factor;
27    let base = f64::max(base, 1.1); // Minimum base.
28    let exp = (x as f64).log(base).ceil();
29    let power = base.powf(exp).ceil() as usize;
30
31    let result = if let Some(max) = max {
32        core::cmp::min(power, max)
33    } else {
34        power
35    };
36
37    if let Some(min) = min {
38        core::cmp::max(result, min)
39    } else {
40        result
41    }
42}
43
44fn load_autotune_level() -> u32 {
45    let autotune_level = AUTOTUNE_LEVEL.load(Ordering::Relaxed);
46    if autotune_level == -1 {
47        let config = GlobalConfig::get();
48        let level = match config.autotune.level {
49            crate::config::autotune::AutotuneLevel::Minimal => 0,
50            crate::config::autotune::AutotuneLevel::Balanced => 1,
51            crate::config::autotune::AutotuneLevel::Extensive => 2,
52            crate::config::autotune::AutotuneLevel::Full => 3,
53        };
54        AUTOTUNE_LEVEL.store(level, Ordering::Relaxed);
55        level as u32
56    } else {
57        autotune_level as u32
58    }
59}
cubecl_runtime/tune/util.rs

cubecl_runtime/tune/
util.rs
