zink: improve descriptor set oom handling

this attempts to dynamically establish an upper bound for per-batch descriptor
use, flushing all batches and resetting the pools on alloc failure in
an attempt to be more robust about it

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

src/gallium/drivers/zink/zink_batch.c

@@ -15,9 +15,16 @@
 #include "util/set.h"
 
 void
-zink_batch_release(struct zink_screen *screen, struct zink_batch *batch)
+zink_reset_batch(struct zink_context *ctx, struct zink_batch *batch)
 {
-   zink_fence_reference(screen, &batch->fence, NULL);
+   struct zink_screen *screen = zink_screen(ctx->base.screen);
+   batch->descs_used = 0;
+
+   // cmdbuf hasn't been submitted before without a fence
+   if (batch->fence) {
+      zink_fence_finish(screen, batch->fence, PIPE_TIMEOUT_INFINITE);
+      zink_fence_reference(screen, &batch->fence, NULL);
+   }
 
    zink_framebuffer_reference(screen, &batch->fb, NULL);
    set_foreach(batch->programs, entry) {
@@ -56,20 +63,6 @@ zink_batch_release(struct zink_screen *screen, struct zink_batch *batch)
    }
    util_dynarray_clear(&batch->zombie_samplers);
    util_dynarray_clear(&batch->persistent_resources);
-}
-
-static void
-reset_batch(struct zink_context *ctx, struct zink_batch *batch)
-{
-   struct zink_screen *screen = zink_screen(ctx->base.screen);
-   batch->descs_left = ZINK_BATCH_DESC_SIZE;
-
-   // cmdbuf hasn't been submitted before
-   if (!batch->fence)
-      return;
-
-   zink_fence_finish(screen, batch->fence, PIPE_TIMEOUT_INFINITE);
-   zink_batch_release(screen, batch);
 
    if (vkResetDescriptorPool(screen->dev, batch->descpool, 0) != VK_SUCCESS)
       fprintf(stderr, "vkResetDescriptorPool failed\n");
@@ -82,7 +75,7 @@ reset_batch(struct zink_context *ctx, struct zink_batch *batch)
 void
 zink_start_batch(struct zink_context *ctx, struct zink_batch *batch)
 {
-   reset_batch(ctx, batch);
+   zink_reset_batch(ctx, batch);
 
    VkCommandBufferBeginInfo cbbi = {};
    cbbi.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
@@ -182,7 +175,7 @@ zink_batch_reference_resource_rw(struct zink_batch *batch, struct zink_resource
    if (res->persistent_maps)
       util_dynarray_append(&batch->persistent_resources, struct zink_resource*, res);
    /* the batch_uses value for this batch is guaranteed to not be in use now because
-    * reset_batch() waits on the fence and removes access before resetting
+    * zink_reset_batch() waits on the fence and removes access before resetting
     */
    res->batch_uses[batch->batch_id] |= mask;
 

src/gallium/drivers/zink/zink_batch.h

@@ -36,7 +36,6 @@ struct zink_fence;
 struct zink_framebuffer;
 struct zink_render_pass;
 struct zink_resource;
-struct zink_screen;
 struct zink_sampler_view;
 struct zink_surface;
 
@@ -47,7 +46,8 @@ struct zink_batch {
    VkCommandPool cmdpool;
    VkCommandBuffer cmdbuf;
    VkDescriptorPool descpool;
-   int descs_left;
+   unsigned short max_descs; //set if the device gives oom when allocating a new desc set
+   unsigned short descs_used; //number of descriptors currently allocated
    struct zink_fence *fence;
 
    struct zink_framebuffer *fb;
@@ -65,10 +65,8 @@ struct zink_batch {
    bool in_rp; //renderpass is currently active
 };
 
-/* release all resources attached to batch */
 void
-zink_batch_release(struct zink_screen *screen, struct zink_batch *batch);
-
+zink_reset_batch(struct zink_context *ctx, struct zink_batch *batch);
 void
 zink_start_batch(struct zink_context *ctx, struct zink_batch *batch);
 

src/gallium/drivers/zink/zink_context.c

@@ -68,7 +68,7 @@ zink_context_destroy(struct pipe_context *pctx)
       pipe_resource_reference(&ctx->null_buffers[i], NULL);
 
    for (int i = 0; i < ARRAY_SIZE(ctx->batches); ++i) {
-      zink_batch_release(screen, &ctx->batches[i]);
+      zink_reset_batch(ctx, &ctx->batches[i]);
       util_dynarray_fini(&ctx->batches[i].zombie_samplers);
       vkDestroyDescriptorPool(screen->dev, ctx->batches[i].descpool, NULL);
 
@@ -79,7 +79,7 @@ zink_context_destroy(struct pipe_context *pctx)
       vkDestroyCommandPool(screen->dev, ctx->batches[i].cmdpool, NULL);
    }
    if (ctx->compute_batch.cmdpool) {
-      zink_batch_release(screen, &ctx->compute_batch);
+      zink_reset_batch(ctx, &ctx->compute_batch);
       util_dynarray_fini(&ctx->compute_batch.zombie_samplers);
       vkDestroyDescriptorPool(screen->dev, ctx->compute_batch.descpool, NULL);
       vkFreeCommandBuffers(screen->dev, ctx->compute_batch.cmdpool, 1, &ctx->compute_batch.cmdbuf);
@@ -1753,7 +1753,7 @@ init_batch(struct zink_context *ctx, struct zink_batch *batch, unsigned idx)
    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 (vkAllocateCommandBuffers(screen->dev, &cbai, &batch->cmdbuf) != VK_SUCCESS)
@@ -1776,6 +1776,7 @@ init_batch(struct zink_context *ctx, struct zink_batch *batch, unsigned idx)
       return false;
 
    batch->batch_id = idx;
+   batch->max_descs = 1500;
    return true;
 }
 

src/gallium/drivers/zink/zink_draw.c

@@ -22,7 +22,6 @@ allocate_descriptor_set(struct zink_screen *screen,
                         VkDescriptorSetLayout dsl,
                         unsigned num_descriptors)
 {
-   assert(batch->descs_left >= num_descriptors);
    VkDescriptorSetAllocateInfo dsai;
    memset((void *)&dsai, 0, sizeof(dsai));
    dsai.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
@@ -33,11 +32,11 @@ allocate_descriptor_set(struct zink_screen *screen,
 
    VkDescriptorSet desc_set;
    if (vkAllocateDescriptorSets(screen->dev, &dsai, &desc_set) != VK_SUCCESS) {
-      debug_printf("ZINK: failed to allocate descriptor set :/");
+      debug_printf("ZINK: failed to allocate descriptor set :/\n");
       return VK_NULL_HANDLE;
    }
 
-   batch->descs_left -= num_descriptors;
+   batch->descs_used += num_descriptors;
    return desc_set;
 }
 
@@ -490,22 +489,45 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
       dsl = ctx->curr_program->dsl;
    }
 
-   if (batch->descs_left < num_descriptors) {
+   if (is_compute)
+      zink_batch_reference_program(batch, &ctx->curr_compute->reference);
+   else
+      zink_batch_reference_program(batch, &ctx->curr_program->reference);
+
+   if (batch->descs_used + num_descriptors >= batch->max_descs) {
       if (is_compute)
          zink_wait_on_batch(ctx, ZINK_COMPUTE_BATCH_ID);
       else {
          ctx->base.flush(&ctx->base, NULL, 0);
          batch = zink_batch_rp(ctx);
       }
-      assert(batch->descs_left >= num_descriptors);
    }
-   if (is_compute)
-      zink_batch_reference_program(batch, &ctx->curr_compute->reference);
-   else
-      zink_batch_reference_program(batch, &ctx->curr_program->reference);
-
    VkDescriptorSet desc_set = allocate_descriptor_set(screen, batch,
                                                       dsl, num_descriptors);
+   /* probably oom, so we need to stall until we free up some descriptors */
+   if (!desc_set) {
+      /* 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++)
+         max_descs += ctx->batches[i].descs_used;
+      if (ctx->compute_batch.descs_used) {
+         max_descs += ctx->compute_batch.descs_used;
+         /* try to split evenly between number of batches */
+         max_descs /= ZINK_COMPUTE_BATCH_ID;
+      }
+      for (int i = 0; i < ZINK_COMPUTE_BATCH_ID; i++)
+         ctx->batches[i].max_descs = MIN2(max_descs, ctx->batches[i].max_descs);
+      ctx->compute_batch.max_descs = MIN2(max_descs, ctx->compute_batch.max_descs);
+
+      zink_wait_on_batch(ctx, batch->batch_id);
+      if (!is_compute) {
+         batch = zink_curr_batch(ctx);
+         for (int i = 0; i < ZINK_COMPUTE_BATCH_ID; i++) {
+            zink_reset_batch(ctx, &ctx->batches[i]);
+         }
+      }
+      desc_set = allocate_descriptor_set(screen, batch, dsl, num_descriptors);
+   }
    assert(desc_set != VK_NULL_HANDLE);
 
    for (int i = 0; i < num_stages; i++) {