nir: Add a nir_foreach_image_variable() iterator

Reviewed-by: Caio Marcelo de Oliveira Filho <caio.oliveira@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/4743>

src/compiler/nir/nir.h

@@ -715,6 +715,27 @@ _nir_shader_variable_has_mode(nir_variable *var, unsigned modes)
 #define nir_foreach_uniform_variable_safe(var, shader) \
    nir_foreach_variable_with_modes_safe(var, shader, nir_var_uniform)
 
+static inline bool
+_nir_variable_is_image(const nir_variable *var)
+{
+   if (!glsl_type_contains_image(var->type))
+      return false;
+
+   /* GL, Vulkan, and OpenCL only allows arrays of arrays of images */
+   assert(glsl_type_is_image(glsl_without_array(var->type)));
+   return true;
+}
+
+#define nir_foreach_image_variable(var, shader) \
+   nir_foreach_variable_with_modes(var, shader, nir_var_uniform | \
+                                                nir_var_mem_image) \
+      if (_nir_variable_is_image(var))
+
+#define nir_foreach_image_variable_safe(var, shader) \
+   nir_foreach_variable_with_modes_safe(var, shader, nir_var_uniform | \
+                                                     nir_var_mem_image) \
+      if (_nir_variable_is_image(var))
+
 static inline bool
 nir_variable_is_global(const nir_variable *var)
 {

src/compiler/nir_types.cpp

@@ -396,6 +396,12 @@ glsl_type_contains_64bit(const struct glsl_type *type)
    return type->contains_64bit();
 }
 
+bool
+glsl_type_contains_image(const struct glsl_type *type)
+{
+   return type->contains_image();
+}
+
 const glsl_type *
 glsl_void_type(void)
 {

src/compiler/nir_types.h

@@ -155,6 +155,7 @@ bool glsl_type_is_numeric(const struct glsl_type *type);
 bool glsl_type_is_boolean(const struct glsl_type *type);
 bool glsl_type_is_integer(const struct glsl_type *type);
 bool glsl_type_contains_64bit(const struct glsl_type *type);
+bool glsl_type_contains_image(const struct glsl_type *type);
 bool glsl_sampler_type_is_shadow(const struct glsl_type *type);
 bool glsl_sampler_type_is_array(const struct glsl_type *type);
 bool glsl_struct_type_is_packed(const struct glsl_type *type);