Skip to main content

burn_std/config/
base.rs

1use cubecl_common::config::RuntimeConfig;
2use cubecl_common::stub::Arc;
3
4use super::autodiff::AutodiffConfig;
5use super::fusion::FusionConfig;
6use super::remote::RemoteConfig;
7
8/// Static mutex holding the global Burn configuration, initialized as `None`.
9static BURN_GLOBAL_CONFIG: spin::Mutex<Option<Arc<BurnConfig>>> = spin::Mutex::new(None);
10
11/// Represents the global configuration for Burn.
12#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)]
13pub struct BurnConfig {
14    /// Configuration for operation fusion.
15    #[serde(default)]
16    fusion: FusionConfig,
17
18    /// Configuration for autodiff.
19    #[serde(default)]
20    autodiff: AutodiffConfig,
21
22    /// Configuration for the remote backend.
23    #[serde(default)]
24    remote: RemoteConfig,
25}
26
27impl BurnConfig {
28    /// Returns a reference to the operation-fusion configuration.
29    pub fn fusion(&self) -> &FusionConfig {
30        &self.fusion
31    }
32
33    /// Returns a reference to the autodiff configuration.
34    pub fn autodiff(&self) -> &AutodiffConfig {
35        &self.autodiff
36    }
37
38    /// Returns a reference to the remote-backend configuration.
39    pub fn remote(&self) -> &RemoteConfig {
40        &self.remote
41    }
42}
43
44impl RuntimeConfig for BurnConfig {
45    fn storage() -> &'static spin::Mutex<Option<Arc<Self>>> {
46        &BURN_GLOBAL_CONFIG
47    }
48
49    fn file_names() -> &'static [&'static str] {
50        &["burn.toml", "Burn.toml"]
51    }
52
53    // Match cubecl-common's `std_io` cfg: only available on platforms where
54    // the trait method exists. See cubecl-common's build.rs.
55    #[cfg(all(
56        feature = "std",
57        any(
58            target_os = "windows",
59            target_os = "linux",
60            target_os = "macos",
61            target_os = "android"
62        )
63    ))]
64    fn override_from_env(mut self) -> Self {
65        use super::fusion::FusionLogLevel;
66        use super::remote::RemoteLogLevel;
67
68        if let Ok(val) = std::env::var("BURN_FUSION_LOG") {
69            let level = match val.to_ascii_lowercase().as_str() {
70                "disabled" | "off" | "0" => FusionLogLevel::Disabled,
71                "basic" => FusionLogLevel::Basic,
72                "medium" => FusionLogLevel::Medium,
73                "full" | "1" => FusionLogLevel::Full,
74                _ => self.fusion.logger.level,
75            };
76            self.fusion.logger.level = level;
77            // Default to stderr so tests can see the output via `cargo test -- --nocapture`.
78            if level != FusionLogLevel::Disabled {
79                self.fusion.logger.stderr = true;
80            }
81        }
82
83        if let Ok(val) = std::env::var("BURN_FUSION_MAX_EXPLORATIONS")
84            && let Ok(n) = val.parse::<usize>()
85        {
86            self.fusion.beam_search.max_explorations = Some(n);
87        }
88
89        if let Ok(val) = std::env::var("BURN_REMOTE_LOG") {
90            let level = match val.to_ascii_lowercase().as_str() {
91                "disabled" | "off" | "0" => RemoteLogLevel::Disabled,
92                "basic" | "1" => RemoteLogLevel::Basic,
93                "full" | "2" => RemoteLogLevel::Full,
94                _ => self.remote.logger.level,
95            };
96            self.remote.logger.level = level;
97            if level != RemoteLogLevel::Disabled {
98                self.remote.logger.stderr = true;
99            }
100        }
101
102        self
103    }
104}