nir: Invalidate live SSA def information when making new SSA defs.
[mesa.git] / src / compiler / glsl / gl_nir_link_uniforms.c
index 7ed484444606cacba97c1511fc0d6b045a34c2ab..05b70a965077a25703ae43f3ffb4a76176078d85 100644 (file)
 
 #define UNMAPPED_UNIFORM_LOC ~0u
 
+struct uniform_array_info {
+   /** List of dereferences of the uniform array. */
+   struct util_dynarray *deref_list;
+
+   /** Set of bit-flags to note which array elements have been accessed. */
+   BITSET_WORD *indices;
+};
+
 /**
  * Built-in / reserved GL variables names start with "gl_"
  */
@@ -45,6 +53,126 @@ is_gl_identifier(const char *s)
    return s && s[0] == 'g' && s[1] == 'l' && s[2] == '_';
 }
 
+static unsigned
+uniform_storage_size(const struct glsl_type *type)
+{
+   switch (glsl_get_base_type(type)) {
+   case GLSL_TYPE_STRUCT:
+   case GLSL_TYPE_INTERFACE: {
+      unsigned size = 0;
+      for (unsigned i = 0; i < glsl_get_length(type); i++)
+         size += uniform_storage_size(glsl_get_struct_field(type, i));
+      return size;
+   }
+   case GLSL_TYPE_ARRAY: {
+      const struct glsl_type *e_type = glsl_get_array_element(type);
+      enum glsl_base_type e_base_type = glsl_get_base_type(e_type);
+      if (e_base_type == GLSL_TYPE_STRUCT ||
+          e_base_type == GLSL_TYPE_INTERFACE ||
+          e_base_type == GLSL_TYPE_ARRAY) {
+         unsigned length = !glsl_type_is_unsized_array(type) ?
+            glsl_get_length(type) : 1;
+         return length * uniform_storage_size(e_type);
+      } else
+         return 1;
+   }
+   default:
+      return 1;
+   }
+}
+
+/**
+ * Update the sizes of linked shader uniform arrays to the maximum
+ * array index used.
+ *
+ * From page 81 (page 95 of the PDF) of the OpenGL 2.1 spec:
+ *
+ *     If one or more elements of an array are active,
+ *     GetActiveUniform will return the name of the array in name,
+ *     subject to the restrictions listed above. The type of the array
+ *     is returned in type. The size parameter contains the highest
+ *     array element index used, plus one. The compiler or linker
+ *     determines the highest index used.  There will be only one
+ *     active uniform reported by the GL per uniform array.
+ */
+static void
+update_array_sizes(struct gl_shader_program *prog, nir_variable *var,
+                   struct hash_table **referenced_uniforms,
+                   unsigned current_var_stage)
+{
+   /* For now we only resize 1D arrays.
+    * TODO: add support for resizing more complex array types ??
+    */
+   if (!glsl_type_is_array(var->type) ||
+       glsl_type_is_array(glsl_get_array_element(var->type)))
+      return;
+
+   /* GL_ARB_uniform_buffer_object says that std140 uniforms
+    * will not be eliminated.  Since we always do std140, just
+    * don't resize arrays in UBOs.
+    *
+    * Atomic counters are supposed to get deterministic
+    * locations assigned based on the declaration ordering and
+    * sizes, array compaction would mess that up.
+    *
+    * Subroutine uniforms are not removed.
+    */
+   if (nir_variable_is_in_block(var) || glsl_contains_atomic(var->type) ||
+       glsl_get_base_type(glsl_without_array(var->type)) == GLSL_TYPE_SUBROUTINE ||
+       var->constant_initializer)
+      return;
+
+   struct uniform_array_info *ainfo = NULL;
+   int words = BITSET_WORDS(glsl_array_size(var->type));
+   int max_array_size = 0;
+   for (unsigned stage = 0; stage < MESA_SHADER_STAGES; stage++) {
+      struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
+      if (!sh)
+         continue;
+
+      struct hash_entry *entry =
+         _mesa_hash_table_search(referenced_uniforms[stage], var->name);
+      if (entry) {
+         ainfo = (struct uniform_array_info *)  entry->data;
+         max_array_size = MAX2(BITSET_LAST_BIT(ainfo->indices, words),
+                               max_array_size);
+      }
+
+      if (max_array_size == glsl_array_size(var->type))
+         return;
+   }
+
+   if (max_array_size != glsl_array_size(var->type)) {
+      /* If this is a built-in uniform (i.e., it's backed by some
+       * fixed-function state), adjust the number of state slots to
+       * match the new array size.  The number of slots per array entry
+       * is not known.  It seems safe to assume that the total number of
+       * slots is an integer multiple of the number of array elements.
+       * Determine the number of slots per array element by dividing by
+       * the old (total) size.
+       */
+      const unsigned num_slots = var->num_state_slots;
+      if (num_slots > 0) {
+         var->num_state_slots =
+            (max_array_size * (num_slots / glsl_array_size(var->type)));
+      }
+
+      var->type = glsl_array_type(glsl_get_array_element(var->type),
+                                  max_array_size, 0);
+
+      /* Update the types of dereferences in case we changed any. */
+      struct hash_entry *entry =
+         _mesa_hash_table_search(referenced_uniforms[current_var_stage], var->name);
+      if (entry) {
+         struct uniform_array_info *ainfo =
+            (struct uniform_array_info *) entry->data;
+         util_dynarray_foreach(ainfo->deref_list, nir_deref_instr *, deref) {
+            (*deref)->type = var->type;
+         }
+      }
+   }
+}
+
 static void
 nir_setup_uniform_remap_tables(struct gl_context *ctx,
                                struct gl_shader_program *prog)
@@ -321,18 +449,23 @@ add_var_use_deref(nir_deref_instr *deref, struct hash_table *live,
 
    nir_deref_path_finish(&path);
 
-   /** Set of bit-flags to note which array elements have been accessed. */
-   BITSET_WORD *bits = NULL;
+
+   struct uniform_array_info *ainfo = NULL;
 
    struct hash_entry *entry =
-      _mesa_hash_table_search(live, deref->var);
+      _mesa_hash_table_search(live, deref->var->name);
    if (!entry && glsl_type_is_array(deref->var->type)) {
+      ainfo = ralloc(live, struct uniform_array_info);
+
       unsigned num_bits = MAX2(1, glsl_get_aoa_size(deref->var->type));
-      bits = rzalloc_array(live, BITSET_WORD, BITSET_WORDS(num_bits));
+      ainfo->indices = rzalloc_array(live, BITSET_WORD, BITSET_WORDS(num_bits));
+
+      ainfo->deref_list = ralloc(live, struct util_dynarray);
+      util_dynarray_init(ainfo->deref_list, live);
    }
 
    if (entry)
-      bits = (BITSET_WORD *) entry->data;
+      ainfo = (struct uniform_array_info *) entry->data;
 
    if (glsl_type_is_array(deref->var->type)) {
       /* Count the "depth" of the arrays-of-arrays. */
@@ -344,11 +477,13 @@ add_var_use_deref(nir_deref_instr *deref, struct hash_table *live,
       }
 
       link_util_mark_array_elements_referenced(*derefs, num_derefs, array_depth,
-                                               bits);
+                                               ainfo->indices);
+
+      util_dynarray_append(ainfo->deref_list, nir_deref_instr *, deref);
    }
 
    assert(deref->mode == deref->var->data.mode);
-   _mesa_hash_table_insert(live, deref->var, bits);
+   _mesa_hash_table_insert(live, deref->var->name, ainfo);
 }
 
 /* Iterate over the shader and collect infomation about uniform use */
@@ -466,11 +601,13 @@ struct nir_link_uniforms_state {
    unsigned num_hidden_uniforms;
    unsigned num_values;
    unsigned max_uniform_location;
-   unsigned next_subroutine;
 
    /* per-shader stage */
+   unsigned next_bindless_image_index;
+   unsigned next_bindless_sampler_index;
    unsigned next_image_index;
    unsigned next_sampler_index;
+   unsigned next_subroutine;
    unsigned num_shader_samplers;
    unsigned num_shader_images;
    unsigned num_shader_uniform_components;
@@ -489,7 +626,7 @@ struct nir_link_uniforms_state {
    int top_level_array_stride;
 
    struct type_tree_entry *current_type;
-   struct hash_table *referenced_uniforms;
+   struct hash_table *referenced_uniforms[MESA_SHADER_STAGES];
    struct hash_table *uniform_hash;
 };
 
@@ -500,7 +637,8 @@ add_parameter(struct gl_uniform_storage *uniform,
               const struct glsl_type *type,
               struct nir_link_uniforms_state *state)
 {
-   if (!state->params || uniform->is_shader_storage || glsl_contains_opaque(type))
+   if (!state->params || uniform->is_shader_storage ||
+       (glsl_contains_opaque(type) && !state->current_var->data.bindless))
       return;
 
    unsigned num_params = glsl_get_aoa_size(type);
@@ -581,6 +719,136 @@ get_next_index(struct nir_link_uniforms_state *state,
    return index;
 }
 
+/* Update the uniforms info for the current shader stage */
+static void
+update_uniforms_shader_info(struct gl_shader_program *prog,
+                            struct nir_link_uniforms_state *state,
+                            struct gl_uniform_storage *uniform,
+                            const struct glsl_type *type,
+                            unsigned stage)
+{
+   unsigned values = glsl_get_component_slots(type);
+   const struct glsl_type *type_no_array = glsl_without_array(type);
+
+   if (glsl_type_is_sampler(type_no_array)) {
+      bool init_idx;
+      unsigned *next_index = state->current_var->data.bindless ?
+         &state->next_bindless_sampler_index :
+         &state->next_sampler_index;
+      int sampler_index = get_next_index(state, uniform, next_index, &init_idx);
+      struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
+
+      if (state->current_var->data.bindless) {
+         if (init_idx) {
+            sh->Program->sh.BindlessSamplers =
+               rerzalloc(sh->Program, sh->Program->sh.BindlessSamplers,
+                         struct gl_bindless_sampler,
+                         sh->Program->sh.NumBindlessSamplers,
+                         state->next_bindless_sampler_index);
+
+            for (unsigned j = sh->Program->sh.NumBindlessSamplers;
+                 j < state->next_bindless_sampler_index; j++) {
+               sh->Program->sh.BindlessSamplers[j].target =
+                  glsl_get_sampler_target(type_no_array);
+            }
+
+            sh->Program->sh.NumBindlessSamplers =
+               state->next_bindless_sampler_index;
+         }
+
+         if (!state->var_is_in_block)
+            state->num_shader_uniform_components += values;
+      } else {
+         /* Samplers (bound or bindless) are counted as two components
+          * as specified by ARB_bindless_texture.
+          */
+         state->num_shader_samplers += values / 2;
+
+         if (init_idx) {
+            const unsigned shadow = glsl_sampler_type_is_shadow(type_no_array);
+            for (unsigned i = sampler_index;
+                 i < MIN2(state->next_sampler_index, MAX_SAMPLERS); i++) {
+               sh->Program->sh.SamplerTargets[i] =
+                  glsl_get_sampler_target(type_no_array);
+               state->shader_samplers_used |= 1U << i;
+               state->shader_shadow_samplers |= shadow << i;
+            }
+         }
+      }
+
+      uniform->opaque[stage].active = true;
+      uniform->opaque[stage].index = sampler_index;
+   } else if (glsl_type_is_image(type_no_array)) {
+      struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
+
+      /* Set image access qualifiers */
+      enum gl_access_qualifier image_access =
+         state->current_var->data.access;
+      const GLenum access =
+         (image_access & ACCESS_NON_WRITEABLE) ?
+         ((image_access & ACCESS_NON_READABLE) ? GL_NONE :
+                                                 GL_READ_ONLY) :
+         ((image_access & ACCESS_NON_READABLE) ? GL_WRITE_ONLY :
+                                                 GL_READ_WRITE);
+
+      int image_index;
+      if (state->current_var->data.bindless) {
+         image_index = state->next_bindless_image_index;
+         state->next_bindless_image_index += MAX2(1, uniform->array_elements);
+
+         sh->Program->sh.BindlessImages =
+            rerzalloc(sh->Program, sh->Program->sh.BindlessImages,
+                      struct gl_bindless_image,
+                      sh->Program->sh.NumBindlessImages,
+                      state->next_bindless_image_index);
+
+         for (unsigned j = sh->Program->sh.NumBindlessImages;
+              j < state->next_bindless_image_index; j++) {
+            sh->Program->sh.BindlessImages[j].access = access;
+         }
+
+         sh->Program->sh.NumBindlessImages = state->next_bindless_image_index;
+
+      } else {
+         image_index = state->next_image_index;
+         state->next_image_index += MAX2(1, uniform->array_elements);
+
+         /* Images (bound or bindless) are counted as two components as
+          * specified by ARB_bindless_texture.
+          */
+         state->num_shader_images += values / 2;
+
+         for (unsigned i = image_index;
+              i < MIN2(state->next_image_index, MAX_IMAGE_UNIFORMS); i++) {
+            sh->Program->sh.ImageAccess[i] = access;
+         }
+      }
+
+      uniform->opaque[stage].active = true;
+      uniform->opaque[stage].index = image_index;
+
+      if (!uniform->is_shader_storage)
+         state->num_shader_uniform_components += values;
+   } else {
+      if (glsl_get_base_type(type_no_array) == GLSL_TYPE_SUBROUTINE) {
+         struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
+
+         uniform->opaque[stage].index = state->next_subroutine;
+         uniform->opaque[stage].active = true;
+
+         sh->Program->sh.NumSubroutineUniforms++;
+
+         /* Increment the subroutine index by 1 for non-arrays and by the
+          * number of array elements for arrays.
+          */
+         state->next_subroutine += MAX2(1, uniform->array_elements);
+      }
+
+      if (!state->var_is_in_block)
+         state->num_shader_uniform_components += values;
+   }
+}
+
 static bool
 find_and_update_named_uniform_storage(struct gl_context *ctx,
                                       struct gl_shader_program *prog,
@@ -659,70 +927,15 @@ find_and_update_named_uniform_storage(struct gl_context *ctx,
             var->data.location = uniform - prog->data->UniformStorage;
          }
 
-         unsigned values = glsl_get_component_slots(type);
-         const struct glsl_type *type_no_array = glsl_without_array(type);
-         if (glsl_type_is_sampler(type_no_array)) {
-            struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
-            bool init_idx;
-            unsigned sampler_index =
-               get_next_index(state, uniform, &state->next_sampler_index,
-                              &init_idx);
-
-            /* Samplers (bound or bindless) are counted as two components as
-             * specified by ARB_bindless_texture.
-             */
-            state->num_shader_samplers += values / 2;
-
-            uniform->opaque[stage].active = true;
-            uniform->opaque[stage].index = sampler_index;
-
-            if (init_idx) {
-               const unsigned shadow =
-                  glsl_sampler_type_is_shadow(type_no_array);
-               for (unsigned i = sampler_index;
-                    i < MIN2(state->next_sampler_index, MAX_SAMPLERS);
-                    i++) {
-                  sh->Program->sh.SamplerTargets[i] =
-                     glsl_get_sampler_target(type_no_array);
-                  state->shader_samplers_used |= 1U << i;
-                  state->shader_shadow_samplers |= shadow << i;
-               }
-            }
-         } else if (glsl_type_is_image(type_no_array)) {
-            struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
-            int image_index = state->next_image_index;
-            /* TODO: handle structs when bindless support is added */
-            state->next_image_index += MAX2(1, uniform->array_elements);
-
-            /* Images (bound or bindless) are counted as two components as
-             * specified by ARB_bindless_texture.
-             */
-            state->num_shader_images += values / 2;
-
-            uniform->opaque[stage].active = true;
-            uniform->opaque[stage].index = image_index;
-
-            /* Set image access qualifiers */
-            enum gl_access_qualifier image_access =
-               state->current_var->data.access;
-            const GLenum access =
-               (image_access & ACCESS_NON_WRITEABLE) ?
-               ((image_access & ACCESS_NON_READABLE) ? GL_NONE :
-                                                       GL_READ_ONLY) :
-               ((image_access & ACCESS_NON_READABLE) ? GL_WRITE_ONLY :
-                                                       GL_READ_WRITE);
-            for (unsigned i = image_index;
-                 i < MIN2(state->next_image_index, MAX_IMAGE_UNIFORMS);
-                 i++) {
-               sh->Program->sh.ImageAccess[i] = access;
-            }
-         }
+         update_uniforms_shader_info(prog, state, uniform, type, stage);
 
-         struct hash_entry *entry =
-            _mesa_hash_table_search(state->referenced_uniforms,
-                                    state->current_var);
+         const struct glsl_type *type_no_array = glsl_without_array(type);
+         struct hash_entry *entry = prog->data->spirv ? NULL :
+            _mesa_hash_table_search(state->referenced_uniforms[stage],
+                                    state->current_var->name);
          if (entry != NULL ||
-             glsl_get_base_type(type_no_array) == GLSL_TYPE_SUBROUTINE)
+             glsl_get_base_type(type_no_array) == GLSL_TYPE_SUBROUTINE ||
+             prog->data->spirv)
             uniform->active_shader_mask |= 1 << stage;
 
          if (!state->var_is_in_block)
@@ -1065,15 +1278,20 @@ nir_link_uniform(struct gl_context *ctx,
 
       return location_count;
    } else {
-      /* Create a new uniform storage entry */
-      prog->data->UniformStorage =
-         reralloc(prog->data,
-                  prog->data->UniformStorage,
-                  struct gl_uniform_storage,
-                  prog->data->NumUniformStorage + 1);
-      if (!prog->data->UniformStorage) {
-         linker_error(prog, "Out of memory during linking.\n");
-         return -1;
+      /* TODO: reallocating storage is slow, we should figure out a way to
+       * allocate storage up front for spirv like we do for GLSL.
+       */
+      if (prog->data->spirv) {
+         /* Create a new uniform storage entry */
+         prog->data->UniformStorage =
+            reralloc(prog->data,
+                     prog->data->UniformStorage,
+                     struct gl_uniform_storage,
+                     prog->data->NumUniformStorage + 1);
+         if (!prog->data->UniformStorage) {
+            linker_error(prog, "Out of memory during linking.\n");
+            return -1;
+         }
       }
 
       uniform = &prog->data->UniformStorage[prog->data->NumUniformStorage];
@@ -1096,11 +1314,12 @@ nir_link_uniform(struct gl_context *ctx,
       uniform->top_level_array_size = state->top_level_array_size;
       uniform->top_level_array_stride = state->top_level_array_stride;
 
-      struct hash_entry *entry =
-         _mesa_hash_table_search(state->referenced_uniforms,
-                                 state->current_var);
+      struct hash_entry *entry = prog->data->spirv ? NULL :
+         _mesa_hash_table_search(state->referenced_uniforms[stage],
+                                 state->current_var->name);
       if (entry != NULL ||
-          glsl_get_base_type(type_no_array) == GLSL_TYPE_SUBROUTINE)
+          glsl_get_base_type(type_no_array) == GLSL_TYPE_SUBROUTINE ||
+          prog->data->spirv)
          uniform->active_shader_mask |= 1 << stage;
 
       if (location >= 0) {
@@ -1115,6 +1334,7 @@ nir_link_uniform(struct gl_context *ctx,
          state->num_hidden_uniforms++;
 
       uniform->is_shader_storage = nir_variable_is_in_ssbo(state->current_var);
+      uniform->is_bindless = state->current_var->data.bindless;
 
       /* Set fields whose default value depend on the variable being inside a
        * block.
@@ -1229,14 +1449,8 @@ nir_link_uniform(struct gl_context *ctx,
       }
 
       uniform->block_index = buffer_block_index;
-
-      /* @FIXME: the initialization of the following will be done as we
-       * implement support for their specific features, like SSBO, atomics,
-       * etc.
-       */
       uniform->builtin = is_gl_identifier(uniform->name);
       uniform->atomic_buffer_index = -1;
-      uniform->is_bindless = false;
 
       /* The following are not for features not supported by ARB_gl_spirv */
       uniform->num_compatible_subroutines = 0;
@@ -1244,81 +1458,7 @@ nir_link_uniform(struct gl_context *ctx,
       unsigned entries = MAX2(1, uniform->array_elements);
       unsigned values = glsl_get_component_slots(type);
 
-      if (glsl_type_is_sampler(type_no_array)) {
-         bool init_idx;
-         int sampler_index =
-            get_next_index(state, uniform, &state->next_sampler_index,
-                           &init_idx);
-
-         /* Samplers (bound or bindless) are counted as two components as
-          * specified by ARB_bindless_texture.
-          */
-         state->num_shader_samplers += values / 2;
-
-         uniform->opaque[stage].active = true;
-         uniform->opaque[stage].index = sampler_index;
-
-         if (init_idx) {
-            const unsigned shadow = glsl_sampler_type_is_shadow(type_no_array);
-            for (unsigned i = sampler_index;
-                 i < MIN2(state->next_sampler_index, MAX_SAMPLERS);
-                 i++) {
-               stage_program->sh.SamplerTargets[i] =
-                  glsl_get_sampler_target(type_no_array);
-               state->shader_samplers_used |= 1U << i;
-               state->shader_shadow_samplers |= shadow << i;
-            }
-         }
-      } else if (glsl_type_is_image(type_no_array)) {
-         /* @FIXME: image_index should match that of the same image
-          * uniform in other shaders. This means we need to match image
-          * uniforms by location (GLSL does it by variable name, but we
-          * want to avoid that).
-          */
-         int image_index = state->next_image_index;
-         state->next_image_index += entries;
-
-         /* Images (bound or bindless) are counted as two components as
-          * specified by ARB_bindless_texture.
-          */
-         state->num_shader_images += values / 2;
-
-         uniform->opaque[stage].active = true;
-         uniform->opaque[stage].index = image_index;
-
-         /* Set image access qualifiers */
-         enum gl_access_qualifier image_access =
-            state->current_var->data.access;
-         const GLenum access =
-            (image_access & ACCESS_NON_WRITEABLE) ?
-            ((image_access & ACCESS_NON_READABLE) ? GL_NONE :
-                                                    GL_READ_ONLY) :
-            ((image_access & ACCESS_NON_READABLE) ? GL_WRITE_ONLY :
-                                                    GL_READ_WRITE);
-         for (unsigned i = image_index;
-              i < MIN2(state->next_image_index, MAX_IMAGE_UNIFORMS);
-              i++) {
-            stage_program->sh.ImageAccess[i] = access;
-         }
-
-         if (!uniform->is_shader_storage)
-            state->num_shader_uniform_components += values;
-      } else {
-         if (glsl_get_base_type(type_no_array) == GLSL_TYPE_SUBROUTINE) {
-            uniform->opaque[stage].index = state->next_subroutine;
-            uniform->opaque[stage].active = true;
-
-            prog->_LinkedShaders[stage]->Program->sh.NumSubroutineUniforms++;
-
-            /* Increment the subroutine index by 1 for non-arrays and by the
-             * number of array elements for arrays.
-             */
-            state->next_subroutine += MAX2(1, uniform->array_elements);
-         }
-
-         if (!state->var_is_in_block)
-            state->num_shader_uniform_components += values;
-      }
+      update_uniforms_shader_info(prog, state, uniform, type, stage);
 
       if (uniform->remap_location != UNMAPPED_UNIFORM_LOC &&
           state->max_uniform_location < uniform->remap_location + entries)
@@ -1353,6 +1493,71 @@ gl_nir_link_uniforms(struct gl_context *ctx,
 
    /* Iterate through all linked shaders */
    struct nir_link_uniforms_state state = {0,};
+
+   if (!prog->data->spirv) {
+      /* Gather information on uniform use */
+      for (unsigned stage = 0; stage < MESA_SHADER_STAGES; stage++) {
+         struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
+         if (!sh)
+            continue;
+
+         state.referenced_uniforms[stage] =
+            _mesa_hash_table_create(NULL, _mesa_hash_string,
+                                    _mesa_key_string_equal);
+
+         nir_shader *nir = sh->Program->nir;
+         add_var_use_shader(nir, state.referenced_uniforms[stage]);
+      }
+
+      /* Resize uniform arrays based on the maximum array index */
+      for (unsigned stage = 0; stage < MESA_SHADER_STAGES; stage++) {
+         struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
+         if (!sh)
+            continue;
+
+         nir_foreach_gl_uniform_variable(var, sh->Program->nir)
+            update_array_sizes(prog, var, state.referenced_uniforms, stage);
+      }
+   }
+
+   /* Count total number of uniforms and allocate storage */
+   unsigned storage_size = 0;
+   if (!prog->data->spirv) {
+      struct set *storage_counted =
+         _mesa_set_create(NULL, _mesa_hash_string, _mesa_key_string_equal);
+      for (unsigned stage = 0; stage < MESA_SHADER_STAGES; stage++) {
+         struct gl_linked_shader *sh = prog->_LinkedShaders[stage];
+         if (!sh)
+            continue;
+
+         nir_foreach_gl_uniform_variable(var, sh->Program->nir) {
+            const struct glsl_type *type = var->type;
+            const char *name = var->name;
+            if (nir_variable_is_in_block(var) &&
+                glsl_without_array(type) == var->interface_type) {
+               type = glsl_without_array(var->type);
+               name = glsl_get_type_name(type);
+            }
+
+            struct set_entry *entry = _mesa_set_search(storage_counted, name);
+            if (!entry) {
+               storage_size += uniform_storage_size(type);
+               _mesa_set_add(storage_counted, name);
+            }
+         }
+      }
+      _mesa_set_destroy(storage_counted, NULL);
+
+      prog->data->UniformStorage = rzalloc_array(prog->data,
+                                                 struct gl_uniform_storage,
+                                                 storage_size);
+      if (!prog->data->UniformStorage) {
+         linker_error(prog, "Out of memory while linking uniforms.\n");
+         return false;
+      }
+   }
+
+   /* Iterate through all linked shaders */
    state.uniform_hash = _mesa_hash_table_create(NULL, _mesa_hash_string,
                                                 _mesa_key_string_equal);
 
@@ -1364,9 +1569,8 @@ gl_nir_link_uniforms(struct gl_context *ctx,
       nir_shader *nir = sh->Program->nir;
       assert(nir);
 
-      state.referenced_uniforms =
-         _mesa_hash_table_create(NULL, _mesa_hash_pointer,
-                                 _mesa_key_pointer_equal);
+      state.next_bindless_image_index = 0;
+      state.next_bindless_sampler_index = 0;
       state.next_image_index = 0;
       state.next_sampler_index = 0;
       state.num_shader_samplers = 0;
@@ -1377,9 +1581,7 @@ gl_nir_link_uniforms(struct gl_context *ctx,
       state.shader_shadow_samplers = 0;
       state.params = fill_parameters ? sh->Program->Parameters : NULL;
 
-      add_var_use_shader(nir, state.referenced_uniforms);
-
-      nir_foreach_variable(var, &nir->uniforms) {
+      nir_foreach_gl_uniform_variable(var, nir) {
          state.current_var = var;
          state.current_ifc_type = NULL;
          state.offset = 0;
@@ -1474,10 +1676,12 @@ gl_nir_link_uniforms(struct gl_context *ctx,
                         buffer_block_index = i;
 
                      struct hash_entry *entry =
-                        _mesa_hash_table_search(state.referenced_uniforms, var);
+                        _mesa_hash_table_search(state.referenced_uniforms[shader_type],
+                                                var->name);
                      if (entry) {
-                        BITSET_WORD *bits = (BITSET_WORD *) entry->data;
-                        if (BITSET_TEST(bits, blocks[i].linearized_array_index))
+                        struct uniform_array_info *ainfo =
+                           (struct uniform_array_info *) entry->data;
+                        if (BITSET_TEST(ainfo->indices, blocks[i].linearized_array_index))
                            blocks[i].stageref |= 1U << shader_type;
                      }
                   }
@@ -1488,7 +1692,8 @@ gl_nir_link_uniforms(struct gl_context *ctx,
                      buffer_block_index = i;
 
                      struct hash_entry *entry =
-                        _mesa_hash_table_search(state.referenced_uniforms, var);
+                        _mesa_hash_table_search(state.referenced_uniforms[shader_type],
+                                                var->name);
                      if (entry)
                         blocks[i].stageref |= 1U << shader_type;
 
@@ -1551,7 +1756,7 @@ gl_nir_link_uniforms(struct gl_context *ctx,
                      location = j;
 
                      struct hash_entry *entry =
-                        _mesa_hash_table_search(state.referenced_uniforms, var);
+                        _mesa_hash_table_search(state.referenced_uniforms[shader_type], var->name);
                      if (entry)
                         blocks[i].stageref |= 1U << shader_type;
 
@@ -1608,7 +1813,10 @@ gl_nir_link_uniforms(struct gl_context *ctx,
             return false;
       }
 
-      _mesa_hash_table_destroy(state.referenced_uniforms, NULL);
+      if (!prog->data->spirv) {
+         _mesa_hash_table_destroy(state.referenced_uniforms[shader_type],
+                                  NULL);
+      }
 
       if (state.num_shader_samplers >
           ctx->Const.Program[shader_type].MaxTextureImageUnits) {
@@ -1639,6 +1847,8 @@ gl_nir_link_uniforms(struct gl_context *ctx,
    prog->data->NumHiddenUniforms = state.num_hidden_uniforms;
    prog->data->NumUniformDataSlots = state.num_values;
 
+   assert(prog->data->spirv || prog->data->NumUniformStorage == storage_size);
+
    if (prog->data->spirv)
       prog->NumUniformRemapTable = state.max_uniform_location;