nir/lower_amul: Add a variable mode check

[mesa.git] / src / compiler / nir / nir_lower_amul.c

diff --git a/src/compiler/nir/nir_lower_amul.c b/src/compiler/nir/nir_lower_amul.c
index 4c7f10859cbcff72e93a76a6b9e455c719f62308..561cf77af33f61163559d4fc8fa3608156680500 100644 (file)
--- a/src/compiler/nir/nir_lower_amul.c
+++ b/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,26 @@ 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_foreach_variable (var, &shader->uniforms) {
+      if (!(var->data.mode & (nir_var_mem_ubo | nir_var_mem_ssbo)))
+         continue;
+
+      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, num_uniforms, 0);
-   NIR_VLA_FILL(bool, large_ssbos, num_uniforms, 0);
+   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 +260,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;
          }
       }
    }