zink: track program usages for each shader

when shaders are created and destroyed in large numbers, the same pointers
get reused for different shaders, which can lead to bad lookups in the
program_cache hash table.

now each shader tracks its program usage to automatically remove itself from
that program in order to avoid hash collisions

fixes mesa/mesa#3053

Reviewed-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5315>

src/gallium/drivers/zink/zink_compiler.c

@@ -22,6 +22,7 @@
  */
 
 #include "zink_compiler.h"
+#include "zink_program.h"
 #include "zink_screen.h"
 #include "nir_to_spirv/nir_to_spirv.h"
 
@@ -136,6 +137,8 @@ zink_compile_nir(struct zink_screen *screen, struct nir_shader *nir)
 {
    struct zink_shader *ret = CALLOC_STRUCT(zink_shader);
 
+   ret->programs = _mesa_pointer_set_create(NULL);
+
    NIR_PASS_V(nir, nir_lower_uniforms_to_ubo, 1);
    NIR_PASS_V(nir, nir_lower_clip_halfz);
    NIR_PASS_V(nir, nir_lower_regs_to_ssa);
@@ -217,5 +220,9 @@ void
 zink_shader_free(struct zink_screen *screen, struct zink_shader *shader)
 {
    vkDestroyShaderModule(screen->dev, shader->shader_module, NULL);
+   set_foreach(shader->programs, entry) {
+      zink_gfx_program_remove_shader((void*)entry->key, shader);
+   }
+   _mesa_set_destroy(shader->programs, NULL);
    FREE(shader);
 }

src/gallium/drivers/zink/zink_compiler.h

@@ -33,10 +33,13 @@
 
 struct pipe_screen;
 struct zink_screen;
+struct zink_gfx_program;
 
 struct nir_shader_compiler_options;
 struct nir_shader;
 
+struct set;
+
 struct tgsi_token;
 
 const void *
@@ -58,6 +61,7 @@ struct zink_shader {
       VkDescriptorType type;
    } bindings[PIPE_MAX_CONSTANT_BUFFERS + PIPE_MAX_SHADER_SAMPLER_VIEWS];
    size_t num_bindings;
+   struct set *programs;
 };
 
 struct zink_shader *

src/gallium/drivers/zink/zink_program.c

@@ -32,6 +32,7 @@
 #include "util/set.h"
 #include "util/u_debug.h"
 #include "util/u_memory.h"
+#include "tgsi/tgsi_from_mesa.h"
 
 static VkDescriptorSetLayout
 create_desc_set_layout(VkDevice dev,
@@ -123,8 +124,11 @@ zink_create_gfx_program(struct zink_screen *screen,
          goto fail;
    }
 
-   for (int i = 0; i < PIPE_SHADER_TYPES - 1; ++i)
+   for (int i = 0; i < PIPE_SHADER_TYPES - 1; ++i) {
       prog->stages[i] = stages[i];
+      if (stages[i])
+         _mesa_set_add(stages[i]->programs, prog);
+   }
 
    prog->dsl = create_desc_set_layout(screen->dev, stages,
                                       &prog->num_descriptors);
@@ -148,6 +152,16 @@ fail:
    return NULL;
 }
 
+void
+zink_gfx_program_remove_shader(struct zink_gfx_program *prog, struct zink_shader *shader)
+{
+   enum pipe_shader_type p_stage = pipe_shader_type_from_mesa(shader->info.stage);
+
+   assert(prog->stages[p_stage] == shader);
+   prog->stages[p_stage] = NULL;
+   _mesa_set_remove_key(shader->programs, prog);
+}
+
 void
 zink_destroy_gfx_program(struct zink_screen *screen,
                          struct zink_gfx_program *prog)
@@ -158,6 +172,11 @@ zink_destroy_gfx_program(struct zink_screen *screen,
    if (prog->dsl)
       vkDestroyDescriptorSetLayout(screen->dev, prog->dsl, NULL);
 
+   for (int i = 0; i < PIPE_SHADER_TYPES - 1; ++i) {
+      if (prog->stages[i])
+         zink_gfx_program_remove_shader(prog, prog->stages[i]);
+   }
+
    /* unref all used render-passes */
    if (prog->render_passes) {
       set_foreach(prog->render_passes, entry) {

src/gallium/drivers/zink/zink_program.h

@@ -58,4 +58,7 @@ zink_get_gfx_pipeline(struct zink_screen *screen,
                       struct zink_gfx_pipeline_state *state,
                       enum pipe_prim_type mode);
 
+void
+zink_gfx_program_remove_shader(struct zink_gfx_program *prog, struct zink_shader *shader);
+
 #endif