zink: refcount zink_gfx_program objects

now that we're tracking these by shader, we want to ensure that they live through
each render pass successfully if there's no flush regardless of the timing when the
shader objects are destroyed. this becomes useful when we split up shader create and compile
functionality in future patches, at which point program refcounts can be changed
during successive draw calls, potentially resulting in a program being destroyed at that
point when it shouldn't be

with this patch, each shader used by the program gets a reference, with the renderpass
batch itself becoming the owner of the program such that it will be deleted
when the draw state gets invalidated and a new program is created

Reviewed-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5970>
diff --git a/src/gallium/drivers/zink/zink_batch.c b/src/gallium/drivers/zink/zink_batch.c
index 2403524..2b92b1a 100644
--- a/src/gallium/drivers/zink/zink_batch.c
+++ b/src/gallium/drivers/zink/zink_batch.c
@@ -4,6 +4,7 @@
 #include "zink_fence.h"
 #include "zink_framebuffer.h"
 #include "zink_query.h"
+#include "zink_program.h"
 #include "zink_render_pass.h"
 #include "zink_resource.h"
 #include "zink_screen.h"
@@ -26,6 +27,11 @@
 
    zink_render_pass_reference(screen, &batch->rp, NULL);
    zink_framebuffer_reference(screen, &batch->fb, NULL);
+   set_foreach(batch->programs, entry) {
+      struct zink_gfx_program *prog = (struct zink_gfx_program*)entry->key;
+      zink_gfx_program_reference(screen, &prog, NULL);
+   }
+   _mesa_set_clear(batch->programs, NULL);
 
    /* unref all used resources */
    set_foreach(batch->resources, entry) {
@@ -118,3 +124,14 @@
       pipe_reference(NULL, &sv->base.reference);
    }
 }
+
+void
+zink_batch_reference_program(struct zink_batch *batch,
+                             struct zink_gfx_program *prog)
+{
+   struct set_entry *entry = _mesa_set_search(batch->programs, prog);
+   if (!entry) {
+      entry = _mesa_set_add(batch->programs, prog);
+      pipe_reference(NULL, &prog->reference);
+   }
+}
diff --git a/src/gallium/drivers/zink/zink_batch.h b/src/gallium/drivers/zink/zink_batch.h
index 950d3a9..7b5ffe9 100644
--- a/src/gallium/drivers/zink/zink_batch.h
+++ b/src/gallium/drivers/zink/zink_batch.h
@@ -32,6 +32,7 @@
 struct zink_context;
 struct zink_fence;
 struct zink_framebuffer;
+struct zink_gfx_program;
 struct zink_render_pass;
 struct zink_resource;
 struct zink_sampler_view;
@@ -46,6 +47,7 @@
 
    struct zink_render_pass *rp;
    struct zink_framebuffer *fb;
+   struct set *programs;
 
    struct set *resources;
    struct set *sampler_views;
@@ -69,4 +71,7 @@
 zink_batch_reference_sampler_view(struct zink_batch *batch,
                                   struct zink_sampler_view *sv);
 
+void
+zink_batch_reference_program(struct zink_batch *batch,
+                             struct zink_gfx_program *prog);
 #endif
diff --git a/src/gallium/drivers/zink/zink_compiler.c b/src/gallium/drivers/zink/zink_compiler.c
index d42379e..db43303 100644
--- a/src/gallium/drivers/zink/zink_compiler.c
+++ b/src/gallium/drivers/zink/zink_compiler.c
@@ -321,7 +321,8 @@
    set_foreach(shader->programs, entry) {
       struct zink_gfx_program *prog = (void*)entry->key;
       _mesa_hash_table_remove_key(ctx->program_cache, prog->stages);
-      zink_destroy_gfx_program(screen, prog);
+      prog->stages[pipe_shader_type_from_mesa(shader->info.stage)] = NULL;
+      zink_gfx_program_reference(screen, &prog, NULL);
    }
    _mesa_set_destroy(shader->programs, NULL
    free(shader->streamout.so_info_slots);
diff --git a/src/gallium/drivers/zink/zink_context.c b/src/gallium/drivers/zink/zink_context.c
index 2f9a385..823c19e 100644
--- a/src/gallium/drivers/zink/zink_context.c
+++ b/src/gallium/drivers/zink/zink_context.c
@@ -1147,6 +1147,9 @@
       ctx->batches[i].sampler_views = _mesa_set_create(NULL,
                                                        _mesa_hash_pointer,
                                                        _mesa_key_pointer_equal);
+      ctx->batches[i].programs = _mesa_set_create(NULL,
+                                                  _mesa_hash_pointer,
+                                                  _mesa_key_pointer_equal);
 
       if (!ctx->batches[i].resources || !ctx->batches[i].sampler_views)
          goto fail;
diff --git a/src/gallium/drivers/zink/zink_draw.c b/src/gallium/drivers/zink/zink_draw.c
index 34953f6..1a351ee 100644
--- a/src/gallium/drivers/zink/zink_draw.c
+++ b/src/gallium/drivers/zink/zink_draw.c
@@ -368,6 +368,7 @@
       batch = zink_batch_rp(ctx);
       assert(batch->descs_left >= gfx_program->num_descriptors);
    }
+   zink_batch_reference_program(batch, ctx->curr_program);
 
    VkDescriptorSet desc_set = allocate_descriptor_set(screen, batch,
                                                       gfx_program);
diff --git a/src/gallium/drivers/zink/zink_program.c b/src/gallium/drivers/zink/zink_program.c
index 4bc1b21..c6b0c1c 100644
--- a/src/gallium/drivers/zink/zink_program.c
+++ b/src/gallium/drivers/zink/zink_program.c
@@ -39,6 +39,12 @@
    VkPipeline pipeline;
 };
 
+void
+debug_describe_zink_gfx_program(char *buf, const struct zink_gfx_program *ptr)
+{
+   sprintf(buf, "zink_gfx_program");
+}
+
 static VkDescriptorSetLayout
 create_desc_set_layout(VkDevice dev,
                        struct zink_shader *stages[PIPE_SHADER_TYPES - 1],
@@ -122,6 +128,8 @@
    if (!prog)
       goto fail;
 
+   pipe_reference_init(&prog->reference, 1);
+
    for (int i = 0; i < ARRAY_SIZE(prog->pipelines); ++i) {
       prog->pipelines[i] = _mesa_hash_table_create(NULL,
                                                    hash_gfx_pipeline_state,
@@ -132,8 +140,10 @@
 
    for (int i = 0; i < PIPE_SHADER_TYPES - 1; ++i) {
       prog->stages[i] = stages[i];
-      if (stages[i])
+      if (stages[i]) {
          _mesa_set_add(stages[i]->programs, prog);
+         zink_gfx_program_reference(screen, NULL, prog);
+      }
    }
 
    prog->dsl = create_desc_set_layout(screen->dev, stages,
diff --git a/src/gallium/drivers/zink/zink_program.h b/src/gallium/drivers/zink/zink_program.h
index aaad9f3..60fc778 100644
--- a/src/gallium/drivers/zink/zink_program.h
+++ b/src/gallium/drivers/zink/zink_program.h
@@ -27,6 +27,7 @@
 #include <vulkan/vulkan.h>
 
 #include "pipe/p_state.h"
+#include "util/u_inlines.h"
 
 struct zink_context;
 struct zink_screen;
@@ -37,6 +38,8 @@
 struct set;
 
 struct zink_gfx_program {
+   struct pipe_reference reference;
+
    struct zink_shader *stages[PIPE_SHADER_TYPES - 1]; // compute stage doesn't belong here
    VkDescriptorSetLayout dsl;
    VkPipelineLayout layout;
@@ -61,4 +64,20 @@
 
 void
 zink_program_init(struct zink_context *ctx);
+
+void
+debug_describe_zink_gfx_program(char* buf, const struct zink_gfx_program *ptr);
+
+static inline void
+zink_gfx_program_reference(struct zink_screen *screen,
+                           struct zink_gfx_program **dst,
+                           struct zink_gfx_program *src)
+{
+   struct zink_gfx_program *old_dst = dst ? *dst : NULL;
+
+   if (pipe_reference_described(old_dst ? &old_dst->reference : NULL, &src->reference,
+                                (debug_reference_descriptor)debug_describe_zink_gfx_program))
+      zink_destroy_gfx_program(screen, old_dst);
+   if (dst) *dst = src;
+}
 #endif