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 Ok(val) = std::env::var("CUBECL_ENVIRONMENT_RECORDS") {
182            use cubecl_environment::records::RecordLevel;
183
184            match val.as_str() {
185                "off" | "0" => self.environment.records.level = RecordLevel::Off,
186                "basic" | "1" => self.environment.records.level = RecordLevel::Basic,
187                "full" | "2" => self.environment.records.level = RecordLevel::Full,
188                _ => {}
189            }
190        }
191
192        if let Some(enabled) = env_bool("CUBECL_AUTOTUNE_CACHE") {
193            self.autotune.disable_cache = !enabled;
194        }
195
196        if let Some(enabled) = env_bool("CUBECL_AUTOTUNE_SHORT_CIRCUIT") {
197            self.autotune.disable_short_circuit = !enabled;
198        }
199
200        if let Some(enabled) = env_bool("CUBECL_AUTOTUNE_BENCH_ADAPTIVE") {
201            self.autotune.bench.adaptive = enabled;
202        }
203
204        self
205    }
206}
207
208/// A boolean environment variable, or `None` when it is unset or unreadable.
209///
210/// An unrecognized value is `None` rather than an error: the variable is an
211/// override, so failing to parse it means leaving the configured value alone.
212#[cfg(std_io)]
213fn env_bool(name: &str) -> Option<bool> {
214    match std::env::var(name).ok()?.as_str() {
215        "true" | "1" | "on" => Some(true),
216        "false" | "0" | "off" => Some(false),
217        _ => None,
218    }
219}
220
221#[derive(Clone, Copy, Debug)]
222/// How to format cubecl type names.
223pub enum TypeNameFormatLevel {
224    /// No formatting apply, full information is included.
225    Full,
226    /// Most information is removed for a small formatted name.
227    Short,
228    /// Balanced info is kept.
229    Balanced,
230}
231
232/// Format a type name with different options.
233pub fn type_name_format(name: &str, level: TypeNameFormatLevel) -> String {
234    match level {
235        TypeNameFormatLevel::Full => name.to_string(),
236        TypeNameFormatLevel::Short => {
237            if let Some(val) = name.split("<").next() {
238                val.split("::").last().unwrap_or(name).to_string()
239            } else {
240                name.to_string()
241            }
242        }
243        TypeNameFormatLevel::Balanced => {
244            let mut split = name.split("<");
245            let before_generic = split.next();
246            let after_generic = split.next();
247
248            let before_generic = match before_generic {
249                None => return name.to_string(),
250                Some(val) => val
251                    .split("::")
252                    .last()
253                    .unwrap_or(val)
254                    .trim()
255                    .replace(">", "")
256                    .to_string(),
257            };
258            let inside_generic = match after_generic {
259                None => return before_generic.to_string(),
260                Some(val) => {
261                    let mut val = val.to_string();
262                    for s in split {
263                        val += "<";
264                        val += s;
265                    }
266                    val
267                }
268            };
269
270            let inside = type_name_list_format(&inside_generic, level);
271
272            format!("{before_generic}{inside}")
273        }
274    }
275}
276
277fn type_name_list_format(name: &str, level: TypeNameFormatLevel) -> String {
278    let mut acc = String::new();
279    let splits = name.split(", ");
280
281    for a in splits {
282        acc += " | ";
283        acc += &type_name_format(a, level);
284    }
285
286    acc
287}
288
289#[cfg(test)]
290mod test {
291    use super::*;
292
293    #[test_log::test]
294    fn test_format_name() {
295        let full_name = "burn_cubecl::kernel::unary_numeric::unary_numeric::UnaryNumeric<f32, burn_cubecl::tensor::base::CubeTensor<_>::copy::Copy, cubecl_cuda::runtime::CudaRuntime>";
296        let name = type_name_format(full_name, TypeNameFormatLevel::Balanced);
297
298        assert_eq!(name, "UnaryNumeric | f32 | CubeTensor | Copy | CudaRuntime");
299    }
300}