nir: Add a find_variable_with_[driver_]location helper

We've hand-rolled this loop 10 places and those are just the ones I
found easily.

Reviewed-by: Eric Anholt <eric@anholt.net>
Reviewed-by: Iago Toral Quiroga <itoral@igalia.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5966>

diff --git a/src/amd/vulkan/radv_shader.c b/src/amd/vulkan/radv_shader.c
index 749167e17236c18a125ac9b3d570206ba2809deb..810d4055186e2979dd2360ac010b1ca1780d7a21 100644 (file)
--- a/src/amd/vulkan/radv_shader.c
+++ b/src/amd/vulkan/radv_shader.c
@@ -586,14 +586,12 @@ type_size_vec4(const struct glsl_type *type, bool bindless)
 static nir_variable *
 find_layer_in_var(nir_shader *nir)
 {
-       nir_foreach_shader_in_variable(var, nir) {
-               if (var->data.location == VARYING_SLOT_LAYER) {
-                       return var;
-               }
-       }
-
        nir_variable *var =
-               nir_variable_create(nir, nir_var_shader_in, glsl_int_type(), "layer id");
+               nir_find_variable_with_location(nir, nir_var_shader_in, VARYING_SLOT_LAYER);
+       if (var != NULL)
+               return var;
+
+       var = nir_variable_create(nir, nir_var_shader_in, glsl_int_type(), "layer id");
        var->data.location = VARYING_SLOT_LAYER;
        var->data.interpolation = INTERP_MODE_FLAT;
        return var;

diff --git a/src/compiler/nir/nir.c b/src/compiler/nir/nir.c
index dc56c4922fcc6c93e15aefe617f38f07f802975d..e17b11cdd4a31b8c1025f9c0151774400dee3a6e 100644 (file)
--- a/src/compiler/nir/nir.c
+++ b/src/compiler/nir/nir.c
@@ -171,6 +171,32 @@ nir_local_variable_create(nir_function_impl *impl,
    return var;
 }
 
+nir_variable *
+nir_find_variable_with_location(nir_shader *shader,
+                                nir_variable_mode mode,
+                                unsigned location)
+{
+   assert(util_bitcount(mode) == 1 && mode != nir_var_function_temp);
+   nir_foreach_variable_with_modes(var, shader, mode) {
+      if (var->data.location == location)
+         return var;
+   }
+   return NULL;
+}
+
+nir_variable *
+nir_find_variable_with_driver_location(nir_shader *shader,
+                                       nir_variable_mode mode,
+                                       unsigned location)
+{
+   assert(util_bitcount(mode) == 1 && mode != nir_var_function_temp);
+   nir_foreach_variable_with_modes(var, shader, mode) {
+      if (var->data.driver_location == location)
+         return var;
+   }
+   return NULL;
+}
+
 nir_function *
 nir_function_create(nir_shader *shader, const char *name)
 {

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index b216d293b3e4ad4ee47495adcf401def80820449..52544b9f06c023829d3c5d6a2620e33d113d6748 100644 (file)
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -3308,6 +3308,14 @@ nir_variable *nir_local_variable_create(nir_function_impl *impl,
                                         const struct glsl_type *type,
                                         const char *name);
 
+nir_variable *nir_find_variable_with_location(nir_shader *shader,
+                                              nir_variable_mode mode,
+                                              unsigned location);
+
+nir_variable *nir_find_variable_with_driver_location(nir_shader *shader,
+                                                     nir_variable_mode mode,
+                                                     unsigned location);
+
 /** creates a function and adds it to the shader's list of functions */
 nir_function *nir_function_create(nir_shader *shader, const char *name);
 

diff --git a/src/compiler/nir/nir_lower_bitmap.c b/src/compiler/nir/nir_lower_bitmap.c
index f9173a64796797220d5a485061bf7b9492cfbe60..4d361f5e2dfc1132f7d1c88bc4c01f03b07dda42 100644 (file)
--- a/src/compiler/nir/nir_lower_bitmap.c
+++ b/src/compiler/nir/nir_lower_bitmap.c
@@ -55,16 +55,9 @@
 static nir_variable *
 get_texcoord(nir_shader *shader)
 {
-   nir_variable *texcoord = NULL;
-
-   /* find gl_TexCoord, if it exists: */
-   nir_foreach_shader_in_variable(var, shader) {
-      if (var->data.location == VARYING_SLOT_TEX0) {
-         texcoord = var;
-         break;
-      }
-   }
-
+   nir_variable *texcoord =
+      nir_find_variable_with_location(shader, nir_var_shader_in,
+                                      VARYING_SLOT_TEX0);
    /* otherwise create it: */
    if (texcoord == NULL) {
       texcoord = nir_variable_create(shader,

diff --git a/src/compiler/nir/nir_lower_input_attachments.c b/src/compiler/nir/nir_lower_input_attachments.c
index 75b9e1faff975b0433fa6d0821fc0fdf439674d3..20fa801b45d4abe169a1210be73fc30f1c4ec18f 100644 (file)
--- a/src/compiler/nir/nir_lower_input_attachments.c
+++ b/src/compiler/nir/nir_lower_input_attachments.c
@@ -27,14 +27,14 @@
 static nir_ssa_def *
 load_frag_coord(nir_builder *b)
 {
-   nir_foreach_shader_in_variable(var, b->shader) {
-      if (var->data.location == VARYING_SLOT_POS)
-         return nir_load_var(b, var);
+   nir_variable *pos =
+      nir_find_variable_with_location(b->shader, nir_var_shader_in,
+                                      VARYING_SLOT_POS);
+   if (pos == NULL) {
+      pos = nir_variable_create(b->shader, nir_var_shader_in,
+                                glsl_vec4_type(), NULL);
+      pos->data.location = VARYING_SLOT_POS;
    }
-
-   nir_variable *pos = nir_variable_create(b->shader, nir_var_shader_in,
-                                           glsl_vec4_type(), NULL);
-   pos->data.location = VARYING_SLOT_POS;
    /**
     * From Vulkan spec:
     *   "The OriginLowerLeft execution mode must not be used; fragment entry

diff --git a/src/compiler/nir/nir_lower_point_size_mov.c b/src/compiler/nir/nir_lower_point_size_mov.c
index efe162c8d8fbd2e7b30f4b8ea7ea5151b6715217..f0526ba55fb6873f8b01bd30e2e6f63796f5b2df 100644 (file)
--- a/src/compiler/nir/nir_lower_point_size_mov.c
+++ b/src/compiler/nir/nir_lower_point_size_mov.c
@@ -71,13 +71,9 @@ nir_lower_point_size_mov(nir_shader *shader,
    assert(shader->info.stage != MESA_SHADER_FRAGMENT &&
           shader->info.stage != MESA_SHADER_COMPUTE);
 
-   nir_variable *out = NULL;
-   nir_foreach_shader_out_variable(var, shader) {
-      if (var->data.location == VARYING_SLOT_PSIZ) {
-         out = var;
-         break;
-      }
-   }
+   nir_variable *out =
+      nir_find_variable_with_location(shader, nir_var_shader_out,
+                                      VARYING_SLOT_PSIZ);
 
    lower_impl(nir_shader_get_entrypoint(shader), pointsize_state_tokens,
               out);

diff --git a/src/compiler/nir/nir_lower_two_sided_color.c b/src/compiler/nir/nir_lower_two_sided_color.c
index aa0d9fe84b20b1bd90eda59630b986aca5fda9d6..0507c4d2317a69c80a35aab74df45c761e4c1b19 100644 (file)
--- a/src/compiler/nir/nir_lower_two_sided_color.c
+++ b/src/compiler/nir/nir_lower_two_sided_color.c
@@ -66,19 +66,20 @@ create_input(nir_shader *shader, gl_varying_slot slot,
 static nir_variable *
 create_face_input(nir_shader *shader)
 {
-   nir_foreach_shader_in_variable(var, shader) {
-      if (var->data.location == VARYING_SLOT_FACE)
-         return var;
+   nir_variable *var =
+      nir_find_variable_with_location(shader, nir_var_shader_in,
+                                      VARYING_SLOT_FACE);
+
+   if (var == NULL) {
+      var = nir_variable_create(shader, nir_var_shader_in,
+                                glsl_bool_type(), "gl_FrontFacing");
+
+      var->data.driver_location = shader->num_inputs++;
+      var->data.index = 0;
+      var->data.location = VARYING_SLOT_FACE;
+      var->data.interpolation = INTERP_MODE_FLAT;
    }
 
-   nir_variable *var = nir_variable_create(shader, nir_var_shader_in,
-                                           glsl_bool_type(), "gl_FrontFacing");
-
-   var->data.driver_location = shader->num_inputs++;
-   var->data.index = 0;
-   var->data.location = VARYING_SLOT_FACE;
-   var->data.interpolation = INTERP_MODE_FLAT;
-
    return var;
 }
 

diff --git a/src/compiler/nir/nir_range_analysis.c b/src/compiler/nir/nir_range_analysis.c
index 3d2ff8b2c84305a92e8d38da048d1641a9bae346..14f930610406ab5a63b136a0d9251be195305851 100644 (file)
--- a/src/compiler/nir/nir_range_analysis.c
+++ b/src/compiler/nir/nir_range_analysis.c
@@ -1141,11 +1141,8 @@ search_phi_bcsel(nir_ssa_scalar scalar, nir_ssa_scalar *buf, unsigned buf_size,
 static nir_variable *
 lookup_input(nir_shader *shader, unsigned driver_location)
 {
-   nir_foreach_shader_in_variable(var, shader) {
-      if (driver_location == var->data.driver_location)
-         return var;
-   }
-   return NULL;
+   return nir_find_variable_with_driver_location(shader, nir_var_shader_in,
+                                                 driver_location);
 }
 
 uint32_t

diff --git a/src/gallium/drivers/iris/iris_program.c b/src/gallium/drivers/iris/iris_program.c
index f6c2ab67518caa065a2db1c23a4ae76834cc9fe4..ce037f7632305b7172b1d40cb5670aaff16a09a9 100644 (file)
--- a/src/gallium/drivers/iris/iris_program.c
+++ b/src/gallium/drivers/iris/iris_program.c
@@ -273,14 +273,8 @@ iris_fix_edge_flags(nir_shader *nir)
       return false;
    }
 
-   nir_variable *var = NULL;
-   nir_foreach_shader_out_variable(v, nir) {
-      if (v->data.location == VARYING_SLOT_EDGE) {
-         var = v;
-         break;
-      }
-   }
-
+   nir_variable *var = nir_find_variable_with_location(nir, nir_var_shader_out,
+                                                       VARYING_SLOT_EDGE);
    if (!var) {
       nir_shader_preserve_all_metadata(nir);
       return false;

diff --git a/src/gallium/drivers/vc4/vc4_nir_lower_io.c b/src/gallium/drivers/vc4/vc4_nir_lower_io.c
index eb9bcfffd66ec57cff2e4ea2f8b4fd94b9f17ef9..119ccf276d5f440d6945a6f416da9a0af46cee89 100644 (file)
--- a/src/gallium/drivers/vc4/vc4_nir_lower_io.c
+++ b/src/gallium/drivers/vc4/vc4_nir_lower_io.c
@@ -239,13 +239,9 @@ vc4_nir_lower_fs_input(struct vc4_compile *c, nir_builder *b,
                 return;
         }
 
-        nir_variable *input_var = NULL;
-        nir_foreach_shader_in_variable(var, c->s) {
-                if (var->data.driver_location == nir_intrinsic_base(intr)) {
-                        input_var = var;
-                        break;
-                }
-        }
+        nir_variable *input_var =
+                nir_find_variable_with_driver_location(c->s, nir_var_shader_in,
+                                                       nir_intrinsic_base(intr));
         assert(input_var);
 
         int comp = nir_intrinsic_component(intr);
@@ -290,13 +286,9 @@ static void
 vc4_nir_lower_output(struct vc4_compile *c, nir_builder *b,
                      nir_intrinsic_instr *intr)
 {
-        nir_variable *output_var = NULL;
-        nir_foreach_shader_out_variable(var, c->s) {
-                if (var->data.driver_location == nir_intrinsic_base(intr)) {
-                        output_var = var;
-                        break;
-                }
-        }
+        nir_variable *output_var =
+                nir_find_variable_with_driver_location(c->s, nir_var_shader_out,
+                                                       nir_intrinsic_base(intr));
         assert(output_var);
 
         if (c->stage == QSTAGE_COORD &&

diff --git a/src/intel/compiler/brw_nir_lower_alpha_to_coverage.c b/src/intel/compiler/brw_nir_lower_alpha_to_coverage.c
index 1eed9cd62b4ce2062ab10f207c49ac9c19ea0e6b..113c2266d9c2f3d107ff0e37f9a14bf829a3c1d3 100644 (file)
--- a/src/intel/compiler/brw_nir_lower_alpha_to_coverage.c
+++ b/src/intel/compiler/brw_nir_lower_alpha_to_coverage.c
@@ -88,15 +88,8 @@ brw_nir_lower_alpha_to_coverage(nir_shader *shader)
    assert(shader->info.stage == MESA_SHADER_FRAGMENT);
 
    /* Bail out early if we don't have gl_SampleMask */
-   bool is_sample_mask = false;
-   nir_foreach_shader_out_variable(var, shader) {
-      if (var->data.location == FRAG_RESULT_SAMPLE_MASK) {
-         is_sample_mask = true;
-         break;
-      }
-   }
-
-   if (!is_sample_mask)
+   if (!nir_find_variable_with_location(shader, nir_var_shader_out,
+                                        FRAG_RESULT_SAMPLE_MASK))
       return;
 
    nir_foreach_function(function, shader) {
@@ -115,16 +108,9 @@ brw_nir_lower_alpha_to_coverage(nir_shader *shader)
 
                switch (intr->intrinsic) {
                case nir_intrinsic_store_output:
-                  nir_foreach_shader_out_variable(var, shader) {
-                     int drvloc = var->data.driver_location;
-                     if (nir_intrinsic_base(intr) == drvloc) {
-                        out = var;
-                        break;
-                     }
-                  }
-
-                  if (out->data.mode != nir_var_shader_out)
-                     continue;
+                  out = nir_find_variable_with_driver_location(shader, nir_var_shader_out,
+                                                               nir_intrinsic_base(intr));
+                  assert(out->data.mode == nir_var_shader_out);
 
                   /* save gl_SampleMask instruction pointer */
                   if (out->data.location == FRAG_RESULT_SAMPLE_MASK) {