cubecl-runtime 0.10.0-pre.4

Crate that helps creating high performance async runtimes for CubeCL.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

use super::{CubeClRuntimeConfig, RuntimeConfig};
use crate::config::{
    autotune::AutotuneLogLevel, compilation::CompilationLogLevel, memory::MemoryLogLevel,
    profiling::ProfilingLogLevel, streaming::StreamingLogLevel,
};
use alloc::{sync::Arc, vec::Vec};
use core::fmt::Display;

use cubecl_common::config::logger::LoggerSinks;
pub(crate) use cubecl_common::config::logger::{LogLevel, LoggerConfig};

/// Central logging utility for `CubeCL`, managing multiple log outputs.
#[derive(Debug)]
pub struct Logger {
    sinks: LoggerSinks,
    compilation_index: Vec<usize>,
    profiling_index: Vec<usize>,
    autotune_index: Vec<usize>,
    streaming_index: Vec<usize>,
    memory_index: Vec<usize>,
    /// Global configuration for logging settings.
    pub config: Arc<CubeClRuntimeConfig>,
}

impl Default for Logger {
    fn default() -> Self {
        Self::new()
    }
}

impl Logger {
    /// Creates a new `Logger` instance based on the global configuration.
    ///
    /// Note that creating a logger is quite expensive.
    pub fn new() -> Self {
        let config = CubeClRuntimeConfig::get();
        let mut sinks = LoggerSinks::new();

        let compilation_index = register_enabled(
            &mut sinks,
            &config.compilation.logger,
            !matches!(
                config.compilation.logger.level,
                CompilationLogLevel::Disabled
            ),
        );
        let profiling_index = register_enabled(
            &mut sinks,
            &config.profiling.logger,
            !matches!(config.profiling.logger.level, ProfilingLogLevel::Disabled),
        );
        let autotune_index = register_enabled(
            &mut sinks,
            &config.autotune.logger,
            !matches!(config.autotune.logger.level, AutotuneLogLevel::Disabled),
        );
        let streaming_index = register_enabled(
            &mut sinks,
            &config.streaming.logger,
            !matches!(config.streaming.logger.level, StreamingLogLevel::Disabled),
        );
        let memory_index = register_enabled(
            &mut sinks,
            &config.memory.logger,
            !matches!(config.memory.logger.level, MemoryLogLevel::Disabled),
        );

        Self {
            sinks,
            compilation_index,
            profiling_index,
            autotune_index,
            streaming_index,
            memory_index,
            config,
        }
    }

    /// Logs a message for streaming, directing it to all configured streaming loggers.
    pub fn log_streaming<S: Display>(&mut self, msg: &S) {
        self.sinks.log(&self.streaming_index, msg);
    }

    /// Logs a message for memory, directing it to all configured memory loggers.
    pub fn log_memory<S: Display>(&mut self, msg: &S) {
        self.sinks.log(&self.memory_index, msg);
    }

    /// Logs a message for compilation, directing it to all configured compilation loggers.
    pub fn log_compilation<S: Display>(&mut self, msg: &S) {
        self.sinks.log(&self.compilation_index, msg);
    }

    /// Logs a message for profiling, directing it to all configured profiling loggers.
    pub fn log_profiling<S: Display>(&mut self, msg: &S) {
        self.sinks.log(&self.profiling_index, msg);
    }

    /// Logs a message for autotuning, directing it to all configured autotuning loggers.
    pub fn log_autotune<S: Display>(&mut self, msg: &S) {
        self.sinks.log(&self.autotune_index, msg);
    }

    /// Returns the current streaming log level from the global configuration.
    pub fn log_level_streaming(&self) -> StreamingLogLevel {
        self.config.streaming.logger.level
    }

    /// Returns the current autotune log level from the global configuration.
    pub fn log_level_autotune(&self) -> AutotuneLogLevel {
        self.config.autotune.logger.level
    }

    /// Returns the current compilation log level from the global configuration.
    pub fn log_level_compilation(&self) -> CompilationLogLevel {
        self.config.compilation.logger.level
    }

    /// Returns the current profiling log level from the global configuration.
    pub fn log_level_profiling(&self) -> ProfilingLogLevel {
        self.config.profiling.logger.level
    }
}

fn register_enabled<L: LogLevel>(
    sinks: &mut LoggerSinks,
    config: &LoggerConfig<L>,
    enabled: bool,
) -> Vec<usize> {
    if enabled {
        sinks.register(config)
    } else {
        Vec::new()
    }
}