Update budget by OS budget change events.

.github/workflows/.patches/dawn.diff

@@ -1,4 +1,4 @@
-From f2724b093745556795b76ad7214b694f3f93ed37 Mon Sep 17 00:00:00 2001
+From e41eefd8175a3cd7125699987d452dbfac72fb44 Mon Sep 17 00:00:00 2001
 From: Bryan Bernhart <bryan.bernhart@intel.com>
 Date: Tue, 15 Feb 2022 17:25:29 -0800
 Subject: [PATCH] Use GPGMM for D3D12 backend.
@@ -495,7 +495,7 @@ index 9747b41c0..50db56dbd 100644
 
  AdapterDiscoveryOptions::AdapterDiscoveryOptions()
 diff --git a/src/dawn/native/d3d12/DeviceD3D12.cpp b/src/dawn/native/d3d12/DeviceD3D12.cpp
-index 90ddc3041..c82b3fb90 100644
+index 90ddc3041..8545685aa 100644
 --- a/src/dawn/native/d3d12/DeviceD3D12.cpp
 +++ b/src/dawn/native/d3d12/DeviceD3D12.cpp
 @@ -130,8 +130,43 @@ MaybeError Device::Initialize(const DeviceDescriptor* descriptor) {
@@ -514,32 +514,32 @@ index 90ddc3041..c82b3fb90 100644
 +        static_cast<D3D12_RESOURCE_HEAP_TIER>(adapter->GetDeviceInfo().resourceHeapTier);
 +    allocatorDesc.PreferredResourceHeapSize = 4ll * 1024ll * 1024ll;  // 4MB
 +
-+    gpgmm::d3d12::ResidencyManager** residencyManager = nullptr;
-+    if (IsToggleEnabled(Toggle::UseD3D12ResidencyManagement)) {
-+        allocatorDesc.MaxVideoMemoryBudget = 0.95;  // Use up to 95%.
-+        residencyManager = mResidencyManager.GetAddressOf();
-+    }
++    gpgmm::d3d12::RESIDENCY_DESC residencyDesc = {};
++    residencyDesc.Adapter = adapter->GetHardwareAdapter();
++    residencyDesc.Device = mD3d12Device;
 +
-+    if (residencyManager != nullptr &&
-+        IsToggleEnabled(Toggle::UseD3D12SmallResidencyBudgetForTesting)) {
-+        allocatorDesc.Budget = 100000000;  // 100MB
++    if (IsToggleEnabled(Toggle::UseD3D12SmallResidencyBudgetForTesting)) {
++        residencyDesc.UpdateBudgetByPolling = true;
++        residencyDesc.Budget = 100000000;  // 100MB
 +        allocatorDesc.Flags |= gpgmm::d3d12::ALLOCATOR_FLAG_DISABLE_MEMORY_PREFETCH;
 +        allocatorDesc.Flags |= gpgmm::d3d12::ALLOCATOR_FLAG_ALWAYS_IN_BUDGET;
 +    }
 +
++    if (IsToggleEnabled(Toggle::UseD3D12ResidencyManagement)) {
++        residencyDesc.VideoMemoryBudget = 0.95;  // Use up to 95%
++        DAWN_TRY(CheckHRESULT(gpgmm::d3d12::ResidencyManager::CreateResidencyManager(
++                                  residencyDesc, &mResidencyManager),
++                              "D3D12 create residency manager"));
++    }
++
 +    if (IsToggleEnabled(Toggle::DumpResourceAllocator)) {
 +        allocatorDesc.RecordOptions.Flags |= gpgmm::d3d12::ALLOCATOR_RECORD_FLAG_ALL_EVENTS;
-+        allocatorDesc.RecordOptions.MinMessageLevel = D3D12_MESSAGE_SEVERITY_MESSAGE;
 +        allocatorDesc.RecordOptions.EventScope = gpgmm::d3d12::ALLOCATOR_RECORD_SCOPE_PER_INSTANCE;
 +        allocatorDesc.RecordOptions.TraceFile = "dawn_resource_allocator_dump.json";
 +    }
 +
-+    if (IsToggleEnabled(Toggle::RecordDetailedTimingInTraceEvents)) {
-+        allocatorDesc.RecordOptions.UseDetailedTimingEvents = true;
-+    }
-+
 +    DAWN_TRY(CheckHRESULT(gpgmm::d3d12::ResourceAllocator::CreateAllocator(
-+                              allocatorDesc, &mResourceAllocator, residencyManager),
++                              allocatorDesc, mResidencyManager.Get(), &mResourceAllocator),
 +                          "D3D12 create resource allocator"));
 
      // ShaderVisibleDescriptorAllocators use the ResidencyManager and must be initialized after.

src/gpgmm/common/EventMessage.h

@@ -30,7 +30,8 @@ namespace gpgmm {
         AlignmentMismatch,
         AllocatorFailed,
         PrefetchFailed,
-        BudgetExceeded
+        BudgetExceeded,
+        BudgetUpdate,
     };
 
     struct EventMessageInfo {

src/gpgmm/d3d12/ResidencyManagerD3D12.cpp

@@ -18,6 +18,7 @@
 #include "gpgmm/common/EventMessage.h"
 #include "gpgmm/common/SizeClass.h"
 #include "gpgmm/common/TraceEvent.h"
+#include "gpgmm/common/WorkerThread.h"
 #include "gpgmm/d3d12/ErrorD3D12.h"
 #include "gpgmm/d3d12/FenceD3D12.h"
 #include "gpgmm/d3d12/HeapD3D12.h"
@@ -33,6 +34,115 @@ namespace gpgmm::d3d12 {
     static constexpr uint32_t kDefaultEvictBatchSize = GPGMM_MB_TO_BYTES(50);
     static constexpr float kDefaultVideoMemoryBudget = 0.95f;               // 95%
 
+    // Creates a long-lived task to recieve and process OS budget change events.
+    class BudgetUpdateTask : public VoidCallback {
+      public:
+        BudgetUpdateTask(ResidencyManager* const residencyManager, ComPtr<IDXGIAdapter3> adapter)
+            : mResidencyManager(residencyManager),
+              mAdapter(std::move(adapter)),
+              mBudgetNotificationUpdateEvent(CreateEventW(NULL, FALSE, FALSE, NULL)),
+              mUnregisterAndExitEvent(CreateEventW(NULL, FALSE, FALSE, NULL)) {
+            ASSERT(mResidencyManager != nullptr);
+            ASSERT(mAdapter != nullptr);
+            mLastError = mAdapter->RegisterVideoMemoryBudgetChangeNotificationEvent(
+                mBudgetNotificationUpdateEvent, &mCookie);
+        }
+
+        void operator()() override {
+            HRESULT hr = GetLastError();
+            bool isExiting = false;
+            while (!isExiting && SUCCEEDED(hr)) {
+                // Wait on two events: one to unblock for OS budget changes, and another to unblock
+                // for shutdown.
+                HANDLE hWaitEvents[2] = {mBudgetNotificationUpdateEvent, mUnregisterAndExitEvent};
+                const DWORD waitedEvent =
+                    WaitForMultipleObjects(2, hWaitEvents, /*bWaitAll*/ false, INFINITE);
+                switch (waitedEvent) {
+                    // mBudgetNotificationUpdateEvent
+                    case (WAIT_OBJECT_0 + 0): {
+                        hr = mResidencyManager->UpdateVideoMemorySegments();
+                        if (SUCCEEDED(hr)) {
+                            gpgmm::InfoEvent("ResidencyManager", EventMessageId::BudgetUpdate)
+                                << "Updated video budget from OS notification.";
+                        }
+                        break;
+                    }
+                    // mUnregisterAndExitEvent
+                    case (WAIT_OBJECT_0 + 1): {
+                        isExiting = true;
+                        break;
+                    }
+                    default: {
+                        UNREACHABLE();
+                        break;
+                    }
+                }
+            }
+
+            SetLastError(hr);
+        }
+
+        HRESULT GetLastError() const {
+            std::lock_guard<std::mutex> lock(mMutex);
+            return mLastError;
+        }
+
+        // Shutdown the event loop.
+        bool UnregisterAndExit() {
+            mAdapter->UnregisterVideoMemoryBudgetChangeNotification(mCookie);
+            return SetEvent(mUnregisterAndExitEvent);
+        }
+
+      private:
+        void SetLastError(HRESULT hr) {
+            std::lock_guard<std::mutex> lock(mMutex);
+            mLastError = hr;
+        }
+
+        ResidencyManager* const mResidencyManager;
+        ComPtr<IDXGIAdapter3> mAdapter;
+
+        HANDLE mBudgetNotificationUpdateEvent = INVALID_HANDLE_VALUE;
+        HANDLE mUnregisterAndExitEvent = INVALID_HANDLE_VALUE;
+
+        DWORD mCookie = 0;  // Used to unregister from notifications.
+
+        mutable std::mutex mMutex;  // Protect access between threads for members below.
+        HRESULT mLastError = S_OK;
+    };
+
+    class BudgetUpdateEvent final : public Event {
+      public:
+        BudgetUpdateEvent(std::shared_ptr<Event> event, std::shared_ptr<BudgetUpdateTask> task)
+            : mTask(task), mEvent(event) {
+        }
+
+        // Event overrides
+        void Wait() override {
+            mEvent->Wait();
+        }
+
+        bool IsSignaled() override {
+            return mEvent->IsSignaled();
+        }
+
+        void Signal() override {
+            return mEvent->Signal();
+        }
+
+        bool UnregisterAndExit() {
+            return mTask->UnregisterAndExit();
+        }
+
+        bool GetLastError() const {
+            return mTask->GetLastError();
+        }
+
+      private:
+        std::shared_ptr<BudgetUpdateTask> mTask;
+        std::shared_ptr<Event> mEvent;
+    };
+
     // static
     HRESULT ResidencyManager::CreateResidencyManager(const RESIDENCY_DESC& descriptor,
                                                      ResidencyManager** residencyManagerOut) {
@@ -54,6 +164,11 @@ namespace gpgmm::d3d12 {
         std::unique_ptr<ResidencyManager> residencyManager = std::unique_ptr<ResidencyManager>(
             new ResidencyManager(descriptor, std::move(residencyFence)));
 
+        // Require automatic video memory budget updates.
+        if (!descriptor.UpdateBudgetByPolling) {
+            ReturnIfFailed(residencyManager->StartBudgetNotificationUpdates());
+        }
+
         // Set the initial video memory limits per segment.
         ReturnIfFailed(residencyManager->UpdateVideoMemorySegments());
 
@@ -94,7 +209,9 @@ namespace gpgmm::d3d12 {
           mIsBudgetRestricted(descriptor.Budget > 0),
           mEvictBatchSize(descriptor.EvictBatchSize == 0 ? kDefaultEvictBatchSize
                                                          : descriptor.EvictBatchSize),
-          mIsUMA(descriptor.IsUMA) {
+          mIsUMA(descriptor.IsUMA),
+          mIsBudgetChangeEventsDisabled(descriptor.UpdateBudgetByPolling),
+          mThreadPool(ThreadPool::Create()) {
         GPGMM_TRACE_EVENT_OBJECT_NEW(this);
 
         ASSERT(mDevice != nullptr);
@@ -104,6 +221,7 @@ namespace gpgmm::d3d12 {
 
     ResidencyManager::~ResidencyManager() {
         GPGMM_TRACE_EVENT_OBJECT_DESTROY(this);
+        StopBudgetNotificationUpdates();
     }
 
     const char* ResidencyManager::GetTypename() const {
@@ -293,6 +411,11 @@ namespace gpgmm::d3d12 {
     }
 
     HRESULT ResidencyManager::UpdateVideoMemorySegments() {
+        std::lock_guard<std::mutex> lock(mMutex);
+        return UpdateVideoMemorySegmentsInternal();
+    }
+
+    HRESULT ResidencyManager::UpdateVideoMemorySegmentsInternal() {
         DXGI_QUERY_VIDEO_MEMORY_INFO* queryVideoMemoryInfo =
             GetVideoMemoryInfo(DXGI_MEMORY_SEGMENT_GROUP_LOCAL);
 
@@ -320,7 +443,10 @@ namespace gpgmm::d3d12 {
 
         DXGI_QUERY_VIDEO_MEMORY_INFO* videoMemorySegmentInfo =
             GetVideoMemoryInfo(memorySegmentGroup);
-        ReturnIfFailed(QueryVideoMemoryInfo(memorySegmentGroup, videoMemorySegmentInfo));
+
+        if (IsBudgetNotificationUpdatesDisabled()) {
+            ReturnIfFailed(QueryVideoMemoryInfo(memorySegmentGroup, videoMemorySegmentInfo));
+        }
 
         const uint64_t currentUsageAfterEvict =
             evictSizeInBytes + videoMemorySegmentInfo->CurrentUsage;
@@ -403,19 +529,15 @@ namespace gpgmm::d3d12 {
 
         std::lock_guard<std::mutex> lock(mMutex);
 
+        if (count == 0) {
+            return E_INVALIDARG;
+        }
+
         // TODO: support multiple command lists.
         if (count > 1) {
             return E_NOTIMPL;
         }
 
-        // Keep video memory segments up-to-date. This must always happen because tests may
-        // want to execute a no-op to get the current budget and usage.
-        ReturnIfFailed(UpdateVideoMemorySegments());
-
-        if (count == 0) {
-            return S_OK;
-        }
-
         ResidencySet* residencySet = residencySets[0];
 
         std::vector<ID3D12Pageable*> localHeapsToMakeResident;
@@ -480,6 +602,13 @@ namespace gpgmm::d3d12 {
             ReturnIfFailed(mFence->Signal(queue));
         }
 
+        // Keep video memory segments up-to-date. This must always happen because if the budget
+        // never changes (ie. not manually updated or through budget change events), the
+        // residency manager wouldn't know what to page in or out.
+        if (IsBudgetNotificationUpdatesDisabled()) {
+            ReturnIfFailed(UpdateVideoMemorySegmentsInternal());
+        }
+
         GPGMM_TRACE_EVENT_OBJECT_CALL("ResidencyManager.ExecuteCommandLists",
                                       (EXECUTE_COMMAND_LISTS_DESC{residencySets, count}));
 
@@ -540,4 +669,36 @@ namespace gpgmm::d3d12 {
         return info;
     }
 
+    // Starts updating video memory budget from OS notifications.
+    // Return True if successfully registered or False if error.
+    HRESULT ResidencyManager::StartBudgetNotificationUpdates() {
+        if (mBudgetNotificationUpdateEvent == nullptr) {
+            std::shared_ptr<BudgetUpdateTask> task =
+                std::make_shared<BudgetUpdateTask>(this, mAdapter);
+            mBudgetNotificationUpdateEvent = std::make_shared<BudgetUpdateEvent>(
+                ThreadPool::PostTask(mThreadPool, task, "GPGMM_ThreadBudgetChangeWorker"), task);
+        }
+
+        ASSERT(mBudgetNotificationUpdateEvent != nullptr);
+        return mBudgetNotificationUpdateEvent->GetLastError();
+    }
+
+    bool ResidencyManager::IsBudgetNotificationUpdatesDisabled() const {
+        return mIsBudgetChangeEventsDisabled ||
+               (mBudgetNotificationUpdateEvent != nullptr &&
+                FAILED(mBudgetNotificationUpdateEvent->GetLastError()));
+    }
+
+    void ResidencyManager::StopBudgetNotificationUpdates() {
+        if (mBudgetNotificationUpdateEvent == nullptr) {
+            return;
+        }
+
+        const bool success = mBudgetNotificationUpdateEvent->UnregisterAndExit();
+        ASSERT(success);
+
+        mBudgetNotificationUpdateEvent->Wait();
+        mBudgetNotificationUpdateEvent = nullptr;
+    }
+
 }  // namespace gpgmm::d3d12