Add device callback for GMM library to notify AUB subcapture

This commit adds a new callback to be called by GMM library
when it allocates/deallocates compressed resources to collect data
about their addresses and sizes and handle their AUB residency.

Change-Id: I075d3ff4cb049cfe626da82892069c4460ea585c

runtime/command_stream/aub_command_stream_receiver_hw.h

@@ -37,6 +37,7 @@ class AUBCommandStreamReceiverHw : public CommandStreamReceiverHw<GfxFamily> {
     typedef CommandStreamReceiverHw<GfxFamily> BaseClass;
     typedef typename AUBFamilyMapper<GfxFamily>::AUB AUB;
     typedef typename AUB::MiContextDescriptorReg MiContextDescriptorReg;
+    using ExternalAllocationsContainer = std::vector<AllocationView>;
 
   public:
     FlushStamp flush(BatchBuffer &batchBuffer, EngineType engineType, ResidencyContainer *allocationsForResidency) override;
@@ -45,7 +46,11 @@ class AUBCommandStreamReceiverHw : public CommandStreamReceiverHw<GfxFamily> {
 
     void processResidency(ResidencyContainer *allocationsForResidency) override;
 
+    void makeResidentExternal(AllocationView &allocationView);
+    void makeNonResidentExternal(uint64_t gpuAddress);
+
     MOCKABLE_VIRTUAL bool writeMemory(GraphicsAllocation &gfxAllocation);
+    MOCKABLE_VIRTUAL bool writeMemory(AllocationView &allocationView);
 
     void activateAubSubCapture(const MultiDispatchInfo &dispatchInfo) override;
 
@@ -114,5 +119,6 @@ class AUBCommandStreamReceiverHw : public CommandStreamReceiverHw<GfxFamily> {
 
   protected:
     bool dumpAubNonWritable = false;
+    ExternalAllocationsContainer externalAllocations;
 };
 } // namespace OCLRT

runtime/command_stream/aub_command_stream_receiver_hw.inl

@@ -518,6 +518,21 @@ void AUBCommandStreamReceiverHw<GfxFamily>::makeResident(GraphicsAllocation &gfx
     gfxAllocation.residencyTaskCount = submissionTaskCount;
 }
 
+template <typename GfxFamily>
+void AUBCommandStreamReceiverHw<GfxFamily>::makeResidentExternal(AllocationView &allocationView) {
+    externalAllocations.push_back(allocationView);
+}
+
+template <typename GfxFamily>
+void AUBCommandStreamReceiverHw<GfxFamily>::makeNonResidentExternal(uint64_t gpuAddress) {
+    for (auto it = externalAllocations.begin(); it != externalAllocations.end(); it++) {
+        if (it->first == gpuAddress) {
+            externalAllocations.erase(it);
+            break;
+        }
+    }
+}
+
 template <typename GfxFamily>
 bool AUBCommandStreamReceiverHw<GfxFamily>::writeMemory(GraphicsAllocation &gfxAllocation) {
     auto cpuAddress = gfxAllocation.getUnderlyingBuffer();
@@ -555,6 +570,12 @@ bool AUBCommandStreamReceiverHw<GfxFamily>::writeMemory(GraphicsAllocation &gfxA
     return true;
 }
 
+template <typename GfxFamily>
+bool AUBCommandStreamReceiverHw<GfxFamily>::writeMemory(AllocationView &allocationView) {
+    GraphicsAllocation gfxAllocation(reinterpret_cast<void *>(allocationView.first), allocationView.second);
+    return writeMemory(gfxAllocation);
+}
+
 template <typename GfxFamily>
 void AUBCommandStreamReceiverHw<GfxFamily>::processResidency(ResidencyContainer *allocationsForResidency) {
     if (subCaptureManager->isSubCaptureMode()) {
@@ -563,6 +584,12 @@ void AUBCommandStreamReceiverHw<GfxFamily>::processResidency(ResidencyContainer
         }
     }
 
+    for (auto &externalAllocation : externalAllocations) {
+        if (!writeMemory(externalAllocation)) {
+            DEBUG_BREAK_IF(externalAllocation.second != 0);
+        }
+    }
+
     auto &residencyAllocations = allocationsForResidency ? *allocationsForResidency : this->getMemoryManager()->getResidencyAllocations();
 
     for (auto &gfxAllocation : residencyAllocations) {

runtime/command_stream/command_stream_receiver.h

@@ -42,6 +42,7 @@ class GraphicsAllocation;
 class IndirectHeap;
 class LinearStream;
 class MemoryManager;
+class GmmPageTableMngr;
 class OSInterface;
 class ExecutionEnvironment;
 
@@ -87,6 +88,8 @@ class CommandStreamReceiver {
     virtual MemoryManager *createMemoryManager(bool enable64kbPages) { return nullptr; }
     void setMemoryManager(MemoryManager *mm);
 
+    virtual GmmPageTableMngr *createPageTableManager() { return nullptr; }
+
     GraphicsAllocation *createAllocationAndHandleResidency(const void *address, size_t size, bool addToDefferFreeList = true);
     void waitForTaskCountAndCleanAllocationList(uint32_t requiredTaskCount, uint32_t allocationType);
 

runtime/enable_gens.cmake

@@ -20,7 +20,7 @@
 
 set(RUNTIME_SRCS_GENX_CPP_WINDOWS
   windows/command_stream_receiver
-  windows/translationtable_callbacks
+  windows/gmm_callbacks
   windows/wddm
 )
 

runtime/execution_environment/execution_environment.cpp

@@ -47,6 +47,9 @@ bool ExecutionEnvironment::initializeCommandStreamReceiver(const HardwareInfo *p
     if (!commandStreamReceiver) {
         return false;
     }
+    if (pHwInfo->capabilityTable.ftrRenderCompressedBuffers || pHwInfo->capabilityTable.ftrRenderCompressedImages) {
+        commandStreamReceiver->createPageTableManager();
+    }
     this->commandStreamReceiver.reset(commandStreamReceiver);
     return true;
 }

runtime/gen10/windows/gmm_callbacks_gen10.cpp

@@ -21,9 +21,10 @@
  */
 
 #include "hw_cmds.h"
-#include "runtime/helpers/translationtable_callbacks.h"
+#include "runtime/helpers/gmm_callbacks.h"
-#include "runtime/helpers/translationtable_callbacks.inl"
+#include "runtime/helpers/gmm_callbacks.inl"
 
 using namespace OCLRT;
 
+template struct DeviceCallbacks<CNLFamily>;
 template struct TTCallbacks<CNLFamily>;

runtime/gen8/windows/gmm_callbacks_gen8.cpp

@@ -21,9 +21,10 @@
  */
 
 #include "hw_cmds.h"
-#include "runtime/helpers/translationtable_callbacks.h"
+#include "runtime/helpers/gmm_callbacks.h"
-#include "runtime/helpers/translationtable_callbacks.inl"
+#include "runtime/helpers/gmm_callbacks.inl"
 
 using namespace OCLRT;
 
+template struct DeviceCallbacks<BDWFamily>;
 template struct TTCallbacks<BDWFamily>;

runtime/gen9/windows/gmm_callbacks_gen9.cpp

@@ -21,9 +21,10 @@
  */
 
 #include "hw_cmds.h"
-#include "runtime/helpers/translationtable_callbacks.h"
+#include "runtime/helpers/gmm_callbacks.h"
-#include "runtime/helpers/translationtable_callbacks.inl"
+#include "runtime/helpers/gmm_callbacks.inl"
 
 using namespace OCLRT;
 
+template struct DeviceCallbacks<SKLFamily>;
 template struct TTCallbacks<SKLFamily>;

runtime/helpers/CMakeLists.txt

@@ -93,8 +93,8 @@ set(RUNTIME_SRCS_HELPERS_BASE
   ${CMAKE_CURRENT_SOURCE_DIR}/validators.h
 )
 set(RUNTIME_SRCS_HELPERS_WINDOWS
-  ${CMAKE_CURRENT_SOURCE_DIR}/translationtable_callbacks.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/gmm_callbacks.h
-  ${CMAKE_CURRENT_SOURCE_DIR}/translationtable_callbacks.inl
+  ${CMAKE_CURRENT_SOURCE_DIR}/gmm_callbacks.inl
   ${CMAKE_CURRENT_SOURCE_DIR}/windows/kmd_notify_properties_windows.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/wddm_helper.h
 )

runtime/helpers/gmm_callbacks.h

@@ -25,6 +25,11 @@
 
 namespace OCLRT {
 
+template <typename GfxFamily>
+struct DeviceCallbacks {
+    static long __stdcall notifyAubCapture(void *csrHandle, uint64_t gfxAddress, size_t gfxSize, bool allocate);
+};
+
 template <typename GfxFamily>
 struct TTCallbacks {
     using MI_LOAD_REGISTER_IMM = typename GfxFamily::MI_LOAD_REGISTER_IMM;

runtime/helpers/gmm_callbacks.inl

@@ -20,15 +20,21 @@
  * OTHER DEALINGS IN THE SOFTWARE.
  */
 
-#include "runtime/helpers/translationtable_callbacks.h"
+#include "runtime/helpers/gmm_callbacks.h"
 #include "runtime/command_stream/linear_stream.h"
 #include "runtime/helpers/hw_helper.h"
 #include <cstdint>
 
 namespace OCLRT {
+
+template <typename GfxFamily>
+long __stdcall DeviceCallbacks<GfxFamily>::notifyAubCapture(void *csrHandle, uint64_t gfxAddress, size_t gfxSize, bool allocate) {
+    return 0;
+}
+
 template <typename GfxFamily>
 int __stdcall TTCallbacks<GfxFamily>::writeL3Address(void *queueHandle, uint64_t l3GfxAddress, uint64_t regOffset) {
     return 0;
-};
+}
 
 } // namespace OCLRT

runtime/memory_manager/residency_container.h

@@ -22,9 +22,11 @@
 
 #pragma once
 #include <vector>
+#include <utility>
 
 namespace OCLRT {
 class GraphicsAllocation;
 using ResidencyContainer = std::vector<GraphicsAllocation *>;
+using AllocationView = std::pair<uint64_t /*address*/, size_t /*size*/>;
 
 } // namespace OCLRT

runtime/os_interface/windows/wddm_device_command_stream.h

@@ -51,11 +51,11 @@ class WddmCommandStreamReceiver : public DeviceCommandStreamReceiver<GfxFamily>
     Wddm *peekWddm() {
         return wddm;
     }
+    GmmPageTableMngr *createPageTableManager() override;
 
   protected:
     void initPageTableManagerRegisters(LinearStream &csr) override;
     void kmDafLockAllocations(ResidencyContainer *allocationsForResidency);
-    GmmPageTableMngr *createPageTableManager();
 
     Wddm *wddm;
     COMMAND_BUFFER_HEADER_REC *commandBufferHeader;

runtime/os_interface/windows/wddm_device_command_stream.inl

@@ -29,8 +29,8 @@
 #include "runtime/command_stream/preemption.h"
 #include "runtime/device/device.h"
 #include "runtime/gmm_helper/page_table_mngr.h"
+#include "runtime/helpers/gmm_callbacks.h"
 #include "runtime/helpers/ptr_math.h"
-#include "runtime/helpers/translationtable_callbacks.h"
 #include "runtime/mem_obj/mem_obj.h"
 #include "runtime/os_interface/windows/wddm/wddm.h"
 #include "runtime/os_interface/windows/wddm_device_command_stream.h"
@@ -77,10 +77,6 @@ WddmCommandStreamReceiver<GfxFamily>::WddmCommandStreamReceiver(const HardwareIn
 
     bool success = this->wddm->init<GfxFamily>();
     DEBUG_BREAK_IF(!success);
-
-    if (hwInfoIn.capabilityTable.ftrRenderCompressedBuffers || hwInfoIn.capabilityTable.ftrRenderCompressedImages) {
-        this->wddm->resetPageTableManager(createPageTableManager());
-    }
 }
 
 template <typename GfxFamily>
@@ -184,6 +180,7 @@ GmmPageTableMngr *WddmCommandStreamReceiver<GfxFamily>::createPageTableManager()
     // clang-format off
     deviceCallbacks.Adapter.KmtHandle         = wddm->getAdapter();
     deviceCallbacks.hDevice.KmtHandle         = wddm->getDevice();
+    deviceCallbacks.hCsr            = static_cast<CommandStreamReceiverHw<GfxFamily> *>(this);
     deviceCallbacks.PagingQueue     = wddm->getPagingQueue();
     deviceCallbacks.PagingFence     = wddm->getPagingQueueSyncObject();
 
@@ -198,11 +195,14 @@ GmmPageTableMngr *WddmCommandStreamReceiver<GfxFamily>::createPageTableManager()
     deviceCallbacks.DevCbPtrs.KmtCbPtrs.pfnLock         = gdi->lock2;
     deviceCallbacks.DevCbPtrs.KmtCbPtrs.pfnUnLock       = gdi->unlock2;
     deviceCallbacks.DevCbPtrs.KmtCbPtrs.pfnEscape       = gdi->escape;
+    deviceCallbacks.DevCbPtrs.KmtCbPtrs.pfnNotifyAubCapture = DeviceCallbacks<GfxFamily>::notifyAubCapture;
 
     ttCallbacks.pfWriteL3Adr        = TTCallbacks<GfxFamily>::writeL3Address;
     // clang-format on
 
-    return GmmPageTableMngr::create(&deviceCallbacks, TT_TYPE::TRTT | TT_TYPE::AUXTT, &ttCallbacks);
+    GmmPageTableMngr *gmmPageTableMngr = GmmPageTableMngr::create(&deviceCallbacks, TT_TYPE::TRTT | TT_TYPE::AUXTT, &ttCallbacks);
+    this->wddm->resetPageTableManager(gmmPageTableMngr);
+    return gmmPageTableMngr;
 }
 
 template <typename GfxFamily>

unit_tests/command_stream/aub_command_stream_receiver_tests.cpp

@@ -764,6 +764,23 @@ HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenGraphic
     EXPECT_FALSE(aubCsr->writeMemory(gfxAllocation));
 }
 
+HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenAllocationDataIsPassedInAllocationViewThenWriteMemoryIsAllowed) {
+    auto aubCsr = std::make_unique<AUBCommandStreamReceiverHw<FamilyType>>(*platformDevices[0], "", true, *pDevice->executionEnvironment);
+    size_t size = 100;
+    auto ptr = std::make_unique<char[]>(size);
+    auto addr = reinterpret_cast<uint64_t>(ptr.get());
+    AllocationView allocationView(addr, size);
+
+    EXPECT_TRUE(aubCsr->writeMemory(allocationView));
+}
+
+HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenAllocationSizeInAllocationViewIsZeroThenWriteMemoryIsNotAllowed) {
+    auto aubCsr = std::make_unique<AUBCommandStreamReceiverHw<FamilyType>>(*platformDevices[0], "", true, *pDevice->executionEnvironment);
+    AllocationView allocationView(0x1234, 0);
+
+    EXPECT_FALSE(aubCsr->writeMemory(allocationView));
+}
+
 HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverInSubCaptureModeWhenAubSubCaptureIsActivatedThenFileIsOpened) {
     DebugManagerStateRestore stateRestore;
     std::unique_ptr<MockAubCsr<FamilyType>> aubCsr(new MockAubCsr<FamilyType>(*platformDevices[0], "", false, *pDevice->executionEnvironment));
@@ -1776,6 +1793,97 @@ HWTEST_F(AubCommandStreamReceiverTests, givenPhysicalAddressWhenSetGttEntryIsCal
     EXPECT_FALSE(entry.pageConfig.LocalMemory);
 }
 
+template <typename GfxFamily>
+struct MockAubCsrToTestExternalAllocations : public AUBCommandStreamReceiverHw<GfxFamily> {
+    using AUBCommandStreamReceiverHw<GfxFamily>::AUBCommandStreamReceiverHw;
+    using AUBCommandStreamReceiverHw<GfxFamily>::externalAllocations;
+
+    bool writeMemory(AllocationView &allocationView) override {
+        writeMemoryParametrization.wasCalled = true;
+        writeMemoryParametrization.receivedAllocationView = allocationView;
+        writeMemoryParametrization.statusToReturn = (0 != allocationView.second) ? true : false;
+        return writeMemoryParametrization.statusToReturn;
+    }
+    struct WriteMemoryParametrization {
+        bool wasCalled = false;
+        AllocationView receivedAllocationView = {};
+        bool statusToReturn = false;
+    } writeMemoryParametrization;
+};
+
+HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenMakeResidentExternalIsCalledThenGivenAllocationViewShouldBeAddedToExternalAllocations) {
+    auto aubCsr = std::make_unique<MockAubCsrToTestExternalAllocations<FamilyType>>(*platformDevices[0], "", true, *pDevice->executionEnvironment);
+    size_t size = 100;
+    auto ptr = std::make_unique<char[]>(size);
+    auto addr = reinterpret_cast<uint64_t>(ptr.get());
+    AllocationView externalAllocation(addr, size);
+
+    ASSERT_EQ(0u, aubCsr->externalAllocations.size());
+    aubCsr->makeResidentExternal(externalAllocation);
+    EXPECT_EQ(1u, aubCsr->externalAllocations.size());
+    EXPECT_EQ(addr, aubCsr->externalAllocations[0].first);
+    EXPECT_EQ(size, aubCsr->externalAllocations[0].second);
+}
+
+HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenMakeNonResidentExternalIsCalledThenMatchingAllocationViewShouldBeRemovedFromExternalAllocations) {
+    auto aubCsr = std::make_unique<MockAubCsrToTestExternalAllocations<FamilyType>>(*platformDevices[0], "", true, *pDevice->executionEnvironment);
+    size_t size = 100;
+    auto ptr = std::make_unique<char[]>(size);
+    auto addr = reinterpret_cast<uint64_t>(ptr.get());
+    AllocationView externalAllocation(addr, size);
+    aubCsr->makeResidentExternal(externalAllocation);
+
+    ASSERT_EQ(1u, aubCsr->externalAllocations.size());
+    aubCsr->makeNonResidentExternal(addr);
+    EXPECT_EQ(0u, aubCsr->externalAllocations.size());
+}
+
+HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenMakeNonResidentExternalIsCalledThenNonMatchingAllocationViewShouldNotBeRemovedFromExternalAllocations) {
+    auto aubCsr = std::make_unique<MockAubCsrToTestExternalAllocations<FamilyType>>(*platformDevices[0], "", true, *pDevice->executionEnvironment);
+    size_t size = 100;
+    auto ptr = std::make_unique<char[]>(size);
+    auto addr = reinterpret_cast<uint64_t>(ptr.get());
+    AllocationView externalAllocation(addr, size);
+    aubCsr->makeResidentExternal(externalAllocation);
+
+    ASSERT_EQ(1u, aubCsr->externalAllocations.size());
+    aubCsr->makeNonResidentExternal(0);
+    EXPECT_EQ(1u, aubCsr->externalAllocations.size());
+}
+
+HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenProcessResidencyIsCalledThenExternalAllocationsShouldBeMadeResident) {
+    auto aubCsr = std::make_unique<MockAubCsrToTestExternalAllocations<FamilyType>>(*platformDevices[0], "", true, *pDevice->executionEnvironment);
+    size_t size = 100;
+    auto ptr = std::make_unique<char[]>(size);
+    auto addr = reinterpret_cast<uint64_t>(ptr.get());
+    AllocationView externalAllocation(addr, size);
+    aubCsr->makeResidentExternal(externalAllocation);
+
+    ASSERT_EQ(1u, aubCsr->externalAllocations.size());
+    ResidencyContainer allocationsForResidency;
+    aubCsr->processResidency(&allocationsForResidency);
+
+    EXPECT_TRUE(aubCsr->writeMemoryParametrization.wasCalled);
+    EXPECT_EQ(addr, aubCsr->writeMemoryParametrization.receivedAllocationView.first);
+    EXPECT_EQ(size, aubCsr->writeMemoryParametrization.receivedAllocationView.second);
+    EXPECT_TRUE(aubCsr->writeMemoryParametrization.statusToReturn);
+}
+
+HWTEST_F(AubCommandStreamReceiverTests, givenAubCommandStreamReceiverWhenProcessResidencyIsCalledThenExternalAllocationWithZeroSizeShouldNotBeMadeResident) {
+    auto aubCsr = std::make_unique<MockAubCsrToTestExternalAllocations<FamilyType>>(*platformDevices[0], "", true, *pDevice->executionEnvironment);
+    AllocationView externalAllocation(0, 0);
+    aubCsr->makeResidentExternal(externalAllocation);
+
+    ASSERT_EQ(1u, aubCsr->externalAllocations.size());
+    ResidencyContainer allocationsForResidency;
+    aubCsr->processResidency(&allocationsForResidency);
+
+    EXPECT_TRUE(aubCsr->writeMemoryParametrization.wasCalled);
+    EXPECT_EQ(0u, aubCsr->writeMemoryParametrization.receivedAllocationView.first);
+    EXPECT_EQ(0u, aubCsr->writeMemoryParametrization.receivedAllocationView.second);
+    EXPECT_FALSE(aubCsr->writeMemoryParametrization.statusToReturn);
+}
+
 #if defined(__clang__)
 #pragma clang diagnostic pop
 #endif

unit_tests/command_stream/command_stream_receiver_flush_task_tests.cpp

@@ -1366,6 +1366,8 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, CreateCommandStreamReceiverHw) {
 
     MemoryManager *mm = csrHw->createMemoryManager(false);
     EXPECT_EQ(nullptr, mm);
+    GmmPageTableMngr *ptm = csrHw->createPageTableManager();
+    EXPECT_EQ(nullptr, ptm);
     delete csrHw;
 
     DebugManager.flags.SetCommandStreamReceiver.set(0);

unit_tests/execution_environment/execution_environment_tests.cpp

@@ -178,3 +178,26 @@ TEST(ExecutionEnvironment, givenMultipleDevicesWhenTheyAreCreatedTheyAllReuseThe
     EXPECT_EQ(&commandStreamReceiver, &device->getCommandStreamReceiver());
     EXPECT_EQ(memoryManager, device2->getMemoryManager());
 }
+
+typedef ::testing::Test ExecutionEnvironmentHw;
+
+HWTEST_F(ExecutionEnvironmentHw, givenExecutionEnvironmentWhenCommandStreamReceiverIsInitializedForCompressedBuffersThenCreatePageTableManagerIsCalled) {
+    ExecutionEnvironment executionEnvironment;
+    HardwareInfo localHwInfo = *platformDevices[0];
+    localHwInfo.capabilityTable.ftrRenderCompressedBuffers = true;
+    executionEnvironment.initializeCommandStreamReceiver(&localHwInfo);
+    auto csr = static_cast<UltCommandStreamReceiver<FamilyType> *>(executionEnvironment.commandStreamReceiver.get());
+    ASSERT_NE(nullptr, csr);
+    EXPECT_TRUE(csr->createPageTableManagerCalled);
+}
+
+HWTEST_F(ExecutionEnvironmentHw, givenExecutionEnvironmentWhenCommandStreamReceiverIsInitializedForCompressedImagesThenCreatePageTableManagerIsCalled) {
+    ExecutionEnvironment executionEnvironment;
+    HardwareInfo localHwInfo = *platformDevices[0];
+    localHwInfo.capabilityTable.ftrRenderCompressedImages = true;
+    executionEnvironment.initializeCommandStreamReceiver(&localHwInfo);
+    EXPECT_NE(nullptr, executionEnvironment.commandStreamReceiver);
+    auto csr = static_cast<UltCommandStreamReceiver<FamilyType> *>(executionEnvironment.commandStreamReceiver.get());
+    ASSERT_NE(nullptr, csr);
+    EXPECT_TRUE(csr->createPageTableManagerCalled);
+}

unit_tests/gen10/windows/CMakeLists.txt

@@ -20,10 +20,10 @@
 
 set(IGDRCL_SRCS_tests_gen10_windows
   ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
-  ${CMAKE_CURRENT_SOURCE_DIR}/os_interface_tests_gen10.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/gmm_callbacks_tests_gen10.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/translationtable_callbacks_tests_gen10.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_mapper_tests_gen10.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/hw_info_config_tests_gen10.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/os_interface_tests_gen10.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_mapper_tests_gen10.cpp
 )
 if(WIN32)
   target_sources(igdrcl_tests PRIVATE ${IGDRCL_SRCS_tests_gen10_windows})

unit_tests/gen10/windows/gmm_callbacks_tests_gen10.cpp

@@ -20,13 +20,17 @@
  * OTHER DEALINGS IN THE SOFTWARE.
  */
 
-#include "runtime/helpers/translationtable_callbacks.h"
+#include "runtime/helpers/gmm_callbacks.h"
 #include "test.h"
 
 using namespace OCLRT;
 
-typedef ::testing::Test Gen10TTCallbacksTests;
+typedef ::testing::Test Gen10GmmCallbacksTests;
 
-GEN10TEST_F(Gen10TTCallbacksTests, notSupported) {
+GEN10TEST_F(Gen10GmmCallbacksTests, notSupportedDeviceCallback) {
+    EXPECT_EQ(0, DeviceCallbacks<FamilyType>::notifyAubCapture(nullptr, 0, 0, false));
+}
+
+GEN10TEST_F(Gen10GmmCallbacksTests, notSupportedTTCallback) {
     EXPECT_EQ(0, TTCallbacks<FamilyType>::writeL3Address(nullptr, 1, 2));
 }

unit_tests/gen8/windows/CMakeLists.txt

@@ -20,8 +20,8 @@
 
 set(IGDRCL_SRCS_tests_gen8_windows
   ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
+  ${CMAKE_CURRENT_SOURCE_DIR}/gmm_callbacks_tests_gen8.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/wddm_mapper_tests_gen8.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/translationtable_callbacks_tests_gen8.cpp
 )
 if(WIN32)
   target_sources(igdrcl_tests PRIVATE ${IGDRCL_SRCS_tests_gen8_windows})

unit_tests/gen8/windows/gmm_callbacks_tests_gen8.cpp

@@ -20,13 +20,17 @@
  * OTHER DEALINGS IN THE SOFTWARE.
  */
 
-#include "runtime/helpers/translationtable_callbacks.h"
+#include "runtime/helpers/gmm_callbacks.h"
 #include "test.h"
 
 using namespace OCLRT;
 
-typedef ::testing::Test Gen9TTCallbacksTests;
+typedef ::testing::Test Gen8GmmCallbacksTests;
 
-GEN9TEST_F(Gen9TTCallbacksTests, notSupported) {
+GEN8TEST_F(Gen8GmmCallbacksTests, notSupportedDeviceCallback) {
+    EXPECT_EQ(0, DeviceCallbacks<FamilyType>::notifyAubCapture(nullptr, 0, 0, false));
+}
+
+GEN8TEST_F(Gen8GmmCallbacksTests, notSupportedTTCallback) {
     EXPECT_EQ(0, TTCallbacks<FamilyType>::writeL3Address(nullptr, 1, 2));
 }

unit_tests/gen9/windows/CMakeLists.txt

@@ -20,9 +20,9 @@
 
 set(IGDRCL_SRCS_tests_gen9_windows
   ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt
+  ${CMAKE_CURRENT_SOURCE_DIR}/gmm_callbacks_tests_gen9.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/os_interface_tests_gen9.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/wddm_mapper_tests_gen9.cpp
-  ${CMAKE_CURRENT_SOURCE_DIR}/translationtable_callbacks_tests_gen9.cpp
 )
 if(WIN32)
   target_sources(igdrcl_tests PRIVATE ${IGDRCL_SRCS_tests_gen9_windows})

unit_tests/gen9/windows/gmm_callbacks_tests_gen9.cpp

@@ -20,13 +20,17 @@
  * OTHER DEALINGS IN THE SOFTWARE.
  */
 
-#include "runtime/helpers/translationtable_callbacks.h"
+#include "runtime/helpers/gmm_callbacks.h"
 #include "test.h"
 
 using namespace OCLRT;
 
-typedef ::testing::Test Gen8TTCallbacksTests;
+typedef ::testing::Test Gen9GmmCallbacksTests;
 
-GEN8TEST_F(Gen8TTCallbacksTests, notSupported) {
+GEN9TEST_F(Gen9GmmCallbacksTests, notSupportedDeviceCallback) {
+    EXPECT_EQ(0, DeviceCallbacks<FamilyType>::notifyAubCapture(nullptr, 0, 0, false));
+}
+
+GEN9TEST_F(Gen9GmmCallbacksTests, notSupportedTTCallback) {
     EXPECT_EQ(0, TTCallbacks<FamilyType>::writeL3Address(nullptr, 1, 2));
 }

unit_tests/libult/ult_command_stream_receiver.h

@@ -28,6 +28,8 @@
 
 namespace OCLRT {
 
+class GmmPageTableMngr;
+
 template <typename GfxFamily>
 class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily> {
     using BaseClass = CommandStreamReceiverHw<GfxFamily>;
@@ -80,6 +82,11 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily> {
         return memoryManager;
     }
 
+    virtual GmmPageTableMngr *createPageTableManager() override {
+        createPageTableManagerCalled = true;
+        return nullptr;
+    }
+
     void overrideCsrSizeReqFlags(CsrSizeRequestFlags &flags) { this->csrSizeRequestFlags = flags; }
     bool isPreambleProgrammed() const { return this->isPreambleSent; }
     bool isGSBAFor32BitProgrammed() const { return this->GSBAFor32BitProgrammed; }
@@ -133,6 +140,7 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily> {
         initProgrammingFlagsCalled = true;
     }
 
+    bool createPageTableManagerCalled = false;
     bool activateAubSubCaptureCalled = false;
     bool flushBatchedSubmissionsCalled = false;
     bool initProgrammingFlagsCalled = false;

unit_tests/os_interface/windows/device_command_stream_tests.cpp

@@ -28,8 +28,8 @@
 #include "runtime/command_stream/preemption.h"
 #include "runtime/gen_common/hw_cmds.h"
 #include "runtime/helpers/built_ins_helper.h"
+#include "runtime/helpers/gmm_callbacks.h"
 #include "runtime/helpers/options.h"
-#include "runtime/helpers/translationtable_callbacks.h"
 #include "runtime/mem_obj/buffer.h"
 #include "runtime/memory_manager/memory_manager.h"
 #include "runtime/os_interface/windows/wddm_device_command_stream.h"
@@ -830,6 +830,7 @@ HWTEST_F(WddmCsrCompressionTests, givenEnabledCompressionWhenInitializedThenCrea
         createMockWddm();
         EXPECT_EQ(nullptr, myMockWddm->getPageTableManager());
         MockWddmCsr<FamilyType> mockWddmCsr(hwInfo, myMockWddm, *executionEnvironment);
+        mockWddmCsr.createPageTableManager();
         ASSERT_NE(nullptr, myMockWddm->getPageTableManager());
 
         auto mockMngr = reinterpret_cast<MockGmmPageTableMngr *>(myMockWddm->getPageTableManager());
@@ -841,6 +842,7 @@ HWTEST_F(WddmCsrCompressionTests, givenEnabledCompressionWhenInitializedThenCrea
         // clang-format off
         expectedDeviceCb.Adapter.KmtHandle = myMockWddm->getAdapter();
         expectedDeviceCb.hDevice.KmtHandle = myMockWddm->getDevice();
+        expectedDeviceCb.hCsr = &mockWddmCsr;
         expectedDeviceCb.PagingQueue = myMockWddm->getPagingQueue();
         expectedDeviceCb.PagingFence = myMockWddm->getPagingQueueSyncObject();
 
@@ -855,6 +857,7 @@ HWTEST_F(WddmCsrCompressionTests, givenEnabledCompressionWhenInitializedThenCrea
         expectedDeviceCb.DevCbPtrs.KmtCbPtrs.pfnLock = myGdi->lock2;
         expectedDeviceCb.DevCbPtrs.KmtCbPtrs.pfnUnLock = myGdi->unlock2;
         expectedDeviceCb.DevCbPtrs.KmtCbPtrs.pfnEscape = myGdi->escape;
+        expectedDeviceCb.DevCbPtrs.KmtCbPtrs.pfnNotifyAubCapture = DeviceCallbacks<FamilyType>::notifyAubCapture;
 
         expectedTTCallbacks.pfWriteL3Adr = TTCallbacks<FamilyType>::writeL3Address;
         // clang-format on
@@ -883,8 +886,8 @@ HWTEST_F(WddmCsrCompressionTests, givenEnabledCompressionWhenFlushingThenInitTra
         ExecutionEnvironment *executionEnvironment = new ExecutionEnvironment;
         setCompressionEnabled(compressionEnabled[i][0], compressionEnabled[i][1]);
         createMockWddm();
-
         auto mockWddmCsr = new MockWddmCsr<FamilyType>(hwInfo, myMockWddm, *executionEnvironment);
+        mockWddmCsr->createPageTableManager();
         mockWddmCsr->overrideDispatchPolicy(DispatchMode::BatchedDispatch);
 
         executionEnvironment->commandStreamReceiver.reset(mockWddmCsr);