cubecl_core/codegen/
dump.rs

1use cubecl_runtime::config::{CubeClRuntimeConfig, RuntimeConfig};
2use pliron::pass::PMConfig;
3
4/// The compiler dump of one kernel: its IR after every pass, plus the files each compiler adds.
5///
6/// Enabled when `compilation.dump_dir` (or `CUBECL_DEBUG_PLIRON`) is set. Otherwise every method
7/// is a no-op, so compilers call them unconditionally.
8#[derive(Clone, Debug, Default)]
9pub struct KernelDump {
10    #[cfg(feature = "std")]
11    dir: Option<std::path::PathBuf>,
12}
13
14impl KernelDump {
15    /// The dump of `kernel_name`, in a subdirectory of the configured dump directory.
16    #[must_use]
17    #[cfg_attr(not(feature = "std"), allow(unused_variables))]
18    pub fn new(kernel_name: &str) -> Self {
19        Self {
20            #[cfg(feature = "std")]
21            dir: kernel_dir(kernel_name),
22        }
23    }
24
25    /// Whether this kernel is dumped.
26    #[must_use]
27    pub fn is_enabled(&self) -> bool {
28        #[cfg(feature = "std")]
29        return self.dir.is_some();
30        #[cfg(not(feature = "std"))]
31        false
32    }
33
34    /// Writes `file_name` to the dump. `contents` runs only when the dump is enabled.
35    #[cfg_attr(not(feature = "std"), allow(unused_variables))]
36    pub fn write<C: AsRef<[u8]>>(&self, file_name: &str, contents: impl FnOnce() -> C) {
37        #[cfg(feature = "std")]
38        if let Some(dir) = &self.dir {
39            let path = dir.join(file_name);
40            if let Err(err) = std::fs::write(&path, contents()) {
41                log::warn!("Can't write the dump file {}: {err}", path.display());
42            }
43        }
44    }
45
46    /// Pass manager settings: print the IR after each pass to the dump, and time the passes
47    /// when `compilation.time_passes` is set.
48    #[must_use]
49    pub fn pass_config(&self) -> PMConfig {
50        PMConfig {
51            print_after_all: self.is_enabled(),
52            #[cfg(feature = "std")]
53            ir_printing_dir: self.dir.clone(),
54            time_all_passes: CubeClRuntimeConfig::get().compilation.time_passes,
55            ..Default::default()
56        }
57    }
58}
59
60/// Creates and returns the dump directory of `kernel_name`, if dumps are enabled.
61#[cfg(feature = "std")]
62fn kernel_dir(kernel_name: &str) -> Option<std::path::PathBuf> {
63    let root = CubeClRuntimeConfig::get().compilation.dump_dir.clone()?;
64    create_kernel_dir(&root, kernel_name)
65}
66
67#[cfg(feature = "std")]
68fn create_kernel_dir(root: &std::path::Path, kernel_name: &str) -> Option<std::path::PathBuf> {
69    let options = sanitize_filename::Options {
70        replacement: "_",
71        ..Default::default()
72    };
73    let dir = root.join(sanitize_filename::sanitize_with_options(
74        kernel_name,
75        options,
76    ));
77    match std::fs::create_dir_all(&dir) {
78        Ok(()) => Some(dir),
79        Err(err) => {
80            log::warn!("Can't create the dump directory {}: {err}", dir.display());
81            None
82        }
83    }
84}
85
86#[cfg(all(test, feature = "std"))]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn disabled_dump_builds_no_contents() {
92        KernelDump::default().write("never.txt", || -> &[u8] {
93            panic!("contents must not be built")
94        });
95    }
96
97    #[test]
98    fn enabled_dump_writes_to_a_sanitized_kernel_dir() {
99        let root = std::env::temp_dir().join(alloc::format!("cubecl-dump-{}", std::process::id()));
100        let dump = KernelDump {
101            dir: create_kernel_dir(&root, "matmul<f32>/tile"),
102        };
103        dump.write("module.cpp", || "kernel");
104
105        let dir = root.join("matmul_f32__tile");
106        assert_eq!(
107            std::fs::read_to_string(dir.join("module.cpp")).unwrap(),
108            "kernel"
109        );
110        let config = dump.pass_config();
111        assert!(config.print_after_all);
112        assert_eq!(config.ir_printing_dir, Some(dir));
113        std::fs::remove_dir_all(root).unwrap();
114    }
115}
cubecl_core/codegen/dump.rs

cubecl_core/codegen/
dump.rs
