cubecl_runtime/config/
base.rs

1use crate::config::memory::MemoryConfig;
2use crate::config::streaming::StreamingConfig;
3
4use super::{
5    autotune::AutotuneConfig, compilation::CompilationConfig, profiling::ProfilingConfig,
6    throughput::ThroughputConfig,
7};
8use alloc::format;
9use alloc::string::{String, ToString};
10use alloc::sync::Arc;
11use cubecl_environment::config::RuntimeConfig;
12use cubecl_environment::sync::Mutex;
13
14/// Static mutex holding the global configuration, initialized as `None`.
15static CUBE_GLOBAL_CONFIG: Mutex<Option<Arc<CubeClRuntimeConfig>>> = Mutex::new(None);
16
17/// Represents the global configuration for `CubeCL`, combining profiling, autotuning, and compilation settings.
18#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
19pub struct CubeClRuntimeConfig {
20    /// Configuration for profiling `CubeCL` operations.
21    #[serde(default)]
22    pub profiling: ProfilingConfig,
23
24    /// Configuration for autotuning performance parameters.
25    #[serde(default)]
26    pub autotune: AutotuneConfig,
27
28    /// Configuration for throughput settings.
29    #[serde(default)]
30    pub throughput: ThroughputConfig,
31
32    /// Configuration for compilation settings.
33    #[serde(default)]
34    pub compilation: CompilationConfig,
35
36    /// Configuration for streaming settings.
37    #[serde(default)]
38    pub streaming: StreamingConfig,
39
40    /// Configuration for memory settings.
41    #[serde(default)]
42    pub memory: MemoryConfig,
43
44    /// Which named environment to warm into.
45    #[serde(default)]
46    pub environment: super::environment::EnvironmentConfig,
47}
48
49impl RuntimeConfig for CubeClRuntimeConfig {
50    fn storage() -> &'static Mutex<Option<Arc<Self>>> {
51        &CUBE_GLOBAL_CONFIG
52    }
53
54    fn on_loaded(&self) {
55        cubecl_environment::stream::set_policy_from_config(self.streaming.policy);
56        // Before any device is initialized, so every cache opened afterwards
57        // lands in the chosen environment.
58        cubecl_environment::environment::activate(&self.environment.name);
59        #[cfg(std_io)]
60        cubecl_environment::environment::set_root(self.environment.path.root());
61    }
62
63    fn file_names() -> &'static [&'static str] {
64        &["cubecl.toml", "CubeCL.toml"]
65    }
66
67    fn section_file_names() -> &'static [(&'static str, &'static str)] {
68        &[("burn.toml", "cubecl"), ("Burn.toml", "cubecl")]
69    }
70
71    #[cfg(std_io)]
72    fn override_from_env(mut self) -> Self {
73        use super::compilation::CompilationLogLevel;
74        use crate::config::{
75            autotune::{AutotuneLevel, AutotuneLogLevel},
76            profiling::ProfilingLogLevel,
77        };
78
79        if let Ok(val) = std::env::var("CUBECL_DEBUG_LOG") {
80            self.compilation.logger.level = CompilationLogLevel::Full;
81            self.profiling.logger.level = ProfilingLogLevel::Medium;
82            self.autotune.logger.level = AutotuneLogLevel::Full;
83
84            match val.as_str() {
85                "stdout" => {
86                    self.compilation.logger.stdout = true;
87                    self.profiling.logger.stdout = true;
88                    self.autotune.logger.stdout = true;
89                }
90                "stderr" => {
91                    self.compilation.logger.stderr = true;
92                    self.profiling.logger.stderr = true;
93                    self.autotune.logger.stderr = true;
94                }
95                "1" | "true" => {
96                    let file_path = "/tmp/cubecl.log";
97                    self.compilation.logger.file = Some(file_path.into());
98                    self.profiling.logger.file = Some(file_path.into());
99                    self.autotune.logger.file = Some(file_path.into());
100                }
101                "0" | "false" => {
102                    self.compilation.logger.level = CompilationLogLevel::Disabled;
103                    self.profiling.logger.level = ProfilingLogLevel::Disabled;
104                    self.autotune.logger.level = AutotuneLogLevel::Disabled;
105                }
106                file_path => {
107                    self.compilation.logger.file = Some(file_path.into());
108                    self.profiling.logger.file = Some(file_path.into());
109                    self.autotune.logger.file = Some(file_path.into());
110                }
111            }
112        };
113
114        if let Ok(val) = std::env::var("CUBECL_DEBUG_OPTION") {
115            match val.as_str() {
116                "debug" => {
117                    self.compilation.logger.level = CompilationLogLevel::Full;
118                    self.profiling.logger.level = ProfilingLogLevel::Medium;
119                    self.autotune.logger.level = AutotuneLogLevel::Full;
120                }
121                "debug-full" => {
122                    self.compilation.logger.level = CompilationLogLevel::Full;
123                    self.profiling.logger.level = ProfilingLogLevel::Full;
124                    self.autotune.logger.level = AutotuneLogLevel::Full;
125                }
126                "profile" => {
127                    self.profiling.logger.level = ProfilingLogLevel::Basic;
128                }
129                "profile-medium" => {
130                    self.profiling.logger.level = ProfilingLogLevel::Medium;
131                }
132                "profile-full" => {
133                    self.profiling.logger.level = ProfilingLogLevel::Full;
134                }
135                _ => {}
136            }
137        };
138
139        if let Ok(val) = std::env::var("CUBECL_AUTOTUNE_LEVEL") {
140            match val.as_str() {
141                "minimal" | "0" => {
142                    self.autotune.level = AutotuneLevel::Minimal;
143                }
144                "balanced" | "1" => {
145                    self.autotune.level = AutotuneLevel::Balanced;
146                }
147                "extensive" | "2" => {
148                    self.autotune.level = AutotuneLevel::Extensive;
149                }
150                "full" | "3" => {
151                    self.autotune.level = AutotuneLevel::Full;
152                }
153                _ => {}
154            }
155        }
156
157        if let Some(enabled) = env_bool("CUBECL_THROUGHPUT_CACHE") {
158            self.throughput.disable_cache = !enabled;
159        }
160
161        if let Ok(val) = std::env::var("CUBECL_ENVIRONMENT") {
162            self.environment.name = val;
163        }
164
165        if let Some(enabled) = env_bool("CUBECL_AUTOTUNE_CACHE") {
166            self.autotune.disable_cache = !enabled;
167        }
168
169        if let Some(enabled) = env_bool("CUBECL_AUTOTUNE_SHORT_CIRCUIT") {
170            self.autotune.disable_short_circuit = !enabled;
171        }
172
173        if let Some(enabled) = env_bool("CUBECL_AUTOTUNE_BENCH_ADAPTIVE") {
174            self.autotune.bench.adaptive = enabled;
175        }
176
177        self
178    }
179}
180
181/// A boolean environment variable, or `None` when it is unset or unreadable.
182///
183/// An unrecognized value is `None` rather than an error: the variable is an
184/// override, so failing to parse it means leaving the configured value alone.
185#[cfg(std_io)]
186fn env_bool(name: &str) -> Option<bool> {
187    match std::env::var(name).ok()?.as_str() {
188        "true" | "1" | "on" => Some(true),
189        "false" | "0" | "off" => Some(false),
190        _ => None,
191    }
192}
193
194#[derive(Clone, Copy, Debug)]
195/// How to format cubecl type names.
196pub enum TypeNameFormatLevel {
197    /// No formatting apply, full information is included.
198    Full,
199    /// Most information is removed for a small formatted name.
200    Short,
201    /// Balanced info is kept.
202    Balanced,
203}
204
205/// Format a type name with different options.
206pub fn type_name_format(name: &str, level: TypeNameFormatLevel) -> String {
207    match level {
208        TypeNameFormatLevel::Full => name.to_string(),
209        TypeNameFormatLevel::Short => {
210            if let Some(val) = name.split("<").next() {
211                val.split("::").last().unwrap_or(name).to_string()
212            } else {
213                name.to_string()
214            }
215        }
216        TypeNameFormatLevel::Balanced => {
217            let mut split = name.split("<");
218            let before_generic = split.next();
219            let after_generic = split.next();
220
221            let before_generic = match before_generic {
222                None => return name.to_string(),
223                Some(val) => val
224                    .split("::")
225                    .last()
226                    .unwrap_or(val)
227                    .trim()
228                    .replace(">", "")
229                    .to_string(),
230            };
231            let inside_generic = match after_generic {
232                None => return before_generic.to_string(),
233                Some(val) => {
234                    let mut val = val.to_string();
235                    for s in split {
236                        val += "<";
237                        val += s;
238                    }
239                    val
240                }
241            };
242
243            let inside = type_name_list_format(&inside_generic, level);
244
245            format!("{before_generic}{inside}")
246        }
247    }
248}
249
250fn type_name_list_format(name: &str, level: TypeNameFormatLevel) -> String {
251    let mut acc = String::new();
252    let splits = name.split(", ");
253
254    for a in splits {
255        acc += " | ";
256        acc += &type_name_format(a, level);
257    }
258
259    acc
260}
261
262#[cfg(test)]
263mod test {
264    use super::*;
265
266    #[test_log::test]
267    fn test_format_name() {
268        let full_name = "burn_cubecl::kernel::unary_numeric::unary_numeric::UnaryNumeric<f32, burn_cubecl::tensor::base::CubeTensor<_>::copy::Copy, cubecl_cuda::runtime::CudaRuntime>";
269        let name = type_name_format(full_name, TypeNameFormatLevel::Balanced);
270
271        assert_eq!(name, "UnaryNumeric | f32 | CubeTensor | Copy | CudaRuntime");
272    }
273}
cubecl_runtime/config/base.rs

cubecl_runtime/config/
base.rs
