From 514e0b0d787f476d58c8976d48ed2f432ed2b66c Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Mon, 25 Nov 2024 12:40:10 +0000
Subject: [PATCH] Reduce time searching for suballocs

---
 renderdoc/driver/d3d12/d3d12_manager.cpp |  8 ++-
 renderdoc/driver/d3d12/d3d12_manager.h   | 63 +++++++++++++++---------
 2 files changed, 45 insertions(+), 26 deletions(-)

diff --git a/renderdoc/driver/d3d12/d3d12_manager.cpp b/renderdoc/driver/d3d12/d3d12_manager.cpp
index 09fb22c08..417c08dba 100644
--- a/renderdoc/driver/d3d12/d3d12_manager.cpp
+++ b/renderdoc/driver/d3d12/d3d12_manager.cpp
@@ -2747,7 +2747,7 @@ bool D3D12GpuBufferAllocator::D3D12GpuBufferResource::ReleaseGpuBufferResource(
 
 D3D12GpuBufferAllocator::D3D12GpuBufferResource::D3D12GpuBufferResource(ID3D12Resource *resource,
                                                                         D3D12_HEAP_TYPE heapType)
-    : m_resource(resource), m_heapType(heapType)
+    : m_resource(resource), m_heapType(heapType), m_subRanges(), m_lastFree(m_subRanges.begin())
 
 {
   if(m_resource)
@@ -2755,6 +2755,7 @@ D3D12GpuBufferAllocator::D3D12GpuBufferResource::D3D12GpuBufferResource(ID3D12Re
     m_resDesc = m_resource->GetDesc();
     m_resourceGpuAddressRange.start = resource->GetGPUVirtualAddress();
     m_resourceGpuAddressRange.realEnd = m_resourceGpuAddressRange.start + m_resDesc.Width;
+    m_bytesFree = m_resDesc.Width;
   }
 }
 
@@ -2774,6 +2775,9 @@ bool D3D12GpuBufferAllocator::D3D12GpuBufferPool::Alloc(WrappedID3D12Device *wra
     D3D12_GPU_VIRTUAL_ADDRESS gpuAddress = 0;
     for(D3D12GpuBufferResource *bufferRes : m_bufferResourceList)
     {
+      if(size > bufferRes->m_bytesFree)
+        continue;
+
       if(bufferRes->SubAlloc(size, alignment, gpuAddress))
       {
         *gpuBuffer = new D3D12GpuBuffer(allocator, m_bufferPoolHeapType,
@@ -2836,7 +2840,7 @@ void D3D12GpuBufferAllocator::D3D12GpuBufferPool::Free(const D3D12GpuBuffer &gpu
       {
         if(bufferRes->SubAllocationInRange(gpuBuffer.Address()))
         {
-          if(!bufferRes->Free(gpuBuffer.Address()))
+          if(!bufferRes->Free(gpuBuffer.Address(), gpuBuffer.Size(), gpuBuffer.Alignment()))
           {
             RDCERR("Invalid address when freeing buffer");
           }
diff --git a/renderdoc/driver/d3d12/d3d12_manager.h b/renderdoc/driver/d3d12/d3d12_manager.h
index f962da22a..de8c107db 100644
--- a/renderdoc/driver/d3d12/d3d12_manager.h
+++ b/renderdoc/driver/d3d12/d3d12_manager.h
@@ -859,28 +859,37 @@ private:
 
     bool SubAllocationInRange(D3D12_GPU_VIRTUAL_ADDRESS gpuAddress) const
     {
-      if(m_resourceGpuAddressRange.start <= gpuAddress &&
-         gpuAddress < m_resourceGpuAddressRange.realEnd)
-
-      {
-        return true;
-      }
-
-      return false;
+      return (m_resourceGpuAddressRange.start <= gpuAddress &&
+              gpuAddress < m_resourceGpuAddressRange.realEnd);
     }
 
-    bool Free(D3D12_GPU_VIRTUAL_ADDRESS gpuAddress)
+    bool Free(D3D12_GPU_VIRTUAL_ADDRESS gpuAddress, uint64_t size, uint64_t alignment)
     {
       uint64_t offset = gpuAddress - m_resourceGpuAddressRange.start;
       auto iter = m_subRanges.find(offset);
       if(iter != m_subRanges.end() && iter->value() == D3D12SubRangeFlag::Used)
       {
+        uint64_t iterOffset = iter->start();
+        uint64_t alignedOffset = iterOffset;
+        if(alignment)
+          alignedOffset = AlignUp(m_resourceGpuAddressRange.start + alignedOffset, alignment) -
+                          m_resourceGpuAddressRange.start;
+
+        uint64_t padding = alignedOffset - iterOffset;
+
+        m_bytesFree += size + padding;
         iter->setValue(D3D12SubRangeFlag::Free);
         // Merging will only occur if the adjacent sub-ranges are also free
         iter->mergeLeft();
+        m_lastFree = iter;
+
         ++iter;
         if(iter != m_subRanges.end())
+        {
           iter->mergeLeft();
+          m_lastFree = iter;
+        }
+
         return true;
       }
       return false;
@@ -890,31 +899,35 @@ private:
     {
       uint64_t resourceWidth = m_resourceGpuAddressRange.realEnd - m_resourceGpuAddressRange.start;
 
-      for(auto iter = m_subRanges.begin(); iter != m_subRanges.end(); ++iter)
+      for(auto iter = m_lastFree; iter != m_subRanges.end(); ++iter)
       {
         if(iter->value() == D3D12SubRangeFlag::Free)
         {
-          uint64_t addr = iter->start() + m_resourceGpuAddressRange.start;
-          uint64_t end = RDCMIN(iter->finish(), resourceWidth) + m_resourceGpuAddressRange.start;
-          uint64_t alignedAddr = alignment != 0 ? AlignUp(addr, alignment) : addr;
+          uint64_t freeRangeStart = iter->start();
+          uint64_t freeRangeEnd = RDCMIN(iter->finish(), resourceWidth);
+          uint64_t alignedStart = freeRangeStart;
 
-          if(alignedAddr < end && size <= (end - alignedAddr))
+          if(alignment)
+            alignedStart = AlignUp(m_resourceGpuAddressRange.start + alignedStart, alignment) -
+                           m_resourceGpuAddressRange.start;
+
+          uint64_t padding = alignedStart - freeRangeStart;
+
+          if(alignedStart < freeRangeEnd && alignedStart + size <= freeRangeEnd)
           {
-            uint64_t offset = alignedAddr - m_resourceGpuAddressRange.start;
-            // Free the extra space from aligning
-            if(alignedAddr > addr)
-            {
-              iter->split(offset);
-            }
-
             iter->setValue(D3D12SubRangeFlag::Used);
-            address = alignedAddr;
+            address = m_resourceGpuAddressRange.start + alignedStart;
             // Split the sub-range if there's extra space beyond this allocation
-            if(size < (end - alignedAddr))
+            if(alignedStart + size < freeRangeEnd)
             {
-              iter->split(offset + size);
+              iter->split(alignedStart + size);
               iter->setValue(D3D12SubRangeFlag::Free);
             }
+
+            m_bytesFree -= size + padding;
+
+            m_lastFree = iter;
+
             return true;
           }
         }
@@ -929,10 +942,12 @@ private:
     };
 
     Intervals<D3D12SubRangeFlag> m_subRanges;
+    Intervals<D3D12SubRangeFlag>::iterator m_lastFree;
     GPUAddressRange m_resourceGpuAddressRange;
     ID3D12Resource *m_resource;
     D3D12_RESOURCE_DESC m_resDesc;
     D3D12_HEAP_TYPE m_heapType;
+    uint64_t m_bytesFree;
   };
 
   class D3D12GpuBufferPool