kk: Complete VK_EXT_memory_budget

Metal provides device properties for the recommended maximum memory usage and
the current amount of memory used. These can be used to provide an estimate
of heap usage and calculate a budget of memory usage by the application before
performance may degrade.

Reviewed-by: Aitor Camacho <aitor@lunarg.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/41523>

docs/features.txt

@@ -660,7 +660,7 @@ Khronos extensions that are not part of any Vulkan version:
   VK_EXT_line_rasterization                             DONE (anv, hasvk, hk, kk, 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_map_memory_placed                              DONE (anv, hk, nvk, panvk, pvr, radv, tu, vn)
-  VK_EXT_memory_budget                                  DONE (anv, hasvk, lvp, nvk, panvk, radv, tu, v3dv, vn)
+  VK_EXT_memory_budget                                  DONE (anv, hasvk, kk, 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, kk, lvp, nvk, radv, tu, vn, v3dv)

src/kosmickrisp/bridge/mtl_device.h

@@ -32,6 +32,8 @@ bool mtl_device_supports_sample_count(mtl_device *dev, uint32_t sample_count);
 struct mtl_size mtl_device_max_threads_per_threadgroup(mtl_device *dev);
 uint32_t mtl_device_max_threadgroup_memory_length(mtl_device *dev);
 uint64_t mtl_device_max_buffer_length(mtl_device *dev);
+uint64_t mtl_device_recommended_max_working_set_size(mtl_device *dev);
+uint64_t mtl_device_current_allocated_size(mtl_device *dev);
 
 /* Timestamp query */
 uint64_t mtl_device_get_gpu_timestamp(mtl_device *dev);

src/kosmickrisp/bridge/mtl_device.m

@@ -156,6 +156,24 @@ mtl_device_max_buffer_length(mtl_device *dev)
    }
 }
 
+uint64_t
+mtl_device_recommended_max_working_set_size(mtl_device *dev)
+{
+   @autoreleasepool {
+      id<MTLDevice> device = (id<MTLDevice>)dev;
+      return device.recommendedMaxWorkingSetSize;
+   }
+}
+
+uint64_t
+mtl_device_current_allocated_size(mtl_device *dev)
+{
+   @autoreleasepool {
+      id<MTLDevice> device = (id<MTLDevice>)dev;
+      return device.currentAllocatedSize;
+   }
+}
+
 /* Timestamp query */
 uint64_t
 mtl_device_get_gpu_timestamp(mtl_device *dev)

src/kosmickrisp/bridge/stubs/mtl_device.c

@@ -77,6 +77,18 @@ mtl_device_max_buffer_length(mtl_device *dev)
    return 0u;
 }
 
+uint64_t
+mtl_device_recommended_max_working_set_size(mtl_device *dev)
+{
+   return 0u;
+}
+
+uint64_t
+mtl_device_current_allocated_size(mtl_device *dev)
+{
+   return 0u;
+}
+
 /* Timestamp query */
 uint64_t
 mtl_device_get_gpu_timestamp(mtl_device *dev)

src/kosmickrisp/vulkan/kk_device_memory.c

@@ -60,12 +60,10 @@ kk_AllocateMemory(VkDevice device, const VkMemoryAllocateInfo *pAllocateInfo,
                   const VkAllocationCallbacks *pAllocator, VkDeviceMemory *pMem)
 {
    VK_FROM_HANDLE(kk_device, dev, device);
-   struct kk_physical_device *pdev = kk_device_physical(dev);
    struct kk_device_memory *mem;
    VkResult result = VK_SUCCESS;
    const VkImportMemoryMetalHandleInfoEXT *metal_info = vk_find_struct_const(
       pAllocateInfo->pNext, IMPORT_MEMORY_METAL_HANDLE_INFO_EXT);
-   const VkMemoryType *type = &pdev->mem_types[pAllocateInfo->memoryTypeIndex];
 
    // TODO_KOSMICKRISP Do the actual memory allocation with alignment requirements
    uint32_t alignment = (1ULL << 12);
@@ -102,9 +100,6 @@ kk_AllocateMemory(VkDevice device, const VkMemoryAllocateInfo *pAllocateInfo,
          goto fail_alloc;
    }
 
-   struct kk_memory_heap *heap = &pdev->mem_heaps[type->heapIndex];
-   p_atomic_add(&heap->used, mem->bo->size_B);
-
    *pMem = kk_device_memory_to_handle(mem);
 
    return VK_SUCCESS;
@@ -120,15 +115,10 @@ kk_FreeMemory(VkDevice device, VkDeviceMemory _mem,
 {
    VK_FROM_HANDLE(kk_device, dev, device);
    VK_FROM_HANDLE(kk_device_memory, mem, _mem);
-   struct kk_physical_device *pdev = kk_device_physical(dev);
 
    if (!mem)
       return;
 
-   const VkMemoryType *type = &pdev->mem_types[mem->vk.memory_type_index];
-   struct kk_memory_heap *heap = &pdev->mem_heaps[type->heapIndex];
-   p_atomic_add(&heap->used, -((int64_t)mem->bo->size_B));
-
    kk_destroy_bo(dev, mem->bo);
 
    vk_device_memory_destroy(&dev->vk, pAllocator, &mem->vk);

src/kosmickrisp/vulkan/kk_physical_device.c

@@ -149,6 +149,7 @@ kk_get_device_extensions(const struct kk_instance *instance,
       .EXT_external_memory_metal = true,
       .EXT_image_2d_view_of_3d = true,
       .EXT_load_store_op_none = true,
+      .EXT_memory_budget = true,
       .EXT_multi_draw = true,
       .EXT_mutable_descriptor_type = true,
       .EXT_post_depth_coverage = true,
@@ -806,25 +807,80 @@ kk_physical_device_free_disk_cache(struct kk_physical_device *pdev)
 static uint64_t
 kk_get_sysmem_heap_size(void)
 {
+   /* Report the total amount of system memory as the actual heap size */
    uint64_t sysmem_size_B = 0;
    if (!os_get_total_physical_memory(&sysmem_size_B))
       return 0;
 
-   /* Use 3/4 of total size to avoid swapping */
-   return ROUND_DOWN_TO(sysmem_size_B * 3 / 4, 1 << 20);
+   return sysmem_size_B;
 }
 
 static uint64_t
-kk_get_sysmem_heap_available(struct kk_physical_device *pdev)
+kk_get_sysmem_heap_budget(struct kk_physical_device *pdev)
 {
+   /* 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."
+    *
+    * From Metal documentation for recommendedMaxWorkingSetSize:
+    *
+    *     An approximation of how much memory, in bytes, this GPU device can
+    *     allocate without affecting its runtime performance.
+    *
+    * From Metal documentation for currentAllocatedSize:
+    *
+    *     The total amount of memory, in bytes, the GPU device is using for all
+    *     of its resources.
+    *
+    * First, determine the total and available system memory to calculate the
+    * amount of used memory. Then, subtract this from the Metal-defined budget,
+    * and add back the current used memory by this device.
+    */
    uint64_t sysmem_size_B = 0;
-   if (!os_get_available_system_memory(&sysmem_size_B)) {
-      vk_loge(VK_LOG_OBJS(pdev), "Failed to query available system memory");
+   uint64_t sysmem_available_B = 0;
+   if (!os_get_total_physical_memory(&sysmem_size_B) ||
+       !os_get_available_system_memory(&sysmem_available_B))
       return 0;
-   }
 
-   /* Use 3/4 of available to avoid swapping */
-   return ROUND_DOWN_TO(sysmem_size_B * 3 / 4, 1 << 20);
+   uint64_t sysmem_used_B = sysmem_size_B - sysmem_available_B;
+   uint64_t sysmem_budget_B =
+      mtl_device_recommended_max_working_set_size(pdev->mtl_dev_handle);
+   uint64_t remaining_budget_B = sysmem_budget_B > sysmem_used_B ?
+                                 sysmem_budget_B - sysmem_used_B : 0u;
+   return remaining_budget_B +
+          mtl_device_current_allocated_size(pdev->mtl_dev_handle);
+}
+
+static uint64_t
+kk_get_sysmem_heap_used(struct kk_physical_device *pdev)
+{
+   /* 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."
+    *
+    * From Metal documentation for currentAllocatedSize:
+    *
+    *     The total amount of memory, in bytes, the GPU device is using for all
+    *     of its resources.
+    *
+    * We can trivially report estimated heap usage using Metal's reported
+    * allocated size
+    */
+   return mtl_device_current_allocated_size(pdev->mtl_dev_handle);
 }
 
 static void
@@ -901,7 +957,8 @@ kk_enumerate_physical_devices(struct vk_instance *_instance)
    pdev->mem_heaps[sysmem_heap_idx] = (struct kk_memory_heap){
       .size = sysmem_size_B,
       .flags = VK_MEMORY_HEAP_DEVICE_LOCAL_BIT,
-      .available = kk_get_sysmem_heap_available,
+      .budget = kk_get_sysmem_heap_budget,
+      .used = kk_get_sysmem_heap_used,
    };
 
    pdev->mem_types[pdev->mem_type_count++] = (VkMemoryType){
@@ -987,46 +1044,8 @@ kk_GetPhysicalDeviceMemoryProperties2(
 
          for (unsigned i = 0; i < pdev->mem_heap_count; i++) {
             const struct kk_memory_heap *heap = &pdev->mem_heaps[i];
-            uint64_t used = p_atomic_read(&heap->used);
-
-            /* 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."
-             *
-             * TODO: Include internal allocations?
-             */
-            p->heapUsage[i] = used;
-
-            uint64_t available = heap->size;
-            if (heap->available)
-               available = heap->available(pdev);
-
-            /* 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 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[i] = ROUND_DOWN_TO(budget * 9 / 10, 1 << 20);
+            p->heapBudget[i] = heap->budget ? heap->budget(pdev) : 0;
+            p->heapUsage[i] = heap->used ? heap->used(pdev) : 0;
          }
 
          /* From the Vulkan 1.3.278 spec:

src/kosmickrisp/vulkan/kk_physical_device.h

@@ -30,9 +30,9 @@ struct kk_queue_family {
 
 struct kk_memory_heap {
    uint64_t size;
-   uint64_t used;
    VkMemoryHeapFlags flags;
-   uint64_t (*available)(struct kk_physical_device *pdev);
+   uint64_t (*budget)(struct kk_physical_device *pdev);
+   uint64_t (*used)(struct kk_physical_device *pdev);
 };
 
 struct kk_device_info {