zink: fix gfx program cache pruning with generated tcs

if the tcs was generated, then the prgram was added to the non-tcs cache,
which means deleting it from the tcs+tes cache will fail and then
context_destroy will explode

Fixes: 4123ee3c714 ("zink: invoke descriptor_program_deinit for programs on context destroy")

Reviewed-by: Dave Airlie <airlied@redhat.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/17866>

2 files changed, 9 insertions, 1 deletions

diff --git a/src/gallium/drivers/zink/zink_compiler.c b/src/gallium/drivers/zink/zink_compiler.c
index dddc82307bc..307e8ae9000 100644
--- a/src/gallium/drivers/zink/zink_compiler.c
+++ b/src/gallium/drivers/zink/zink_compiler.c
@@ -3139,7 +3139,13 @@ zink_shader_free(struct zink_context *ctx, struct zink_shader *shader)
          enum pipe_shader_type pstage = pipe_shader_type_from_mesa(shader->nir->info.stage);
          assert(pstage < ZINK_SHADER_COUNT);
          if (!prog->base.removed && (shader->nir->info.stage != MESA_SHADER_TESS_CTRL || !shader->is_generated)) {
-            _mesa_hash_table_remove_key(&ctx->program_cache[prog->stages_present >> 2], prog->shaders);
+            unsigned stages_present = prog->stages_present;
+            if (prog->shaders[PIPE_SHADER_TESS_CTRL] && prog->shaders[PIPE_SHADER_TESS_CTRL]->is_generated)
+               stages_present &= ~BITFIELD_BIT(PIPE_SHADER_TESS_CTRL);
+            struct hash_table *ht = &ctx->program_cache[stages_present >> 2];
+            struct hash_entry *he = _mesa_hash_table_search(ht, prog->shaders);
+            assert(he);
+            _mesa_hash_table_remove(ht, he);
             prog->base.removed = true;
          }
          if (shader->nir->info.stage != MESA_SHADER_TESS_CTRL || !shader->is_generated)
diff --git a/src/gallium/drivers/zink/zink_context.c b/src/gallium/drivers/zink/zink_context.c
index 74fb0d978e9..5ba56c1473e 100644
--- a/src/gallium/drivers/zink/zink_context.c
+++ b/src/gallium/drivers/zink/zink_context.c
@@ -108,11 +108,13 @@ zink_context_destroy(struct pipe_context *pctx)
    for (unsigned i = 0; i < ARRAY_SIZE(ctx->program_cache); i++) {
       hash_table_foreach(&ctx->program_cache[i], entry) {
          struct zink_program *pg = entry->data;
+         pg->removed = true;
          screen->descriptor_program_deinit(ctx, pg);
       }
    }
    hash_table_foreach(&ctx->compute_program_cache, entry) {
       struct zink_program *pg = entry->data;
+      pg->removed = true;
       screen->descriptor_program_deinit(ctx, pg);
    }