Register trim callback after creating monitored fence

Change-Id: Ib0510a0aed180c7bfe883a60632fd642664e1085
Signed-off-by: Mateusz Jablonski <[email protected]>

Author: JablonskiMateusz

runtime/os_interface/windows/os_context_win.cpp

@@ -24,6 +24,7 @@ OsContextWin::OsContextImpl(Wddm &wddm, uint32_t osContextId) : wddm(wddm), resi
         }
     }
     initialized = wddmInterface->createMonitoredFence(this->residencyController);
+    this->residencyController.registerCallback();
 };
 OsContextWin::~OsContextImpl() {
     wddm.getWddmInterface()->destroyHwQueue(hwQueueHandle);

runtime/os_interface/windows/wddm/wddm.h

@@ -78,7 +78,7 @@ class Wddm {
     MOCKABLE_VIRTUAL bool waitFromCpu(uint64_t lastFenceValue, const MonitoredFence &monitoredFence);
 
     NTSTATUS escape(D3DKMT_ESCAPE &escapeCommand);
-    VOID *registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmResidencyController &residencyController);
+    MOCKABLE_VIRTUAL VOID *registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmResidencyController &residencyController);
     void unregisterTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, VOID *trimCallbackHandle);
     MOCKABLE_VIRTUAL void releaseReservedAddress(void *reservedAddress);
     MOCKABLE_VIRTUAL bool reserveValidAddressRange(size_t size, void *&reservedMem);

runtime/os_interface/windows/wddm_residency_controller.cpp

@@ -16,6 +16,9 @@
 namespace OCLRT {
 
 WddmResidencyController::WddmResidencyController(Wddm &wddm, uint32_t osContextId) : wddm(wddm), osContextId(osContextId) {
+}
+
+void WddmResidencyController::registerCallback() {
     this->trimCallbackHandle = wddm.registerTrimCallback(WddmMemoryManager::trimCallback, *this);
 }
 

runtime/os_interface/windows/wddm_residency_controller.h

@@ -46,6 +46,8 @@ class WddmResidencyController {
     MonitoredFence &getMonitoredFence() { return monitoredFence; }
     void resetMonitoredFenceParams(D3DKMT_HANDLE &handle, uint64_t *cpuAddress, D3DGPU_VIRTUAL_ADDRESS &gpuAddress);
 
+    void registerCallback();
+
     void trimResidency(D3DDDI_TRIMRESIDENCYSET_FLAGS flags, uint64_t bytes);
     bool trimResidencyToBudget(uint64_t bytes);
 

unit_tests/mocks/mock_wddm.cpp

@@ -220,6 +220,10 @@ bool WddmMock::reserveValidAddressRange(size_t size, void *&reservedMem) {
     }
     return ret;
 }
+VOID *WddmMock::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmResidencyController &residencyController) {
+    registerTrimCallbackResult.called++;
+    return Wddm::registerTrimCallback(callback, residencyController);
+}
 
 GmmMemory *WddmMock::getGmmMemory() const {
     return gmmMemory.get();

unit_tests/mocks/mock_wddm.h

@@ -77,6 +77,7 @@ class WddmMock : public Wddm {
     void *virtualAlloc(void *inPtr, size_t size, unsigned long flags, unsigned long type) override;
     int virtualFree(void *ptr, size_t size, unsigned long flags) override;
     void releaseReservedAddress(void *reservedAddress) override;
+    VOID *registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, WddmResidencyController &residencyController) override;
     bool reserveValidAddressRange(size_t size, void *&reservedMem);
     GmmMemory *getGmmMemory() const;
     PLATFORM *getGfxPlatform() { return gfxPlatform.get(); }
@@ -110,6 +111,7 @@ class WddmMock : public Wddm {
     WddmMockHelpers::CallResult waitFromCpuResult;
     WddmMockHelpers::CallResult releaseReservedAddressResult;
     WddmMockHelpers::CallResult reserveValidAddressRangeResult;
+    WddmMockHelpers::CallResult registerTrimCallbackResult;
 
     NTSTATUS createAllocationStatus;
     bool mapGpuVaStatus;

unit_tests/os_interface/windows/os_interface_win_tests.cpp

@@ -32,15 +32,17 @@ TEST(OsContextTest, givenWddmWhenCreateOsContextBeforeInitWddmThenOsContextIsNot
     EXPECT_THROW(auto osContext = std::make_unique<OsContext>(&osInterface, 0u), std::exception);
 }
 
-TEST(OsContextTest, givenWddmWhenCreateOsContextAfterInitWddmThenOsContextIsInitialized) {
+TEST(OsContextTest, givenWddmWhenCreateOsContextAfterInitWddmThenOsContextIsInitializedAndTrimCallbackIsRegistered) {
     auto wddm = new WddmMock;
     OSInterface osInterface;
     osInterface.get()->setWddm(wddm);
     wddm->init();
+    EXPECT_EQ(0u, wddm->registerTrimCallbackResult.called);
     auto osContext = std::make_unique<OsContext>(&osInterface, 0u);
     EXPECT_NE(nullptr, osContext->get());
     EXPECT_TRUE(osContext->get()->isInitialized());
     EXPECT_EQ(osContext->get()->getWddm(), wddm);
+    EXPECT_EQ(1u, wddm->registerTrimCallbackResult.called);
 }
 
 TEST(OsContextTest, whenCreateOsContextWithoutOsInterfaceThenOsContextImplIsNotAvailable) {

unit_tests/os_interface/windows/wddm_residency_controller_tests.cpp

@@ -63,6 +63,7 @@ struct WddmResidencyControllerWithGdiTest : ::testing::Test {
 
         residencyController = std::make_unique<MockWddmResidencyController>(*wddm, osContextId);
         wddm->getWddmInterface()->createMonitoredFence(*residencyController);
+        residencyController->registerCallback();
     }
 
     std::unique_ptr<WddmMock> wddm;
@@ -142,16 +143,37 @@ struct WddmResidencyControllerWithGdiAndMemoryManagerTest : ::testing::Test {
     WddmResidencyController *residencyController = nullptr;
 };
 
-TEST(WddmResidencyController, givenWddmResidencyControllerWhenItIsConstructedThenRegisterTrimCallback) {
+TEST(WddmResidencyController, givenWddmResidencyControllerWhenItIsConstructedThenDoNotRegisterTrimCallback) {
     ExecutionEnvironment executionEnvironment;
     auto gdi = new MockGdi();
     auto wddm = std::unique_ptr<WddmMock>{static_cast<WddmMock *>(Wddm::createWddm())};
     wddm->gdi.reset(gdi);
     wddm->init();
 
     std::memset(&gdi->getRegisterTrimNotificationArg(), 0, sizeof(D3DKMT_REGISTERTRIMNOTIFICATION));
+    MockWddmResidencyController residencyController{*wddm, 0u};
+
+    EXPECT_EQ(0u, wddm->registerTrimCallbackResult.called);
+    EXPECT_EQ(nullptr, residencyController.trimCallbackHandle);
+
+    EXPECT_EQ(nullptr, gdi->getRegisterTrimNotificationArg().Callback);
+    EXPECT_EQ(nullptr, gdi->getRegisterTrimNotificationArg().Context);
+    EXPECT_EQ(0u, gdi->getRegisterTrimNotificationArg().hDevice);
+}
+
+TEST(WddmResidencyController, givenWddmResidencyControllerWhenRegisterCallbackThenCallbackIsSetUpProperly) {
+    ExecutionEnvironment executionEnvironment;
+    auto gdi = new MockGdi();
+    auto wddm = std::unique_ptr<WddmMock>{static_cast<WddmMock *>(Wddm::createWddm())};
+    wddm->gdi.reset(gdi);
+    wddm->init();
+
+    std::memset(&gdi->getRegisterTrimNotificationArg(), 0, sizeof(D3DKMT_REGISTERTRIMNOTIFICATION));
+
     WddmResidencyController residencyController{*wddm, 0u};
+    residencyController.registerCallback();
 
+    EXPECT_EQ(1u, wddm->registerTrimCallbackResult.called);
     EXPECT_EQ(reinterpret_cast<PFND3DKMT_TRIMNOTIFICATIONCALLBACK>(WddmMemoryManager::trimCallback), gdi->getRegisterTrimNotificationArg().Callback);
     EXPECT_EQ(reinterpret_cast<void *>(&residencyController), gdi->getRegisterTrimNotificationArg().Context);
     EXPECT_EQ(wddm->getDevice(), gdi->getRegisterTrimNotificationArg().hDevice);