nak: Re-plumb compute shader info

Put the stuff in nak_shader_info in nak.h into a union so we can have
other stage info exported as well.  Then plumb local_size and smem_size
through ShaderInfo like we do for a bunch of the FS things.  While we're
shuffling things around, pull nak_shader_info::cs into a union.

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

src/nouveau/compiler/nak.h

@@ -49,13 +49,18 @@ struct nak_shader_info {
    /** Size of thread-local storage */
    uint32_t tls_size;
 
-   struct {
-      /* Local workgroup size */
-      uint16_t local_size[3];
+   union {
+      struct {
+         /* Local workgroup size */
+         uint16_t local_size[3];
 
-      /* Shared memory size */
-      uint16_t smem_size;
-   } cs;
+         /* Shared memory size */
+         uint16_t smem_size;
+      } cs;
+
+      /* Used to initialize the union for other stages */
+      uint32_t dummy;
+   };
 
    /** Shader header for 3D stages */
    uint32_t hdr[32];

src/nouveau/compiler/nak.rs

@@ -496,13 +496,22 @@ pub extern "C" fn nak_compile_shader(
             0
         },
         tls_size: s.info.tls_size,
-        cs: nak_shader_info__bindgen_ty_1 {
-            local_size: [
-                nir.info.workgroup_size[0].into(),
-                nir.info.workgroup_size[1].into(),
-                nir.info.workgroup_size[2].into(),
-            ],
-            smem_size: nir.info.shared_size.try_into().unwrap(),
+        __bindgen_anon_1: match &s.info.stage {
+            ShaderStageInfo::Compute(cs_info) => {
+                nak_shader_info__bindgen_ty_1 {
+                    cs: nak_shader_info__bindgen_ty_1__bindgen_ty_1 {
+                        local_size: [
+                            cs_info.local_size[0],
+                            cs_info.local_size[1],
+                            cs_info.local_size[2],
+                        ],
+                        smem_size: cs_info.smem_size,
+                    },
+                }
+            }
+            _ => nak_shader_info__bindgen_ty_1 {
+                dummy: 0,
+            },
         },
         hdr: encode_hdr_for_nir(nir, &s.info, fs_key),
     };
@@ -521,15 +530,6 @@ pub extern "C" fn nak_compile_shader(
         eprintln!("Stage: {}", stage_name);
         eprintln!("Num GPRs: {}", info.num_gprs);
         eprintln!("TLS size: {}", info.tls_size);
-        if info.stage == MESA_SHADER_COMPUTE {
-            eprintln!(
-                "Local size: {}x{}x{}",
-                info.cs.local_size[0],
-                info.cs.local_size[1],
-                info.cs.local_size[2],
-            );
-            eprintln!("Shared memory size: {:#x}", info.cs.smem_size);
-        }
 
         if info.stage != MESA_SHADER_COMPUTE {
             eprint_hex("Header", &info.hdr);

src/nouveau/compiler/nak_from_nir.rs

@@ -22,7 +22,16 @@ fn init_info_from_nir(nir: &nir_shader, sm: u8) -> ShaderInfo {
         num_gprs: 0,
         tls_size: nir.scratch_size,
         stage: match nir.info.stage() {
-            MESA_SHADER_COMPUTE => ShaderStageInfo::Compute,
+            MESA_SHADER_COMPUTE => {
+                ShaderStageInfo::Compute(ComputeShaderInfo {
+                    local_size: [
+                        nir.info.workgroup_size[0].into(),
+                        nir.info.workgroup_size[1].into(),
+                        nir.info.workgroup_size[2].into(),
+                    ],
+                    smem_size: nir.info.shared_size.try_into().unwrap(),
+                })
+            }
             MESA_SHADER_FRAGMENT => {
                 ShaderStageInfo::Fragment(Default::default())
             }

src/nouveau/compiler/nak_ir.rs

@@ -4330,6 +4330,12 @@ impl fmt::Display for Function {
     }
 }
 
+#[derive(Debug)]
+pub struct ComputeShaderInfo {
+    pub local_size: [u16; 3],
+    pub smem_size: u16,
+}
+
 #[derive(Debug, Default)]
 pub struct FragmentShaderInfo {
     pub writes_color: u32,
@@ -4339,7 +4345,7 @@ pub struct FragmentShaderInfo {
 
 #[derive(Debug)]
 pub enum ShaderStageInfo {
-    Compute,
+    Compute(ComputeShaderInfo),
     Fragment(FragmentShaderInfo),
 }