nak: Add a NAK_DEBUG environment variable

Also, hide printing behind it and clean up the print code a bit.

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

src/nouveau/compiler/meson.build

@@ -56,6 +56,7 @@ nak_bindings_rs = rust.bindgen(
     '--allowlist-type', 'mesa_scope',
     '--allowlist-type', 'mesa_prim',
     '--allowlist-var', 'nir_.*_infos',
+    '--allowlist-function', '_mesa_shader_stage_to_string',
     '--allowlist-function', 'nak_.*',
     '--allowlist-function', 'nir_.*',
     '--allowlist-function', 'glsl_.*',

src/nouveau/compiler/nak.h

@@ -6,6 +6,8 @@
 #ifndef NAK_H
 #define NAK_H
 
+#include "compiler/shader_enums.h"
+
 #include <assert.h>
 #include <stdbool.h>
 #include <stdint.h>
@@ -22,6 +24,8 @@ struct nv_device_info;
 struct nak_compiler *nak_compiler_create(const struct nv_device_info *dev);
 void nak_compiler_destroy(struct nak_compiler *nak);
 
+uint64_t nak_debug_flags(const struct nak_compiler *nak);
+
 const struct nir_shader_compiler_options *
 nak_nir_options(const struct nak_compiler *nak);
 
@@ -30,6 +34,8 @@ void nak_preprocess_nir(nir_shader *nir, const struct nak_compiler *nak);
 void nak_postprocess_nir(nir_shader *nir, const struct nak_compiler *nak);
 
 struct nak_shader_info {
+   gl_shader_stage stage;
+
    /** Number of GPRs used */
    uint8_t num_gprs;
 

src/nouveau/compiler/nak.rs

@@ -24,9 +24,63 @@ use crate::nir::NirShader;
 use bitview::*;
 use nak_bindings::*;
 use nak_from_nir::*;
+use std::env;
+use std::ffi::CStr;
 use std::os::raw::c_void;
+use std::sync::OnceLock;
 use util::NextMultipleOf;
 
+#[repr(u8)]
+enum DebugFlags {
+    Print,
+}
+
+struct Debug {
+    flags: u32,
+}
+
+impl Debug {
+    fn new() -> Debug {
+        let debug_var = "NAK_DEBUG";
+        let debug_str = match env::var(debug_var) {
+            Ok(s) => s,
+            Err(_) => {
+                return Debug { flags: 0 };
+            }
+        };
+
+        let mut flags = 0;
+        for flag in debug_str.split(',') {
+            match flag.trim() {
+                "print" => flags |= 1 << DebugFlags::Print as u8,
+                unk => eprintln!("Unknown NAK_DEBUG flag \"{}\"", unk),
+            }
+        }
+        Debug { flags: flags }
+    }
+}
+
+trait GetDebugFlags {
+    fn debug_flags(&self) -> u32;
+
+    fn print(&self) -> bool {
+        self.debug_flags() & (1 << DebugFlags::Print as u8) != 0
+    }
+}
+
+static DEBUG: OnceLock<Debug> = OnceLock::new();
+
+impl GetDebugFlags for OnceLock<Debug> {
+    fn debug_flags(&self) -> u32 {
+        self.get().unwrap().flags
+    }
+}
+
+#[no_mangle]
+pub extern "C" fn nak_should_print_nir() -> bool {
+    DEBUG.print()
+}
+
 fn nir_options(dev: &nv_device_info) -> nir_shader_compiler_options {
     let mut op: nir_shader_compiler_options = unsafe { std::mem::zeroed() };
 
@@ -69,6 +123,8 @@ pub extern "C" fn nak_compiler_create(
     assert!(!dev.is_null());
     let dev = unsafe { &*dev };
 
+    DEBUG.get_or_init(|| Debug::new());
+
     let nak = Box::new(nak_compiler {
         sm: dev.sm,
         nir_options: nir_options(dev),
@@ -82,6 +138,11 @@ pub extern "C" fn nak_compiler_destroy(nak: *mut nak_compiler) {
     unsafe { Box::from_raw(nak) };
 }
 
+#[no_mangle]
+pub extern "C" fn nak_debug_flags(nak: *const nak_compiler) -> u64 {
+    DEBUG.debug_flags().into()
+}
+
 #[no_mangle]
 pub extern "C" fn nak_nir_options(
     nak: *const nak_compiler,
@@ -335,6 +396,18 @@ fn encode_hdr_for_nir(nir: &nir_shader, tls_size: u32) -> [u32; 32] {
     hdr
 }
 
+fn print_hex(label: &str, data: &[u32]) {
+    print!("{}:", label);
+    for i in 0..data.len() {
+        if (i % 8) == 0 {
+            println!("");
+            print!(" ");
+        }
+        print!(" {:08x}", data[i]);
+    }
+    println!("");
+}
+
 #[no_mangle]
 pub extern "C" fn nak_compile_shader(
     nir: *mut nir_shader,
@@ -346,32 +419,43 @@ pub extern "C" fn nak_compile_shader(
 
     let mut s = nak_shader_from_nir(nir, nak.sm);
 
-    println!("NAK IR:\n{}", &s);
+    if DEBUG.print() {
+        println!("NAK IR:\n{}", &s);
+    }
 
     s.opt_copy_prop();
-
-    println!("NAK IR:\n{}", &s);
+    if DEBUG.print() {
+        println!("NAK IR:\n{}", &s);
+    }
 
     s.opt_dce();
-
-    println!("NAK IR:\n{}", &s);
+    if DEBUG.print() {
+        println!("NAK IR:\n{}", &s);
+    }
 
     s.legalize();
-
-    println!("NAK IR:\n{}", &s);
+    if DEBUG.print() {
+        println!("NAK IR:\n{}", &s);
+    }
 
     s.assign_regs();
     //s.assign_regs_trivial();
-    println!("NAK IR:\n{}", &s);
+    if DEBUG.print() {
+        println!("NAK IR:\n{}", &s);
+    }
+
     s.lower_vec_split();
     s.lower_par_copies();
     s.lower_swap();
     s.lower_mov_predicate();
     s.calc_instr_deps();
 
-    println!("NAK IR:\n{}", &s);
+    if DEBUG.print() {
+        println!("NAK IR:\n{}", &s);
+    }
 
     let info = nak_shader_info {
+        stage: nir.info.stage(),
         num_gprs: 255,
         tls_size: 0,
         cs: nak_shader_info__bindgen_ty_1 {
@@ -391,14 +475,30 @@ pub extern "C" fn nak_compile_shader(
         panic!("Unsupported shader model");
     };
 
-    print!("Encoded shader:");
-    for i in 0..code.len() {
-        if (i % 8) == 0 {
-            print!("\n  ");
+    if DEBUG.print() {
+        let stage_name = unsafe {
+            let c_name = _mesa_shader_stage_to_string(info.stage as u32);
+            CStr::from_ptr(c_name).to_str().expect("Invalid UTF-8")
+        };
+        println!("Stage: {}", stage_name);
+        println!("Num GPRs: {}", info.num_gprs);
+        println!("TLS size: {}", info.tls_size);
+        if info.stage == MESA_SHADER_COMPUTE {
+            println!(
+                "Local size: {}x{}x{}",
+                info.cs.local_size[0],
+                info.cs.local_size[1],
+                info.cs.local_size[2],
+            );
+            println!("Shared memory size: {:#x}", info.cs.smem_size);
         }
-        print!("  {:#x}", code[i]);
+
+        if info.stage != MESA_SHADER_COMPUTE {
+            print_hex("Header", &info.hdr);
+        }
+
+        print_hex("Encoded shader", &code);
     }
-    print!("\n\n");
 
     Box::into_raw(Box::new(ShaderBin::new(info, code))) as *mut nak_shader_bin
 }

src/nouveau/compiler/nak_nir.c

@@ -578,7 +578,8 @@ nak_postprocess_nir(nir_shader *nir, const struct nak_compiler *nak)
       }
    }
 
-   nir_print_shader(nir, stderr);
+   if (nak_should_print_nir())
+      nir_print_shader(nir, stderr);
 }
 
 static bool

src/nouveau/compiler/nak_private.h

@@ -14,6 +14,8 @@
 extern "C" {
 #endif
 
+bool nak_should_print_nir(void);
+
 struct nak_compiler {
    uint8_t sm;