rusticl: add nir debugging option

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30152>

docs/envvars.rst

@@ -1153,6 +1153,8 @@ Rusticl environment variables
 
    - ``allow_invalid_spirv`` disables validation of any input SPIR-V
    - ``clc`` dumps all OpenCL C source being compiled
+   - ``nir`` dumps nirs in various compilation stages. Might print nothing if shader caching is
+             enabled.
    - ``no_reuse_context`` pipe_contexts are not recycled
    - ``no_variants`` disable kernel variants (e.g. specialized binaries for offsets == 0)
    - ``perf`` prints a warning when hitting slow paths once

src/gallium/frontends/rusticl/core/kernel.rs

@@ -22,6 +22,8 @@ use spirv::SpirvKernelInfo;
 use std::cmp;
 use std::collections::HashMap;
 use std::convert::TryInto;
+use std::fmt::Debug;
+use std::fmt::Display;
 use std::ops::Index;
 use std::os::raw::c_void;
 use std::ptr;
@@ -369,6 +371,13 @@ enum NirKernelVariant {
     Optimized,
 }
 
+impl Display for NirKernelVariant {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        // this simply prints the enum name, so that's fine
+        Debug::fmt(self, f)
+    }
+}
+
 pub struct NirKernelBuilds {
     default_build: NirKernelBuild,
     optimized: Option<NirKernelBuild>,
@@ -723,6 +732,7 @@ fn compile_nir_variant(
     dev: &Device,
     variant: NirKernelVariant,
     args: &[KernelArg],
+    name: &str,
 ) {
     let mut lower_state = rusticl_lower_state::default();
     let compiled_args = &mut res.compiled_args;
@@ -954,7 +964,20 @@ fn compile_nir_variant(
             false
         }
     }));
-    dev.screen.finalize_nir(nir);
+
+    if Platform::dbg().nir {
+        eprintln!("=== Printing nir variant '{variant}' for '{name}' before driver finalization");
+        nir.print();
+    }
+
+    if dev.screen.finalize_nir(nir) {
+        if Platform::dbg().nir {
+            eprintln!(
+                "=== Printing nir variant '{variant}' for '{name}' after driver finalization"
+            );
+            nir.print();
+        }
+    }
 
     nir_pass!(nir, nir_opt_dce);
     nir.sweep_mem();
@@ -964,6 +987,7 @@ fn compile_nir_remaining(
     dev: &Device,
     mut nir: NirShader,
     args: &[KernelArg],
+    name: &str,
 ) -> (CompilationResult, Option<CompilationResult>) {
     // add all API kernel args
     let mut compiled_args: Vec<_> = (0..args.len())
@@ -975,6 +999,11 @@ fn compile_nir_remaining(
         .collect();
 
     compile_nir_prepare_for_variants(dev, &mut nir, &mut compiled_args);
+    if Platform::dbg().nir {
+        eprintln!("=== Printing nir for '{name}' before specialization");
+        nir.print();
+    }
+
     let mut default_build = CompilationResult {
         nir: nir,
         compiled_args: compiled_args,
@@ -990,9 +1019,15 @@ fn compile_nir_remaining(
     let mut optimized = (!Platform::dbg().no_variants && (has_offsets || has_wgs_hint))
         .then(|| default_build.clone());
 
-    compile_nir_variant(&mut default_build, dev, NirKernelVariant::Default, args);
+    compile_nir_variant(
+        &mut default_build,
+        dev,
+        NirKernelVariant::Default,
+        args,
+        name,
+    );
     if let Some(optimized) = &mut optimized {
-        compile_nir_variant(optimized, dev, NirKernelVariant::Optimized, args);
+        compile_nir_variant(optimized, dev, NirKernelVariant::Optimized, args, name);
     }
 
     (default_build, optimized)
@@ -1097,8 +1132,14 @@ pub(super) fn convert_spirv_to_nir(
         .and_then(|entry| SPIRVToNirResult::deserialize(&entry, dev, spirv_info))
         .unwrap_or_else(|| {
             let nir = build.to_nir(name, dev);
+
+            if Platform::dbg().nir {
+                eprintln!("=== Printing nir for '{name}' after spirv_to_nir");
+                nir.print();
+            }
+
             let (mut args, nir) = compile_nir_to_args(dev, nir, args, &dev.lib_clc);
-            let (default_build, optimized) = compile_nir_remaining(dev, nir, &args);
+            let (default_build, optimized) = compile_nir_remaining(dev, nir, &args, name);
 
             for build in [Some(&default_build), optimized.as_ref()].into_iter() {
                 let Some(build) = build else {

src/gallium/frontends/rusticl/core/platform.rs

@@ -31,6 +31,7 @@ pub struct PlatformDebug {
     pub allow_invalid_spirv: bool,
     pub clc: bool,
     pub max_grid_size: u64,
+    pub nir: bool,
     pub no_variants: bool,
     pub perf: PerfDebugLevel,
     pub program: bool,
@@ -80,6 +81,7 @@ static mut PLATFORM_DBG: PlatformDebug = PlatformDebug {
     allow_invalid_spirv: false,
     clc: false,
     max_grid_size: 0,
+    nir: false,
     no_variants: false,
     perf: PerfDebugLevel::None,
     program: false,
@@ -100,6 +102,7 @@ fn load_env() {
             match flag {
                 "allow_invalid_spirv" => debug.allow_invalid_spirv = true,
                 "clc" => debug.clc = true,
+                "nir" => debug.nir = true,
                 "no_reuse_context" => debug.reuse_context = false,
                 "no_variants" => debug.no_variants = true,
                 "perf" => debug.perf = PerfDebugLevel::Once,

src/gallium/frontends/rusticl/mesa/pipe/screen.rs

@@ -414,11 +414,15 @@ impl PipeScreen {
         DiskCacheBorrowed::from_ptr(ptr)
     }
 
-    pub fn finalize_nir(&self, nir: &NirShader) {
+    /// returns true if finalize_nir was called
+    pub fn finalize_nir(&self, nir: &NirShader) -> bool {
         if let Some(func) = self.screen().finalize_nir {
             unsafe {
                 func(self.screen.as_ptr(), nir.get_nir().cast());
             }
+            true
+        } else {
+            false
         }
     }