nir/nir_lower_uniforms_to_ubo: Use the nir_shader_instructions_pass() helper

Reviewed-by: Alyssa Rosenzweig <alyssa.rosenzweig@collabora.com>
Reviewed-by: Emma Anholt <emma@anholt.net>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/11683>

diff --git a/src/compiler/nir/nir_lower_uniforms_to_ubo.c b/src/compiler/nir/nir_lower_uniforms_to_ubo.c
index 8fefcfd..6424229 100644 (file)
--- a/src/compiler/nir/nir_lower_uniforms_to_ubo.c
+++ b/src/compiler/nir/nir_lower_uniforms_to_ubo.c
@@ -38,43 +38,55 @@
 #include "nir.h"
 #include "nir_builder.h"
 
+struct nir_lower_uniforms_to_ubo_state {
+   bool dword_packed;
+   bool load_vec4;
+};
+
 static bool
-lower_instr(nir_intrinsic_instr *instr, nir_builder *b, bool dword_packed, bool load_vec4)
+nir_lower_uniforms_to_ubo_instr(nir_builder *b, nir_instr *instr, void *data)
 {
-   b->cursor = nir_before_instr(&instr->instr);
+   struct nir_lower_uniforms_to_ubo_state *state = data;
+
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+
+   b->cursor = nir_before_instr(&intr->instr);
 
    /* Increase all UBO binding points by 1. */
-   if (instr->intrinsic == nir_intrinsic_load_ubo &&
+   if (intr->intrinsic == nir_intrinsic_load_ubo &&
        !b->shader->info.first_ubo_is_default_ubo) {
-      nir_ssa_def *old_idx = nir_ssa_for_src(b, instr->src[0], 1);
+      nir_ssa_def *old_idx = nir_ssa_for_src(b, intr->src[0], 1);
       nir_ssa_def *new_idx = nir_iadd(b, old_idx, nir_imm_int(b, 1));
-      nir_instr_rewrite_src(&instr->instr, &instr->src[0],
+      nir_instr_rewrite_src(&intr->instr, &intr->src[0],
                             nir_src_for_ssa(new_idx));
       return true;
    }
 
-   if (instr->intrinsic == nir_intrinsic_load_uniform) {
+   if (intr->intrinsic == nir_intrinsic_load_uniform) {
       nir_ssa_def *ubo_idx = nir_imm_int(b, 0);
-      nir_ssa_def *uniform_offset = nir_ssa_for_src(b, instr->src[0], 1);
+      nir_ssa_def *uniform_offset = nir_ssa_for_src(b, intr->src[0], 1);
 
-      assert(instr->dest.ssa.bit_size >= 8);
+      assert(intr->dest.ssa.bit_size >= 8);
       nir_ssa_def *load_result;
-      if (load_vec4) {
+      if (state->load_vec4) {
          /* No asking us to generate load_vec4 when you've packed your uniforms
           * as dwords instead of vec4s.
           */
-         assert(!dword_packed);
-         load_result = nir_load_ubo_vec4(b, instr->num_components, instr->dest.ssa.bit_size,
-                                         ubo_idx, uniform_offset, .base=nir_intrinsic_base(instr));
+         assert(!state->dword_packed);
+         load_result = nir_load_ubo_vec4(b, intr->num_components, intr->dest.ssa.bit_size,
+                                         ubo_idx, uniform_offset, .base=nir_intrinsic_base(intr));
       } else {
          /* For PIPE_CAP_PACKED_UNIFORMS, the uniforms are packed with the
           * base/offset in dword units instead of vec4 units.
           */
-         int multiplier = dword_packed ? 4 : 16;
-         load_result = nir_load_ubo(b, instr->num_components, instr->dest.ssa.bit_size,
+         int multiplier = state->dword_packed ? 4 : 16;
+         load_result = nir_load_ubo(b, intr->num_components, intr->dest.ssa.bit_size,
                              ubo_idx,
                              nir_iadd_imm(b, nir_imul_imm(b, uniform_offset, multiplier),
-                                          nir_intrinsic_base(instr) * multiplier));
+                                          nir_intrinsic_base(intr) * multiplier));
          nir_intrinsic_instr *load = nir_instr_as_intrinsic(load_result->parent_instr);
 
          /* If it's const, set the alignment to our known constant offset.  If
@@ -85,22 +97,22 @@ lower_instr(nir_intrinsic_instr *instr, nir_builder *b, bool dword_packed, bool
           * knowing what features are enabled in the APIs (see comment in
           * nir_lower_ubo_vec4.c)
           */
-         if (nir_src_is_const(instr->src[0])) {
+         if (nir_src_is_const(intr->src[0])) {
             nir_intrinsic_set_align(load, NIR_ALIGN_MUL_MAX,
-                                    (nir_src_as_uint(instr->src[0]) +
-                                    nir_intrinsic_base(instr) * multiplier) %
+                                    (nir_src_as_uint(intr->src[0]) +
+                                    nir_intrinsic_base(intr) * multiplier) %
                                     NIR_ALIGN_MUL_MAX);
          } else {
             nir_intrinsic_set_align(load, MAX2(multiplier,
-                                             instr->dest.ssa.bit_size / 8), 0);
+                                             intr->dest.ssa.bit_size / 8), 0);
          }
 
-         nir_intrinsic_set_range_base(load, nir_intrinsic_base(instr) * multiplier);
-         nir_intrinsic_set_range(load, nir_intrinsic_range(instr) * multiplier);
+         nir_intrinsic_set_range_base(load, nir_intrinsic_base(intr) * multiplier);
+         nir_intrinsic_set_range(load, nir_intrinsic_range(intr) * multiplier);
       }
-      nir_ssa_def_rewrite_uses(&instr->dest.ssa, load_result);
+      nir_ssa_def_rewrite_uses(&intr->dest.ssa, load_result);
 
-      nir_instr_remove(&instr->instr);
+      nir_instr_remove(&intr->instr);
       return true;
    }
 
@@ -112,23 +124,16 @@ nir_lower_uniforms_to_ubo(nir_shader *shader, bool dword_packed, bool load_vec4)
 {
    bool progress = false;
 
-   nir_foreach_function(function, shader) {
-      if (function->impl) {
-         nir_builder builder;
-         nir_builder_init(&builder, function->impl);
-         nir_foreach_block(block, function->impl) {
-            nir_foreach_instr_safe(instr, block) {
-               if (instr->type == nir_instr_type_intrinsic)
-                  progress |= lower_instr(nir_instr_as_intrinsic(instr),
-                                          &builder,
-                                          dword_packed, load_vec4);
-            }
-         }
+   struct nir_lower_uniforms_to_ubo_state state = {
+      .dword_packed = dword_packed,
+      .load_vec4 = load_vec4,
+   };
 
-         nir_metadata_preserve(function->impl, nir_metadata_block_index |
-                                               nir_metadata_dominance);
-      }
-   }
+   progress = nir_shader_instructions_pass(shader,
+                                           nir_lower_uniforms_to_ubo_instr,
+                                           nir_metadata_block_index |
+                                           nir_metadata_dominance,
+                                           &state);
 
    if (progress) {
       if (!shader->info.first_ubo_is_default_ubo) {