From 653393ee70f110719bde1f51ae71ac624c17dced Mon Sep 17 00:00:00 2001
From: "Bernhart, Bryan" <bryan.bernhart@intel.com>
Date: Mon, 21 Aug 2023 15:32:23 -0700
Subject: [PATCH] Never evict mapped resource allocations.

Tracks number of Map/Unmap calls to never evict the underlying resource heap.
---
 src/gpgmm/d3d12/ResourceAllocationD3D12.cpp   | 11 +++-
 src/gpgmm/d3d12/ResourceAllocationD3D12.h     |  3 +
 .../end2end/D3D12ResidencyManagerTests.cpp    | 62 +++++++++++++++++++
 3 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/src/gpgmm/d3d12/ResourceAllocationD3D12.cpp b/src/gpgmm/d3d12/ResourceAllocationD3D12.cpp
index b1d6b75c..bd0b4e13 100644
--- a/src/gpgmm/d3d12/ResourceAllocationD3D12.cpp
+++ b/src/gpgmm/d3d12/ResourceAllocationD3D12.cpp
@@ -63,6 +63,12 @@ namespace gpgmm::d3d12 {
 
     ResourceAllocation::~ResourceAllocation() {
         GPGMM_TRACE_EVENT_OBJECT_DESTROY(this);
+        if (mMappedCount.GetRefCount() > 0) {
+            WarnLog(MessageId::kPerformanceWarning, this)
+                << "Destroying a mapped resource allocation is allowed but discouraged. Please "
+                   "call Unmap the same number of times as Map before releasing the resource "
+                   "allocation.";
+        }
     }
 
     void ResourceAllocation::DeleteThis() {
@@ -105,6 +111,8 @@ namespace gpgmm::d3d12 {
         GPGMM_RETURN_IF_FAILED(mResource->Map(subresource, newReadRangePtr, &mappedData),
                                GetDevice(mResource.Get()));
 
+        mMappedCount.Ref();
+
         if (ppDataOut != nullptr) {
             *ppDataOut = static_cast<uint8_t*>(mappedData) + mOffsetFromResource;
         }
@@ -122,7 +130,8 @@ namespace gpgmm::d3d12 {
             return;
         }
 
-        if (mResidencyManager != nullptr) {
+        // Underlying heap cannot be evicted until the last Unmap.
+        if (mResidencyManager != nullptr && mMappedCount.Unref()) {
             mResidencyManager->UnlockHeap(GetMemory());
         }
 
diff --git a/src/gpgmm/d3d12/ResourceAllocationD3D12.h b/src/gpgmm/d3d12/ResourceAllocationD3D12.h
index caa6b357..d68d5ca3 100644
--- a/src/gpgmm/d3d12/ResourceAllocationD3D12.h
+++ b/src/gpgmm/d3d12/ResourceAllocationD3D12.h
@@ -82,6 +82,9 @@ namespace gpgmm::d3d12 {
         ComPtr<ID3D12Resource> mResource;
 
         const uint64_t mOffsetFromResource;
+
+        // Keeps track of the number of outstanding calls to Map to avoid paging-out those heaps.
+        RefCounted mMappedCount = RefCounted{0u};
     };
 
 }  // namespace gpgmm::d3d12
diff --git a/src/tests/end2end/D3D12ResidencyManagerTests.cpp b/src/tests/end2end/D3D12ResidencyManagerTests.cpp
index bb50ad0b..a4ba17af 100644
--- a/src/tests/end2end/D3D12ResidencyManagerTests.cpp
+++ b/src/tests/end2end/D3D12ResidencyManagerTests.cpp
@@ -703,6 +703,68 @@ TEST_F(D3D12ResidencyManagerTests, OverBudgetWithLockedHeaps) {
         bufferAllocationDesc, bufferDesc, D3D12_RESOURCE_STATE_COMMON, nullptr, nullptr));
 }
 
+// Keeps creating mapped resources until it reaches the restricted budget then over-commits to
+// ensure mapped resources cannot be evicted.
+TEST_F(D3D12ResidencyManagerTests, OverBudgetWithMappedResources) {
+    RESIDENCY_MANAGER_DESC residencyDesc = CreateBasicResidencyDesc(kDefaultBudget);
+
+    ComPtr<IResidencyManager> residencyManager;
+    ASSERT_SUCCEEDED(
+        CreateResidencyManager(residencyDesc, mDevice.Get(), mAdapter.Get(), &residencyManager));
+
+    ComPtr<IResourceAllocator> resourceAllocator;
+    ASSERT_SUCCEEDED(CreateResourceAllocator(CreateBasicAllocatorDesc(), mDevice.Get(),
+                                             mAdapter.Get(), residencyManager.Get(),
+                                             &resourceAllocator));
+
+    constexpr uint64_t kBufferMemorySize = GPGMM_MB_TO_BYTES(1);
+    const D3D12_RESOURCE_DESC bufferDesc = CreateBasicBufferDesc(kBufferMemorySize);
+
+    RESOURCE_ALLOCATION_DESC bufferAllocationDesc = {};
+    bufferAllocationDesc.HeapType = D3D12_HEAP_TYPE_UPLOAD;
+
+    // Keep allocating until we reach the budget.
+    std::vector<ComPtr<IResourceAllocation>> mappedAllocationsBelowBudget = {};
+    while (GetStats(resourceAllocator).UsedHeapUsage + kBufferMemorySize <= kDefaultBudget) {
+        ComPtr<IResourceAllocation> allocation;
+        ASSERT_SUCCEEDED(resourceAllocator->CreateResource(bufferAllocationDesc, bufferDesc,
+                                                           D3D12_RESOURCE_STATE_GENERIC_READ,
+                                                           nullptr, &allocation));
+
+        ASSERT_SUCCEEDED(allocation->Map(0, nullptr, nullptr));
+        mappedAllocationsBelowBudget.push_back(std::move(allocation));
+    }
+
+    // Mapped allocations should stay locked.
+    for (auto& allocation : mappedAllocationsBelowBudget) {
+        EXPECT_EQ(allocation->GetMemory()->GetInfo().Status, RESIDENCY_HEAP_STATUS_RESIDENT);
+        EXPECT_EQ(allocation->GetMemory()->GetInfo().IsLocked, true);
+    }
+
+    // Budget updates are not occuring frequently enough to detect going over budget will evict the
+    // same amount.
+    if (GetBudgetLeft(residencyManager.Get(), GetHeapSegment(bufferAllocationDesc.HeapType)) > 0) {
+        return;
+    }
+
+    // Since mapped resources are ineligable for eviction and RESIDENCY_HEAP_FLAG_CREATE_IN_BUDGET
+    // is true, CreateResource should always fail since there is not enough budget.
+    ASSERT_FAILED(resourceAllocator->CreateResource(
+        bufferAllocationDesc, bufferDesc, D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, nullptr));
+
+    // Unmapped allocations should be always eligable for eviction.
+    for (auto& allocation : mappedAllocationsBelowBudget) {
+        allocation->Unmap(0, nullptr);
+    }
+
+    for (auto& allocation : mappedAllocationsBelowBudget) {
+        EXPECT_EQ(allocation->GetMemory()->GetInfo().IsLocked, false);
+    }
+
+    ASSERT_SUCCEEDED(resourceAllocator->CreateResource(
+        bufferAllocationDesc, bufferDesc, D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, nullptr));
+}
+
 // Creates two sets of heaps, first set is below the budget, second set is above the budget, then
 // swaps the residency status using ExecuteCommandList: first set gets paged-in again, second set
 // gets paged-out.