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