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 <[email protected]>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5315>

4 files changed, 34 insertions, 1 deletions

diff --git a/src/gallium/drivers/zink/zink_compiler.c b/src/gallium/drivers/zink/zink_compiler.c
index d3efcb987a6..6fbbe8800a2 100644
--- a/src/gallium/drivers/zink/zink_compiler.c
+++ b/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);
 }
diff --git a/src/gallium/drivers/zink/zink_compiler.h b/src/gallium/drivers/zink/zink_compiler.h
index 47e5b4b7c48..f392c636cd3 100644
--- a/src/gallium/drivers/zink/zink_compiler.h
+++ b/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 *
diff --git a/src/gallium/drivers/zink/zink_program.c b/src/gallium/drivers/zink/zink_program.c
index 95b47a69d93..b4827fa56c3 100644
--- a/src/gallium/drivers/zink/zink_program.c
+++ b/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);
@@ -149,6 +153,16 @@ fail:
 }
 
 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) {
diff --git a/src/gallium/drivers/zink/zink_program.h b/src/gallium/drivers/zink/zink_program.h
index 8807f044ae3..da52f6009b2 100644
--- a/src/gallium/drivers/zink/zink_program.h
+++ b/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