Skip to main content

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        cubecl_environment::records::configure(self.environment.records);
57        // Before any device is initialized, so every cache opened afterwards
58        // lands in the chosen environment.
59        cubecl_environment::environment::activate(&self.environment.name);
60        #[cfg(std_io)]
61        cubecl_environment::environment::set_root(self.environment.path.root());
62    }
63
64    fn file_names() -> &'static [&'static str] {
65        &["cubecl.toml", "CubeCL.toml"]
66    }
67
68    fn section_file_names() -> &'static [(&'static str, &'static str)] {
69        &[("burn.toml", "cubecl"), ("Burn.toml", "cubecl")]
70    }
71
72    #[cfg(std_io)]
73    fn override_from_env(mut self) -> Self {
74        use super::compilation::{CompilationLogLevel, F16Evaluation};
75        use crate::config::{
76            autotune::{AutotuneLevel, AutotuneLogLevel},
77            profiling::ProfilingLogLevel,
78        };
79
80        if let Ok(val) = std::env::var("CUBECL_DEBUG_LOG") {
81            self.compilation.logger.level = CompilationLogLevel::Full;
82            self.profiling.logger.level = ProfilingLogLevel::Medium;
83            self.autotune.logger.level = AutotuneLogLevel::Full;
84
85            match val.as_str() {
86                "stdout" => {
87                    self.compilation.logger.stdout = true;
88                    self.profiling.logger.stdout = true;
89                    self.autotune.logger.stdout = true;
90                }
91                "stderr" => {
92                    self.compilation.logger.stderr = true;
93                    self.profiling.logger.stderr = true;
94                    self.autotune.logger.stderr = true;
95                }
96                "1" | "true" => {
97                    let file_path = "/tmp/cubecl.log";
98                    self.compilation.logger.file = Some(file_path.into());
99                    self.profiling.logger.file = Some(file_path.into());
100                    self.autotune.logger.file = Some(file_path.into());
101                }
102                "0" | "false" => {
103                    self.compilation.logger.level = CompilationLogLevel::Disabled;
104                    self.profiling.logger.level = ProfilingLogLevel::Disabled;
105                    self.autotune.logger.level = AutotuneLogLevel::Disabled;
106                }
107                file_path => {
108                    self.compilation.logger.file = Some(file_path.into());
109                    self.profiling.logger.file = Some(file_path.into());
110                    self.autotune.logger.file = Some(file_path.into());
111                }
112            }
113        };
114
115        if let Ok(val) = std::env::var("CUBECL_DEBUG_OPTION") {
116            match val.as_str() {
117                "debug" => {
118                    self.compilation.logger.level = CompilationLogLevel::Full;
119                    self.profiling.logger.level = ProfilingLogLevel::Medium;
120                    self.autotune.logger.level = AutotuneLogLevel::Full;
121                }
122                "debug-full" => {
123                    self.compilation.logger.level = CompilationLogLevel::Full;
124                    self.profiling.logger.level = ProfilingLogLevel::Full;
125                    self.autotune.logger.level = AutotuneLogLevel::Full;
126                }
127                "profile" => {
128                    self.profiling.logger.level = ProfilingLogLevel::Basic;
129                }
130                "profile-medium" => {
131                    self.profiling.logger.level = ProfilingLogLevel::Medium;
132                }
133                "profile-full" => {
134                    self.profiling.logger.level = ProfilingLogLevel::Full;
135                }
136                _ => {}
137            }
138        };
139
140        if let Ok(val) = std::env::var("CUBECL_AUTOTUNE_LEVEL") {
141            match val.as_str() {
142                "minimal" | "0" => {
143                    self.autotune.level = AutotuneLevel::Minimal;
144                }
145                "balanced" | "1" => {
146                    self.autotune.level = AutotuneLevel::Balanced;
147                }
148                "extensive" | "2" => {
149                    self.autotune.level = AutotuneLevel::Extensive;
150                }
151                "full" | "3" => {
152                    self.autotune.level = AutotuneLevel::Full;
153                }
154                _ => {}
155            }
156        }
157
158        if let Ok(val) = std::env::var("CUBECL_CPU_F16_EVAL") {
159            match val.as_str() {
160                "per-operation" => {
161                    self.compilation.f16_evaluation = Some(F16Evaluation::PerOperation);
162                }
163                "chain" => {
164                    self.compilation.f16_evaluation = Some(F16Evaluation::Chain);
165                }
166                "accumulators" => {
167                    self.compilation.f16_evaluation = Some(F16Evaluation::Accumulators);
168                }
169                _ => {}
170            }
171        }
172
173        if let Some(enabled) = env_bool("CUBECL_THROUGHPUT_CACHE") {
174            self.throughput.disable_cache = !enabled;
175        }
176
177        if let Ok(val) = std::env::var("CUBECL_ENVIRONMENT") {
178            self.environment.name = val;
179        }
180
181        if let Some(enabled) = env_bool("CUBECL_AUTOTUNE_CACHE") {
182            self.autotune.disable_cache = !enabled;
183        }
184
185        if let Some(enabled) = env_bool("CUBECL_AUTOTUNE_SHORT_CIRCUIT") {
186            self.autotune.disable_short_circuit = !enabled;
187        }
188
189        if let Some(enabled) = env_bool("CUBECL_AUTOTUNE_BENCH_ADAPTIVE") {
190            self.autotune.bench.adaptive = enabled;
191        }
192
193        self
194    }
195}
196
197/// A boolean environment variable, or `None` when it is unset or unreadable.
198///
199/// An unrecognized value is `None` rather than an error: the variable is an
200/// override, so failing to parse it means leaving the configured value alone.
201#[cfg(std_io)]
202fn env_bool(name: &str) -> Option<bool> {
203    match std::env::var(name).ok()?.as_str() {
204        "true" | "1" | "on" => Some(true),
205        "false" | "0" | "off" => Some(false),
206        _ => None,
207    }
208}
209
210#[derive(Clone, Copy, Debug)]
211/// How to format cubecl type names.
212pub enum TypeNameFormatLevel {
213    /// No formatting apply, full information is included.
214    Full,
215    /// Most information is removed for a small formatted name.
216    Short,
217    /// Balanced info is kept.
218    Balanced,
219}
220
221/// Format a type name with different options.
222pub fn type_name_format(name: &str, level: TypeNameFormatLevel) -> String {
223    match level {
224        TypeNameFormatLevel::Full => name.to_string(),
225        TypeNameFormatLevel::Short => {
226            if let Some(val) = name.split("<").next() {
227                val.split("::").last().unwrap_or(name).to_string()
228            } else {
229                name.to_string()
230            }
231        }
232        TypeNameFormatLevel::Balanced => {
233            let mut split = name.split("<");
234            let before_generic = split.next();
235            let after_generic = split.next();
236
237            let before_generic = match before_generic {
238                None => return name.to_string(),
239                Some(val) => val
240                    .split("::")
241                    .last()
242                    .unwrap_or(val)
243                    .trim()
244                    .replace(">", "")
245                    .to_string(),
246            };
247            let inside_generic = match after_generic {
248                None => return before_generic.to_string(),
249                Some(val) => {
250                    let mut val = val.to_string();
251                    for s in split {
252                        val += "<";
253                        val += s;
254                    }
255                    val
256                }
257            };
258
259            let inside = type_name_list_format(&inside_generic, level);
260
261            format!("{before_generic}{inside}")
262        }
263    }
264}
265
266fn type_name_list_format(name: &str, level: TypeNameFormatLevel) -> String {
267    let mut acc = String::new();
268    let splits = name.split(", ");
269
270    for a in splits {
271        acc += " | ";
272        acc += &type_name_format(a, level);
273    }
274
275    acc
276}
277
278#[cfg(test)]
279mod test {
280    use super::*;
281
282    #[test_log::test]
283    fn test_format_name() {
284        let full_name = "burn_cubecl::kernel::unary_numeric::unary_numeric::UnaryNumeric<f32, burn_cubecl::tensor::base::CubeTensor<_>::copy::Copy, cubecl_cuda::runtime::CudaRuntime>";
285        let name = type_name_format(full_name, TypeNameFormatLevel::Balanced);
286
287        assert_eq!(name, "UnaryNumeric | f32 | CubeTensor | Copy | CudaRuntime");
288    }
289}