nir: Add a nir_foreach_uniform_variable helper

This one's a bit more complex because it filters off only those
variables with mode == nir_var_uniform.  As such, it's not exactly a
drop-in replacement for nir_foreach_variable(var, &nir->uniforms).

Reviewed-by: Gert Wollny <gert.wollny@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5966>

src/broadcom/compiler/vir.c

@@ -570,7 +570,7 @@ v3d_lower_nir(struct v3d_compile *c)
         }
 
         /* CS textures may not have return_size reflecting the shadow state. */
-        nir_foreach_variable(var, &c->s->uniforms) {
+        nir_foreach_uniform_variable(var, c->s) {
                 const struct glsl_type *type = glsl_without_array(var->type);
                 unsigned array_len = MAX2(glsl_get_length(var->type), 1);
 

src/compiler/glsl/gl_nir_link_atomics.c

@@ -148,7 +148,7 @@ find_active_atomic_counters(struct gl_context *ctx,
 
       nir_shader *nir = sh->Program->nir;
 
-      nir_foreach_variable(var, &nir->uniforms) {
+      nir_foreach_uniform_variable(var, nir) {
          if (!glsl_contains_atomic(var->type))
             continue;
 

src/compiler/nir/nir.h

@@ -636,6 +636,14 @@ typedef struct nir_variable {
 #define nir_foreach_shader_out_variable_safe(var, shader) \
    nir_foreach_variable_safe(var, &(shader)->outputs)
 
+#define nir_foreach_uniform_variable(var, shader) \
+   nir_foreach_variable(var, &(shader)->uniforms) \
+      if (var->data.mode == nir_var_uniform)
+
+#define nir_foreach_uniform_variable_safe(var, shader) \
+   nir_foreach_variable_safe(var, &(shader)->uniforms) \
+      if (var->data.mode == nir_var_uniform)
+
 static inline bool
 nir_variable_is_global(const nir_variable *var)
 {

src/compiler/nir/nir_gather_info.c

@@ -601,7 +601,7 @@ nir_shader_gather_info(nir_shader *shader, nir_function_impl *entrypoint)
    shader->info.image_buffers = 0;
    shader->info.msaa_images = 0;
 
-   nir_foreach_variable(var, &shader->uniforms) {
+   nir_foreach_uniform_variable(var, shader) {
       /* Bindless textures and images don't use non-bindless slots.
        * Interface blocks imply inputs, outputs, UBO, or SSBO, which can only
        * mean bindless.

src/compiler/nir/nir_lower_atomics_to_ssbo.c

@@ -188,7 +188,7 @@ nir_lower_atomics_to_ssbo(nir_shader *shader)
    if (progress) {
       /* replace atomic_uint uniforms with ssbo's: */
       unsigned replaced = 0;
-      nir_foreach_variable_safe(var, &shader->uniforms) {
+      nir_foreach_uniform_variable_safe(var, shader) {
          if (is_atomic_uint(var->type)) {
             exec_node_remove(&var->node);
 

src/gallium/auxiliary/nir/nir_draw_helpers.c

@@ -123,7 +123,7 @@ nir_lower_pstipple_fs(struct nir_shader *shader,
       return;
 
    int binding = 0;
-   nir_foreach_variable(var, &shader->uniforms) {
+   nir_foreach_uniform_variable(var, shader) {
       if (glsl_type_is_sampler(var->type)) {
          if (var->data.binding >= binding)
             binding = var->data.binding + 1;

src/gallium/drivers/r600/sfn/sfn_instruction_tex.cpp

@@ -212,7 +212,7 @@ bool r600_nir_lower_int_tg4(nir_shader *shader)
    std::vector<bool> lower_sampler(shader->uniforms.length(), false);
    auto is = lower_sampler.begin();
 
-   nir_foreach_variable(var, &shader->uniforms) {
+   nir_foreach_uniform_variable(var, shader) {
       if (var->type->is_sampler()) {
          if (glsl_base_type_is_integer(var->type->sampled_type)) {
             need_lowering = *is = true;

src/intel/vulkan/anv_nir_apply_pipeline_layout.c

@@ -1286,7 +1286,7 @@ anv_nir_apply_pipeline_layout(const struct anv_physical_device *pdevice,
       }
    }
 
-   nir_foreach_variable(var, &shader->uniforms) {
+   nir_foreach_uniform_variable(var, shader) {
       const struct glsl_type *glsl_type = glsl_without_array(var->type);
 
       if (!glsl_type_is_image(glsl_type))

src/mesa/drivers/dri/i965/brw_link.cpp

@@ -346,7 +346,7 @@ brw_link_shader(struct gl_context *ctx, struct gl_shader_program *shProg)
        * too late.  At that point, the values for the built-in uniforms won't
        * get sent to the shader.
        */
-      nir_foreach_variable(var, &prog->nir->uniforms) {
+      nir_foreach_uniform_variable(var, prog->nir) {
          const nir_state_slot *const slots = var->state_slots;
          for (unsigned int i = 0; i < var->num_state_slots; i++) {
             assert(slots != NULL);

src/mesa/drivers/dri/i965/brw_nir_uniforms.cpp

@@ -222,7 +222,7 @@ brw_nir_setup_glsl_uniforms(void *mem_ctx, nir_shader *shader,
    stage_prog_data->nr_params = nr_params;
    stage_prog_data->param = rzalloc_array(mem_ctx, uint32_t, nr_params);
 
-   nir_foreach_variable(var, &shader->uniforms) {
+   nir_foreach_uniform_variable(var, shader) {
       /* UBO's, atomics and samplers don't take up space in the
          uniform file */
       if (var->interface_type != NULL || var->type->contains_atomic())
@@ -306,7 +306,7 @@ brw_nir_lower_gl_images(nir_shader *shader,
                         const struct gl_program *prog)
 {
    /* We put image uniforms at the end */
-   nir_foreach_variable(var, &shader->uniforms) {
+   nir_foreach_uniform_variable(var, shader) {
       if (!var->type->contains_image())
          continue;
 

src/mesa/state_tracker/st_glsl_to_nir.cpp

@@ -195,21 +195,14 @@ st_nir_lookup_parameter_index(struct gl_program *prog, nir_variable *var)
 static void
 st_nir_assign_uniform_locations(struct gl_context *ctx,
                                 struct gl_program *prog,
-                                struct exec_list *uniform_list)
+                                nir_shader *nir)
 {
    int shaderidx = 0;
    int imageidx = 0;
 
-   nir_foreach_variable(uniform, uniform_list) {
+   nir_foreach_uniform_variable(uniform, nir) {
       int loc;
 
-      /*
-       * UBO's have their own address spaces, so don't count them towards the
-       * number of global uniforms
-       */
-      if (uniform->data.mode == nir_var_mem_ubo || uniform->data.mode == nir_var_mem_ssbo)
-         continue;
-
       const struct glsl_type *type = glsl_without_array(uniform->type);
       if (!uniform->data.bindless && (type->is_sampler() || type->is_image())) {
          if (type->is_sampler()) {
@@ -448,7 +441,7 @@ st_glsl_to_nir_post_opts(struct st_context *st, struct gl_program *prog,
     * too late.  At that point, the values for the built-in uniforms won't
     * get sent to the shader.
     */
-   nir_foreach_variable(var, &nir->uniforms) {
+   nir_foreach_uniform_variable(var, nir) {
       const nir_state_slot *const slots = var->state_slots;
       if (slots != NULL) {
          const struct glsl_type *type = glsl_without_array(var->type);
@@ -949,8 +942,7 @@ st_finalize_nir(struct st_context *st, struct gl_program *prog,
    NIR_PASS_V(nir, nir_lower_var_copies);
 
    st_nir_assign_varying_locations(st, nir);
-   st_nir_assign_uniform_locations(st->ctx, prog,
+   st_nir_assign_uniform_locations(st->ctx, prog, nir);
-                                   &nir->uniforms);
 
    /* Set num_uniforms in number of attribute slots (vec4s) */
    nir->num_uniforms = DIV_ROUND_UP(prog->Parameters->NumParameterValues, 4);

src/mesa/state_tracker/st_nir_lower_builtin.c

@@ -126,7 +126,7 @@ get_variable(lower_builtin_state *state, nir_deref_path *path,
 
    char *name = _mesa_program_state_string(tokens);
 
-   nir_foreach_variable(var, &shader->uniforms) {
+   nir_foreach_uniform_variable(var, shader) {
       if (strcmp(var->name, name) == 0) {
          free(name);
          return var;

src/mesa/state_tracker/st_nir_lower_tex_src_plane.c

@@ -50,7 +50,7 @@ static nir_variable *
 find_sampler(lower_tex_src_state *state, unsigned samp)
 {
    /* NOTE: arrays of samplerExternalOES do not appear to be allowed: */
-   nir_foreach_variable(var, &state->shader->uniforms)
+   nir_foreach_uniform_variable(var, state->shader)
       if (var->data.binding == samp)
          return var;
    return NULL;