panvk: Implement VK_EXT_memory_budget support

Signed-off-by: Christian Gmeiner <cgmeiner@igalia.com>
Reviewed-by: Christoph Pillmayer <christoph.pillmayer@arm.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/40246>

docs/features.txt

@@ -651,7 +651,7 @@ Khronos extensions that are not part of any Vulkan version:
   VK_EXT_legacy_vertex_attributes                       DONE (anv, lvp, nvk, radv, tu, vn)
   VK_EXT_line_rasterization                             DONE (anv, hasvk, hk, nvk, panvk, pvr, lvp, radv, tu, v3dv, vn)
   VK_EXT_load_store_op_none                             DONE (anv, hk, kk, lvp, nvk, panvk, radv, tu, v3dv, vn)
-  VK_EXT_memory_budget                                  DONE (anv, hasvk, lvp, nvk, radv, tu, v3dv, vn)
+  VK_EXT_memory_budget                                  DONE (anv, hasvk, lvp, nvk, panvk, radv, tu, v3dv, vn)
   VK_EXT_memory_priority                                DONE (lvp, radv)
   VK_EXT_mesh_shader                                    DONE (anv/gfx12.5+, lvp, radv, vn)
   VK_EXT_multi_draw                                     DONE (anv, hasvk, hk, lvp, nvk, radv, tu, vn, v3dv)

src/panfrost/vulkan/panvk_device_memory.c

@@ -14,6 +14,7 @@
 
 #include "pan_props.h"
 
+#include "util/u_atomic.h"
 #include "vk_debug_utils.h"
 #include "vk_log.h"
 
@@ -235,6 +236,9 @@ panvk_AllocateMemory(VkDevice _device,
 
    panvk_memory_emit_report(device, mem, pAllocateInfo, VK_SUCCESS);
 
+   p_atomic_add(&physical_device->memory.heap_used,
+                (uint64_t)pan_kmod_bo_size(mem->bo));
+
    *pMem = panvk_device_memory_to_handle(mem);
 
    return VK_SUCCESS;
@@ -264,6 +268,12 @@ panvk_FreeMemory(VkDevice _device, VkDeviceMemory _mem,
    if (mem == NULL)
       return;
 
+   struct panvk_physical_device *physical_device =
+      to_panvk_physical_device(device->vk.physical);
+
+   p_atomic_add(&physical_device->memory.heap_used,
+                -((int64_t)pan_kmod_bo_size(mem->bo)));
+
    if (device->debug.decode_ctx) {
       pandecode_inject_free(device->debug.decode_ctx, mem->addr.dev,
                             pan_kmod_bo_size(mem->bo));

src/panfrost/vulkan/panvk_physical_device.c

@@ -14,6 +14,7 @@
 
 #include "util/disk_cache.h"
 #include "util/os_misc.h"
+#include "util/u_atomic.h"
 #include "git_sha1.h"
 
 #include "vk_android.h"
@@ -620,6 +621,71 @@ panvk_GetPhysicalDeviceMemoryProperties2(
       pMemoryProperties->memoryProperties.memoryTypes[i] =
           physical_device->memory.types[i];
    }
+
+   vk_foreach_struct(ext, pMemoryProperties->pNext) {
+      switch (ext->sType) {
+      case VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_BUDGET_PROPERTIES_EXT: {
+         VkPhysicalDeviceMemoryBudgetPropertiesEXT *p = (void *)ext;
+
+         uint64_t used = p_atomic_read(&physical_device->memory.heap_used);
+         uint64_t heap_size = physical_device->memory.heaps[0].size;
+         uint64_t available;
+
+         if (!os_get_available_system_memory(&available))
+            available = heap_size;
+
+         /* From the Vulkan 1.3.278 spec:
+          *
+          *    "heapUsage is an array of VK_MAX_MEMORY_HEAPS VkDeviceSize
+          *    values in which memory usages are returned, with one element
+          *    for each memory heap. A heap’s usage is an estimate of how
+          *    much memory the process is currently using in that heap."
+          */
+         p->heapUsage[0] = used;
+
+         /* From the Vulkan 1.3.278 spec:
+          *
+          *    "heapBudget is an array of VK_MAX_MEMORY_HEAPS VkDeviceSize
+          *    values in which memory budgets are returned, with one
+          *    element for each memory heap. A heap’s budget is a rough
+          *    estimate of how much memory the process can allocate from
+          *    that heap before allocations may fail or cause performance
+          *    degradation. The budget includes any currently allocated
+          *    device memory."
+          *
+          * and
+          *
+          *    "The heapBudget value must be less than or equal to
+          *    VkMemoryHeap::size for each heap."
+          *
+          * available (queried above) is the total amount of free memory
+          * system-wide and does not include our allocations so we need
+          * to add that in.
+          */
+         uint64_t budget = MIN2(available + used, heap_size);
+
+         /* Set the budget at 90% of available to avoid thrashing */
+         p->heapBudget[0] = ROUND_DOWN_TO(budget * 9 / 10, 1 << 20);
+
+         /* From the Vulkan 1.3.278 spec:
+          *
+          *    "The heapBudget and heapUsage values must be zero for array
+          *    elements greater than or equal to
+          *    VkPhysicalDeviceMemoryProperties::memoryHeapCount. The
+          *    heapBudget value must be non-zero for array elements less than
+          *    VkPhysicalDeviceMemoryProperties::memoryHeapCount."
+          */
+         for (unsigned i = 1; i < VK_MAX_MEMORY_HEAPS; i++) {
+            p->heapBudget[i] = 0;
+            p->heapUsage[i] = 0;
+         }
+         break;
+      }
+      default:
+         vk_debug_ignored_stype(ext->sType);
+         break;
+      }
+   }
 }
 
 VKAPI_ATTR VkResult VKAPI_CALL

src/panfrost/vulkan/panvk_physical_device.h

@@ -65,6 +65,7 @@ struct panvk_physical_device {
       uint32_t type_count;
 
       uint64_t max_supported_va;
+      alignas(8) uint64_t heap_used;
    } memory;
 
    struct vk_sync_type drm_syncobj_type;

src/panfrost/vulkan/panvk_vX_physical_device.c

@@ -165,6 +165,7 @@ panvk_per_arch(get_physical_device_extensions)(
       .EXT_line_rasterization = true,
       .EXT_load_store_op_none = true,
       .EXT_nested_command_buffer = PAN_ARCH >= 10,
+      .EXT_memory_budget = true,
       .EXT_non_seamless_cube_map = true,
       .EXT_mutable_descriptor_type = PAN_ARCH >= 9,
       .EXT_multisampled_render_to_single_sampled = true,