Return trimCallbackHandle from register method instead of storing in Wddm

Change-Id: Iddfb6a926311480e355d52e5bf3e6379ec96f7a9
Signed-off-by: Maciej Dziuban <maciej.dziuban@intel.com>

Author: DziubanMaciejIntel

runtime/os_interface/windows/wddm/wddm.cpp

@@ -811,9 +811,9 @@ uint64_t Wddm::getHeap32Size() {
     return alignDown(gfxPartition.Heap32[0].Limit, MemoryConstants::pageSize);
 }
 
-void Wddm::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmMemoryManager *memoryManager) {
+VOID *Wddm::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmMemoryManager *memoryManager) {
     if (DebugManager.flags.DoNotRegisterTrimCallback.get()) {
-        return;
+        return nullptr;
     }
     D3DKMT_REGISTERTRIMNOTIFICATION registerTrimNotification;
     registerTrimNotification.Callback = callback;
@@ -823,17 +823,19 @@ void Wddm::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, Wdd
 
     NTSTATUS status = gdi->registerTrimNotification(&registerTrimNotification);
     if (status == STATUS_SUCCESS) {
-        trimCallbackHandle = registerTrimNotification.Handle;
+        return registerTrimNotification.Handle;
     }
+    return nullptr;
 }
 
-void Wddm::unregisterTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback) {
-    if (callback == nullptr) {
+void Wddm::unregisterTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, VOID *trimCallbackHandle) {
+    DEBUG_BREAK_IF(callback == nullptr);
+    if (trimCallbackHandle == nullptr) {
         return;
     }
     D3DKMT_UNREGISTERTRIMNOTIFICATION unregisterTrimNotification;
     unregisterTrimNotification.Callback = callback;
-    unregisterTrimNotification.Handle = this->trimCallbackHandle;
+    unregisterTrimNotification.Handle = trimCallbackHandle;
 
     NTSTATUS status = gdi->unregisterTrimNotification(&unregisterTrimNotification);
     DEBUG_BREAK_IF(status != STATUS_SUCCESS);

runtime/os_interface/windows/wddm/wddm.h

@@ -79,8 +79,8 @@ class Wddm {
     MOCKABLE_VIRTUAL bool waitFromCpu(uint64_t lastFenceValue, OsContextWin &osContext);
 
     NTSTATUS escape(D3DKMT_ESCAPE &escapeCommand);
-    void registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmMemoryManager *memoryManager);
-    void unregisterTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback);
+    VOID *registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmMemoryManager *memoryManager);
+    void unregisterTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, VOID *trimCallbackHandle);
     MOCKABLE_VIRTUAL void releaseReservedAddress(void *reservedAddress);
     MOCKABLE_VIRTUAL bool reserveValidAddressRange(size_t size, void *&reservedMem);
 
@@ -177,7 +177,6 @@ class Wddm {
 
     unsigned long hwContextId = 0;
     LUID adapterLuid;
-    void *trimCallbackHandle = nullptr;
     uintptr_t maximumApplicationAddress = 0;
     GPUNODE_ORDINAL node = GPUNODE_3D;
     PreemptionMode preemptionMode = PreemptionMode::Disabled;

runtime/os_interface/windows/wddm_memory_manager.cpp

@@ -32,7 +32,7 @@ WddmMemoryManager::~WddmMemoryManager() {
         if (osContext) {
             auto &residencyController = osContext->get()->getResidencyController();
             residencyController.acquireTrimCallbackLock();
-            wddm->unregisterTrimCallback(trimCallback);
+            wddm->unregisterTrimCallback(trimCallback, this->trimCallbackHandle);
             residencyController.releaseTrimCallbackLock();
 
             // Wait for lock to ensure trimCallback ended
@@ -46,7 +46,7 @@ WddmMemoryManager::WddmMemoryManager(bool enable64kbPages, bool enableLocalMemor
     DEBUG_BREAK_IF(wddm == nullptr);
     this->wddm = wddm;
     allocator32Bit = std::unique_ptr<Allocator32bit>(new Allocator32bit(wddm->getHeap32Base(), wddm->getHeap32Size()));
-    wddm->registerTrimCallback(trimCallback, this);
+    this->trimCallbackHandle = wddm->registerTrimCallback(trimCallback, this);
     asyncDeleterEnabled = DebugManager.flags.EnableDeferredDeleter.get();
     if (asyncDeleterEnabled)
         deferredDeleter = createDeferredDeleter();

runtime/os_interface/windows/wddm_memory_manager.h

@@ -81,6 +81,7 @@ class WddmMemoryManager : public MemoryManager {
   protected:
     void trimResidency(D3DDDI_TRIMRESIDENCYSET_FLAGS flags, uint64_t bytes);
     bool trimResidencyToBudget(uint64_t bytes);
+    VOID *trimCallbackHandle = nullptr;
 
     GraphicsAllocation *createAllocationFromHandle(osHandle handle, bool requireSpecificBitness, bool ntHandle);
     static bool validateAllocation(WddmAllocation *alloc);

unit_tests/mocks/mock_wddm.h

@@ -44,7 +44,6 @@ class WddmMock : public Wddm {
     using Wddm::pagingFenceAddress;
     using Wddm::pagingQueue;
     using Wddm::preemptionMode;
-    using Wddm::trimCallbackHandle;
     using Wddm::wddmInterface;
 
     WddmMock() : Wddm(){};

unit_tests/os_interface/windows/mock_wddm_memory_manager.h

@@ -15,6 +15,7 @@ class MockWddmMemoryManager : public WddmMemoryManager {
 
   public:
     using BaseClass::createWddmAllocation;
+    using BaseClass::trimCallbackHandle;
     using BaseClass::trimResidency;
     using BaseClass::trimResidencyToBudget;
     using BaseClass::WddmMemoryManager;

unit_tests/os_interface/windows/wddm20_tests.cpp

@@ -778,3 +778,34 @@ TEST_F(Wddm20Tests, whenContextIsInitializedThenApplyAdditionalContextFlagsIsCal
     EXPECT_TRUE(result);
     EXPECT_EQ(1u, wddm->applyAdditionalContextFlagsResult.called);
 }
+
+TEST_F(Wddm20Tests, givenTrimCallbackRegistrationIsDisabledInDebugVariableWhenRegisteringCallbackThenReturnNullptr) {
+    DebugManagerStateRestore stateRestore;
+    DebugManager.flags.DoNotRegisterTrimCallback.set(true);
+    EXPECT_EQ(nullptr, wddm->registerTrimCallback([](D3DKMT_TRIMNOTIFICATION *) {}, nullptr));
+}
+
+TEST_F(Wddm20Tests, givenSuccessWhenRegisteringTrimCallbackThenReturnTrimCallbackHandle) {
+    auto trimCallbackHandle = wddm->registerTrimCallback([](D3DKMT_TRIMNOTIFICATION *) {}, nullptr);
+    EXPECT_NE(nullptr, trimCallbackHandle);
+}
+
+TEST_F(Wddm20Tests, givenCorrectArgumentsWhenUnregisteringTrimCallbackThenPassArgumentsToGdiCall) {
+    PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback = [](D3DKMT_TRIMNOTIFICATION *) {};
+    auto trimCallbackHandle = reinterpret_cast<VOID *>(0x9876);
+
+    wddm->unregisterTrimCallback(callback, trimCallbackHandle);
+    EXPECT_EQ(callback, gdi->getUnregisterTrimNotificationArg().Callback);
+    EXPECT_EQ(trimCallbackHandle, gdi->getUnregisterTrimNotificationArg().Handle);
+}
+
+TEST_F(Wddm20Tests, givenNullTrimCallbackHandleWhenUnregisteringTrimCallbackThenDoNotDoGdiCall) {
+    PFND3DKMT_TRIMNOTIFICATIONCALLBACK callbackBefore = [](D3DKMT_TRIMNOTIFICATION *) {};
+    auto trimCallbackHandleBefore = reinterpret_cast<VOID *>(0x9876);
+    gdi->getUnregisterTrimNotificationArg().Callback = callbackBefore;
+    gdi->getUnregisterTrimNotificationArg().Handle = trimCallbackHandleBefore;
+
+    wddm->unregisterTrimCallback([](D3DKMT_TRIMNOTIFICATION *) {}, nullptr);
+    EXPECT_EQ(callbackBefore, gdi->getUnregisterTrimNotificationArg().Callback);
+    EXPECT_EQ(trimCallbackHandleBefore, gdi->getUnregisterTrimNotificationArg().Handle);
+}

unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp

@@ -946,7 +946,7 @@ TEST_F(WddmMemoryManagerResidencyTest, trimCallbackIsRegisteredInWddmMemoryManag
 }
 
 TEST_F(WddmMemoryManagerResidencyTest, givenWddmMemoryManagerWhenCallingDestructorThenUnregisterTrimCallback) {
-    auto trimCallbackHandle = wddm->trimCallbackHandle;
+    auto trimCallbackHandle = memoryManager->trimCallbackHandle;
     auto trimCallbackAddress = reinterpret_cast<PFND3DKMT_TRIMNOTIFICATIONCALLBACK>(memoryManager->trimCallback);
     memoryManager.reset();
 
@@ -955,14 +955,6 @@ TEST_F(WddmMemoryManagerResidencyTest, givenWddmMemoryManagerWhenCallingDestruct
     EXPECT_EQ(trimCallbackHandle, unregisterNotification.Handle);
 }
 
-TEST(WddmDebugModesTests, givenDebugModeWhenItIsActiveThenTrimCallbackIsNotRegistred) {
-    DebugManagerStateRestore stateRestore;
-    DebugManager.flags.DoNotRegisterTrimCallback.set(true);
-    WddmMock wddm;
-    wddm.init();
-    EXPECT_EQ(nullptr, wddm.trimCallbackHandle);
-}
-
 TEST_F(WddmMemoryManagerResidencyTest, givenNotUsedAllocationsFromPreviousPeriodicTrimWhenTrimResidencyPeriodicTrimIsCalledThenAllocationsAreEvictedMarkedAndRemovedFromTrimCandidateList) {
     D3DKMT_TRIMNOTIFICATION trimNotification = {0};
     trimNotification.Flags.PeriodicTrim = 1;