From da3318676f037c5bb80d4f2904709ad496e767ab Mon Sep 17 00:00:00 2001
From: Bryan Bernhart <bryan.bernhart@intel.com>
Date: Wed, 4 May 2022 11:31:38 -0700
Subject: [PATCH] Keep video memory budget up-to-date when tracing.

Budget change events were not being captured because video memory was only updated on residency creation or page-out.

This change keeps the video memory budget always up-to-date by calling it per ExecuteCommandLists when tracing is enabled.
---
 src/gpgmm/d3d12/ResidencyManagerD3D12.cpp | 36 ++++++++++++++---------
 src/gpgmm/d3d12/ResidencyManagerD3D12.h   |  4 +++
 2 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/src/gpgmm/d3d12/ResidencyManagerD3D12.cpp b/src/gpgmm/d3d12/ResidencyManagerD3D12.cpp
index 979a9a3c2..6a4896264 100644
--- a/src/gpgmm/d3d12/ResidencyManagerD3D12.cpp
+++ b/src/gpgmm/d3d12/ResidencyManagerD3D12.cpp
@@ -46,19 +46,8 @@ namespace gpgmm { namespace d3d12 {
         std::unique_ptr<ResidencyManager> residencyManager = std::unique_ptr<ResidencyManager>(
             new ResidencyManager(descriptor, std::move(residencyFence)));
 
-        // Query and set the video memory limits per segment.
-        DXGI_QUERY_VIDEO_MEMORY_INFO* queryVideoMemoryInfo =
-            residencyManager->GetVideoMemorySegmentInfo(DXGI_MEMORY_SEGMENT_GROUP_LOCAL);
-
-        ReturnIfFailed(residencyManager->QueryVideoMemoryInfo(DXGI_MEMORY_SEGMENT_GROUP_LOCAL,
-                                                              queryVideoMemoryInfo));
-        if (!descriptor.IsUMA) {
-            queryVideoMemoryInfo =
-                residencyManager->GetVideoMemorySegmentInfo(DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL);
-
-            ReturnIfFailed(residencyManager->QueryVideoMemoryInfo(
-                DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL, queryVideoMemoryInfo));
-        }
+        // Set the initial video memory limits per segment.
+        ReturnIfFailed(residencyManager->UpdateVideoMemorySegments());
 
         *residencyManagerOut = residencyManager.release();
 
@@ -73,7 +62,8 @@ namespace gpgmm { namespace d3d12 {
           mVideoMemoryBudget(descriptor.VideoMemoryBudget == 0 ? kDefaultVideoMemoryBudget
                                                                : descriptor.VideoMemoryBudget),
           mBudget(descriptor.Budget),
-          mEvictLimit(descriptor.EvictLimit == 0 ? kDefaultEvictLimit : descriptor.EvictLimit) {
+          mEvictLimit(descriptor.EvictLimit == 0 ? kDefaultEvictLimit : descriptor.EvictLimit),
+          mIsUMA(descriptor.IsUMA) {
         GPGMM_TRACE_EVENT_OBJECT_NEW(this);
 
         ASSERT(mDevice != nullptr);
@@ -290,6 +280,19 @@ namespace gpgmm { namespace d3d12 {
         return S_OK;
     }
 
+    HRESULT ResidencyManager::UpdateVideoMemorySegments() {
+        DXGI_QUERY_VIDEO_MEMORY_INFO* queryVideoMemoryInfo =
+            GetVideoMemorySegmentInfo(DXGI_MEMORY_SEGMENT_GROUP_LOCAL);
+
+        ReturnIfFailed(QueryVideoMemoryInfo(DXGI_MEMORY_SEGMENT_GROUP_LOCAL, queryVideoMemoryInfo));
+        if (!mIsUMA) {
+            queryVideoMemoryInfo = GetVideoMemorySegmentInfo(DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL);
+            ReturnIfFailed(
+                QueryVideoMemoryInfo(DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL, queryVideoMemoryInfo));
+        }
+        return S_OK;
+    }
+
     HRESULT ResidencyManager::Evict(uint64_t evictSizeInBytes,
                                     const DXGI_MEMORY_SEGMENT_GROUP& memorySegmentGroup) {
         std::lock_guard<std::mutex> lock(mMutex);
@@ -384,6 +387,11 @@ namespace gpgmm { namespace d3d12 {
             return E_NOTIMPL;
         }
 
+        // Ensure changes to video memory budget is recorded during tracing.
+        if (IsEventTraceEnabled()) {
+            ReturnIfFailed(UpdateVideoMemorySegments());
+        }
+
         ID3D12CommandList* commandList = commandLists[0];
         ResidencySet* residencySet = residencySets[0];
 
diff --git a/src/gpgmm/d3d12/ResidencyManagerD3D12.h b/src/gpgmm/d3d12/ResidencyManagerD3D12.h
index ff22332db..8fcf8cc8a 100644
--- a/src/gpgmm/d3d12/ResidencyManagerD3D12.h
+++ b/src/gpgmm/d3d12/ResidencyManagerD3D12.h
@@ -189,6 +189,9 @@ namespace gpgmm { namespace d3d12 {
         HRESULT QueryVideoMemoryInfo(const DXGI_MEMORY_SEGMENT_GROUP& memorySegmentGroup,
                                      DXGI_QUERY_VIDEO_MEMORY_INFO* pVideoMemoryInfo) const;
 
+        // Query and set the video memory limits for all segments.
+        HRESULT UpdateVideoMemorySegments();
+
         ComPtr<ID3D12Device> mDevice;
         ComPtr<IDXGIAdapter3> mAdapter;
         ComPtr<ID3D12Device3> mDevice3;
@@ -198,6 +201,7 @@ namespace gpgmm { namespace d3d12 {
         const float mVideoMemoryBudget;
         const uint64_t mBudget;
         const uint64_t mEvictLimit;
+        const bool mIsUMA;
 
         VideoMemorySegment mLocalVideoMemorySegment;
         VideoMemorySegment mNonLocalVideoMemorySegment;