refactor: Move residency controller to wddm

Resolves: NEO-13315

Signed-off-by: Lukasz Jobczyk <lukasz.jobczyk@intel.com>

shared/source/os_interface/windows/os_context_win.cpp

@@ -24,11 +24,7 @@ OsContext *OsContextWin::create(OSInterface *osInterface, uint32_t rootDeviceInd
     return new OsContext(rootDeviceIndex, contextId, engineDescriptor);
 }
 
-OsContextWin::OsContextWin(Wddm &wddm, uint32_t rootDeviceIndex, uint32_t contextId, const EngineDescriptor &engineDescriptor)
-    : OsContext(rootDeviceIndex, contextId, engineDescriptor),
-      residencyController(wddm),
-      wddm(wddm) {
-}
+OsContextWin::OsContextWin(Wddm &wddm, uint32_t rootDeviceIndex, uint32_t contextId, const EngineDescriptor &engineDescriptor) : OsContext(rootDeviceIndex, contextId, engineDescriptor), wddm(wddm) {}
 
 bool OsContextWin::initializeContext(bool allocateInterrupt) {
 
@@ -45,9 +41,6 @@ bool OsContextWin::initializeContext(bool allocateInterrupt) {
             UNRECOVERABLE_IF(!wddmInterface->createHwQueue(*this));
         }
         UNRECOVERABLE_IF(!wddmInterface->createMonitoredFence(*this));
-
-        residencyController.registerCallback();
-        UNRECOVERABLE_IF(!residencyController.isInitialized());
     }
 
     return true;
@@ -101,6 +94,8 @@ uint32_t OsContextWin::getDeviceNodeMask() {
     return hwDeviceID->getAdapterNodeMask();
 }
 
+WddmResidencyController &OsContextWin::getResidencyController() { return wddm.getResidencyController(); }
+
 uint64_t OsContextWin::getOfflineDumpContextId(uint32_t deviceIndex) const {
     return 0;
 }

shared/source/os_interface/windows/os_context_win.h

@@ -33,7 +33,7 @@ class OsContextWin : public OsContext {
     void setHwQueue(HardwareQueue hardwareQueue) { this->hardwareQueue = hardwareQueue; }
     bool isDirectSubmissionSupported() const override;
     Wddm *getWddm() const { return &wddm; }
-    MOCKABLE_VIRTUAL WddmResidencyController &getResidencyController() { return residencyController; }
+    MOCKABLE_VIRTUAL WddmResidencyController &getResidencyController();
     static OsContext *create(OSInterface *osInterface, uint32_t rootDeviceIndex, uint32_t contextId, const EngineDescriptor &engineDescriptor);
     MonitoredFence &getMonitoredFence() { return monitoredFence; }
     void resetMonitoredFenceParams(D3DKMT_HANDLE &handle, uint64_t *cpuAddress, D3DGPU_VIRTUAL_ADDRESS &gpuAddress);
@@ -48,8 +48,6 @@ class OsContextWin : public OsContext {
   protected:
     bool initializeContext(bool allocateInterrupt) override;
 
-    WddmResidencyController residencyController;
-
     HardwareQueue hardwareQueue;
 
     MonitoredFence monitoredFence = {};

shared/source/os_interface/windows/os_interface_win.cpp

@@ -34,4 +34,11 @@ uint32_t OSInterface::getAggregatedProcessCount() const {
     return 0;
 }
 
+void OSInterface::registerTrimCallback() {
+    if (driverModel && driverModel->getDriverModelType() == DriverModelType::wddm) {
+        driverModel->as<NEO::Wddm>()->getResidencyController().registerCallback();
+        UNRECOVERABLE_IF(!driverModel->as<NEO::Wddm>()->getResidencyController().isInitialized());
+    }
+}
+
 } // namespace NEO

shared/source/os_interface/windows/trim_callback.cpp

@@ -15,7 +15,11 @@
 
 namespace NEO {
 
-VOID *Wddm::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmResidencyController &residencyController) {
+VOID *Wddm::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback) {
+    if (!hwDeviceId || !this->rootDeviceEnvironment.executionEnvironment.osEnvironment || !getGdi()) {
+        return nullptr;
+    }
+
     if (debugManager.flags.DoNotRegisterTrimCallback.get()) {
         return nullptr;
     }
@@ -33,6 +37,10 @@ VOID *Wddm::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, Wd
 }
 
 void Wddm::unregisterTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, VOID *trimCallbackHandle) {
+    if (!hwDeviceId || !this->rootDeviceEnvironment.executionEnvironment.osEnvironment || !getGdi()) {
+        return;
+    }
+
     DEBUG_BREAK_IF(callback == nullptr);
     if (trimCallbackHandle == nullptr || isShutdownInProgress()) {
         return;
@@ -48,7 +56,6 @@ void Wddm::unregisterTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, V
 void APIENTRY WddmResidencyController::trimCallback(_Inout_ D3DKMT_TRIMNOTIFICATION *trimNotification) {
     auto residencyController = static_cast<WddmResidencyController *>(trimNotification->Context);
     DEBUG_BREAK_IF(residencyController == nullptr);
-    UNRECOVERABLE_IF(residencyController->csr == nullptr);
 
     auto lock = residencyController->acquireTrimCallbackLock();
     residencyController->trimResidency(trimNotification->Flags, trimNotification->NumBytesToTrim);
@@ -57,74 +64,94 @@ void APIENTRY WddmResidencyController::trimCallback(_Inout_ D3DKMT_TRIMNOTIFICAT
 void WddmResidencyController::trimResidency(const D3DDDI_TRIMRESIDENCYSET_FLAGS &flags, uint64_t bytes) {
     std::chrono::high_resolution_clock::time_point callbackStart;
     perfLogResidencyTrimCallbackBegin(wddm.getResidencyLogger(), callbackStart);
-    auto &osContext = static_cast<NEO::OsContextWin &>(this->csr->getOsContext());
-    uint32_t osContextId = osContext.getContextId();
 
     if (flags.PeriodicTrim) {
         uint64_t sizeToTrim = 0;
         auto lock = this->acquireLock();
         WddmAllocation *wddmAllocation = nullptr;
         std::vector<D3DKMT_HANDLE> handlesToEvict;
+        std::vector<WddmAllocation *> allocationsToRemove;
+        allocationsToRemove.reserve(evictionAllocations.size());
         handlesToEvict.reserve(evictionAllocations.size());
-        for (auto allocationIter = evictionAllocations.begin(); allocationIter != evictionAllocations.end();) {
-            wddmAllocation = reinterpret_cast<WddmAllocation *>(*allocationIter);
 
-            DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "lastPeriodicTrimFenceValue = ", osContext.getLastTrimFenceValue());
+        this->wddm.forEachContextWithinWddm<true>([&](const EngineControl &engine) {
+            for (auto allocationIter = evictionAllocations.begin(); allocationIter != evictionAllocations.end(); ++allocationIter) {
+                auto &osContext = static_cast<NEO::OsContextWin &>(*engine.osContext);
+                uint32_t osContextId = osContext.getContextId();
 
-            if (osContext.wasAllocationUsedSinceLastTrim(wddmAllocation->getResidencyData().getFenceValueForContextId(osContextId))) {
-                allocationIter++;
-                continue;
-            }
+                wddmAllocation = reinterpret_cast<WddmAllocation *>(*allocationIter);
 
-            if (wddmAllocation->isAlwaysResident(osContextId)) {
-                allocationIter = evictionAllocations.erase(allocationIter);
-                continue;
-            }
+                DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "lastPeriodicTrimFenceValue = ", osContext.getLastTrimFenceValue());
 
-            if (wddmAllocation->fragmentsStorage.fragmentCount == 0) {
-                for (auto i = 0u; i < wddmAllocation->getNumGmms(); i++) {
-                    handlesToEvict.push_back(wddmAllocation->getHandles()[i]);
+                if (osContext.wasAllocationUsedSinceLastTrim(wddmAllocation->getResidencyData().getFenceValueForContextId(osContextId))) {
+                    continue;
+                }
+
+                if (wddmAllocation->isAlwaysResident(osContextId)) {
+                    continue;
+                }
+
+                if (wddmAllocation->fragmentsStorage.fragmentCount == 0) {
+                    if (!wddmAllocation->getResidencyData().resident[osContextId]) {
+                        continue;
+                    }
+                    for (auto i = 0u; i < wddmAllocation->getNumGmms(); i++) {
+                        handlesToEvict.push_back(wddmAllocation->getHandles()[i]);
+                    }
+                } else {
+                    for (uint32_t allocationId = 0; allocationId < wddmAllocation->fragmentsStorage.fragmentCount; allocationId++) {
+                        AllocationStorageData &fragmentStorageData = wddmAllocation->fragmentsStorage.fragmentStorageData[allocationId];
+                        if (!osContext.wasAllocationUsedSinceLastTrim(fragmentStorageData.residency->getFenceValueForContextId(osContextId))) {
+                            auto osHandle = static_cast<OsHandleWin *>(wddmAllocation->fragmentsStorage.fragmentStorageData[allocationId].osHandleStorage);
+                            DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "Evict fragment: handle =", osHandle->handle, "lastFence =",
+                                    wddmAllocation->fragmentsStorage.fragmentStorageData[allocationId].residency->getFenceValueForContextId(osContextId));
+
+                            handlesToEvict.push_back(static_cast<OsHandleWin *>(fragmentStorageData.osHandleStorage)->handle);
+                            fragmentStorageData.residency->resident[osContextId] = false;
+                        }
+                    }
+                }
+
+                DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "Evict allocation, gpu address = ", std::hex, wddmAllocation->getGpuAddress(), "lastFence =", wddmAllocation->getResidencyData().getFenceValueForContextId(osContextId));
+                wddmAllocation->getResidencyData().resident[osContextId] = false;
+
+                if (std::none_of(wddmAllocation->getResidencyData().resident.begin(), wddmAllocation->getResidencyData().resident.end(), [](bool resident) { return resident; })) {
+                    allocationsToRemove.push_back(wddmAllocation);
                 }
             }
 
-            for (uint32_t allocationId = 0; allocationId < wddmAllocation->fragmentsStorage.fragmentCount; allocationId++) {
-                AllocationStorageData &fragmentStorageData = wddmAllocation->fragmentsStorage.fragmentStorageData[allocationId];
-                if (!osContext.wasAllocationUsedSinceLastTrim(fragmentStorageData.residency->getFenceValueForContextId(osContextId))) {
-                    auto osHandle = static_cast<OsHandleWin *>(wddmAllocation->fragmentsStorage.fragmentStorageData[allocationId].osHandleStorage);
-                    DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "Evict fragment: handle =", osHandle->handle, "lastFence =",
-                            wddmAllocation->fragmentsStorage.fragmentStorageData[allocationId].residency->getFenceValueForContextId(osContextId));
+            this->wddm.evict(handlesToEvict.data(), static_cast<uint32_t>(handlesToEvict.size()), sizeToTrim, false);
+            handlesToEvict.clear();
+        });
 
-                    handlesToEvict.push_back(static_cast<OsHandleWin *>(fragmentStorageData.osHandleStorage)->handle);
-                    fragmentStorageData.residency->resident[osContextId] = false;
-                }
-            }
-
-            DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "Evict allocation, gpu address = ", std::hex, wddmAllocation->getGpuAddress(), "lastFence =", wddmAllocation->getResidencyData().getFenceValueForContextId(osContextId));
-            wddmAllocation->getResidencyData().resident[osContextId] = false;
-            allocationIter = evictionAllocations.erase(allocationIter);
-        }
-
-        this->wddm.evict(handlesToEvict.data(), static_cast<uint32_t>(handlesToEvict.size()), sizeToTrim, false);
+        std::erase_if(evictionAllocations, [&allocationsToRemove](GraphicsAllocation *graphicsAllocation) {
+            return std::find_if(allocationsToRemove.begin(), allocationsToRemove.end(), [&graphicsAllocation](WddmAllocation *graphicsAllocationToRemove) {
+                       return graphicsAllocationToRemove == graphicsAllocation;
+                   }) != allocationsToRemove.end();
+        });
     }
 
     if (flags.TrimToBudget) {
-        auto csrLock = csr->obtainUniqueOwnership();
+        StackVec<std::unique_lock<CommandStreamReceiver::MutexType>, 5> csrLocks;
+        this->wddm.forEachContextWithinWddm<true>([&](const EngineControl &engine) {
+            csrLocks.push_back(engine.commandStreamReceiver->obtainUniqueOwnership());
+        });
         auto lock = this->acquireLock();
         trimResidencyToBudget(bytes);
     }
 
     if (flags.PeriodicTrim || flags.RestartPeriodicTrim) {
-        osContext.updateLastTrimFenceValue();
-        DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "updated lastPeriodicTrimFenceValue =", osContext.getLastTrimFenceValue());
+        this->wddm.forEachContextWithinWddm<true>([&](const EngineControl &engine) {
+            auto osContext = static_cast<NEO::OsContextWin *>(engine.osContext);
+            osContext->updateLastTrimFenceValue();
+            DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "updated lastPeriodicTrimFenceValue =", osContext->getLastTrimFenceValue());
+        });
     }
 
     perfLogResidencyTrimCallbackEnd(wddm.getResidencyLogger(), flags.Value, this, callbackStart);
 }
 
 bool WddmResidencyController::trimResidencyToBudget(uint64_t bytes) {
-    this->csr->drainPagingFenceQueue();
-    auto &osContext = static_cast<NEO::OsContextWin &>(this->csr->getOsContext());
-    uint32_t osContextId = osContext.getContextId();
     uint64_t sizeToTrim = 0;
     uint64_t numberOfBytesToTrim = bytes;
     WddmAllocation *wddmAllocation = nullptr;
@@ -133,55 +160,78 @@ bool WddmResidencyController::trimResidencyToBudget(uint64_t bytes) {
     auto allocationIter = evictionAllocations.begin();
     std::vector<D3DKMT_HANDLE> handlesToEvict;
     handlesToEvict.reserve(evictionAllocations.size());
-    while (numberOfBytesToTrim > 0 && allocationIter != evictionAllocations.end()) {
-        wddmAllocation = reinterpret_cast<WddmAllocation *>(*allocationIter);
-        uint64_t lastFence = wddmAllocation->getResidencyData().getFenceValueForContextId(osContextId);
-        auto &monitoredFence = osContext.getMonitoredFence();
+    std::vector<WddmAllocation *> allocationsToRemove;
+    allocationsToRemove.reserve(evictionAllocations.size());
 
-        if (lastFence > monitoredFence.lastSubmittedFence) {
-            allocationIter++;
-            continue;
-        }
+    this->wddm.forEachContextWithinWddm<true>([&](const EngineControl &engine) {
+        auto csr = engine.commandStreamReceiver;
+        csr->drainPagingFenceQueue();
+        auto &osContext = static_cast<NEO::OsContextWin &>(csr->getOsContext());
+        uint32_t osContextId = osContext.getContextId();
 
-        if (wddmAllocation->isAlwaysResident(osContextId)) {
-            allocationIter = evictionAllocations.erase(allocationIter);
-            continue;
-        }
+        while (numberOfBytesToTrim > 0 && allocationIter != evictionAllocations.end()) {
+            wddmAllocation = reinterpret_cast<WddmAllocation *>(*allocationIter);
+            uint64_t lastFence = wddmAllocation->getResidencyData().getFenceValueForContextId(osContextId);
+            auto &monitoredFence = osContext.getMonitoredFence();
 
-        uint64_t sizeEvicted = 0;
-
-        if (lastFence > *monitoredFence.cpuAddress) {
-            this->wddm.waitFromCpu(lastFence, osContext.getMonitoredFence(), false);
-        }
-
-        if (wddmAllocation->fragmentsStorage.fragmentCount == 0) {
-            for (auto i = 0u; i < wddmAllocation->getNumGmms(); i++) {
-                handlesToEvict.push_back(wddmAllocation->getHandles()[i]);
+            if (lastFence > monitoredFence.lastSubmittedFence) {
+                ++allocationIter;
+                continue;
             }
-            sizeEvicted = wddmAllocation->getAlignedSize();
-        } else {
-            auto &fragmentStorageData = wddmAllocation->fragmentsStorage.fragmentStorageData;
-            for (uint32_t allocationId = 0; allocationId < wddmAllocation->fragmentsStorage.fragmentCount; allocationId++) {
-                if (fragmentStorageData[allocationId].residency->getFenceValueForContextId(osContextId) <= monitoredFence.lastSubmittedFence) {
-                    handlesToEvict.push_back(static_cast<OsHandleWin *>(fragmentStorageData[allocationId].osHandleStorage)->handle);
-                    fragmentStorageData[allocationId].residency->resident[osContextId] = false;
-                    sizeEvicted += fragmentStorageData[allocationId].fragmentSize;
+
+            if (wddmAllocation->isAlwaysResident(osContextId)) {
+                ++allocationIter;
+                continue;
+            }
+
+            uint64_t sizeEvicted = 0;
+
+            if (lastFence > *monitoredFence.cpuAddress) {
+                this->wddm.waitFromCpu(lastFence, osContext.getMonitoredFence(), false);
+            }
+
+            if (wddmAllocation->fragmentsStorage.fragmentCount == 0) {
+                for (auto i = 0u; i < wddmAllocation->getNumGmms(); i++) {
+                    handlesToEvict.push_back(wddmAllocation->getHandles()[i]);
+                }
+                sizeEvicted = wddmAllocation->getAlignedSize();
+            } else {
+                auto &fragmentStorageData = wddmAllocation->fragmentsStorage.fragmentStorageData;
+                for (uint32_t allocationId = 0; allocationId < wddmAllocation->fragmentsStorage.fragmentCount; allocationId++) {
+                    if (fragmentStorageData[allocationId].residency->getFenceValueForContextId(osContextId) <= monitoredFence.lastSubmittedFence) {
+                        handlesToEvict.push_back(static_cast<OsHandleWin *>(fragmentStorageData[allocationId].osHandleStorage)->handle);
+                        fragmentStorageData[allocationId].residency->resident[osContextId] = false;
+                        sizeEvicted += fragmentStorageData[allocationId].fragmentSize;
+                    }
                 }
             }
-        }
-        DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "Evict allocation, gpu address = ", std::hex, wddmAllocation->getGpuAddress(), "lastFence =", wddmAllocation->getResidencyData().getFenceValueForContextId(osContextId));
+            DBG_LOG(ResidencyDebugEnable, "Residency:", __FUNCTION__, "Evict allocation, gpu address = ", std::hex, wddmAllocation->getGpuAddress(), "lastFence =", wddmAllocation->getResidencyData().getFenceValueForContextId(osContextId));
 
-        if (sizeEvicted >= numberOfBytesToTrim) {
-            numberOfBytesToTrim = 0;
-        } else {
-            numberOfBytesToTrim -= sizeEvicted;
+            if (sizeEvicted >= numberOfBytesToTrim) {
+                numberOfBytesToTrim = 0;
+            } else {
+                numberOfBytesToTrim -= sizeEvicted;
+            }
+
+            wddmAllocation->getResidencyData().resident[osContextId] = false;
+
+            if (std::none_of(wddmAllocation->getResidencyData().resident.begin(), wddmAllocation->getResidencyData().resident.end(), [](bool resident) { return resident; })) {
+                allocationsToRemove.push_back(wddmAllocation);
+            }
+
+            ++allocationIter;
         }
 
-        wddmAllocation->getResidencyData().resident[osContextId] = false;
-        allocationIter = evictionAllocations.erase(allocationIter);
-    }
+        this->wddm.evict(handlesToEvict.data(), static_cast<uint32_t>(handlesToEvict.size()), sizeToTrim, true);
+        handlesToEvict.clear();
+    });
+
+    std::erase_if(evictionAllocations, [&allocationsToRemove](GraphicsAllocation *graphicsAllocation) {
+        return std::find_if(allocationsToRemove.begin(), allocationsToRemove.end(), [&graphicsAllocation](WddmAllocation *graphicsAllocationToRemove) {
+                   return graphicsAllocationToRemove == graphicsAllocation;
+               }) != allocationsToRemove.end();
+    });
 
-    this->wddm.evict(handlesToEvict.data(), static_cast<uint32_t>(handlesToEvict.size()), sizeToTrim, true);
     return numberOfBytesToTrim == 0;
 }
 

shared/source/os_interface/windows/trim_callback_stub.cpp

@@ -11,7 +11,7 @@
 
 namespace NEO {
 
-VOID *Wddm::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmResidencyController &residencyController) {
+VOID *Wddm::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback) {
     return nullptr;
 }
 

shared/source/os_interface/windows/wddm/wddm.cpp

@@ -54,7 +54,7 @@ Wddm::CreateDXGIFactoryFcn Wddm::createDxgiFactory = getCreateDxgiFactory();
 Wddm::GetSystemInfoFcn Wddm::getSystemInfo = getGetSystemInfo();
 
 Wddm::Wddm(std::unique_ptr<HwDeviceIdWddm> &&hwDeviceIdIn, RootDeviceEnvironment &rootDeviceEnvironment)
-    : DriverModel(DriverModelType::wddm), hwDeviceId(std::move(hwDeviceIdIn)), rootDeviceEnvironment(rootDeviceEnvironment) {
+    : DriverModel(DriverModelType::wddm), residencyController(*this), hwDeviceId(std::move(hwDeviceIdIn)), rootDeviceEnvironment(rootDeviceEnvironment) {
     UNRECOVERABLE_IF(!hwDeviceId);
     featureTable.reset(new FeatureTable());
     workaroundTable.reset(new WorkaroundTable());
@@ -76,6 +76,7 @@ Wddm::Wddm(std::unique_ptr<HwDeviceIdWddm> &&hwDeviceIdIn, RootDeviceEnvironment
 }
 
 Wddm::~Wddm() {
+    this->residencyController.unregisterCallback();
     temporaryResources.reset();
     destroyPagingQueue();
     destroyDevice();
@@ -610,7 +611,7 @@ bool Wddm::mapGpuVirtualAddress(Gmm *gmm, D3DKMT_HANDLE handle, D3DGPU_VIRTUAL_A
     bool ret = true;
     auto &productHelper = rootDeviceEnvironment.getHelper<ProductHelper>();
     if (gmm->isCompressionEnabled() && productHelper.isPageTableManagerSupported(*rootDeviceEnvironment.getHardwareInfo())) {
-        this->forEachContextWithinWddm([&](const EngineControl &engine) {
+        this->forEachContextWithinWddm<false>([&](const EngineControl &engine) {
             if (engine.commandStreamReceiver->pageTableManager.get()) {
                 ret &= engine.commandStreamReceiver->pageTableManager->updateAuxTable(gpuPtr, gmm, true);
             }
@@ -1196,7 +1197,7 @@ bool Wddm::waitFromCpu(uint64_t lastFenceValue, const MonitoredFence &monitoredF
 
     if (!skipResourceCleanup() && lastFenceValue > *monitoredFence.cpuAddress) {
         CommandStreamReceiver *csr = nullptr;
-        this->forEachContextWithinWddm([&monitoredFence, &csr](const EngineControl &engine) {
+        this->forEachContextWithinWddm<false>([&monitoredFence, &csr](const EngineControl &engine) {
             auto &contextMonitoredFence = static_cast<OsContextWin *>(engine.osContext)->getMonitoredFence();
             if (contextMonitoredFence.cpuAddress == monitoredFence.cpuAddress) {
                 csr = engine.commandStreamReceiver;
@@ -1439,7 +1440,7 @@ void Wddm::setNewResourceBoundToPageTable() {
     if (!this->rootDeviceEnvironment.getProductHelper().isTlbFlushRequired()) {
         return;
     }
-    this->forEachContextWithinWddm([](const EngineControl &engine) { engine.osContext->setNewResourceBound(); });
+    this->forEachContextWithinWddm<false>([](const EngineControl &engine) { engine.osContext->setNewResourceBound(); });
 }
 
 bool Wddm::needsNotifyAubCaptureCallback() const {

shared/source/os_interface/windows/wddm/wddm.h

@@ -13,6 +13,7 @@
 #include "shared/source/os_interface/windows/sharedata_wrapper.h"
 #include "shared/source/os_interface/windows/wddm/wddm_defs.h"
 #include "shared/source/os_interface/windows/wddm_memory_manager.h"
+#include "shared/source/os_interface/windows/wddm_residency_controller.h"
 
 #include <atomic>
 
@@ -30,7 +31,6 @@ class ProductHelper;
 class SettingsReader;
 class WddmAllocation;
 class WddmInterface;
-class WddmResidencyController;
 class WddmResidencyLogger;
 class WddmResidentAllocationsContainer;
 
@@ -95,7 +95,8 @@ class Wddm : public DriverModel {
     MOCKABLE_VIRTUAL bool waitFromCpu(uint64_t lastFenceValue, const MonitoredFence &monitoredFence, bool busyWait);
 
     MOCKABLE_VIRTUAL NTSTATUS escape(D3DKMT_ESCAPE &escapeCommand);
-    MOCKABLE_VIRTUAL VOID *registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmResidencyController &residencyController);
+    WddmResidencyController &getResidencyController() { return residencyController; }
+    MOCKABLE_VIRTUAL VOID *registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback);
     void unregisterTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, VOID *trimCallbackHandle);
     MOCKABLE_VIRTUAL void releaseReservedAddress(void *reservedAddress);
     MOCKABLE_VIRTUAL bool reserveValidAddressRange(size_t size, void *&reservedMem);
@@ -217,12 +218,18 @@ class Wddm : public DriverModel {
         return additionalAdapterInfoOptions;
     }
 
-    template <typename F>
+    template <bool onlyInitialized, typename F>
     void forEachContextWithinWddm(F func) {
         for (auto rootDeviceIndex = 0u; rootDeviceIndex < rootDeviceEnvironment.executionEnvironment.rootDeviceEnvironments.size(); rootDeviceIndex++) {
             if (rootDeviceEnvironment.executionEnvironment.rootDeviceEnvironments[rootDeviceIndex].get() == &rootDeviceEnvironment) {
                 for (auto &engine : rootDeviceEnvironment.executionEnvironment.memoryManager->getRegisteredEngines(rootDeviceIndex)) {
-                    func(engine);
+                    if constexpr (onlyInitialized) {
+                        if (engine.osContext->isInitialized()) {
+                            func(engine);
+                        }
+                    } else {
+                        func(engine);
+                    }
                 }
             }
         }
@@ -262,6 +269,8 @@ class Wddm : public DriverModel {
     bool isReadOnlyFlagFallbackSupported() const;
     bool isReadOnlyFlagFallbackAvailable(const D3DKMT_CREATEALLOCATION &createAllocation) const;
 
+    WddmResidencyController residencyController;
+
     GMM_GFX_PARTITIONING gfxPartition{};
     ADAPTER_BDF adapterBDF{};
 

shared/source/os_interface/windows/wddm_device_command_stream.inl

@@ -147,7 +147,7 @@ SubmissionStatus WddmCommandStreamReceiver<GfxFamily>::flush(BatchBuffer &batchB
 
 template <typename GfxFamily>
 SubmissionStatus WddmCommandStreamReceiver<GfxFamily>::processResidency(ResidencyContainer &allocationsForResidency, uint32_t handleId) {
-    return static_cast<OsContextWin *>(this->osContext)->getResidencyController().makeResidentResidencyAllocations(allocationsForResidency, this->requiresBlockingResidencyHandling, *static_cast<OsContextWin *>(this->osContext)) ? SubmissionStatus::success : SubmissionStatus::outOfMemory;
+    return static_cast<OsContextWin *>(this->osContext)->getResidencyController().makeResidentResidencyAllocations(allocationsForResidency, this->requiresBlockingResidencyHandling, this) ? SubmissionStatus::success : SubmissionStatus::outOfMemory;
 }
 
 template <typename GfxFamily>
@@ -212,7 +212,6 @@ CommandStreamReceiver *createWddmDeviceCommandStreamReceiver(bool withAubDump,
 template <typename GfxFamily>
 void WddmCommandStreamReceiver<GfxFamily>::setupContext(OsContext &osContext) {
     this->osContext = &osContext;
-    static_cast<OsContextWin *>(this->osContext)->getResidencyController().setCommandStreamReceiver(this);
 }
 
 template <typename GfxFamily>

shared/source/os_interface/windows/wddm_residency_controller.cpp

@@ -22,10 +22,10 @@ WddmResidencyController::WddmResidencyController(Wddm &wddm) : wddm(wddm) {
 }
 
 void WddmResidencyController::registerCallback() {
-    this->trimCallbackHandle = wddm.registerTrimCallback(WddmResidencyController::trimCallback, *this);
+    this->trimCallbackHandle = wddm.registerTrimCallback(WddmResidencyController::trimCallback);
 }
 
-WddmResidencyController::~WddmResidencyController() {
+void WddmResidencyController::unregisterCallback() {
     auto lock = this->acquireTrimCallbackLock();
     wddm.unregisterTrimCallback(WddmResidencyController::trimCallback, this->trimCallbackHandle);
     lock.unlock();
@@ -54,7 +54,8 @@ std::unique_lock<SpinLock> WddmResidencyController::acquireTrimCallbackLock() {
  *
  * @return returns true if all allocations either succeeded or are pending to be resident
  */
-bool WddmResidencyController::makeResidentResidencyAllocations(ResidencyContainer &allocationsForResidency, bool &requiresBlockingResidencyHandling, OsContextWin &osContext) {
+bool WddmResidencyController::makeResidentResidencyAllocations(ResidencyContainer &allocationsForResidency, bool &requiresBlockingResidencyHandling, CommandStreamReceiver *csr) {
+    auto &osContext = static_cast<OsContextWin &>(csr->getOsContext());
     auto osContextId = osContext.getContextId();
     const size_t residencyCount = allocationsForResidency.size();
     requiresBlockingResidencyHandling = false;
@@ -159,17 +160,13 @@ size_t WddmResidencyController::fillHandlesContainer(ResidencyContainer &allocat
 
 bool WddmResidencyController::isInitialized() const {
     bool requiresTrimCallbacks = OSInterface::requiresSupportForWddmTrimNotification;
-    requiresTrimCallbacks = requiresTrimCallbacks && (false == debugManager.flags.DoNotRegisterTrimCallback.get());
+    requiresTrimCallbacks = requiresTrimCallbacks && (false == debugManager.flags.DoNotRegisterTrimCallback.get()) && wddm.getHwDeviceId() && wddm.getRootDeviceEnvironment().executionEnvironment.osEnvironment && wddm.getGdi();
     if (requiresTrimCallbacks) {
         return trimCallbackHandle != nullptr;
     }
     return true;
 }
 
-void WddmResidencyController::setCommandStreamReceiver(CommandStreamReceiver *csr) {
-    this->csr = csr;
-}
-
 void WddmResidencyController::removeAllocation(ResidencyContainer &container, GraphicsAllocation *gfxAllocation) {
     std::unique_lock<std::mutex> lock1(this->lock, std::defer_lock);
     std::unique_lock<std::mutex> lock2(this->trimCallbackLock, std::defer_lock);

shared/source/os_interface/windows/wddm_residency_controller.h

@@ -30,7 +30,7 @@ class OsContextWin;
 class WddmResidencyController {
   public:
     WddmResidencyController(Wddm &wddm);
-    MOCKABLE_VIRTUAL ~WddmResidencyController();
+    MOCKABLE_VIRTUAL ~WddmResidencyController() = default;
 
     static void APIENTRY trimCallback(_Inout_ D3DKMT_TRIMNOTIFICATION *trimNotification);
 
@@ -38,6 +38,7 @@ class WddmResidencyController {
     [[nodiscard]] std::unique_lock<SpinLock> acquireTrimCallbackLock();
 
     void registerCallback();
+    void unregisterCallback();
 
     void trimResidency(const D3DDDI_TRIMRESIDENCYSET_FLAGS &flags, uint64_t bytes);
     bool trimResidencyToBudget(uint64_t bytes);
@@ -45,12 +46,10 @@ class WddmResidencyController {
     bool isMemoryBudgetExhausted() const { return memoryBudgetExhausted; }
     void setMemoryBudgetExhausted() { memoryBudgetExhausted = true; }
 
-    bool makeResidentResidencyAllocations(ResidencyContainer &allocationsForResidency, bool &requiresBlockingResidencyHandling, OsContextWin &osContext);
+    bool makeResidentResidencyAllocations(ResidencyContainer &allocationsForResidency, bool &requiresBlockingResidencyHandling, CommandStreamReceiver *csr);
 
     bool isInitialized() const;
 
-    void setCommandStreamReceiver(CommandStreamReceiver *csr);
-
     void removeAllocation(ResidencyContainer &container, GraphicsAllocation *gfxAllocation);
 
     ResidencyContainer &getEvictionAllocations() {
@@ -68,8 +67,6 @@ class WddmResidencyController {
 
     bool memoryBudgetExhausted = false;
 
-    CommandStreamReceiver *csr = nullptr;
-
     ResidencyContainer evictionAllocations;
 
     ResidencyContainer backupResidencyContainer;    // Stores allocations which should be resident