zink: store and reuse descriptorsets after batch completion

since we know that the layout is going to match, we can store descriptorsets
in the program and then overwrite them instead of needing to free sets or reset
the pool

Reviewed-by: Bas Nieuwenhuizen <bas@basnieuwenhuizen.nl>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/9348>

src/gallium/drivers/zink/zink_batch.c

@@ -54,8 +54,12 @@ zink_reset_batch(struct zink_context *ctx, struct zink_batch *batch)
       struct zink_program *pg = (struct zink_program*)entry->key;
       struct set *desc_sets = (struct set*)entry->data;
       set_foreach(desc_sets, sentry) {
-         VkDescriptorSet desc_set = (VkDescriptorSet)sentry->key;
-         vkFreeDescriptorSets(screen->dev, pg->descpool, 1, &desc_set);
+         struct zink_descriptor_set *zds = (void*)sentry->key;
+         /* reset descriptor pools when no batch is using this program to avoid
+          * having some inactive program hogging a billion descriptors
+          */
+         pipe_reference(&zds->reference, NULL);
+         zink_program_invalidate_desc_set(pg, zds);
       }
       _mesa_set_destroy(desc_sets, NULL);
       if (batch->batch_id == ZINK_COMPUTE_BATCH_ID) {
@@ -231,13 +235,14 @@ zink_batch_reference_program(struct zink_batch *batch,
 }
 
 bool
-zink_batch_add_desc_set(struct zink_batch *batch, struct zink_program *pg, VkDescriptorSet desc_set)
+zink_batch_add_desc_set(struct zink_batch *batch, struct zink_program *pg, struct zink_descriptor_set *zds)
 {
    struct hash_entry *entry = _mesa_hash_table_search(batch->programs, pg);
    assert(entry);
    struct set *desc_sets = (void*)entry->data;
-   if (!_mesa_set_search(desc_sets, desc_set)) {
-      _mesa_set_add(desc_sets, desc_set);
+   if (!_mesa_set_search(desc_sets, zds)) {
+      pipe_reference(NULL, &zds->reference);
+      _mesa_set_add(desc_sets, zds);
       return true;
    }
    return false;

src/gallium/drivers/zink/zink_batch.h

@@ -100,5 +100,5 @@ zink_batch_reference_surface(struct zink_batch *batch,
                              struct zink_surface *surface);
 
 bool
-zink_batch_add_desc_set(struct zink_batch *batch, struct zink_program *pg, VkDescriptorSet desc_set);
+zink_batch_add_desc_set(struct zink_batch *batch, struct zink_program *pg, struct zink_descriptor_set *zds);
 #endif

src/gallium/drivers/zink/zink_draw.c

@@ -17,11 +17,19 @@
 #include "util/u_prim_restart.h"
 
 
-static VkDescriptorSet
+static struct zink_descriptor_set *
 allocate_descriptor_set(struct zink_screen *screen,
                         struct zink_batch *batch,
                         struct zink_program *pg)
 {
+   struct zink_descriptor_set *zds;
+
+   if (util_dynarray_num_elements(&pg->alloc_desc_sets, struct zink_descriptor_set *)) {
+      /* grab one off the allocated array */
+      zds = util_dynarray_pop(&pg->alloc_desc_sets, struct zink_descriptor_set *);
+      goto out;
+   }
+
    VkDescriptorSetAllocateInfo dsai;
    memset((void *)&dsai, 0, sizeof(dsai));
    dsai.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
@@ -35,10 +43,15 @@ allocate_descriptor_set(struct zink_screen *screen,
       debug_printf("ZINK: %p failed to allocate descriptor set :/\n", pg);
       return VK_NULL_HANDLE;
    }
-   if (zink_batch_add_desc_set(batch, pg, desc_set))
+   zds = ralloc_size(NULL, sizeof(struct zink_descriptor_set));
+   assert(zds);
+   pipe_reference_init(&zds->reference, 1);
+   zds->desc_set = desc_set;
+out:
+   if (zink_batch_add_desc_set(batch, pg, zds))
       batch->descs_used += pg->num_descriptors;
 
-   return desc_set;
+   return zds;
 }
 
 static void
@@ -530,9 +543,9 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
    struct zink_program *pg = is_compute ? &ctx->curr_compute->base : &ctx->curr_program->base;
    zink_batch_reference_program(batch, pg);
    assert(pg->num_descriptors == num_descriptors);
-   VkDescriptorSet desc_set = allocate_descriptor_set(screen, batch, pg);
+   struct zink_descriptor_set *zds = allocate_descriptor_set(screen, batch, pg);
    /* probably oom, so we need to stall until we free up some descriptors */
-   if (!desc_set) {
+   if (!zds) {
       /* update our max descriptor count so we can try and avoid this happening again */
       unsigned short max_descs = 0;
       for (int i = 0; i < ZINK_COMPUTE_BATCH_ID; i++)
@@ -554,15 +567,15 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
          }
       }
       zink_batch_reference_program(batch, pg);
-      desc_set = allocate_descriptor_set(screen, batch, pg);
+      zds = allocate_descriptor_set(screen, batch, pg);
    }
-   assert(desc_set != VK_NULL_HANDLE);
+   assert(zds != VK_NULL_HANDLE);
 
    unsigned check_flush_id = is_compute ? 0 : ZINK_COMPUTE_BATCH_ID;
    bool need_flush = false;
    if (num_wds > 0) {
       for (int i = 0; i < num_wds; ++i) {
-         wds[i].dstSet = desc_set;
+         wds[i].dstSet = zds->desc_set;
          struct zink_resource *res = resources[i].res;
          if (res) {
             need_flush |= zink_batch_reference_resource_rw(batch, res, resources[i].write) == check_flush_id;
@@ -577,10 +590,10 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
 
    if (is_compute)
       vkCmdBindDescriptorSets(batch->cmdbuf, VK_PIPELINE_BIND_POINT_COMPUTE,
-                              ctx->curr_compute->layout, 0, 1, &desc_set, 0, NULL);
+                              ctx->curr_compute->layout, 0, 1, &zds->desc_set, 0, NULL);
    else
       vkCmdBindDescriptorSets(batch->cmdbuf, VK_PIPELINE_BIND_POINT_GRAPHICS,
-                              ctx->curr_program->layout, 0, 1, &desc_set, 0, NULL);
+                              ctx->curr_program->layout, 0, 1, &zds->desc_set, 0, NULL);
 
    for (int i = 0; i < num_stages; i++) {
       struct zink_shader *shader = stages[i];

src/gallium/drivers/zink/zink_program.c

@@ -165,7 +165,7 @@ create_desc_set_layout(VkDevice dev,
    dpci.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    dpci.pPoolSizes = sizes;
    dpci.poolSizeCount = ARRAY_SIZE(sizes);
-   dpci.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;
+   dpci.flags = 0;
    dpci.maxSets = ZINK_BATCH_DESC_SIZE;
    if (vkCreateDescriptorPool(dev, &dpci, 0, descpool) != VK_SUCCESS) {
       vkDestroyDescriptorSetLayout(dev, dsl, NULL);
@@ -670,6 +670,15 @@ fail:
    return NULL;
 }
 
+void
+zink_program_invalidate_desc_set(struct zink_program *pg, struct zink_descriptor_set *zds)
+{
+   uint32_t refcount = p_atomic_read(&zds->reference.count);
+   /* refcount > 1 means this is currently in use, so we can't recycle it yet */
+   if (refcount == 1)
+      util_dynarray_append(&pg->alloc_desc_sets, struct zink_descriptor_set *, zds);
+}
+
 static void
 gfx_program_remove_shader(struct zink_gfx_program *prog, struct zink_shader *shader)
 {

src/gallium/drivers/zink/zink_program.h

@@ -63,6 +63,11 @@ struct zink_shader_cache {
    struct hash_table *shader_cache;
 };
 
+struct zink_descriptor_set {
+   struct pipe_reference reference; //incremented for batch usage
+   VkDescriptorSet desc_set;
+};
+
 struct zink_program {
    struct pipe_reference reference;
 
@@ -161,4 +166,7 @@ zink_get_compute_pipeline(struct zink_screen *screen,
                       struct zink_compute_program *comp,
                       struct zink_compute_pipeline_state *state);
 
+void
+zink_program_invalidate_desc_set(struct zink_program *pg, struct zink_descriptor_set *zds);
+
 #endif