From 53df391d27c4d8410c148c490a03f72b39c20aa5 Mon Sep 17 00:00:00 2001
From: Karol Herbst <kherbst@redhat.com>
Date: Thu, 15 Feb 2024 14:29:20 +0100
Subject: [PATCH] rusticl/kernel: recalculate scratch and shared memory after
 opts

Turns out we have to do it regardless as opts after explicit_types could
get rid of even more scratch and shared memory.

Fixes: ea023ff5cdd ("rusticl/kernel: no need to reset the scratch size anymore")
Signed-off-by: Karol Herbst <kherbst@redhat.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27634>
---
 src/gallium/frontends/rusticl/core/kernel.rs   | 18 ++++++++++++++++++
 .../frontends/rusticl/mesa/compiler/nir.rs     | 11 +++++++++++
 2 files changed, 29 insertions(+)

diff --git a/src/gallium/frontends/rusticl/core/kernel.rs b/src/gallium/frontends/rusticl/core/kernel.rs
index fc4c387dce0..6c4b477fb18 100644
--- a/src/gallium/frontends/rusticl/core/kernel.rs
+++ b/src/gallium/frontends/rusticl/core/kernel.rs
@@ -634,6 +634,24 @@ fn lower_and_optimize_nir(
     opt_nir(nir, dev, true);
     nir_pass!(nir, nir_lower_memcpy);
 
+    // we might have got rid of more function_temp or shared memory
+    nir.reset_scratch_size();
+    nir.reset_shared_size();
+    nir_pass!(
+        nir,
+        nir_remove_dead_variables,
+        nir_variable_mode::nir_var_function_temp | nir_variable_mode::nir_var_mem_shared,
+        &dv_opts,
+    );
+    nir_pass!(
+        nir,
+        nir_lower_vars_to_explicit_types,
+        nir_variable_mode::nir_var_function_temp
+            | nir_variable_mode::nir_var_mem_shared
+            | nir_variable_mode::nir_var_mem_generic,
+        Some(glsl_get_cl_type_size_align),
+    );
+
     nir_pass!(
         nir,
         nir_lower_explicit_io,
diff --git a/src/gallium/frontends/rusticl/mesa/compiler/nir.rs b/src/gallium/frontends/rusticl/mesa/compiler/nir.rs
index 5845d017d13..5fadf1810c5 100644
--- a/src/gallium/frontends/rusticl/mesa/compiler/nir.rs
+++ b/src/gallium/frontends/rusticl/mesa/compiler/nir.rs
@@ -308,10 +308,21 @@ impl NirShader {
         unsafe { (*self.nir.as_ptr()).info.num_textures }
     }
 
+    pub fn reset_scratch_size(&mut self) {
+        unsafe {
+            (*self.nir.as_ptr()).scratch_size = 0;
+        }
+    }
+
     pub fn scratch_size(&self) -> u32 {
         unsafe { (*self.nir.as_ptr()).scratch_size }
     }
 
+    pub fn reset_shared_size(&mut self) {
+        unsafe {
+            (*self.nir.as_ptr()).info.shared_size = 0;
+        }
+    }
     pub fn shared_size(&self) -> u32 {
         unsafe { (*self.nir.as_ptr()).info.shared_size }
     }