nir/nir_lower_vec3_to_vec4: Use the nir_shader_instructions_pass() helper

Extracts some per-impl code to nir_lower_vec3_to_vec4 and then
converts to 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_vec3_to_vec4.c b/src/compiler/nir/nir_lower_vec3_to_vec4.c
index d5f67ab..a0ab162 100644 (file)
--- a/src/compiler/nir/nir_lower_vec3_to_vec4.c
+++ b/src/compiler/nir/nir_lower_vec3_to_vec4.c
@@ -24,121 +24,94 @@
 #include "nir_builder.h"
 
 static bool
-lower_vec3_to_vec4_impl(nir_function_impl *impl, nir_variable_mode modes)
+lower_vec3_to_vec4_instr(nir_builder *b, nir_instr *instr, void *data)
 {
-   bool progress = false;
-
-   if (modes & nir_var_function_temp) {
-      nir_foreach_function_temp_variable(var, impl) {
-         const struct glsl_type *vec4_type =
-            glsl_type_replace_vec3_with_vec4(var->type);
-         if (var->type != vec4_type) {
-            var->type = vec4_type;
-            progress = true;
-         }
+   nir_variable_mode modes = *((nir_variable_mode*) data);
+
+   switch (instr->type) {
+   case nir_instr_type_deref: {
+      nir_deref_instr *deref = nir_instr_as_deref(instr);
+      if (!nir_deref_mode_is_in_set(deref, modes))
+         return false;
+
+      const struct glsl_type *vec4_type =
+         glsl_type_replace_vec3_with_vec4(deref->type);
+      if (deref->type != vec4_type) {
+         deref->type = vec4_type;
+         return true;
       }
+      break;
    }
 
-   nir_builder b;
-   nir_builder_init(&b, impl);
-
-   nir_foreach_block(block, impl) {
-      nir_foreach_instr_safe(instr, block) {
-         switch (instr->type) {
-         case nir_instr_type_deref: {
-            nir_deref_instr *deref = nir_instr_as_deref(instr);
-            if (!nir_deref_mode_is_in_set(deref, modes))
-               continue;
+   case nir_instr_type_intrinsic: {
+      nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+      switch (intrin->intrinsic) {
+      case nir_intrinsic_load_deref: {
+         if (intrin->num_components != 3)
+            break;
 
-            const struct glsl_type *vec4_type =
-               glsl_type_replace_vec3_with_vec4(deref->type);
-            if (deref->type != vec4_type) {
-               deref->type = vec4_type;
-               progress = true;
-            }
+         nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
+         if (!nir_deref_mode_is_in_set(deref, modes))
             break;
-         }
 
-         case nir_instr_type_intrinsic: {
-            nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
-            switch (intrin->intrinsic) {
-            case nir_intrinsic_load_deref: {
-               if (intrin->num_components != 3)
-                  break;
-
-               nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
-               if (!nir_deref_mode_is_in_set(deref, modes))
-                  break;
-
-               assert(intrin->dest.is_ssa);
-               intrin->num_components = 4;
-               intrin->dest.ssa.num_components = 4;
-
-               b.cursor = nir_after_instr(&intrin->instr);
-               nir_ssa_def *vec3 = nir_channels(&b, &intrin->dest.ssa, 0x7);
-               nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
-                                              vec3,
-                                              vec3->parent_instr);
-               progress = true;
-               break;
-            }
+         assert(intrin->dest.is_ssa);
+         intrin->num_components = 4;
+         intrin->dest.ssa.num_components = 4;
 
-            case nir_intrinsic_store_deref: {
-               if (intrin->num_components != 3)
-                  break;
+         b->cursor = nir_after_instr(&intrin->instr);
+         nir_ssa_def *vec3 = nir_channels(b, &intrin->dest.ssa, 0x7);
+         nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
+                                        vec3,
+                                        vec3->parent_instr);
+         return true;
+      }
 
-               nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
-               if (!nir_deref_mode_is_in_set(deref, modes))
-                  break;
+      case nir_intrinsic_store_deref: {
+         if (intrin->num_components != 3)
+            break;
 
-               assert(intrin->src[1].is_ssa);
-               nir_ssa_def *data = intrin->src[1].ssa;
+         nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
+         if (!nir_deref_mode_is_in_set(deref, modes))
+            break;
 
-               b.cursor = nir_before_instr(&intrin->instr);
-               unsigned swiz[] = { 0, 1, 2, 2 };
-               data = nir_swizzle(&b, data, swiz, 4);
+         assert(intrin->src[1].is_ssa);
+         nir_ssa_def *data = intrin->src[1].ssa;
 
-               intrin->num_components = 4;
-               nir_instr_rewrite_src(&intrin->instr, &intrin->src[1],
-                                     nir_src_for_ssa(data));
-               progress = true;
-               break;
-            }
+         b->cursor = nir_before_instr(&intrin->instr);
+         unsigned swiz[] = { 0, 1, 2, 2 };
+         data = nir_swizzle(b, data, swiz, 4);
 
-            case nir_intrinsic_copy_deref: {
-               nir_deref_instr *dst = nir_src_as_deref(intrin->src[0]);
-               nir_deref_instr *src = nir_src_as_deref(intrin->src[0]);
-               /* If we convert once side of a copy and not the other, that
-                * would be very bad.
-                */
-               if (nir_deref_mode_may_be(dst, modes) ||
-                   nir_deref_mode_may_be(src, modes)) {
-                  assert(nir_deref_mode_must_be(dst, modes));
-                  assert(nir_deref_mode_must_be(src, modes));
-               }
-               break;
-            }
+         intrin->num_components = 4;
+         nir_instr_rewrite_src(&intrin->instr, &intrin->src[1],
+                               nir_src_for_ssa(data));
+         return true;
+      }
 
-            default:
-               break;
-            }
-            break;
+      case nir_intrinsic_copy_deref: {
+         nir_deref_instr *dst = nir_src_as_deref(intrin->src[0]);
+         nir_deref_instr *src = nir_src_as_deref(intrin->src[0]);
+         /* If we convert once side of a copy and not the other, that
+         * would be very bad.
+         */
+         if (nir_deref_mode_may_be(dst, modes) ||
+             nir_deref_mode_may_be(src, modes)) {
+             assert(nir_deref_mode_must_be(dst, modes));
+             assert(nir_deref_mode_must_be(src, modes));
          }
+         break;
+      }
 
-         default:
-            break;
-         }
+      default:
+         break;
       }
+      break;
    }
 
-   if (progress) {
-      nir_metadata_preserve(impl, nir_metadata_block_index |
-                                  nir_metadata_dominance);
-   } else {
-      nir_metadata_preserve(impl, nir_metadata_all);
+   default:
+      break;
    }
 
-   return progress;
+   return false;
 }
 
 bool
@@ -160,10 +133,27 @@ nir_lower_vec3_to_vec4(nir_shader *shader, nir_variable_mode modes)
       }
    }
 
-   nir_foreach_function(function, shader) {
-      if (function->impl && lower_vec3_to_vec4_impl(function->impl, modes))
-         progress = true;
+   if (modes & nir_var_function_temp) {
+      nir_foreach_function(function, shader) {
+         if (!function->impl)
+            continue;
+
+         nir_foreach_function_temp_variable(var, function->impl) {
+            const struct glsl_type *vec4_type =
+               glsl_type_replace_vec3_with_vec4(var->type);
+            if (var->type != vec4_type) {
+               var->type = vec4_type;
+               progress = true;
+            }
+         }
+      }
    }
 
+   progress |= nir_shader_instructions_pass(shader,
+                                            lower_vec3_to_vec4_instr,
+                                            nir_metadata_block_index |
+                                            nir_metadata_dominance,
+                                            &modes);
+
    return progress;
 }