zink: improve descriptor cache invalidation

we can pass the offset of the 'invalid' flag directly to the resources
to let them run through and invalidate their sets in time for us to detect
it when we recycle the set during batch reset and throw it onto our allocation array

additionally, by adding refs for the actual objects used in a descriptor set, we can
ensure that our cache is as accurate as possible

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

src/gallium/drivers/zink/zink_context.c

@@ -449,6 +449,7 @@ zink_create_sampler_state(struct pipe_context *pctx,
       FREE(sampler);
       return NULL;
    }
+   util_dynarray_init(&sampler->desc_set_refs.refs, NULL);
    calc_descriptor_hash_sampler_state(sampler);
 
    return sampler;
@@ -480,6 +481,7 @@ zink_delete_sampler_state(struct pipe_context *pctx,
 {
    struct zink_sampler_state *sampler = sampler_state;
    struct zink_batch *batch = zink_curr_batch(zink_context(pctx));
+   zink_descriptor_set_refs_clear(&sampler->desc_set_refs, sampler_state);
    util_dynarray_append(&batch->zombie_samplers, VkSampler,
                         sampler->sampler);
    if (sampler->custom_border_color)
@@ -487,7 +489,6 @@ zink_delete_sampler_state(struct pipe_context *pctx,
    FREE(sampler);
 }
 
-
 static VkImageViewType
 image_view_type(enum pipe_texture_target target)
 {
@@ -599,6 +600,7 @@ zink_create_sampler_view(struct pipe_context *pctx, struct pipe_resource *pres,
       FREE(sampler_view);
       return NULL;
    }
+   util_dynarray_init(&sampler_view->desc_set_refs.refs, NULL);
    calc_descriptor_hash_sampler_view(zink_context(pctx), sampler_view);
    return &sampler_view->base;
 }
@@ -608,6 +610,7 @@ zink_sampler_view_destroy(struct pipe_context *pctx,
                           struct pipe_sampler_view *pview)
 {
    struct zink_sampler_view *view = zink_sampler_view(pview);
+   zink_descriptor_set_refs_clear(&view->desc_set_refs, view);
    if (pview->texture->target == PIPE_BUFFER)
       vkDestroyBufferView(zink_screen(pctx->screen)->dev, view->buffer_view, NULL);
    else
@@ -829,6 +832,7 @@ unbind_shader_image(struct zink_context *ctx, enum pipe_shader_type stage, unsig
    if (!image_view->base.resource)
       return;
 
+   zink_descriptor_set_refs_clear(&image_view->desc_set_refs, image_view);
    if (image_view->base.resource->target == PIPE_BUFFER)
       vkDestroyBufferView(zink_screen(ctx->base.screen)->dev, image_view->buffer_view, NULL);
    else
@@ -850,6 +854,7 @@ zink_set_shader_images(struct pipe_context *pctx,
    for (unsigned i = 0; i < count; i++) {
       struct zink_image_view *image_view = &ctx->image_views[p_stage][start_slot + i];
       if (images && images[i].resource) {
+         util_dynarray_init(&image_view->desc_set_refs.refs, NULL);
          struct zink_resource *res = zink_resource(images[i].resource);
          util_copy_image_view(&image_view->base, images + i);
          if (images[i].resource->target == PIPE_BUFFER) {

src/gallium/drivers/zink/zink_context.h

@@ -32,6 +32,7 @@
 #include "pipe/p_context.h"
 #include "pipe/p_state.h"
 #include "util/u_rect.h"
+#include "util/u_dynarray.h"
 
 #include "util/slab.h"
 #include "util/list.h"
@@ -50,6 +51,11 @@ struct zink_resource;
 struct zink_surface;
 struct zink_vertex_elements_state;
 
+
+struct zink_descriptor_refs {
+   struct util_dynarray refs;
+};
+
 enum zink_blit_flags {
    ZINK_BLIT_NORMAL = 1 << 0,
    ZINK_BLIT_SAVE_FS = 1 << 1,
@@ -61,11 +67,13 @@ enum zink_blit_flags {
 struct zink_sampler_state {
    VkSampler sampler;
    uint32_t hash;
+   struct zink_descriptor_refs desc_set_refs;
    bool custom_border_color;
 };
 
 struct zink_sampler_view {
    struct pipe_sampler_view base;
+   struct zink_descriptor_refs desc_set_refs;
    union {
       VkImageView image_view;
       VkBufferView buffer_view;
@@ -75,6 +83,7 @@ struct zink_sampler_view {
 
 struct zink_image_view {
    struct pipe_image_view base;
+   struct zink_descriptor_refs desc_set_refs;
    union {
       struct zink_surface *surface;
       VkBufferView buffer_view;

src/gallium/drivers/zink/zink_draw.c

@@ -16,6 +16,49 @@
 #include "util/u_prim.h"
 #include "util/u_prim_restart.h"
 
+
+static void
+desc_set_res_add(struct zink_descriptor_set *zds, struct zink_resource *res, unsigned int i, bool cache_hit)
+{
+   /* if we got a cache hit, we have to verify that the cached set is still valid;
+    * we store the vk resource to the set here to avoid a more complex and costly mechanism of maintaining a
+    * hash table on every resource with the associated descriptor sets that then needs to be iterated through
+    * whenever a resource is destroyed
+    */
+   assert(!cache_hit || zds->resources[i] == res);
+   if (!cache_hit)
+      zink_resource_desc_set_add(res, zds, i);
+}
+
+static void
+desc_set_sampler_add(struct zink_descriptor_set *zds, struct zink_sampler_view *sv, struct zink_sampler_state *state, unsigned int i, bool cache_hit)
+{
+   /* if we got a cache hit, we have to verify that the cached set is still valid;
+    * we store the vk resource to the set here to avoid a more complex and costly mechanism of maintaining a
+    * hash table on every resource with the associated descriptor sets that then needs to be iterated through
+    * whenever a resource is destroyed
+    */
+   assert(!cache_hit || zds->sampler_views[i] == sv);
+   assert(!cache_hit || zds->sampler_states[i] == state);
+   if (!cache_hit) {
+      zink_sampler_view_desc_set_add(sv, zds, i);
+      zink_sampler_state_desc_set_add(state, zds, i);
+   }
+}
+
+static void
+desc_set_image_add(struct zink_descriptor_set *zds, struct zink_image_view *image_view, unsigned int i, bool cache_hit)
+{
+   /* if we got a cache hit, we have to verify that the cached set is still valid;
+    * we store the vk resource to the set here to avoid a more complex and costly mechanism of maintaining a
+    * hash table on every resource with the associated descriptor sets that then needs to be iterated through
+    * whenever a resource is destroyed
+    */
+   assert(!cache_hit || zds->image_views[i] == image_view);
+   if (!cache_hit)
+      zink_image_view_desc_set_add(image_view, zds, i);
+}
+
 static void
 zink_emit_xfb_counter_barrier(struct zink_context *ctx)
 {
@@ -356,7 +399,6 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
    int num_transitions = 0;
    struct set *ht = _mesa_set_create(NULL, transition_hash, transition_equals);
 
-
    for (int h = 0; h < ZINK_DESCRIPTOR_TYPES; h++) {
       for (int i = 0; i < num_stages; i++) {
          struct zink_shader *shader = stages[i];
@@ -373,6 +415,7 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
                assert(!res || ctx->ubos[stage][index].buffer_size > 0);
                assert(!res || ctx->ubos[stage][index].buffer);
                assert(num_resources[h] < num_bindings);
+               desc_set_res_add(zds[h], res, num_resources[h], cache_hit[h]);
                read_descriptor_resource(&resources[h][num_resources[h]], res, &num_resources[h]);
                assert(num_buffer_info[h] < num_bindings);
                buffer_infos[h][num_buffer_info[h]].buffer = res ? res->buffer :
@@ -393,6 +436,7 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
                ++num_buffer_info[h];
             } else if (shader->bindings[h][j].type == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER) {
                struct zink_resource *res = zink_resource(ctx->ssbos[stage][index].buffer);
+               desc_set_res_add(zds[h], res, num_resources[h], cache_hit[h]);
                if (res) {
                   assert(ctx->ssbos[stage][index].buffer_size > 0);
                   assert(ctx->ssbos[stage][index].buffer_size <= screen->info.props.limits.maxStorageBufferRange);
@@ -431,6 +475,7 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
                   case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER: {
                      struct pipe_sampler_view *psampler_view = ctx->sampler_views[stage][index + k];
                      struct zink_sampler_view *sampler_view = zink_sampler_view(psampler_view);
+                     struct zink_sampler_state *sampler_state = NULL;
                      res = psampler_view ? zink_resource(psampler_view->texture) : NULL;
                      if (!res)
                         break;
@@ -440,9 +485,11 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
                         imageview = sampler_view->image_view;
                         layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
                         sampler = ctx->samplers[stage][index + k];
+                        sampler_state = ctx->sampler_states[stage][index + k];
                      }
                      add_transition(res, layout, VK_ACCESS_SHADER_READ_BIT, stage, &transitions[num_transitions], &num_transitions, ht);
                      assert(num_resources[h] < num_bindings);
+                     desc_set_sampler_add(zds[h], sampler_view, sampler_state, num_resources[h], cache_hit[h]);
                      read_descriptor_resource(&resources[h][num_resources[h]], res, &num_resources[h]);
                   }
                   break;
@@ -470,6 +517,7 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
                         flags |= VK_ACCESS_SHADER_WRITE_BIT;
                      add_transition(res, layout, flags, stage, &transitions[num_transitions], &num_transitions, ht);
                      assert(num_resources[h] < num_bindings);
+                     desc_set_image_add(zds[h], image_view, num_resources[h], cache_hit[h]);
                      if (image_view->base.access & PIPE_IMAGE_ACCESS_WRITE)
                         write_descriptor_resource(&resources[h][num_resources[h]], res, &num_resources[h]);
                      else
@@ -562,14 +610,6 @@ update_descriptors(struct zink_context *ctx, struct zink_screen *screen, bool is
          if (res) {
             need_flush |= zink_batch_reference_resource_rw(batch, res, resources[h][i].write) == check_flush_id;
          }
-         /* if we got a cache hit, we have to verify that the cached set is still valid;
-          * we store the vk resource to the set here to avoid a more complex and costly mechanism of maintaining a
-          * hash table on every resource with the associated descriptor sets that then needs to be iterated through
-          * whenever a resource is destroyed
-          */
-         cache_hit[h] = cache_hit[h] && zds[h]->resources[i] == res;
-         if (zds[h]->resources[i] != res)
-            zink_resource_desc_set_add(res, zds[h], i);
       }
       if (!cache_hit[h] && num_wds[h])
          vkUpdateDescriptorSets(screen->dev, num_wds[h], wds[h], 0, NULL);

src/gallium/drivers/zink/zink_program.c

@@ -877,8 +877,13 @@ allocate_desc_set(struct zink_screen *screen, struct zink_program *pg, enum zink
    struct zink_descriptor_set *alloc = ralloc_array(pg, struct zink_descriptor_set, bucket_size);
    assert(alloc);
    unsigned num_resources = zink_program_num_bindings_typed(pg, type, is_compute);
-   struct zink_resource **resources = rzalloc_array(pg, struct zink_resource*, num_resources * bucket_size);
+   void **resources = rzalloc_array(pg, void*, num_resources * bucket_size);
    assert(resources);
+   void **samplers = NULL;
+   if (type == ZINK_DESCRIPTOR_TYPE_SAMPLER_VIEW) {
+      samplers = rzalloc_array(pg, void*, num_resources * bucket_size);
+      assert(samplers);
+   }
    for (unsigned i = 0; i < bucket_size; i ++) {
       struct zink_descriptor_set *zds = &alloc[i];
       pipe_reference_init(&zds->reference, 1);
@@ -889,7 +894,11 @@ allocate_desc_set(struct zink_screen *screen, struct zink_program *pg, enum zink
 #ifndef NDEBUG
       zds->num_resources = num_resources;
 #endif
-      zds->resources = &resources[i * pg->num_descriptors[type]];
+      if (type == ZINK_DESCRIPTOR_TYPE_SAMPLER_VIEW) {
+         zds->sampler_views = (struct zink_sampler_view**)&resources[i * pg->num_descriptors[type]];
+         zds->sampler_states = (struct zink_sampler_state**)&samplers[i * pg->num_descriptors[type]];
+      } else
+         zds->resources = (struct zink_resource**)&resources[i * pg->num_descriptors[type]];
       zds->desc_set = desc_set[i];
       if (i > 0)
          util_dynarray_append(&pg->alloc_desc_sets[type], struct zink_descriptor_set *, zds);
@@ -932,7 +941,7 @@ zink_program_allocate_desc_set(struct zink_context *ctx,
    if (pg->last_set[type] && pg->last_set[type]->hash == hash &&
        desc_state_equal(&pg->last_set[type]->key, &key)) {
       zds = pg->last_set[type];
-      *cache_hit = true;
+      *cache_hit = !zds->invalid;
       if (pg->num_descriptors[type]) {
          struct hash_entry *he = _mesa_hash_table_search_pre_hashed(pg->free_desc_sets[type], hash, &key);
          if (he)
@@ -1039,7 +1048,11 @@ zink_program_recycle_desc_set(struct zink_program *pg, struct zink_descriptor_se
       return;
 
    _mesa_hash_table_remove(pg->desc_sets[zds->type], he);
-   _mesa_hash_table_insert_pre_hashed(pg->free_desc_sets[zds->type], zds->hash, &zds->key, zds);
+   if (zds->invalid) {
+      desc_set_invalidate_resources(pg, zds);
+      util_dynarray_append(&pg->alloc_desc_sets[zds->type], struct zink_descriptor_set *, zds);
+   } else
+      _mesa_hash_table_insert_pre_hashed(pg->free_desc_sets[zds->type], zds->hash, &zds->key, zds);
 }
 
 static void
@@ -1518,3 +1531,49 @@ zink_program_init(struct zink_context *ctx)
    ctx->base.bind_compute_state = zink_bind_cs_state;
    ctx->base.delete_compute_state = zink_delete_shader_state;
 }
+
+
+static void
+desc_set_ref_add(struct zink_descriptor_set *zds, struct zink_descriptor_refs *refs, void **ref_ptr, void *ptr)
+{
+   struct zink_descriptor_reference ref = {ref_ptr, &zds->invalid};
+   *ref_ptr = ptr;
+   if (ptr)
+      util_dynarray_append(&refs->refs, struct zink_descriptor_reference, ref);
+}
+
+void
+zink_image_view_desc_set_add(struct zink_image_view *image_view, struct zink_descriptor_set *zds, unsigned idx)
+{
+   desc_set_ref_add(zds, &image_view->desc_set_refs, (void**)&zds->image_views[idx], image_view);
+}
+
+void
+zink_sampler_state_desc_set_add(struct zink_sampler_state *sampler_state, struct zink_descriptor_set *zds, unsigned idx)
+{
+   desc_set_ref_add(zds, &sampler_state->desc_set_refs, (void**)&zds->sampler_states[idx], sampler_state);
+}
+
+void
+zink_sampler_view_desc_set_add(struct zink_sampler_view *sampler_view, struct zink_descriptor_set *zds, unsigned idx)
+{
+   desc_set_ref_add(zds, &sampler_view->desc_set_refs, (void**)&zds->sampler_views[idx], sampler_view);
+}
+
+void
+zink_resource_desc_set_add(struct zink_resource *res, struct zink_descriptor_set *zds, unsigned idx)
+{
+   desc_set_ref_add(zds, &res->desc_set_refs, (void**)&zds->resources[idx], res);
+}
+
+void
+zink_descriptor_set_refs_clear(struct zink_descriptor_refs *refs, void *ptr)
+{
+   util_dynarray_foreach(&refs->refs, struct zink_descriptor_reference, ref) {
+      if (*ref->ref == ptr) {
+         *ref->invalid = true;
+         *ref->ref = NULL;
+      }
+   }
+   util_dynarray_fini(&refs->refs);
+}

src/gallium/drivers/zink/zink_program.h

@@ -83,9 +83,33 @@ struct zink_descriptor_set {
    /* for extra debug asserts */
    unsigned num_resources;
 #endif
-   struct zink_resource **resources;
+   union {
+      struct zink_resource **resources;
+      struct zink_image_view **image_views;
+      struct {
+         struct zink_sampler_view **sampler_views;
+         struct zink_sampler_state **sampler_states;
+      };
+   };
 };
 
+
+struct zink_descriptor_reference {
+   void **ref;
+   bool *invalid;
+};
+void
+zink_descriptor_set_refs_clear(struct zink_descriptor_refs *refs, void *ptr);
+
+void
+zink_image_view_desc_set_add(struct zink_image_view *image_view, struct zink_descriptor_set *zds, unsigned idx);
+void
+zink_sampler_state_desc_set_add(struct zink_sampler_state *sampler_state, struct zink_descriptor_set *zds, unsigned idx);
+void
+zink_sampler_view_desc_set_add(struct zink_sampler_view *sv, struct zink_descriptor_set *zds, unsigned idx);
+void
+zink_resource_desc_set_add(struct zink_resource *res, struct zink_descriptor_set *zds, unsigned idx);
+
 struct zink_program {
    struct pipe_reference reference;
 

src/gallium/drivers/zink/zink_resource.c

@@ -79,11 +79,7 @@ zink_resource_destroy(struct pipe_screen *pscreen,
    } else
       vkDestroyImage(screen->dev, res->image, NULL);
 
-   util_dynarray_foreach(&res->desc_set_refs, struct zink_resource **, ref) {
-      if (**ref == res)
-         **ref = NULL;
-   }
-   util_dynarray_fini(&res->desc_set_refs);
+   zink_descriptor_set_refs_clear(&res->desc_set_refs, res);
 
    vkFreeMemory(screen->dev, res->mem, NULL);
    FREE(res);
@@ -398,7 +394,7 @@ resource_create(struct pipe_screen *pscreen,
                                              &res->dt_stride);
    }
 
-   util_dynarray_init(&res->desc_set_refs, NULL);
+   util_dynarray_init(&res->desc_set_refs.refs, NULL);
    return &res->base;
 
 fail:
@@ -835,13 +831,6 @@ zink_resource_setup_transfer_layouts(struct zink_context *ctx, struct zink_resou
    }
 }
 
-void
-zink_resource_desc_set_add(struct zink_resource *res, struct zink_descriptor_set *zds, unsigned idx)
-{
-   zds->resources[idx] = res;
-   util_dynarray_append(&res->desc_set_refs, struct zink_resource**, &zds->resources[idx]);
-}
-
 void
 zink_get_depth_stencil_resources(struct pipe_resource *res,
                                  struct zink_resource **out_z,

src/gallium/drivers/zink/zink_resource.h

@@ -28,12 +28,13 @@ struct pipe_screen;
 struct sw_displaytarget;
 struct zink_batch;
 struct zink_context;
-struct zink_descriptor_set;
 
 #include "util/u_transfer.h"
 #include "util/u_range.h"
 #include "util/u_dynarray.h"
 
+#include "zink_context.h"
+
 #include <vulkan/vulkan.h>
 
 #define ZINK_RESOURCE_ACCESS_READ 1
@@ -67,7 +68,7 @@ struct zink_resource {
    unsigned dt_stride;
    unsigned persistent_maps; //if nonzero, requires vkFlushMappedMemoryRanges during batch use
 
-   struct util_dynarray desc_set_refs;
+   struct zink_descriptor_refs desc_set_refs;
 
    /* this has to be atomic for fence access, so we can't use a bitmask and make everything neat */
    uint8_t batch_uses[5]; //ZINK_NUM_BATCHES