nir/lower_amul: fix slot calculation

Fixes incorrect indexing in
dEQP-GLES31.functional.ssbo.layout.instance_array_basic_type.packed.mat2x3

Signed-off-by: Rob Clark <robdclark@chromium.org>
Reviewed-by: Eric Anholt <eric@anholt.net>
Tested-by: Marge Bot <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/4455>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/4455>

src/compiler/nir/nir_lower_amul.c

@@ -61,6 +61,8 @@ typedef struct {
     */
    bool has_large_ubo;
    bool has_large_ssbo;
+
+   unsigned max_slot;
 } lower_state;
 
 /* Lower 'amul's in offset src of large variables to 'imul': */
@@ -99,7 +101,9 @@ large_ubo(lower_state *state, nir_src src)
 {
    if (!nir_src_is_const(src))
       return state->has_large_ubo;
-   return state->large_ubos[nir_src_as_uint(src)];
+   unsigned idx = nir_src_as_uint(src);
+   assert(idx < state->max_slot);
+   return state->large_ubos[idx];
 }
 
 static bool
@@ -107,7 +111,9 @@ large_ssbo(lower_state *state, nir_src src)
 {
    if (!nir_src_is_const(src))
       return state->has_large_ssbo;
-   return state->large_ssbos[nir_src_as_uint(src)];
+   unsigned idx = nir_src_as_uint(src);
+   assert(idx < state->max_slot);
+   return state->large_ssbos[idx];
 }
 
 static bool
@@ -209,7 +215,8 @@ lower_instr(lower_state *state, nir_instr *instr)
 static bool
 is_large(lower_state *state, nir_variable *var)
 {
-   unsigned size = state->type_size(var->type, false);
+   const struct glsl_type *type = glsl_without_array(var->type);
+   unsigned size = state->type_size(type, false);
 
    /* if size is not known (ie. VLA) then assume the worst: */
    if (!size)
@@ -226,15 +233,23 @@ nir_lower_amul(nir_shader *shader,
    assert(type_size);
 
    /* uniforms list actually includes ubo's and ssbo's: */
-   int num_uniforms = exec_list_length(&shader->uniforms);
+   int max_slot = 0;
 
-   NIR_VLA_FILL(bool, large_ubos, num_uniforms, 0);
-   NIR_VLA_FILL(bool, large_ssbos, num_uniforms, 0);
+   nir_foreach_variable (var, &shader->uniforms) {
+      int base = var->data.binding;
+      int size = MAX2(1, glsl_array_size(var->type));
+
+      max_slot = MAX2(max_slot, base + size);
+   }
+
+   NIR_VLA_FILL(bool, large_ubos, max_slot, 0);
+   NIR_VLA_FILL(bool, large_ssbos, max_slot, 0);
 
    lower_state state = {
          .type_size = type_size,
          .large_ubos = large_ubos,
          .large_ssbos = large_ssbos,
+         .max_slot = max_slot,
    };
 
    /* Figure out which UBOs or SSBOs are large enough to be
@@ -242,16 +257,18 @@ nir_lower_amul(nir_shader *shader,
     */
    nir_foreach_variable(var, &shader->uniforms) {
       if (var->data.mode == nir_var_mem_ubo) {
-         assert(var->data.driver_location < num_uniforms);
          if (is_large(&state, var)) {
             state.has_large_ubo = true;
-            state.large_ubos[var->data.driver_location] = true;
+            unsigned size = MAX2(1, glsl_array_size(var->type));
+            for (unsigned i = 0; i < size; i++)
+               state.large_ubos[var->data.binding + i] = true;
          }
       } else if (var->data.mode == nir_var_mem_ssbo) {
-         assert(var->data.driver_location < num_uniforms);
          if (is_large(&state, var)) {
             state.has_large_ssbo = true;
-            state.large_ssbos[var->data.driver_location] = true;
+            unsigned size = MAX2(1, glsl_array_size(var->type));
+            for (unsigned i = 0; i < size; i++)
+               state.large_ssbos[var->data.binding + i] = true;
          }
       }
    }

src/compiler/nir_types.cpp

@@ -34,6 +34,12 @@ glsl_get_type_name(const glsl_type *type)
    return type->name;
 }
 
+int
+glsl_array_size(const struct glsl_type *type)
+{
+   return type->array_size();
+}
+
 const glsl_type *
 glsl_get_array_element(const glsl_type* type)
 {

src/compiler/nir_types.h

@@ -63,6 +63,7 @@ unsigned glsl_get_std430_base_alignment(const struct glsl_type *type,
                                         bool row_major);
 unsigned glsl_get_std430_size(const struct glsl_type *type, bool row_major);
 unsigned glsl_get_explicit_stride(const struct glsl_type *type);
+int glsl_array_size(const struct glsl_type *type);
 const struct glsl_type *glsl_get_array_element(const struct glsl_type *type);
 const struct glsl_type *glsl_without_array(const struct glsl_type *type);
 const struct glsl_type *glsl_without_array_or_matrix(const struct glsl_type *type);