zink: fix binding-usage

Rewriting the variable bindings is nasty and error-prone, and this code
triggered an assert when trying to resolve API bindings into Vulkan
bindings.

This code still needs some tweaks, but this makes things much better,
and fixes a few bugs where we incorrectly accounted for the
array-indexes.

Fixes: 1c3f4c07047 ("zink: fixup sampler-usage")
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/merge_requests/3826>

3 files changed, 57 insertions, 40 deletions

diff --git a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
index 9e516eae56d..b53a25feba5 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
+++ b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
@@ -384,6 +384,31 @@ type_to_dim(enum glsl_sampler_dim gdim, bool *is_ms)
    return SpvDim2D;
 }
 
+uint32_t
+zink_binding(gl_shader_stage stage, VkDescriptorType type, int index)
+{
+   if (stage == MESA_SHADER_NONE ||
+       stage >= MESA_SHADER_COMPUTE) {
+      unreachable("not supported");
+   } else {
+      uint32_t stage_offset = (uint32_t)stage * (PIPE_MAX_CONSTANT_BUFFERS +
+                                                 PIPE_MAX_SHADER_SAMPLER_VIEWS);
+
+      switch (type) {
+      case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER:
+         assert(index < PIPE_MAX_CONSTANT_BUFFERS);
+         return stage_offset + index;
+
+      case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
+         assert(index < PIPE_MAX_SHADER_SAMPLER_VIEWS);
+         return stage_offset + PIPE_MAX_CONSTANT_BUFFERS + index;
+
+      default:
+         unreachable("unexpected type");
+      }
+   }
+}
+
 static void
 emit_sampler(struct ntv_context *ctx, struct nir_variable *var)
 {
@@ -416,7 +441,7 @@ emit_sampler(struct ntv_context *ctx, struct nir_variable *var)
             spirv_builder_emit_name(&ctx->builder, var_id, var->name);
          }
 
-         int index = var->data.driver_location + i;
+         int index = var->data.binding + i;
          assert(!(ctx->samplers_used & (1 << index)));
          assert(!ctx->image_types[index]);
          ctx->image_types[index] = image_type;
@@ -425,7 +450,10 @@ emit_sampler(struct ntv_context *ctx, struct nir_variable *var)
 
          spirv_builder_emit_descriptor_set(&ctx->builder, var_id,
                                            var->data.descriptor_set);
-         spirv_builder_emit_binding(&ctx->builder, var_id, var->data.binding);
+         int binding = zink_binding(ctx->stage,
+                                    VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
+                                    var->data.binding + i);
+         spirv_builder_emit_binding(&ctx->builder, var_id, binding);
       }
    } else {
       SpvId var_id = spirv_builder_emit_var(&ctx->builder, pointer_type,
@@ -434,7 +462,7 @@ emit_sampler(struct ntv_context *ctx, struct nir_variable *var)
       if (var->name)
          spirv_builder_emit_name(&ctx->builder, var_id, var->name);
 
-      int index = var->data.driver_location;
+      int index = var->data.binding;
       assert(!(ctx->samplers_used & (1 << index)));
       assert(!ctx->image_types[index]);
       ctx->image_types[index] = image_type;
@@ -443,7 +471,10 @@ emit_sampler(struct ntv_context *ctx, struct nir_variable *var)
 
       spirv_builder_emit_descriptor_set(&ctx->builder, var_id,
                                         var->data.descriptor_set);
-      spirv_builder_emit_binding(&ctx->builder, var_id, var->data.binding);
+      int binding = zink_binding(ctx->stage,
+                                 VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
+                                 var->data.binding);
+      spirv_builder_emit_binding(&ctx->builder, var_id, binding);
    }
 }
 
@@ -484,7 +515,10 @@ emit_ubo(struct ntv_context *ctx, struct nir_variable *var)
 
    spirv_builder_emit_descriptor_set(&ctx->builder, var_id,
                                      var->data.descriptor_set);
-   spirv_builder_emit_binding(&ctx->builder, var_id, var->data.binding);
+   int binding = zink_binding(ctx->stage,
+                              VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
+                              var->data.binding);
+   spirv_builder_emit_binding(&ctx->builder, var_id, binding);
 }
 
 static void
diff --git a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.h b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.h
index 758ca901a2d..e709515da82 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.h
+++ b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.h
@@ -28,6 +28,8 @@
 #include <stdint.h>
 #include <vulkan/vulkan.h>
 
+#include "compiler/shader_enums.h"
+
 struct spirv_shader {
    uint32_t *words;
    size_t num_words;
@@ -41,4 +43,7 @@ nir_to_spirv(struct nir_shader *s);
 void
 spirv_shader_delete(struct spirv_shader *s);
 
+uint32_t
+zink_binding(gl_shader_stage stage, VkDescriptorType type, int index);
+
 #endif
diff --git a/src/gallium/drivers/zink/zink_compiler.c b/src/gallium/drivers/zink/zink_compiler.c
index f1981dbbd15..c8ad749ba74 100644
--- a/src/gallium/drivers/zink/zink_compiler.c
+++ b/src/gallium/drivers/zink/zink_compiler.c
@@ -210,30 +210,6 @@ optimize_nir(struct nir_shader *s)
    } while (progress);
 }
 
-static uint32_t
-zink_binding(enum pipe_shader_type stage, VkDescriptorType type, int index)
-{
-   if (stage == PIPE_SHADER_COMPUTE) {
-      unreachable("not supported");
-   } else {
-      uint32_t stage_offset = (uint32_t)stage * (PIPE_MAX_CONSTANT_BUFFERS +
-                                                 PIPE_MAX_SHADER_SAMPLER_VIEWS);
-
-      switch (type) {
-      case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER:
-         assert(index < PIPE_MAX_CONSTANT_BUFFERS);
-         return stage_offset + index;
-
-      case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
-         assert(index < PIPE_MAX_SHADER_SAMPLER_VIEWS);
-         return stage_offset + PIPE_MAX_CONSTANT_BUFFERS + index;
-
-      default:
-         unreachable("unexpected type");
-      }
-   }
-}
-
 struct zink_shader *
 zink_compile_nir(struct zink_screen *screen, struct nir_shader *nir)
 {
@@ -253,14 +229,14 @@ zink_compile_nir(struct zink_screen *screen, struct nir_shader *nir)
       fprintf(stderr, "---8<---\n");
    }
 
-   enum pipe_shader_type stage = pipe_shader_type_from_mesa(nir->info.stage);
-
    ret->num_bindings = 0;
    nir_foreach_variable(var, &nir->uniforms) {
       if (var->data.mode == nir_var_mem_ubo) {
+         int binding = zink_binding(nir->info.stage,
+                                    VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
+                                    var->data.binding);
          ret->bindings[ret->num_bindings].index = var->data.binding;
-         var->data.binding = zink_binding(stage, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, var->data.binding);
-         ret->bindings[ret->num_bindings].binding = var->data.binding;
+         ret->bindings[ret->num_bindings].binding = binding;
          ret->bindings[ret->num_bindings].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
          ret->num_bindings++;
       } else {
@@ -268,18 +244,20 @@ zink_compile_nir(struct zink_screen *screen, struct nir_shader *nir)
          if (glsl_type_is_array(var->type) &&
              glsl_type_is_sampler(glsl_get_array_element(var->type))) {
             for (int i = 0; i < glsl_get_length(var->type); ++i) {
-               ret->bindings[ret->num_bindings].index = var->data.binding;
-               var->data.driver_location = var->data.binding + i;
-               var->data.binding = zink_binding(stage, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, var->data.driver_location);
-               ret->bindings[ret->num_bindings].binding = var->data.binding;
+               int binding = zink_binding(nir->info.stage,
+                                          VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
+                                          var->data.binding + i);
+               ret->bindings[ret->num_bindings].index = var->data.binding + i;
+               ret->bindings[ret->num_bindings].binding = binding;
                ret->bindings[ret->num_bindings].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
                ret->num_bindings++;
             }
          } else if (glsl_type_is_sampler(var->type)) {
+            int binding = zink_binding(nir->info.stage,
+                                       VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
+                                       var->data.binding);
             ret->bindings[ret->num_bindings].index = var->data.binding;
-            var->data.driver_location = var->data.binding;
-            var->data.binding = zink_binding(stage, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, var->data.driver_location);
-            ret->bindings[ret->num_bindings].binding = var->data.binding;
+            ret->bindings[ret->num_bindings].binding = binding;
             ret->bindings[ret->num_bindings].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
             ret->num_bindings++;
          }