From 1110fcccc2636cbe7aa571a9e53ccef3c4fb0076 Mon Sep 17 00:00:00 2001
From: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Date: Mon, 7 Nov 2022 20:51:36 -0500
Subject: [PATCH] agx: Split off NIR preprocessing from compiling

So we can specialize after lowering I/O.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/19811>
---
 src/asahi/compiler/agx_compile.c      | 49 ++++++++++++++++++---------
 src/asahi/compiler/agx_compile.h      |  3 ++
 src/gallium/drivers/asahi/agx_blit.c  |  1 +
 src/gallium/drivers/asahi/agx_state.c |  1 +
 4 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c
index 1ade3f4d1fa..b5522af93b1 100644
--- a/src/asahi/compiler/agx_compile.c
+++ b/src/asahi/compiler/agx_compile.c
@@ -1883,24 +1883,13 @@ agx_compile_function_nir(nir_shader *nir, nir_function_impl *impl,
    return offset;
 }
 
+/*
+ * Preprocess NIR. In particular, this lowers I/O. Drivers should call this
+ * as soon as they don't need unlowered I/O.
+ */
 void
-agx_compile_shader_nir(nir_shader *nir,
-      struct agx_shader_key *key,
-      struct util_debug_callback *debug,
-      struct util_dynarray *binary,
-      struct agx_shader_info *out)
+agx_preprocess_nir(nir_shader *nir)
 {
-   agx_debug = debug_get_option_agx_debug();
-
-   memset(out, 0, sizeof *out);
-
-   if (nir->info.stage == MESA_SHADER_VERTEX) {
-      out->writes_psiz = nir->info.outputs_written &
-         BITFIELD_BIT(VARYING_SLOT_PSIZ);
-   } else if (nir->info.stage == MESA_SHADER_FRAGMENT) {
-      out->no_colour_output = !(nir->info.outputs_written >> FRAG_RESULT_DATA0);
-   }
-
    NIR_PASS_V(nir, nir_lower_vars_to_ssa);
 
    /* Lower large arrays to scratch and small arrays to csel */
@@ -1934,6 +1923,9 @@ agx_compile_shader_nir(nir_shader *nir,
       NIR_PASS_V(nir, nir_lower_io_to_scalar, nir_var_shader_out);
    }
 
+   /* Clean up deref gunk after lowering I/O */
+   NIR_PASS_V(nir, nir_opt_dce);
+
    nir_lower_tex_options lower_tex_options = {
       .lower_txp = ~0,
       .lower_invalid_implicit_lod = true,
@@ -1954,6 +1946,31 @@ agx_compile_shader_nir(nir_shader *nir,
    NIR_PASS_V(nir, agx_lower_resinfo);
    NIR_PASS_V(nir, nir_legalize_16bit_sampler_srcs, tex_constraints);
 
+   nir->info.io_lowered = true;
+}
+
+void
+agx_compile_shader_nir(nir_shader *nir,
+      struct agx_shader_key *key,
+      struct util_debug_callback *debug,
+      struct util_dynarray *binary,
+      struct agx_shader_info *out)
+{
+   agx_debug = debug_get_option_agx_debug();
+
+   memset(out, 0, sizeof *out);
+
+   assert(nir->info.io_lowered &&
+          "agx_preprocess_nir is called first, then the shader is specalized,"
+          "then the specialized shader is compiled");
+
+   if (nir->info.stage == MESA_SHADER_VERTEX) {
+      out->writes_psiz = nir->info.outputs_written &
+         BITFIELD_BIT(VARYING_SLOT_PSIZ);
+   } else if (nir->info.stage == MESA_SHADER_FRAGMENT) {
+      out->no_colour_output = !(nir->info.outputs_written >> FRAG_RESULT_DATA0);
+   }
+
    agx_optimize_nir(nir, &out->push_count);
 
    /* Implement conditional discard with real control flow like Metal */
diff --git a/src/asahi/compiler/agx_compile.h b/src/asahi/compiler/agx_compile.h
index f3e6ac4644f..738890b3795 100644
--- a/src/asahi/compiler/agx_compile.h
+++ b/src/asahi/compiler/agx_compile.h
@@ -276,6 +276,9 @@ struct agx_shader_key {
    };
 };
 
+void
+agx_preprocess_nir(nir_shader *nir);
+
 void
 agx_compile_shader_nir(nir_shader *nir,
       struct agx_shader_key *key,
diff --git a/src/gallium/drivers/asahi/agx_blit.c b/src/gallium/drivers/asahi/agx_blit.c
index 06a904d671a..bbb95c52b25 100644
--- a/src/gallium/drivers/asahi/agx_blit.c
+++ b/src/gallium/drivers/asahi/agx_blit.c
@@ -70,6 +70,7 @@ agx_build_reload_shader(struct agx_device *dev)
          .fs.ignore_tib_dependencies = true,
       };
 
+      agx_preprocess_nir(s);
       agx_compile_shader_nir(s, &key, NULL, &binary, &info);
 
       assert(offset + binary.size < bo_size);
diff --git a/src/gallium/drivers/asahi/agx_state.c b/src/gallium/drivers/asahi/agx_state.c
index 9e4803efcdd..d463922864c 100644
--- a/src/gallium/drivers/asahi/agx_state.c
+++ b/src/gallium/drivers/asahi/agx_state.c
@@ -1078,6 +1078,7 @@ agx_compile_variant(struct agx_device *dev,
       }
    }
 
+   agx_preprocess_nir(nir);
    agx_compile_shader_nir(nir, &key->base, debug, &binary, &compiled->info);
 
    if (binary.size) {