From 513312885e9f6dfc78abdb084c915dda0b11ce48 Mon Sep 17 00:00:00 2001
From: Philip Rebohle <philip.rebohle@tu-dortmund.de>
Date: Thu, 3 Oct 2024 01:02:26 +0200
Subject: [PATCH] [dxvk] Use actual memory budget if provided by the
 implementation

For now, this is only used for the heuristic on when to aggressively
free empty chunks and replaces the 80% heap size heuristic.

Periodically updates the memory budget from the worker thread.
---
 src/dxvk/dxvk_memory.cpp | 27 +++++++++++++++++++++++++--
 src/dxvk/dxvk_memory.h   |  3 +++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/dxvk/dxvk_memory.cpp b/src/dxvk/dxvk_memory.cpp
index c23664ff9..ec28b14ff 100644
--- a/src/dxvk/dxvk_memory.cpp
+++ b/src/dxvk/dxvk_memory.cpp
@@ -435,7 +435,7 @@ namespace dxvk {
 
   DxvkMemoryAllocator::DxvkMemoryAllocator(DxvkDevice* device)
   : m_device(device) {
-    VkPhysicalDeviceMemoryProperties memInfo = device->adapter()->memoryProperties();
+    VkPhysicalDeviceMemoryProperties memInfo = m_device->adapter()->memoryProperties();
 
     m_memTypeCount = memInfo.memoryTypeCount;
     m_memHeapCount = memInfo.memoryHeapCount;
@@ -444,6 +444,7 @@ namespace dxvk {
       auto& heap = m_memHeaps[i];
 
       heap.index = i;
+      heap.memoryBudget = memInfo.memoryHeaps[i].size;
       heap.properties = memInfo.memoryHeaps[i];
     }
 
@@ -466,6 +467,8 @@ namespace dxvk {
 
     determineBufferUsageFlagsPerMemoryType();
 
+    updateMemoryHeapBudgets();
+
     // Start worker after setting up everything else
     m_worker = dxvk::thread([this] { runWorker(); });
   }
@@ -1248,7 +1251,7 @@ namespace dxvk {
       maxUnusedMemory *= 4u;
 
     // Factor current memory allocation into the decision to free chunks
-    VkDeviceSize heapBudget = (type.heap->properties.size * 4) / 5;
+    VkDeviceSize heapBudget = type.heap->memoryBudget;
     VkDeviceSize heapAllocated = getMemoryStats(type.heap->index).memoryAllocated;
 
     VkDeviceSize unusedMemory = 0u;
@@ -1836,6 +1839,23 @@ namespace dxvk {
   }
 
 
+  void DxvkMemoryAllocator::updateMemoryHeapBudgets() {
+    if (!m_device->features().extMemoryBudget)
+      return;
+
+    VkPhysicalDeviceMemoryBudgetPropertiesEXT memBudget = { VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_BUDGET_PROPERTIES_EXT };
+    VkPhysicalDeviceMemoryProperties2 memInfo = { VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_PROPERTIES_2, &memBudget };
+
+    auto vki = m_device->adapter()->vki();
+    vki->vkGetPhysicalDeviceMemoryProperties2(m_device->adapter()->handle(), &memInfo);
+
+    for (uint32_t i = 0; i < m_memHeapCount; i++) {
+      if (memBudget.heapBudget[i])
+        m_memHeaps[i].memoryBudget = std::min(memBudget.heapBudget[i], m_memHeaps[i].properties.size);
+    }
+  }
+
+
   void DxvkMemoryAllocator::runWorker() {
     env::setThreadName("dxvk-memory");
 
@@ -1851,6 +1871,9 @@ namespace dxvk {
       if (m_stopWorker)
         break;
 
+      // Re-query current memory budgets
+      updateMemoryHeapBudgets();
+
       // Periodically free unused memory chunks and update
       // memory allocation statistics for the adapter.
       auto currentTime = high_resolution_clock::now();
diff --git a/src/dxvk/dxvk_memory.h b/src/dxvk/dxvk_memory.h
index d723c1ac2..c7807989d 100644
--- a/src/dxvk/dxvk_memory.h
+++ b/src/dxvk/dxvk_memory.h
@@ -138,6 +138,7 @@ namespace dxvk {
   struct DxvkMemoryHeap {
     uint32_t          index         = 0u;
     uint32_t          memoryTypes   = 0u;
+    VkDeviceSize      memoryBudget  = 0u;
     VkMemoryHeap      properties    = { };
   };
 
@@ -1239,6 +1240,8 @@ namespace dxvk {
     uint32_t findGlobalBufferMemoryTypeMask(
             VkBufferUsageFlags    usage) const;
 
+    void updateMemoryHeapBudgets();
+
     void runWorker();
 
   };