Move isMemoryBudgetExhausted to WddmResidencyController

Change-Id: Ic9da29ab954835e93cfbcd6690c1764d99939613
Signed-off-by: Maciej Dziuban <[email protected]>

Author: DziubanMaciejIntel

runtime/os_interface/windows/wddm_memory_manager.cpp

@@ -342,6 +342,15 @@ bool WddmMemoryManager::tryDeferDeletions(D3DKMT_HANDLE *handles, uint32_t alloc
     return status;
 }
 
+bool WddmMemoryManager::isMemoryBudgetExhausted() const {
+    for (auto osContext : this->registeredOsContexts) {
+        if (osContext != nullptr && osContext->get()->getResidencyController().isMemoryBudgetExhausted()) {
+            return true;
+        }
+    }
+    return false;
+}
+
 bool WddmMemoryManager::validateAllocation(WddmAllocation *alloc) {
     if (alloc == nullptr)
         return false;
@@ -491,7 +500,7 @@ bool WddmMemoryManager::makeResidentResidencyAllocations(ResidencyContainer &all
     if (totalHandlesCount) {
         uint64_t bytesToTrim = 0;
         while ((result = wddm->makeResident(handlesForResidency.get(), totalHandlesCount, false, &bytesToTrim)) == false) {
-            this->memoryBudgetExhausted = true;
+            osContext.get()->getResidencyController().setMemoryBudgetExhausted();
             bool trimmingDone = this->getRegisteredOsContext(0u)->get()->getResidencyController().trimResidencyToBudget(bytesToTrim);
             bool cantTrimFurther = !trimmingDone;
             if (cantTrimFurther) {

runtime/os_interface/windows/wddm_memory_manager.h

@@ -72,7 +72,7 @@ class WddmMemoryManager : public MemoryManager {
 
     bool tryDeferDeletions(D3DKMT_HANDLE *handles, uint32_t allocationCount, D3DKMT_HANDLE resourceHandle);
 
-    bool isMemoryBudgetExhausted() const override { return memoryBudgetExhausted; }
+    bool isMemoryBudgetExhausted() const override;
 
     bool mapAuxGpuVA(GraphicsAllocation *graphicsAllocation) override;
 
@@ -82,7 +82,6 @@ class WddmMemoryManager : public MemoryManager {
     GraphicsAllocation *createAllocationFromHandle(osHandle handle, bool requireSpecificBitness, bool ntHandle);
     static bool validateAllocation(WddmAllocation *alloc);
     bool createWddmAllocation(WddmAllocation *allocation, AllocationOrigin origin);
-    bool memoryBudgetExhausted = false;
     AlignedMallocRestrictions mallocRestrictions;
 
   private:

runtime/os_interface/windows/wddm_residency_controller.h

@@ -49,6 +49,9 @@ class WddmResidencyController {
     void trimResidency(D3DDDI_TRIMRESIDENCYSET_FLAGS flags, uint64_t bytes);
     bool trimResidencyToBudget(uint64_t bytes);
 
+    bool isMemoryBudgetExhausted() const { return memoryBudgetExhausted; }
+    void setMemoryBudgetExhausted() { memoryBudgetExhausted = true; }
+
   protected:
     Wddm &wddm;
     uint32_t osContextId;
@@ -57,6 +60,7 @@ class WddmResidencyController {
     SpinLock lock;
     SpinLock trimCallbackLock;
 
+    bool memoryBudgetExhausted = false;
     uint64_t lastTrimFenceValue = 0u;
     ResidencyContainer trimCandidateList;
     uint32_t trimCandidatesCount = 0;

unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp

@@ -21,6 +21,7 @@
 #include "unit_tests/mocks/mock_device.h"
 #include "unit_tests/os_interface/windows/mock_wddm_allocation.h"
 #include "unit_tests/os_interface/windows/wddm_memory_manager_tests.h"
+#include "unit_tests/utilities/base_object_utils.h"
 
 using namespace OCLRT;
 using namespace ::testing;
@@ -1373,10 +1374,24 @@ TEST_F(MockWddmMemoryManagerTest, givenDefaultMemoryManagerWhenItIsCreatedThenAs
     EXPECT_NE(nullptr, memoryManager.getDeferredDeleter());
 }
 
-TEST_F(MockWddmMemoryManagerTest, givenDefaultWddmMemoryManagerWhenItIsCreatedThenMemoryBudgetIsNotExhausted) {
-    auto wddm = std::make_unique<WddmMock>();
-    WddmMemoryManager memoryManager(false, false, wddm.get(), executionEnvironment);
-    EXPECT_FALSE(memoryManager.isMemoryBudgetExhausted());
+TEST_F(WddmMemoryManagerTest, givenWddmMemoryManagerWithNoRegisteredOsContextsWhenCallingIsMemoryBudgetExhaustedThenReturnFalse) {
+    ASSERT_EQ(0u, memoryManager->getOsContextCount());
+    EXPECT_FALSE(memoryManager->isMemoryBudgetExhausted());
+}
+
+TEST_F(WddmMemoryManagerTest, givenWddmMemoryManagerWithRegisteredOsContextWhenCallingIsMemoryBudgetExhaustedThenReturnFalse) {
+    memoryManager->registerOsContext(new OsContext(osInterface.get(), 0u));
+    memoryManager->registerOsContext(new OsContext(osInterface.get(), 1u));
+    memoryManager->registerOsContext(new OsContext(osInterface.get(), 2u));
+    EXPECT_FALSE(memoryManager->isMemoryBudgetExhausted());
+}
+
+TEST_F(WddmMemoryManagerTest, givenWddmMemoryManagerWithRegisteredOsContextWithExhaustedMemoryBudgetWhenCallingIsMemoryBudgetExhaustedThenReturnTrue) {
+    memoryManager->registerOsContext(new OsContext(osInterface.get(), 0u));
+    memoryManager->registerOsContext(new OsContext(osInterface.get(), 1u));
+    memoryManager->registerOsContext(new OsContext(osInterface.get(), 2u));
+    memoryManager->getRegisteredOsContext(1)->get()->getResidencyController().setMemoryBudgetExhausted();
+    EXPECT_TRUE(memoryManager->isMemoryBudgetExhausted());
 }
 
 TEST_F(MockWddmMemoryManagerTest, givenEnabledAsyncDeleterFlagWhenMemoryManagerIsCreatedThenAsyncDeleterEnabledIsTrueAndDeleterIsNotNullptr) {

unit_tests/os_interface/windows/wddm_residency_controller_tests.cpp

@@ -130,6 +130,12 @@ TEST_F(WddmResidencyControllerTest, givenUsedAllocationWhenCallingRemoveFromTrim
     EXPECT_EQ(trimListUnusedPosition, allocation.getTrimCandidateListPosition(osContextId));
 }
 
+TEST_F(WddmResidencyControllerTest, givenWddmResidencyControllerWhenIsMemoryExhaustedIsCalledThenReturnCorrectResult) {
+    EXPECT_FALSE(residencyController->isMemoryBudgetExhausted());
+    residencyController->setMemoryBudgetExhausted();
+    EXPECT_TRUE(residencyController->isMemoryBudgetExhausted());
+}
+
 TEST_F(WddmResidencyControllerTest, givenUnusedAllocationWhenCallingRemoveFromTrimCandidateListIfUsedThenIgnore) {
     MockWddmAllocation allocation;
     residencyController->removeFromTrimCandidateListIfUsed(&allocation, false);