nvk/descriptor_table: Add support for requesting a specific index

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30580>

src/nouveau/vulkan/nvk_descriptor_table.c

@@ -15,6 +15,7 @@ nvk_descriptor_table_grow_locked(struct nvk_device *dev,
                                  uint32_t new_alloc)
 {
    struct nvkmd_mem *new_mem;
+   BITSET_WORD *new_in_use;
    uint32_t *new_free_table;
    VkResult result;
 
@@ -35,6 +36,23 @@ nvk_descriptor_table_grow_locked(struct nvk_device *dev,
    }
    table->mem = new_mem;
 
+   assert((table->alloc % BITSET_WORDBITS) == 0);
+   assert((new_alloc % BITSET_WORDBITS) == 0);
+   const size_t old_in_use_size =
+      BITSET_WORDS(table->alloc) * sizeof(BITSET_WORD);
+   const size_t new_in_use_size =
+      BITSET_WORDS(new_alloc) * sizeof(BITSET_WORD);
+   new_in_use = vk_realloc(&dev->vk.alloc, table->in_use,
+                           new_in_use_size, sizeof(BITSET_WORD),
+                           VK_SYSTEM_ALLOCATION_SCOPE_OBJECT);
+   if (new_in_use == NULL) {
+      return vk_errorf(dev, VK_ERROR_OUT_OF_HOST_MEMORY,
+                       "Failed to allocate image descriptor in-use set");
+   }
+   memset((char *)new_in_use + old_in_use_size, 0,
+          new_in_use_size - old_in_use_size);
+   table->in_use = new_in_use;
+
    const size_t new_free_table_size = new_alloc * sizeof(uint32_t);
    new_free_table = vk_realloc(&dev->vk.alloc, table->free_table,
                                new_free_table_size, 4,
@@ -99,29 +117,55 @@ nvk_descriptor_table_alloc_locked(struct nvk_device *dev,
 {
    VkResult result;
 
-   if (table->free_count > 0) {
-      *index_out = table->free_table[--table->free_count];
+   while (1) {
+      uint32_t index;
+      if (table->free_count > 0) {
+         index = table->free_table[--table->free_count];
+      } else if (table->next_desc < table->alloc) {
+         index = table->next_desc++;
+      } else {
+         if (table->next_desc >= table->max_alloc) {
+            return vk_errorf(dev, VK_ERROR_OUT_OF_HOST_MEMORY,
+                             "Descriptor table not large enough");
+         }
+
+         result = nvk_descriptor_table_grow_locked(dev, table,
+                                                   table->alloc * 2);
+         if (result != VK_SUCCESS)
+            return result;
+
+         assert(table->next_desc < table->alloc);
+         index = table->next_desc++;
+      }
+
+      if (!BITSET_TEST(table->in_use, index)) {
+         BITSET_SET(table->in_use, index);
+         *index_out = index;
+         return VK_SUCCESS;
+      }
+   }
+}
+
+static VkResult
+nvk_descriptor_table_take_locked(struct nvk_device *dev,
+                                 struct nvk_descriptor_table *table,
+                                 uint32_t index)
+{
+   VkResult result;
+
+   while (index >= table->alloc) {
+      result = nvk_descriptor_table_grow_locked(dev, table, table->alloc * 2);
+      if (result != VK_SUCCESS)
+         return result;
+   }
+
+   if (BITSET_TEST(table->in_use, index)) {
+      return vk_errorf(dev, VK_ERROR_INVALID_OPAQUE_CAPTURE_ADDRESS,
+                       "Descriptor %u is already in use", index);
+   } else {
+      BITSET_SET(table->in_use, index);
       return VK_SUCCESS;
    }
-
-   if (table->next_desc < table->alloc) {
-      *index_out = table->next_desc++;
-      return VK_SUCCESS;
-   }
-
-   if (table->next_desc >= table->max_alloc) {
-      return vk_errorf(dev, VK_ERROR_OUT_OF_HOST_MEMORY,
-                       "Descriptor table not large enough");
-   }
-
-   result = nvk_descriptor_table_grow_locked(dev, table, table->alloc * 2);
-   if (result != VK_SUCCESS)
-      return result;
-
-   assert(table->next_desc < table->alloc);
-   *index_out = table->next_desc++;
-
-   return VK_SUCCESS;
 }
 
 static VkResult
@@ -157,6 +201,66 @@ nvk_descriptor_table_add(struct nvk_device *dev,
    return result;
 }
 
+static VkResult
+nvk_descriptor_table_insert_locked(struct nvk_device *dev,
+                                   struct nvk_descriptor_table *table,
+                                   uint32_t index,
+                                   const void *desc_data, size_t desc_size)
+{
+   VkResult result = nvk_descriptor_table_take_locked(dev, table, index);
+   if (result != VK_SUCCESS)
+      return result;
+
+   void *map = (char *)table->mem->map + (index * table->desc_size);
+
+   assert(desc_size == table->desc_size);
+   memcpy(map, desc_data, table->desc_size);
+
+   return result;
+}
+
+VkResult
+nvk_descriptor_table_insert(struct nvk_device *dev,
+                            struct nvk_descriptor_table *table,
+                            uint32_t index,
+                            const void *desc_data, size_t desc_size)
+{
+   simple_mtx_lock(&table->mutex);
+   VkResult result = nvk_descriptor_table_insert_locked(dev, table, index,
+                                                        desc_data, desc_size);
+   simple_mtx_unlock(&table->mutex);
+
+   return result;
+}
+
+static int
+compar_u32(const void *_a, const void *_b)
+{
+   const uint32_t *a = _a, *b = _b;
+   return *a - *b;
+}
+
+static void
+nvk_descriptor_table_compact_free_table(struct nvk_descriptor_table *table)
+{
+   if (table->free_count <= 1)
+      return;
+
+   qsort(table->free_table, table->free_count,
+         sizeof(*table->free_table), compar_u32);
+
+   uint32_t j = 1;
+   for (uint32_t i = 1; i < table->free_count; i++) {
+      if (table->free_table[i] == table->free_table[j - 1])
+         continue;
+
+      assert(table->free_table[i] > table->free_table[j - 1]);
+      table->free_table[j++] = table->free_table[i];
+   }
+
+   table->free_count = j;
+}
+
 void
 nvk_descriptor_table_remove(struct nvk_device *dev,
                             struct nvk_descriptor_table *table,
@@ -167,11 +271,21 @@ nvk_descriptor_table_remove(struct nvk_device *dev,
    void *map = (char *)table->mem->map + (index * table->desc_size);
    memset(map, 0, table->desc_size);
 
-   /* Sanity check for double-free */
-   assert(table->free_count < table->alloc);
-   for (uint32_t i = 0; i < table->free_count; i++)
-      assert(table->free_table[i] != index);
+   assert(BITSET_TEST(table->in_use, index));
 
+   /* There may be duplicate entries in the free table.  For most operations,
+    * this is fine as we always consult nvk_descriptor_table::in_use when
+    * allocating.  However, it does mean that there's nothing preventing our
+    * free table from growing larger than the memory we allocated for it.  In
+    * the unlikely event that we end up with more entries than we can fit in
+    * the allocated space, compact the table to ensure that the new entry
+    * we're about to add fits.
+    */
+   if (table->free_count >= table->alloc)
+      nvk_descriptor_table_compact_free_table(table);
+   assert(table->free_count < table->alloc);
+
+   BITSET_CLEAR(table->in_use, index);
    table->free_table[table->free_count++] = index;
 
    simple_mtx_unlock(&table->mutex);

src/nouveau/vulkan/nvk_descriptor_table.h

@@ -7,6 +7,7 @@
 
 #include "nvk_private.h"
 
+#include "util/bitset.h"
 #include "util/simple_mtx.h"
 #include "nvkmd/nvkmd.h"
 
@@ -23,6 +24,14 @@ struct nvk_descriptor_table {
 
    struct nvkmd_mem *mem;
 
+   /* Bitset of all descriptors currently in use.  This is the single source
+    * of truth for what is and isn't free.  The free_table and next_desc are
+    * simply hints to make finding a free descrptor fast.  Every free
+    * descriptor will either be above next_desc or in free_table but not
+    * everything which satisfies those two criteria is actually free.
+    */
+   BITSET_WORD *in_use;
+
    /* Stack for free descriptor elements */
    uint32_t *free_table;
 };
@@ -41,6 +50,11 @@ VkResult nvk_descriptor_table_add(struct nvk_device *dev,
                                   const void *desc_data, size_t desc_size,
                                   uint32_t *index_out);
 
+VkResult nvk_descriptor_table_insert(struct nvk_device *dev,
+                                     struct nvk_descriptor_table *table,
+                                     uint32_t index,
+                                     const void *desc_data, size_t desc_size);
+
 void nvk_descriptor_table_remove(struct nvk_device *dev,
                                  struct nvk_descriptor_table *table,
                                  uint32_t index);