burn_std/config/
logger.rs

1use alloc::{string::String, vec::Vec};
2use core::fmt::Display;
3use cubecl_common::{
4    config::{
5        RuntimeConfig,
6        logger::{LogLevel, LoggerConfig, LoggerSinks},
7    },
8    stub::Arc,
9};
10
11use super::{
12    autodiff::AutodiffLogLevel, base::BurnConfig, fusion::FusionLogLevel, remote::RemoteLogLevel,
13};
14
15static BURN_LOGGER: spin::Mutex<Option<Logger>> = spin::Mutex::new(None);
16
17#[cfg(feature = "std")]
18std::thread_local! {
19    static LOCAL_CONFIG: std::cell::OnceCell<Arc<BurnConfig>> =
20        const { std::cell::OnceCell::new() };
21}
22
23/// Returns the current [`BurnConfig`], cached in thread-local storage on native targets.
24///
25/// On the first call from a given thread this fetches the global config via
26/// [`BurnConfig::get`] (which locks a spin mutex) and caches the `Arc` thread-locally.
27/// Subsequent calls on the same thread only pay an `Arc` clone. On `no_std` builds this
28/// is equivalent to [`BurnConfig::get`].
29///
30/// Safe because [`BurnConfig::set`] panics after the first read, so the cached snapshot
31/// matches the global singleton for the whole program lifetime.
32pub fn config() -> Arc<BurnConfig> {
33    #[cfg(feature = "std")]
34    {
35        LOCAL_CONFIG.with(|cell| cell.get_or_init(BurnConfig::get).clone())
36    }
37    #[cfg(not(feature = "std"))]
38    {
39        BurnConfig::get()
40    }
41}
42
43/// Central logging utility for Burn, managing one sink registry shared across subsystems.
44#[derive(Debug)]
45pub struct Logger {
46    sinks: LoggerSinks,
47    fusion_index: Vec<usize>,
48    autodiff_index: Vec<usize>,
49    remote_index: Vec<usize>,
50    /// The configuration snapshot the logger was initialized with.
51    pub config: Arc<BurnConfig>,
52}
53
54impl Default for Logger {
55    fn default() -> Self {
56        Self::new()
57    }
58}
59
60impl Logger {
61    /// Creates a new `Logger` from the current global `BurnConfig`.
62    ///
63    /// Note that creating a logger is somewhat expensive because it opens file handles for any
64    /// sink configured with a file path.
65    pub fn new() -> Self {
66        let config = BurnConfig::get();
67        let mut sinks = LoggerSinks::new();
68
69        let fusion_index = register_enabled(
70            &mut sinks,
71            &config.fusion().logger,
72            config.fusion().logger.level != FusionLogLevel::Disabled,
73        );
74        let autodiff_index = register_enabled(
75            &mut sinks,
76            &config.autodiff().logger,
77            config.autodiff().logger.level != AutodiffLogLevel::Disabled,
78        );
79        let remote_index = register_enabled(
80            &mut sinks,
81            &config.remote().logger,
82            config.remote().logger.level != RemoteLogLevel::Disabled,
83        );
84
85        Self {
86            sinks,
87            fusion_index,
88            autodiff_index,
89            remote_index,
90            config,
91        }
92    }
93
94    /// Writes `msg` to all configured fusion sinks.
95    pub fn log_fusion<S: Display>(&mut self, msg: &S) {
96        self.sinks.log(&self.fusion_index, "burn::fusion", msg);
97    }
98
99    /// Writes `msg` to all configured autodiff sinks.
100    pub fn log_autodiff<S: Display>(&mut self, msg: &S) {
101        self.sinks.log(&self.autodiff_index, "burn::autodiff", msg);
102    }
103
104    /// Writes `msg` to all configured remote-backend sinks.
105    pub fn log_remote<S: Display>(&mut self, msg: &S) {
106        self.sinks.log(&self.remote_index, "burn::remote", msg);
107    }
108
109    /// Returns the current fusion log level.
110    pub fn log_level_fusion(&self) -> FusionLogLevel {
111        self.config.fusion().logger.level
112    }
113
114    /// Returns the current remote-backend log level.
115    pub fn log_level_remote(&self) -> RemoteLogLevel {
116        self.config.remote().logger.level
117    }
118
119    /// Returns the current autodiff log level.
120    pub fn log_level_autodiff(&self) -> AutodiffLogLevel {
121        self.config.autodiff().logger.level
122    }
123}
124
125fn register_enabled<L: LogLevel>(
126    sinks: &mut LoggerSinks,
127    config: &LoggerConfig<L>,
128    enabled: bool,
129) -> Vec<usize> {
130    if enabled {
131        sinks.register(config)
132    } else {
133        Vec::new()
134    }
135}
136
137/// Emit a fusion log message when the configured level is at least `level`.
138///
139/// The message is only constructed when logging is enabled.
140pub fn log_fusion<F>(level: FusionLogLevel, f: F)
141where
142    F: FnOnce() -> String,
143{
144    let current = config().fusion().logger.level;
145    if current < level {
146        return;
147    }
148    let msg = f();
149    let mut guard = BURN_LOGGER.lock();
150    let logger = guard.get_or_insert_with(Logger::new);
151    logger.log_fusion(&msg);
152}
153
154/// Emit an autodiff log message when the configured level is at least `level`.
155///
156/// The message is only constructed when logging is enabled.
157pub fn log_autodiff<F>(level: AutodiffLogLevel, f: F)
158where
159    F: FnOnce() -> String,
160{
161    let current = config().autodiff().logger.level;
162    if current < level {
163        return;
164    }
165    let msg = f();
166    let mut guard = BURN_LOGGER.lock();
167    let logger = guard.get_or_insert_with(Logger::new);
168    logger.log_autodiff(&msg);
169}
170
171/// Emit a remote-backend log message when the configured level is at least `level`.
172///
173/// The message is only constructed when logging is enabled.
174pub fn log_remote<F>(level: RemoteLogLevel, f: F)
175where
176    F: FnOnce() -> String,
177{
178    let current = config().remote().logger.level;
179    if current < level {
180        return;
181    }
182    let msg = f();
183    let mut guard = BURN_LOGGER.lock();
184    let logger = guard.get_or_insert_with(Logger::new);
185    logger.log_remote(&msg);
186}
burn_std/config/logger.rs

burn_std/config/
logger.rs
