zink: fixup sampler-usage

It seems I got this stuff all wrong, and looked at driver_location
rather than the binding. But since we mess with the binding, we need to
adjust things a bit to get things right.

This still isn't great as-is, but it seems to work. In the future, we
should move to having samplers always at bindings 0 and up, and just
update the bindings that are used by either of the stages. But this
band-aid should be OK for now.

This fixes 0AD for me.

Tested-by: Marge Bot <https://gitlab.freedesktop.org/mesa/mesa/merge_requests/3668>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/merge_requests/3668>

2 files changed, 19 insertions, 15 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 043b210caf9..455bb562a47 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
@@ -40,7 +40,7 @@ struct ntv_context {
    size_t num_ubos;
    SpvId image_types[PIPE_MAX_SAMPLERS];
    SpvId samplers[PIPE_MAX_SAMPLERS];
-   size_t num_samplers;
+   unsigned samplers_used : PIPE_MAX_SAMPLERS;
    SpvId entry_ifaces[PIPE_MAX_SHADER_INPUTS * 4 + PIPE_MAX_SHADER_OUTPUTS * 4];
    size_t num_entry_ifaces;
 
@@ -415,11 +415,12 @@ emit_sampler(struct ntv_context *ctx, struct nir_variable *var)
             spirv_builder_emit_name(&ctx->builder, var_id, var->name);
          }
 
-         assert(ctx->num_samplers < ARRAY_SIZE(ctx->image_types));
-         ctx->image_types[ctx->num_samplers] = image_type;
-
-         assert(ctx->num_samplers < ARRAY_SIZE(ctx->samplers));
-         ctx->samplers[ctx->num_samplers++] = var_id;
+         int index = var->data.driver_location + i;
+         assert(!(ctx->samplers_used & (1 << index)));
+         assert(!ctx->image_types[index]);
+         ctx->image_types[index] = image_type;
+         ctx->samplers[index] = var_id;
+         ctx->samplers_used |= 1 << index;
 
          spirv_builder_emit_descriptor_set(&ctx->builder, var_id,
                                            var->data.descriptor_set);
@@ -432,11 +433,12 @@ emit_sampler(struct ntv_context *ctx, struct nir_variable *var)
       if (var->name)
          spirv_builder_emit_name(&ctx->builder, var_id, var->name);
 
-      assert(ctx->num_samplers < ARRAY_SIZE(ctx->image_types));
-      ctx->image_types[ctx->num_samplers] = image_type;
-
-      assert(ctx->num_samplers < ARRAY_SIZE(ctx->samplers));
-      ctx->samplers[ctx->num_samplers++] = var_id;
+      int index = var->data.driver_location;
+      assert(!(ctx->samplers_used & (1 << index)));
+      assert(!ctx->image_types[index]);
+      ctx->image_types[index] = image_type;
+      ctx->samplers[index] = var_id;
+      ctx->samplers_used |= 1 << index;
 
       spirv_builder_emit_descriptor_set(&ctx->builder, var_id,
                                         var->data.descriptor_set);
@@ -1500,7 +1502,7 @@ emit_tex(struct ntv_context *ctx, nir_tex_instr *tex)
    SpvId sampled_type = spirv_builder_type_sampled_image(&ctx->builder,
                                                          image_type);
 
-   assert(tex->texture_index < ctx->num_samplers);
+   assert(ctx->samplers_used & (1u << tex->texture_index));
    SpvId load = spirv_builder_emit_load(&ctx->builder, sampled_type,
                                         ctx->samplers[tex->texture_index]);
 
diff --git a/src/gallium/drivers/zink/zink_compiler.c b/src/gallium/drivers/zink/zink_compiler.c
index 5c20c7b297f..f1981dbbd15 100644
--- a/src/gallium/drivers/zink/zink_compiler.c
+++ b/src/gallium/drivers/zink/zink_compiler.c
@@ -268,14 +268,16 @@ 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.driver_location + i;
-               var->data.binding = zink_binding(stage, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, var->data.driver_location + 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;
                ret->bindings[ret->num_bindings].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
                ret->num_bindings++;
             }
          } else if (glsl_type_is_sampler(var->type)) {
-            ret->bindings[ret->num_bindings].index = var->data.driver_location;
+            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].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;