radv: fix memory leak of descriptor set layout

We need to be able to track the descriptor sets explicity to unref the
descriptor sets, otherwise these descriptor sets will not unref the
descriptor set layout it holds.

Closes: https://gitlab.freedesktop.org/mesa/mesa/-/issues/6222
Fixes: 66f7289d568
("radv: add reference counting for descriptor set layouts")

Tested-by: Jakob Bornecrantz <jakob@collabora.com>
Reviewed-by: Bas Nieuwenhuizen <bas@basnieuwenhuizen.nl>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/15741>
(cherry picked from commit 96a240e176701f9b305c4bd273da9a8aee78e280)

2 files changed, 15 insertions, 18 deletions

diff --git a/.pick_status.json b/.pick_status.json
index 0fa871c70c2..48991e2513d 100644
--- a/.pick_status.json
+++ b/.pick_status.json
@@ -175,7 +175,7 @@
         "description": "radv: fix memory leak of descriptor set layout",
         "nominated": true,
         "nomination_type": 1,
-        "resolution": 0,
+        "resolution": 1,
         "main_sha": null,
         "because_sha": "66f7289d568db8711adb885acc56622e2aff252a"
     },
diff --git a/src/amd/vulkan/radv_descriptor_set.c b/src/amd/vulkan/radv_descriptor_set.c
index 815340c16ee..249c123bf08 100644
--- a/src/amd/vulkan/radv_descriptor_set.c
+++ b/src/amd/vulkan/radv_descriptor_set.c
@@ -667,9 +667,8 @@ radv_descriptor_set_create(struct radv_device *device, struct radv_descriptor_po
       if (!pool->host_memory_base) {
          pool->entries[pool->entry_count].offset = pool->current_offset;
          pool->entries[pool->entry_count].size = layout_size;
-         pool->entries[pool->entry_count].set = set;
-         pool->entry_count++;
       }
+      pool->entries[pool->entry_count].set = set;
       pool->current_offset += layout_size;
    } else if (!pool->host_memory_base) {
       uint64_t offset = 0;
@@ -693,7 +692,6 @@ radv_descriptor_set_create(struct radv_device *device, struct radv_descriptor_po
       pool->entries[index].offset = offset;
       pool->entries[index].size = layout_size;
       pool->entries[index].set = set;
-      pool->entry_count++;
    } else
       return VK_ERROR_OUT_OF_POOL_MEMORY;
 
@@ -716,6 +714,7 @@ radv_descriptor_set_create(struct radv_device *device, struct radv_descriptor_po
       }
    }
 
+   pool->entry_count++;
    radv_descriptor_set_layout_ref(layout);
    *out_set = set;
    return VK_SUCCESS;
@@ -725,10 +724,11 @@ static void
 radv_descriptor_set_destroy(struct radv_device *device, struct radv_descriptor_pool *pool,
                             struct radv_descriptor_set *set, bool free_bo)
 {
-   assert(!pool->host_memory_base);
-
    radv_descriptor_set_layout_unref(device, set->header.layout);
 
+   if (pool->host_memory_base)
+      return;
+
    if (free_bo && !pool->host_memory_base) {
       for (int i = 0; i < pool->entry_count; ++i) {
          if (pool->entries[i].set == set) {
@@ -747,10 +747,8 @@ static void
 radv_destroy_descriptor_pool(struct radv_device *device, const VkAllocationCallbacks *pAllocator,
                              struct radv_descriptor_pool *pool)
 {
-   if (!pool->host_memory_base) {
-      for (int i = 0; i < pool->entry_count; ++i) {
-         radv_descriptor_set_destroy(device, pool, pool->entries[i].set, false);
-      }
+   for (int i = 0; i < pool->entry_count; ++i) {
+      radv_descriptor_set_destroy(device, pool, pool->entries[i].set, false);
    }
 
    if (pool->bo)
@@ -844,13 +842,14 @@ radv_CreateDescriptorPool(VkDevice _device, const VkDescriptorPoolCreateInfo *pC
       }
    }
 
+   uint64_t entries_size = sizeof(struct radv_descriptor_pool_entry) * pCreateInfo->maxSets;
+   size += entries_size;
+
    if (!(pCreateInfo->flags & VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT)) {
       uint64_t host_size = pCreateInfo->maxSets * sizeof(struct radv_descriptor_set);
       host_size += sizeof(struct radeon_winsys_bo *) * bo_count;
       host_size += sizeof(struct radv_descriptor_range) * range_count;
       size += host_size;
-   } else {
-      size += sizeof(struct radv_descriptor_pool_entry) * pCreateInfo->maxSets;
    }
 
    pool = vk_alloc2(&device->vk.alloc, pAllocator, size, 8, VK_SYSTEM_ALLOCATION_SCOPE_OBJECT);
@@ -862,7 +861,7 @@ radv_CreateDescriptorPool(VkDevice _device, const VkDescriptorPoolCreateInfo *pC
    vk_object_base_init(&device->vk, &pool->base, VK_OBJECT_TYPE_DESCRIPTOR_POOL);
 
    if (!(pCreateInfo->flags & VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT)) {
-      pool->host_memory_base = (uint8_t *)pool + sizeof(struct radv_descriptor_pool);
+      pool->host_memory_base = (uint8_t *)pool + sizeof(struct radv_descriptor_pool) + entries_size;
       pool->host_memory_ptr = pool->host_memory_base;
       pool->host_memory_end = (uint8_t *)pool + size;
    }
@@ -924,12 +923,10 @@ radv_ResetDescriptorPool(VkDevice _device, VkDescriptorPool descriptorPool,
    RADV_FROM_HANDLE(radv_device, device, _device);
    RADV_FROM_HANDLE(radv_descriptor_pool, pool, descriptorPool);
 
-   if (!pool->host_memory_base) {
-      for (int i = 0; i < pool->entry_count; ++i) {
-         radv_descriptor_set_destroy(device, pool, pool->entries[i].set, false);
-      }
-      pool->entry_count = 0;
+   for (int i = 0; i < pool->entry_count; ++i) {
+      radv_descriptor_set_destroy(device, pool, pool->entries[i].set, false);
    }
+   pool->entry_count = 0;
 
    pool->current_offset = 0;
    pool->host_memory_ptr = pool->host_memory_base;