From dc77542ed2226173286f4c9b37425efd101dc51c Mon Sep 17 00:00:00 2001
From: Caio Oliveira <caio.oliveira@intel.com>
Date: Tue, 1 Mar 2022 11:29:41 -0800
Subject: [PATCH] intel/compiler: Use pass helper in
 brw_nir_adjust_offset_for_arrayed_indices
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Also change the code to preserve certain metadata: control flow is not changed
so both block indices and dominance information is preserved.

Reviewed-by: Marcin Ślusarz <marcin.slusarz@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/15206>
---
 src/intel/compiler/brw_mesh.cpp | 118 +++++++++++++++-----------------
 1 file changed, 57 insertions(+), 61 deletions(-)

diff --git a/src/intel/compiler/brw_mesh.cpp b/src/intel/compiler/brw_mesh.cpp
index d90dac3e0cf..521f89a2f77 100644
--- a/src/intel/compiler/brw_mesh.cpp
+++ b/src/intel/compiler/brw_mesh.cpp
@@ -457,70 +457,66 @@ brw_nir_lower_mue_outputs(nir_shader *nir, const struct brw_mue_map *map)
                 nir_lower_io_lower_64bit_to_32);
 }
 
+static bool
+brw_nir_adjust_offset_for_arrayed_indices_instr(nir_builder *b, nir_instr *instr, void *data)
+{
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+
+   const struct brw_mue_map *map = (const struct brw_mue_map *) data;
+
+   /* Remap per_vertex and per_primitive offsets using the extra source and
+    * the pitch.
+    */
+   switch (intrin->intrinsic) {
+   case nir_intrinsic_load_per_vertex_output:
+   case nir_intrinsic_store_per_vertex_output: {
+      const bool is_load = intrin->intrinsic == nir_intrinsic_load_per_vertex_output;
+      nir_src *index_src = &intrin->src[is_load ? 0 : 1];
+      nir_src *offset_src = &intrin->src[is_load ? 1 : 2];
+
+      assert(index_src->is_ssa);
+      b->cursor = nir_before_instr(&intrin->instr);
+      nir_ssa_def *offset =
+         nir_iadd(b,
+                  offset_src->ssa,
+                  nir_imul_imm(b, index_src->ssa, map->per_vertex_pitch_dw));
+      nir_instr_rewrite_src(&intrin->instr, offset_src, nir_src_for_ssa(offset));
+      return true;
+   }
+
+   case nir_intrinsic_load_per_primitive_output:
+   case nir_intrinsic_store_per_primitive_output: {
+      const bool is_load = intrin->intrinsic == nir_intrinsic_load_per_primitive_output;
+      nir_src *index_src = &intrin->src[is_load ? 0 : 1];
+      nir_src *offset_src = &intrin->src[is_load ? 1 : 2];
+
+      assert(index_src->is_ssa);
+      b->cursor = nir_before_instr(&intrin->instr);
+
+      assert(index_src->is_ssa);
+      nir_ssa_def *offset =
+         nir_iadd(b,
+                  offset_src->ssa,
+                  nir_imul_imm(b, index_src->ssa, map->per_primitive_pitch_dw));
+      nir_instr_rewrite_src(&intrin->instr, offset_src, nir_src_for_ssa(offset));
+      return true;
+   }
+
+   default:
+      return false;
+   }
+}
+
 static void
 brw_nir_adjust_offset_for_arrayed_indices(nir_shader *nir, const struct brw_mue_map *map)
 {
-   /* TODO(mesh): Check if we need to inject extra vertex header / primitive
-    * setup.  If so, we should add them together some required value for
-    * vertex/primitive.
-    */
-
-   /* Remap per_vertex and per_primitive offsets using the extra source and the pitch. */
-   nir_foreach_function(function, nir) {
-      if (function->impl) {
-         nir_builder b;
-         nir_builder_init(&b, function->impl);
-
-         nir_foreach_block(block, function->impl) {
-            nir_foreach_instr(instr, block) {
-               if (instr->type != nir_instr_type_intrinsic)
-                  continue;
-               nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
-
-               switch (intrin->intrinsic) {
-               case nir_intrinsic_load_per_vertex_output:
-               case nir_intrinsic_store_per_vertex_output: {
-                  const bool is_load = intrin->intrinsic == nir_intrinsic_load_per_vertex_output;
-                  nir_src *index_src = &intrin->src[is_load ? 0 : 1];
-                  nir_src *offset_src = &intrin->src[is_load ? 1 : 2];
-
-                  assert(index_src->is_ssa);
-                  b.cursor = nir_before_instr(&intrin->instr);
-                  nir_ssa_def *offset =
-                     nir_iadd(&b,
-                              offset_src->ssa,
-                              nir_imul_imm(&b, index_src->ssa, map->per_vertex_pitch_dw));
-                  nir_instr_rewrite_src(&intrin->instr, offset_src, nir_src_for_ssa(offset));
-                  break;
-               }
-
-               case nir_intrinsic_load_per_primitive_output:
-               case nir_intrinsic_store_per_primitive_output: {
-                  const bool is_load = intrin->intrinsic == nir_intrinsic_load_per_primitive_output;
-                  nir_src *index_src = &intrin->src[is_load ? 0 : 1];
-                  nir_src *offset_src = &intrin->src[is_load ? 1 : 2];
-
-                  assert(index_src->is_ssa);
-                  b.cursor = nir_before_instr(&intrin->instr);
-
-                  assert(index_src->is_ssa);
-                  nir_ssa_def *offset =
-                     nir_iadd(&b,
-                              offset_src->ssa,
-                              nir_imul_imm(&b, index_src->ssa, map->per_primitive_pitch_dw));
-                  nir_instr_rewrite_src(&intrin->instr, offset_src, nir_src_for_ssa(offset));
-                  break;
-               }
-
-               default:
-                  /* Nothing to do. */
-                  break;
-               }
-            }
-         }
-         nir_metadata_preserve(function->impl, nir_metadata_none);
-      }
-   }
+   nir_shader_instructions_pass(nir, brw_nir_adjust_offset_for_arrayed_indices_instr,
+                                nir_metadata_block_index |
+                                nir_metadata_dominance,
+                                (void *)map);
 }
 
 const unsigned *