fix: aub residency in HW+AUB mode

Related-To: NEO-9230
Signed-off-by: Fabian Zwolinski <fabian.zwolinski@intel.com>

opencl/test/unit_test/mt_tests/os_interface/linux/drm_memory_manager_mt_tests.cpp

@@ -40,7 +40,7 @@ TEST(DrmMemoryManagerTest, givenDrmMemoryManagerWhenSharedAllocationIsCreatedFro
     auto mock = new MockDrm(0, *executionEnvironment.rootDeviceEnvironments[0]);
     mock->setupIoctlHelper(defaultHwInfo->platform.eProductFamily);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
     auto memoryManager = std::make_unique<TestedDrmMemoryManager>(executionEnvironment);
 
@@ -112,7 +112,7 @@ TEST(DrmMemoryManagerTest, givenMultipleThreadsWhenSharedAllocationIsCreatedThen
     auto mock = new MockDrm(0, *executionEnvironment.rootDeviceEnvironments[0]);
     mock->setupIoctlHelper(defaultHwInfo->platform.eProductFamily);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
     auto memoryManager = std::make_unique<TestedDrmMemoryManager>(executionEnvironment);
 

opencl/test/unit_test/os_interface/linux/linux_create_command_queue_with_properties_tests.cpp

@@ -29,7 +29,7 @@ struct ClCreateCommandQueueWithPropertiesLinux : public UltCommandStreamReceiver
         auto osInterface = new OSInterface();
         osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface.reset(osInterface);
-        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, rootDeviceIndex);
+        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, rootDeviceIndex, false);
         executionEnvironment->memoryManager.reset(new TestedDrmMemoryManager(*executionEnvironment));
         mdevice = std::make_unique<MockClDevice>(MockDevice::create<MockDevice>(executionEnvironment, rootDeviceIndex));
 

shared/source/command_stream/command_stream_receiver.h

@@ -16,6 +16,8 @@
 #include "shared/source/helpers/options.h"
 #include "shared/source/utilities/spinlock.h"
 
+#include "aubstream/allocation_params.h"
+
 #include <atomic>
 #include <cstddef>
 #include <cstdint>
@@ -245,6 +247,7 @@ class CommandStreamReceiver {
     virtual bool expectMemory(const void *gfxAddress, const void *srcAddress, size_t length, uint32_t compareOperation);
     MOCKABLE_VIRTUAL bool writeMemory(GraphicsAllocation &gfxAllocation) { return writeMemory(gfxAllocation, false, 0, 0); }
     virtual bool writeMemory(GraphicsAllocation &gfxAllocation, bool isChunkCopy, uint64_t gpuVaChunkOffset, size_t chunkSize) { return false; }
+    virtual void writeMemoryAub(aub_stream::AllocationParams &allocationParams){};
 
     virtual bool isMultiOsContextCapable() const = 0;
 

shared/source/command_stream/command_stream_receiver_simulated_common_hw.h

@@ -7,6 +7,7 @@
 
 #pragma once
 #include "shared/source/command_stream/command_stream_receiver_hw.h"
+#include "shared/source/helpers/hardware_context_controller.h"
 #include "shared/source/memory_manager/memory_banks.h"
 
 #include "aub_mapper_common.h"
@@ -20,7 +21,6 @@ struct AubStream;
 namespace NEO {
 class AddressMapper;
 class GraphicsAllocation;
-class HardwareContextController;
 template <typename GfxFamily>
 class CommandStreamReceiverSimulatedCommonHw : public CommandStreamReceiverHw<GfxFamily> {
   protected:
@@ -59,6 +59,9 @@ class CommandStreamReceiverSimulatedCommonHw : public CommandStreamReceiverHw<Gf
     virtual void writeMemory(uint64_t gpuAddress, void *cpuAddress, size_t size, uint32_t memoryBank, uint64_t entryBits) = 0;
     virtual void writeMemoryWithAubManager(GraphicsAllocation &graphicsAllocation, bool isChunkCopy, uint64_t gpuVaChunkOffset, size_t chunkSize) = 0;
     virtual void writeMMIO(uint32_t offset, uint32_t value) = 0;
+    void writeMemoryAub(aub_stream::AllocationParams &allocationParams) override {
+        hardwareContextController->writeMemory(allocationParams);
+    }
 
     virtual void setAubWritable(bool writable, GraphicsAllocation &graphicsAllocation) = 0;
     virtual bool isAubWritable(GraphicsAllocation &graphicsAllocation) const = 0;

shared/source/os_interface/aub_memory_operations_handler.cpp

@@ -7,9 +7,13 @@
 
 #include "shared/source/os_interface/aub_memory_operations_handler.h"
 
+#include "shared/source/aub/aub_helper.h"
 #include "shared/source/aub_mem_dump/aub_mem_dump.h"
+#include "shared/source/command_stream/command_stream_receiver.h"
+#include "shared/source/device/device.h"
 #include "shared/source/gmm_helper/cache_settings_helper.h"
 #include "shared/source/gmm_helper/gmm.h"
+#include "shared/source/gmm_helper/gmm_helper.h"
 #include "shared/source/memory_manager/graphics_allocation.h"
 
 #include "aubstream/allocation_params.h"
@@ -29,7 +33,12 @@ MemoryOperationsStatus AubMemoryOperationsHandler::makeResident(Device *device,
     auto lock = acquireLock(resourcesLock);
     int hint = AubMemDump::DataTypeHintValues::TraceNotype;
     for (const auto &allocation : gfxAllocations) {
-        aub_stream::AllocationParams params(allocation->getGpuAddress(),
+        if (!isAubWritable(*allocation, device)) {
+            continue;
+        }
+
+        uint64_t gpuAddress = device ? device->getGmmHelper()->decanonize(allocation->getGpuAddress()) : allocation->getGpuAddress();
+        aub_stream::AllocationParams params(gpuAddress,
                                             allocation->getUnderlyingBuffer(),
                                             allocation->getUnderlyingBufferSize(),
                                             allocation->storageInfo.getMemoryBanks(),
@@ -43,10 +52,19 @@ MemoryOperationsStatus AubMemoryOperationsHandler::makeResident(Device *device,
             params.additionalParams.uncached = CacheSettingsHelper::isUncachedType(gmm->resourceParams.Usage);
         }
 
-        aubManager->writeMemory2(params);
+        if (allocation->storageInfo.cloningOfPageTables || !allocation->isAllocatedInLocalMemoryPool()) {
+            aubManager->writeMemory2(params);
+        } else {
+            device->getDefaultEngine().commandStreamReceiver->writeMemoryAub(params);
+        }
+
         if (!allocation->getAubInfo().writeMemoryOnly) {
             residentAllocations.push_back(allocation);
         }
+
+        if (AubHelper::isOneTimeAubWritableAllocationType(allocation->getAllocationType())) {
+            setAubWritable(false, *allocation, device);
+        }
     }
     return MemoryOperationsStatus::SUCCESS;
 }
@@ -79,8 +97,44 @@ MemoryOperationsStatus AubMemoryOperationsHandler::isResident(Device *device, Gr
         return MemoryOperationsStatus::SUCCESS;
     }
 }
+
 void AubMemoryOperationsHandler::setAubManager(aub_stream::AubManager *aubManager) {
     this->aubManager = aubManager;
 }
 
+bool AubMemoryOperationsHandler::isAubWritable(GraphicsAllocation &graphicsAllocation, Device *device) const {
+    if (!device) {
+        return false;
+    }
+    auto bank = static_cast<uint32_t>(getMemoryBanksBitfield(&graphicsAllocation, device).to_ulong());
+    if (bank == 0u || graphicsAllocation.storageInfo.cloningOfPageTables) {
+        bank = GraphicsAllocation::defaultBank;
+    }
+    return graphicsAllocation.isAubWritable(bank);
+}
+
+void AubMemoryOperationsHandler::setAubWritable(bool writable, GraphicsAllocation &graphicsAllocation, Device *device) {
+    if (!device) {
+        return;
+    }
+    auto bank = static_cast<uint32_t>(getMemoryBanksBitfield(&graphicsAllocation, device).to_ulong());
+    if (bank == 0u || graphicsAllocation.storageInfo.cloningOfPageTables) {
+        bank = GraphicsAllocation::defaultBank;
+    }
+    graphicsAllocation.setAubWritable(writable, bank);
+}
+
+DeviceBitfield AubMemoryOperationsHandler::getMemoryBanksBitfield(GraphicsAllocation *allocation, Device *device) const {
+    if (allocation->getMemoryPool() == MemoryPool::LocalMemory) {
+        if (allocation->storageInfo.memoryBanks.any()) {
+            if (allocation->storageInfo.cloningOfPageTables ||
+                device->getDefaultEngine().commandStreamReceiver->isMultiOsContextCapable()) {
+                return allocation->storageInfo.memoryBanks;
+            }
+        }
+        return device->getDeviceBitfield();
+    }
+    return {};
+}
+
 } // namespace NEO

shared/source/os_interface/aub_memory_operations_handler.h

@@ -1,11 +1,12 @@
 /*
- * Copyright (C) 2019-2022 Intel Corporation
+ * Copyright (C) 2019-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
  */
 
 #pragma once
+#include "shared/source/helpers/device_bitfield.h"
 #include "shared/source/memory_manager/memory_operations_handler.h"
 #include "shared/source/utilities/spinlock.h"
 
@@ -30,7 +31,11 @@ class AubMemoryOperationsHandler : public MemoryOperationsHandler {
 
     void setAubManager(aub_stream::AubManager *aubManager);
 
+    bool isAubWritable(GraphicsAllocation &graphicsAllocation, Device *device) const;
+    void setAubWritable(bool writable, GraphicsAllocation &graphicsAllocation, Device *device);
+
   protected:
+    DeviceBitfield getMemoryBanksBitfield(GraphicsAllocation *allocation, Device *device) const;
     [[nodiscard]] MOCKABLE_VIRTUAL std::unique_lock<SpinLock> acquireLock(SpinLock &lock) {
         return std::unique_lock<SpinLock>{lock};
     }

shared/source/os_interface/linux/CMakeLists.txt

@@ -38,6 +38,7 @@ set(NEO_CORE_OS_INTERFACE_LINUX
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_create.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_default.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_default.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_with_aub_dump.h
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_create_multi_host_allocation.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_version.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_wrappers_checks.cpp

shared/source/os_interface/linux/drm_memory_operations_handler.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2022 Intel Corporation
+ * Copyright (C) 2019-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -25,7 +25,7 @@ class DrmMemoryOperationsHandler : public MemoryOperationsHandler {
 
     virtual MemoryOperationsStatus evictUnusedAllocations(bool waitForCompletion, bool isLockNeeded) = 0;
 
-    static std::unique_ptr<DrmMemoryOperationsHandler> create(Drm &drm, uint32_t rootDeviceIndex);
+    static std::unique_ptr<DrmMemoryOperationsHandler> create(Drm &drm, uint32_t rootDeviceIndex, bool withAubDump);
 
     uint32_t getRootDeviceIndex() const {
         return this->rootDeviceIndex;

shared/source/os_interface/linux/drm_memory_operations_handler_create.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2022 Intel Corporation
+ * Copyright (C) 2020-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -8,18 +8,28 @@
 #include "shared/source/execution_environment/root_device_environment.h"
 #include "shared/source/os_interface/linux/drm_memory_operations_handler_bind.h"
 #include "shared/source/os_interface/linux/drm_memory_operations_handler_default.h"
+#include "shared/source/os_interface/linux/drm_memory_operations_handler_with_aub_dump.h"
 #include "shared/source/os_interface/linux/drm_neo.h"
 
 namespace NEO {
 
-std::unique_ptr<DrmMemoryOperationsHandler> DrmMemoryOperationsHandler::create(Drm &drm, uint32_t rootDeviceIndex) {
+std::unique_ptr<DrmMemoryOperationsHandler> DrmMemoryOperationsHandler::create(Drm &drm, uint32_t rootDeviceIndex, bool withAubDump) {
     bool useVmBind = drm.isVmBindAvailable();
+    auto rootDeviceEnv = drm.getRootDeviceEnvironment().executionEnvironment.rootDeviceEnvironments[rootDeviceIndex].get();
 
     if (useVmBind) {
-        return std::make_unique<DrmMemoryOperationsHandlerBind>(drm.getRootDeviceEnvironment(), rootDeviceIndex);
+        if (withAubDump) {
+            return std::make_unique<DrmMemoryOperationsHandlerWithAubDump<DrmMemoryOperationsHandlerBind>>(*rootDeviceEnv, rootDeviceIndex);
+        } else {
+            return std::make_unique<DrmMemoryOperationsHandlerBind>(drm.getRootDeviceEnvironment(), rootDeviceIndex);
+        }
     }
 
-    return std::make_unique<DrmMemoryOperationsHandlerDefault>(rootDeviceIndex);
+    if (withAubDump) {
+        return std::make_unique<DrmMemoryOperationsHandlerWithAubDump<DrmMemoryOperationsHandlerDefault>>(*rootDeviceEnv, rootDeviceIndex);
+    } else {
+        return std::make_unique<DrmMemoryOperationsHandlerDefault>(rootDeviceIndex);
+    }
 }
 
 } // namespace NEO

shared/source/os_interface/linux/drm_memory_operations_handler_default.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2022 Intel Corporation
+ * Copyright (C) 2020-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -12,9 +12,11 @@
 
 namespace NEO {
 class OsContextLinux;
+struct RootDeviceEnvironment;
 class DrmMemoryOperationsHandlerDefault : public DrmMemoryOperationsHandler {
   public:
     DrmMemoryOperationsHandlerDefault(uint32_t rootDeviceIndex);
+    DrmMemoryOperationsHandlerDefault(const RootDeviceEnvironment &rootDeviceEnvironment, uint32_t rootDeviceIndex) : DrmMemoryOperationsHandlerDefault(rootDeviceIndex) {}
     ~DrmMemoryOperationsHandlerDefault() override;
 
     MemoryOperationsStatus makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) override;

shared/source/os_interface/linux/drm_memory_operations_handler_with_aub_dump.h

@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+#include "shared/source/aub/aub_center.h"
+#include "shared/source/command_stream/command_stream_receiver.h"
+#include "shared/source/execution_environment/execution_environment.h"
+#include "shared/source/execution_environment/root_device_environment.h"
+#include "shared/source/helpers/gfx_core_helper.h"
+#include "shared/source/os_interface/aub_memory_operations_handler.h"
+
+namespace NEO {
+class Drm;
+
+template <typename BaseOperationsHandler>
+class DrmMemoryOperationsHandlerWithAubDump : public BaseOperationsHandler {
+  public:
+    DrmMemoryOperationsHandlerWithAubDump(RootDeviceEnvironment &rootDeviceEnvironment, uint32_t rootDeviceIndex)
+        : BaseOperationsHandler(rootDeviceEnvironment, rootDeviceIndex) {
+
+        if (!rootDeviceEnvironment.aubCenter) {
+            auto &gfxCoreHelper = rootDeviceEnvironment.getHelper<GfxCoreHelper>();
+            auto hardwareInfo = rootDeviceEnvironment.getMutableHardwareInfo();
+            auto localMemoryEnabled = gfxCoreHelper.getEnableLocalMemory(*hardwareInfo);
+            rootDeviceEnvironment.initGmm();
+            rootDeviceEnvironment.initAubCenter(localMemoryEnabled, "", static_cast<CommandStreamReceiverType>(CommandStreamReceiverType::CSR_HW_WITH_AUB));
+        }
+
+        const auto aubCenter = rootDeviceEnvironment.aubCenter.get();
+        aubMemoryOperationsHandler = std::make_unique<AubMemoryOperationsHandler>(aubCenter->getAubManager());
+    };
+
+    ~DrmMemoryOperationsHandlerWithAubDump() override = default;
+
+    MemoryOperationsStatus makeResident(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override {
+        aubMemoryOperationsHandler->makeResident(device, gfxAllocations);
+        return BaseOperationsHandler::makeResident(device, gfxAllocations);
+    }
+
+    MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override {
+        aubMemoryOperationsHandler->evict(device, gfxAllocation);
+        return BaseOperationsHandler::evict(device, gfxAllocation);
+    }
+
+    MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) override {
+        aubMemoryOperationsHandler->isResident(device, gfxAllocation);
+        return BaseOperationsHandler::isResident(device, gfxAllocation);
+    }
+
+    MemoryOperationsStatus makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) override {
+        aubMemoryOperationsHandler->makeResidentWithinOsContext(osContext, gfxAllocations, evictable);
+        return BaseOperationsHandler::makeResidentWithinOsContext(osContext, gfxAllocations, evictable);
+    }
+
+    MemoryOperationsStatus evictWithinOsContext(OsContext *osContext, GraphicsAllocation &gfxAllocation) override {
+        aubMemoryOperationsHandler->evictWithinOsContext(osContext, gfxAllocation);
+        return BaseOperationsHandler::evictWithinOsContext(osContext, gfxAllocation);
+    }
+
+  protected:
+    std::unique_ptr<AubMemoryOperationsHandler> aubMemoryOperationsHandler;
+};
+} // namespace NEO

shared/source/os_interface/linux/os_interface_linux.cpp

@@ -5,6 +5,7 @@
  *
  */
 
+#include "shared/source/debug_settings/debug_settings_manager.h"
 #include "shared/source/execution_environment/root_device_environment.h"
 #include "shared/source/os_interface/linux/drm_memory_operations_handler.h"
 #include "shared/source/os_interface/linux/drm_neo.h"
@@ -50,7 +51,9 @@ bool initDrmOsInterface(std::unique_ptr<HwDeviceId> &&hwDeviceId, uint32_t rootD
     if (productHelper.configureHwInfoDrm(hardwareInfo, hardwareInfo, *rootDeviceEnv)) {
         return false;
     }
-    rootDeviceEnv->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, rootDeviceIndex);
+
+    const bool isCsrHwWithAub = DebugManager.flags.SetCommandStreamReceiver.get() == CommandStreamReceiverType::CSR_HW_WITH_AUB;
+    rootDeviceEnv->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, rootDeviceIndex, isCsrHwWithAub);
 
     [[maybe_unused]] bool result = rootDeviceEnv->initAilConfiguration();
     DEBUG_BREAK_IF(!result);

shared/source/os_interface/windows/CMakeLists.txt

@@ -111,6 +111,8 @@ set(NEO_CORE_OS_INTERFACE_WDDM
     ${CMAKE_CURRENT_SOURCE_DIR}/wddm_engine_mapper.h
     ${CMAKE_CURRENT_SOURCE_DIR}/wddm_memory_operations_handler.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/wddm_memory_operations_handler.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/wddm_memory_operations_handler_create.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/wddm_memory_operations_handler_with_aub_dump.h
     ${CMAKE_CURRENT_SOURCE_DIR}/wddm_residency_allocations_container.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/wddm_residency_allocations_container.h
     ${CMAKE_CURRENT_SOURCE_DIR}/wddm_residency_controller.cpp

shared/source/os_interface/windows/init_wddm_os_interface.cpp

@@ -21,7 +21,10 @@ bool initWddmOsInterface(std::unique_ptr<HwDeviceId> &&hwDeviceId, RootDeviceEnv
     if (!wddm->init()) {
         return false;
     }
-    rootDeviceEnv->memoryOperationsInterface = std::make_unique<WddmMemoryOperationsHandler>(wddm);
+
+    const bool isCsrHwWithAub = DebugManager.flags.SetCommandStreamReceiver.get() == CommandStreamReceiverType::CSR_HW_WITH_AUB;
+    rootDeviceEnv->memoryOperationsInterface = WddmMemoryOperationsHandler::create(wddm, rootDeviceEnv, isCsrHwWithAub);
+
     return true;
 }
 } // namespace NEO

shared/source/os_interface/windows/wddm_memory_operations_handler.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2019-2021 Intel Corporation
+ * Copyright (C) 2019-2023 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -12,6 +12,7 @@
 
 namespace NEO {
 
+struct RootDeviceEnvironment;
 class Wddm;
 class WddmResidentAllocationsContainer;
 
@@ -20,6 +21,8 @@ class WddmMemoryOperationsHandler : public MemoryOperationsHandler {
     WddmMemoryOperationsHandler(Wddm *wddm);
     ~WddmMemoryOperationsHandler() override;
 
+    static std::unique_ptr<WddmMemoryOperationsHandler> create(Wddm *wddm, RootDeviceEnvironment *rootDeviceEnvironment, bool withAubDump);
+
     MemoryOperationsStatus makeResident(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override;
     MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override;
     MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) override;

shared/source/os_interface/windows/wddm_memory_operations_handler_create.cpp

@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/execution_environment/root_device_environment.h"
+#include "shared/source/os_interface/windows/wddm/wddm.h"
+#include "shared/source/os_interface/windows/wddm_memory_operations_handler.h"
+#include "shared/source/os_interface/windows/wddm_memory_operations_handler_with_aub_dump.h"
+
+namespace NEO {
+
+std::unique_ptr<WddmMemoryOperationsHandler> WddmMemoryOperationsHandler::create(Wddm *wddm, RootDeviceEnvironment *rootDeviceEnvironment, bool withAubDump) {
+    if (withAubDump) {
+        return std::make_unique<WddmMemoryOperationsHandlerWithAubDump<WddmMemoryOperationsHandler>>(wddm, *rootDeviceEnvironment);
+    } else {
+        return std::make_unique<WddmMemoryOperationsHandler>(wddm);
+    }
+}
+
+} // namespace NEO

shared/source/os_interface/windows/wddm_memory_operations_handler_with_aub_dump.h

@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+#include "shared/source/aub/aub_center.h"
+#include "shared/source/command_stream/command_stream_receiver.h"
+#include "shared/source/execution_environment/execution_environment.h"
+#include "shared/source/execution_environment/root_device_environment.h"
+#include "shared/source/helpers/gfx_core_helper.h"
+#include "shared/source/os_interface/aub_memory_operations_handler.h"
+
+namespace NEO {
+class Wddm;
+
+template <typename BaseOperationsHandler>
+class WddmMemoryOperationsHandlerWithAubDump : public BaseOperationsHandler {
+  public:
+    WddmMemoryOperationsHandlerWithAubDump(Wddm *wddm, RootDeviceEnvironment &rootDeviceEnvironment)
+        : BaseOperationsHandler(wddm) {
+
+        if (!rootDeviceEnvironment.aubCenter) {
+            auto &gfxCoreHelper = rootDeviceEnvironment.getHelper<GfxCoreHelper>();
+            auto hardwareInfo = rootDeviceEnvironment.getMutableHardwareInfo();
+            auto localMemoryEnabled = gfxCoreHelper.getEnableLocalMemory(*hardwareInfo);
+            rootDeviceEnvironment.initGmm();
+            rootDeviceEnvironment.initAubCenter(localMemoryEnabled, "", static_cast<CommandStreamReceiverType>(CommandStreamReceiverType::CSR_HW_WITH_AUB));
+        }
+
+        const auto aubCenter = rootDeviceEnvironment.aubCenter.get();
+        aubMemoryOperationsHandler = std::make_unique<AubMemoryOperationsHandler>(aubCenter->getAubManager());
+    };
+
+    ~WddmMemoryOperationsHandlerWithAubDump() override = default;
+
+    MemoryOperationsStatus makeResident(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override {
+        aubMemoryOperationsHandler->makeResident(device, gfxAllocations);
+        return BaseOperationsHandler::makeResident(device, gfxAllocations);
+    }
+
+    MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override {
+        aubMemoryOperationsHandler->evict(device, gfxAllocation);
+        return BaseOperationsHandler::evict(device, gfxAllocation);
+    }
+
+    MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) override {
+        aubMemoryOperationsHandler->isResident(device, gfxAllocation);
+        return BaseOperationsHandler::isResident(device, gfxAllocation);
+    }
+
+    MemoryOperationsStatus makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) override {
+        aubMemoryOperationsHandler->makeResidentWithinOsContext(osContext, gfxAllocations, evictable);
+        return BaseOperationsHandler::makeResidentWithinOsContext(osContext, gfxAllocations, evictable);
+    }
+
+    MemoryOperationsStatus evictWithinOsContext(OsContext *osContext, GraphicsAllocation &gfxAllocation) override {
+        aubMemoryOperationsHandler->evictWithinOsContext(osContext, gfxAllocation);
+        return BaseOperationsHandler::evictWithinOsContext(osContext, gfxAllocation);
+    }
+
+  protected:
+    std::unique_ptr<AubMemoryOperationsHandler> aubMemoryOperationsHandler;
+};
+} // namespace NEO

shared/test/common/mocks/CMakeLists.txt

@@ -23,6 +23,7 @@ set(NEO_CORE_tests_mocks
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_aub_csr.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_aub_file_stream.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_aub_manager.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/mock_aub_memory_operations_handler.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_aub_subcapture_manager.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_bindless_heaps_helper.h
     ${CMAKE_CURRENT_SOURCE_DIR}/mock_builtins.h

shared/test/common/mocks/mock_aub_memory_operations_handler.h

@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+
+#include "shared/source/os_interface/aub_memory_operations_handler.h"
+#include "shared/source/utilities/arrayref.h"
+
+namespace NEO {
+
+class Device;
+class GraphicsAllocation;
+class OsContext;
+
+struct MockAubMemoryOperationsHandler : public AubMemoryOperationsHandler {
+    using AubMemoryOperationsHandler::AubMemoryOperationsHandler;
+    using AubMemoryOperationsHandler::getMemoryBanksBitfield;
+    using AubMemoryOperationsHandler::residentAllocations;
+
+    MemoryOperationsStatus makeResident(Device *device, ArrayRef<GraphicsAllocation *> gfxAllocations) override {
+        makeResidentCalled = true;
+        return AubMemoryOperationsHandler::makeResident(device, gfxAllocations);
+    }
+
+    MemoryOperationsStatus evict(Device *device, GraphicsAllocation &gfxAllocation) override {
+        evictCalled = true;
+        return AubMemoryOperationsHandler::evict(device, gfxAllocation);
+    }
+
+    MemoryOperationsStatus isResident(Device *device, GraphicsAllocation &gfxAllocation) override {
+        isResidentCalled = true;
+        return AubMemoryOperationsHandler::isResident(device, gfxAllocation);
+    }
+
+    MemoryOperationsStatus makeResidentWithinOsContext(OsContext *osContext, ArrayRef<GraphicsAllocation *> gfxAllocations, bool evictable) override {
+        makeResidentWithinOsContextCalled = true;
+        return AubMemoryOperationsHandler::makeResidentWithinOsContext(osContext, gfxAllocations, evictable);
+    }
+
+    MemoryOperationsStatus evictWithinOsContext(OsContext *osContext, GraphicsAllocation &gfxAllocation) override {
+        evictWithinOsContextCalled = true;
+        return AubMemoryOperationsHandler::evictWithinOsContext(osContext, gfxAllocation);
+    }
+
+    bool makeResidentCalled = false;
+    bool evictCalled = false;
+    bool isResidentCalled = false;
+    bool makeResidentWithinOsContextCalled = false;
+    bool evictWithinOsContextCalled = false;
+};
+
+} // namespace NEO

shared/test/common/mocks/mock_command_stream_receiver.h

@@ -43,6 +43,7 @@ class MockCommandStreamReceiver : public CommandStreamReceiver {
     using CommandStreamReceiver::internalAllocationStorage;
     using CommandStreamReceiver::latestFlushedTaskCount;
     using CommandStreamReceiver::latestSentTaskCount;
+    using CommandStreamReceiver::localMemoryEnabled;
     using CommandStreamReceiver::newResources;
     using CommandStreamReceiver::numClients;
     using CommandStreamReceiver::osContext;
@@ -78,6 +79,10 @@ class MockCommandStreamReceiver : public CommandStreamReceiver {
     void updateTagFromWait() override{};
     bool isUpdateTagFromWaitEnabled() override { return false; };
 
+    void writeMemoryAub(aub_stream::AllocationParams &allocationParams) override {
+        writeMemoryAubCalled++;
+    }
+
     bool isMultiOsContextCapable() const override { return multiOsContextCapable; }
 
     bool isGpuHangDetected() const override {
@@ -206,6 +211,7 @@ class MockCommandStreamReceiver : public CommandStreamReceiver {
     int *flushBatchedSubmissionsCallCounter = nullptr;
     uint32_t waitForCompletionWithTimeoutCalled = 0;
     uint32_t fillReusableAllocationsListCalled = 0;
+    uint32_t writeMemoryAubCalled = 0;
     uint32_t makeResidentCalledTimes = 0;
     uint32_t downloadAllocationsCalledCount = 0;
     int hostPtrSurfaceCreationMutexLockCount = 0;

shared/test/common/os_interface/linux/drm_buffer_object_fixture.h

@@ -78,7 +78,7 @@ class DrmBufferObjectFixture {
     void setUp() {
         this->mock = std::make_unique<DrmClass>(*executionEnvironment.rootDeviceEnvironments[0]);
         ASSERT_NE(nullptr, this->mock);
-        executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u);
+        executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u, false);
         osContext.reset(new OsContextLinux(*this->mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor()));
         this->mock->reset();
         bo = new TestedBufferObject(0, this->mock.get());

shared/test/common/os_interface/linux/drm_command_stream_fixture.h

@@ -46,7 +46,7 @@ class DrmCommandStreamTest : public ::testing::Test {
 
         executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
         executionEnvironment.rootDeviceEnvironments[0]->initGmm();
         gmmHelper = executionEnvironment.rootDeviceEnvironments[0]->getGmmHelper();
 
@@ -126,7 +126,7 @@ class DrmCommandStreamEnhancedTemplate : public ::testing::Test {
         mock = new DrmType(*executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]);
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, rootDeviceIndex);
+        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, rootDeviceIndex, false);
 
         csr = new TestedDrmCommandStreamReceiver<GfxFamily>(*executionEnvironment, rootDeviceIndex, 1);
         ASSERT_NE(nullptr, csr);
@@ -207,7 +207,7 @@ class DrmCommandStreamEnhancedWithFailingExecTemplate : public ::testing::Test {
         mock = new T(*executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]);
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, rootDeviceIndex);
+        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, rootDeviceIndex, false);
 
         csr = new TestedDrmCommandStreamReceiverWithFailingExec<GfxFamily>(*executionEnvironment, rootDeviceIndex, 1);
         ASSERT_NE(nullptr, csr);

shared/test/common/os_interface/linux/drm_memory_manager_fixture.cpp

@@ -25,7 +25,7 @@ void DrmMemoryManagerBasic::SetUp() {
         executionEnvironment.rootDeviceEnvironments[i]->osInterface = std::make_unique<OSInterface>();
         auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[i]);
         executionEnvironment.rootDeviceEnvironments[i]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-        executionEnvironment.rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, i);
+        executionEnvironment.rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, i, false);
         executionEnvironment.rootDeviceEnvironments[i]->initGmm();
     }
 }
@@ -51,7 +51,7 @@ void DrmMemoryManagerFixture::setUp(DrmMockCustom *mock, bool localMemoryEnabled
         rootDeviceEnvironment->setHwInfoAndInitHelpers(defaultHwInfo.get());
         rootDeviceEnvironment->osInterface = std::make_unique<OSInterface>();
         rootDeviceEnvironment->osInterface->setDriverModel(std::unique_ptr<DriverModel>(new DrmMockCustom(*rootDeviceEnvironment)));
-        rootDeviceEnvironment->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*rootDeviceEnvironment->osInterface->getDriverModel()->as<Drm>(), i);
+        rootDeviceEnvironment->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*rootDeviceEnvironment->osInterface->getDriverModel()->as<Drm>(), i, false);
         rootDeviceEnvironment->builtins.reset(new MockBuiltins);
         rootDeviceEnvironment->initGmm();
     }
@@ -135,7 +135,7 @@ void DrmMemoryManagerFixtureWithoutQuietIoctlExpectation::setUp(bool enableLocal
         rootDeviceEnvironment->setHwInfoAndInitHelpers(defaultHwInfo.get());
         rootDeviceEnvironment->osInterface = std::make_unique<OSInterface>();
         rootDeviceEnvironment->osInterface->setDriverModel(std::unique_ptr<DriverModel>(new DrmMockCustom(*rootDeviceEnvironment)));
-        rootDeviceEnvironment->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*rootDeviceEnvironment->osInterface->getDriverModel()->as<Drm>(), i);
+        rootDeviceEnvironment->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*rootDeviceEnvironment->osInterface->getDriverModel()->as<Drm>(), i, false);
         rootDeviceEnvironment->initGmm();
         i++;
     }
@@ -146,7 +146,7 @@ void DrmMemoryManagerFixtureWithoutQuietIoctlExpectation::setUp(bool enableLocal
     regionInfo[1].region = {drm_i915_gem_memory_class::I915_MEMORY_CLASS_DEVICE, 0};
     regionInfo[1].probedSize = 16 * GB;
     mock->memoryInfo.reset(new MockedMemoryInfo(regionInfo, *mock));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
     memoryManager.reset(new TestedDrmMemoryManager(enableLocalMem, false, false, *executionEnvironment));
 
     ASSERT_NE(nullptr, memoryManager);

shared/test/common/os_interface/linux/drm_memory_manager_prelim_fixtures.h

@@ -36,7 +36,7 @@ class DrmMemoryManagerWithSubDevicesPrelimTest : public ::testing::Test {
 
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
         memoryManager = new TestedDrmMemoryManager(true, false, false, *executionEnvironment);
         executionEnvironment->memoryManager.reset(memoryManager);
@@ -75,7 +75,7 @@ class DrmMemoryManagerLocalMemoryPrelimTest : public ::testing::Test {
 
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
         memoryManager = new TestedDrmMemoryManager(localMemoryEnabled, false, false, *executionEnvironment);
         executionEnvironment->memoryManager.reset(memoryManager);

shared/test/unit_test/os_interface/aub_memory_operations_handler_tests.cpp

@@ -7,14 +7,17 @@
 
 #include "shared/test/unit_test/os_interface/aub_memory_operations_handler_tests.h"
 
+#include "shared/source/aub/aub_helper.h"
 #include "shared/source/aub_mem_dump/aub_mem_dump.h"
+#include "shared/test/common/helpers/engine_descriptor_helper.h"
 #include "shared/test/common/mocks/mock_aub_manager.h"
+#include "shared/test/common/mocks/mock_command_stream_receiver.h"
 #include "shared/test/common/mocks/mock_gmm.h"
 
 TEST_F(AubMemoryOperationsHandlerTests, givenNullPtrAsAubManagerWhenMakeResidentCalledThenFalseReturned) {
     getMemoryOperationsHandler()->setAubManager(nullptr);
     auto memoryOperationsInterface = getMemoryOperationsHandler();
-    auto result = memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    auto result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     EXPECT_EQ(result, MemoryOperationsStatus::DEVICE_UNINITIALIZED);
 }
 
@@ -22,7 +25,7 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledThe
     MockAubManager aubManager;
     getMemoryOperationsHandler()->setAubManager(&aubManager);
     auto memoryOperationsInterface = getMemoryOperationsHandler();
-    auto result = memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    auto result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
     EXPECT_TRUE(aubManager.writeMemory2Called);
 
@@ -32,7 +35,7 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledThe
 
     aubManager.writeMemory2Called = false;
 
-    result = memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
     EXPECT_TRUE(aubManager.writeMemory2Called);
 
@@ -41,6 +44,32 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledThe
     EXPECT_EQ(2u, memoryOperationsInterface->residentAllocations.size());
 }
 
+TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerAndAllocationOfOneTimeAubWritableAllocationTypeWhenMakeResidentCalledTwoTimesThenWriteMemoryOnce) {
+    ASSERT_TRUE(AubHelper::isOneTimeAubWritableAllocationType(AllocationType::BUFFER));
+    allocPtr->setAllocationType(AllocationType::BUFFER);
+
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+    auto result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
+    EXPECT_TRUE(aubManager.writeMemory2Called);
+
+    auto itor = std::find(memoryOperationsInterface->residentAllocations.begin(), memoryOperationsInterface->residentAllocations.end(), allocPtr);
+    EXPECT_NE(memoryOperationsInterface->residentAllocations.end(), itor);
+    EXPECT_EQ(1u, memoryOperationsInterface->residentAllocations.size());
+
+    aubManager.writeMemory2Called = false;
+
+    result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
+    EXPECT_FALSE(aubManager.writeMemory2Called);
+
+    itor = std::find(memoryOperationsInterface->residentAllocations.begin(), memoryOperationsInterface->residentAllocations.end(), allocPtr);
+    EXPECT_NE(memoryOperationsInterface->residentAllocations.end(), itor);
+    EXPECT_EQ(1u, memoryOperationsInterface->residentAllocations.size());
+}
+
 TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledOnWriteOnlyAllocationThenTrueReturnedAndWriteCalled) {
     MockAubManager aubManager;
     getMemoryOperationsHandler()->setAubManager(&aubManager);
@@ -48,7 +77,7 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledOnW
 
     allocPtr->setWriteMemoryOnly(true);
 
-    auto result = memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    auto result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
     EXPECT_TRUE(aubManager.writeMemory2Called);
 
@@ -57,7 +86,7 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledOnW
     EXPECT_EQ(0u, memoryOperationsInterface->residentAllocations.size());
 
     aubManager.writeMemory2Called = false;
-    result = memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
     EXPECT_TRUE(aubManager.writeMemory2Called);
 
@@ -66,6 +95,34 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledOnW
     EXPECT_EQ(0u, memoryOperationsInterface->residentAllocations.size());
 }
 
+TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerAndAllocationInLocalMemoryPoolWithoutPageTablesCloningWhenMakeResidentCalledThenWriteMemoryAubIsCalled) {
+    MemoryAllocation allocation(0, AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000u,
+                                MemoryConstants::pageSize, 0, MemoryPool::LocalMemory, false, false, MemoryManager::maxOsContextCount);
+    allocPtr = &allocation;
+    allocPtr->storageInfo.cloningOfPageTables = false;
+
+    DeviceBitfield deviceBitfield(1);
+    const auto csr = std::make_unique<MockCommandStreamReceiver>(*device->getExecutionEnvironment(), 0, deviceBitfield);
+    csr->localMemoryEnabled = true;
+
+    auto oldCsr = device->getDefaultEngine().commandStreamReceiver;
+    device->getDefaultEngine().commandStreamReceiver = csr.get();
+
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+    auto result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
+    EXPECT_FALSE(aubManager.writeMemory2Called);
+    EXPECT_EQ(1u, csr->writeMemoryAubCalled);
+
+    auto itor = std::find(memoryOperationsInterface->residentAllocations.begin(), memoryOperationsInterface->residentAllocations.end(), allocPtr);
+    EXPECT_NE(memoryOperationsInterface->residentAllocations.end(), itor);
+    EXPECT_EQ(1u, memoryOperationsInterface->residentAllocations.size());
+
+    device->getDefaultEngine().commandStreamReceiver = oldCsr;
+}
+
 TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledOnCompressedAllocationThenPassCorrectParams) {
     MockAubManager aubManager;
     aubManager.storeAllocationParams = true;
@@ -78,7 +135,7 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledOnC
     gmm.isCompressionEnabled = true;
     allocPtr->setDefaultGmm(&gmm);
 
-    auto result = memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    auto result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
 
     EXPECT_TRUE(aubManager.writeMemory2Called);
@@ -102,7 +159,7 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAubManagerWhenMakeResidentCalledOnU
     gmm.isCompressionEnabled = false;
     allocPtr->setDefaultGmm(&gmm);
 
-    auto result = memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    auto result = memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
 
     EXPECT_TRUE(aubManager.writeMemory2Called);
@@ -115,7 +172,7 @@ TEST_F(AubMemoryOperationsHandlerTests, givenAllocationWhenMakeResidentCalledThe
     MockAubManager aubManager;
     getMemoryOperationsHandler()->setAubManager(&aubManager);
     auto memoryOperationsInterface = getMemoryOperationsHandler();
-    memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     EXPECT_EQ(aubManager.hintToWriteMemory, AubMemDump::DataTypeHintValues::TraceNotype);
 }
 TEST_F(AubMemoryOperationsHandlerTests, givenNonResidentAllocationWhenIsResidentCalledThenFalseReturned) {
@@ -129,7 +186,7 @@ TEST_F(AubMemoryOperationsHandlerTests, givenResidentAllocationWhenIsResidentCal
     MockAubManager aubManager;
     getMemoryOperationsHandler()->setAubManager(&aubManager);
     auto memoryOperationsInterface = getMemoryOperationsHandler();
-    memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     auto result = memoryOperationsInterface->isResident(nullptr, allocation);
     EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
 }
@@ -144,7 +201,203 @@ TEST_F(AubMemoryOperationsHandlerTests, givenResidentAllocationWhenEvictCalledTh
     MockAubManager aubManager;
     getMemoryOperationsHandler()->setAubManager(&aubManager);
     auto memoryOperationsInterface = getMemoryOperationsHandler();
-    memoryOperationsInterface->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
+    memoryOperationsInterface->makeResident(device.get(), ArrayRef<GraphicsAllocation *>(&allocPtr, 1));
     auto result = memoryOperationsInterface->evict(nullptr, allocation);
     EXPECT_EQ(result, MemoryOperationsStatus::SUCCESS);
 }
+
+TEST_F(AubMemoryOperationsHandlerTests, givenLocalMemoryAndNonLocalMemoryAllocationWithStorageInfoNonZeroWhenGetMemoryBanksBitfieldThenBanksBitfieldForSystemMemoryIsReturned) {
+    auto executionEnvironment = device->getExecutionEnvironment();
+    executionEnvironment->prepareRootDeviceEnvironments(1);
+    executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&hardwareInfo);
+    executionEnvironment->initializeMemoryManager();
+
+    MemoryAllocation allocation(0, AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000u,
+                                MemoryConstants::pageSize, 0, MemoryPool::System64KBPages, false, false, MemoryManager::maxOsContextCount);
+    allocation.storageInfo.memoryBanks = 0x3u;
+
+    DeviceBitfield deviceBitfield(1);
+    const auto csr = std::make_unique<MockCommandStreamReceiver>(*device->getExecutionEnvironment(), 0, deviceBitfield);
+    csr->localMemoryEnabled = true;
+    auto osContext = executionEnvironment->memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor());
+    csr->setupContext(*osContext);
+
+    auto oldCsr = device->getDefaultEngine().commandStreamReceiver;
+    device->getDefaultEngine().commandStreamReceiver = csr.get();
+
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+
+    auto banksBitfield = memoryOperationsInterface->getMemoryBanksBitfield(&allocation, device.get());
+
+    EXPECT_TRUE(banksBitfield.none());
+
+    device->getDefaultEngine().commandStreamReceiver = oldCsr;
+}
+
+TEST_F(AubMemoryOperationsHandlerTests, givenLocalMemoryNoncloneableAllocationWithManyBanksWhenGetMemoryBanksBitfieldThenSingleMemoryBankIsReturned) {
+    auto executionEnvironment = device->getExecutionEnvironment();
+    executionEnvironment->prepareRootDeviceEnvironments(1);
+    executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&hardwareInfo);
+    executionEnvironment->initializeMemoryManager();
+
+    MemoryAllocation allocation(0, AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000u,
+                                MemoryConstants::pageSize, 0, MemoryPool::LocalMemory, false, false, MemoryManager::maxOsContextCount);
+    allocation.storageInfo.memoryBanks = 0x3u;
+    allocation.storageInfo.cloningOfPageTables = false;
+
+    DeviceBitfield deviceBitfield(0x1u);
+    const auto csr = std::make_unique<MockCommandStreamReceiver>(*device->getExecutionEnvironment(), 0, deviceBitfield);
+    csr->localMemoryEnabled = true;
+    auto osContext = executionEnvironment->memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor());
+    csr->setupContext(*osContext);
+    EXPECT_FALSE(csr->isMultiOsContextCapable());
+
+    auto oldCsr = device->getDefaultEngine().commandStreamReceiver;
+    device->getDefaultEngine().commandStreamReceiver = csr.get();
+
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+
+    if (csr->isLocalMemoryEnabled()) {
+        auto banksBitfield = memoryOperationsInterface->getMemoryBanksBitfield(&allocation, device.get());
+        EXPECT_EQ(0x1lu, banksBitfield.to_ulong());
+    }
+
+    device->getDefaultEngine().commandStreamReceiver = oldCsr;
+}
+
+TEST_F(AubMemoryOperationsHandlerTests, givenLocalMemoryCloneableAllocationWithManyBanksWhenGetMemoryBanksBitfieldThenAllMemoryBanksAreReturned) {
+    auto executionEnvironment = device->getExecutionEnvironment();
+    executionEnvironment->prepareRootDeviceEnvironments(1);
+    executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&hardwareInfo);
+    executionEnvironment->initializeMemoryManager();
+
+    MemoryAllocation allocation(0, AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000u,
+                                MemoryConstants::pageSize, 0, MemoryPool::LocalMemory, false, false, MemoryManager::maxOsContextCount);
+    allocation.storageInfo.memoryBanks = 0x3u;
+    allocation.storageInfo.cloningOfPageTables = true;
+
+    DeviceBitfield deviceBitfield(1);
+    const auto csr = std::make_unique<MockCommandStreamReceiver>(*device->getExecutionEnvironment(), 0, deviceBitfield);
+    csr->localMemoryEnabled = true;
+    auto osContext = executionEnvironment->memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor());
+    csr->setupContext(*osContext);
+    EXPECT_FALSE(csr->isMultiOsContextCapable());
+
+    auto oldCsr = device->getDefaultEngine().commandStreamReceiver;
+    device->getDefaultEngine().commandStreamReceiver = csr.get();
+
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+
+    if (csr->isLocalMemoryEnabled()) {
+        auto banksBitfield = memoryOperationsInterface->getMemoryBanksBitfield(&allocation, device.get());
+        EXPECT_EQ(0x3lu, banksBitfield.to_ulong());
+    }
+
+    device->getDefaultEngine().commandStreamReceiver = oldCsr;
+}
+
+TEST_F(AubMemoryOperationsHandlerTests, givenLocalMemoryNoncloneableAllocationWithManyBanksWhenGetMemoryBanksBitfieldThenAllMemoryBanksAreReturned) {
+    auto executionEnvironment = device->getExecutionEnvironment();
+    executionEnvironment->prepareRootDeviceEnvironments(1);
+    executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(&hardwareInfo);
+    executionEnvironment->initializeMemoryManager();
+
+    MemoryAllocation allocation(0, AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000u,
+                                MemoryConstants::pageSize, 0, MemoryPool::LocalMemory, false, false, MemoryManager::maxOsContextCount);
+    allocation.storageInfo.memoryBanks = 0x3u;
+    allocation.storageInfo.cloningOfPageTables = false;
+
+    DeviceBitfield deviceBitfield(0b11);
+    const auto csr = std::make_unique<MockCommandStreamReceiver>(*device->getExecutionEnvironment(), 0, deviceBitfield);
+    csr->localMemoryEnabled = true;
+    csr->multiOsContextCapable = true;
+    auto osContext = executionEnvironment->memoryManager->createAndRegisterOsContext(csr.get(), EngineDescriptorHelper::getDefaultDescriptor());
+    csr->setupContext(*osContext);
+    EXPECT_TRUE(csr->isMultiOsContextCapable());
+
+    auto oldCsr = device->getDefaultEngine().commandStreamReceiver;
+    device->getDefaultEngine().commandStreamReceiver = csr.get();
+
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+
+    if (csr->isLocalMemoryEnabled()) {
+        auto banksBitfield = memoryOperationsInterface->getMemoryBanksBitfield(&allocation, device.get());
+        EXPECT_EQ(0x3lu, banksBitfield.to_ulong());
+    }
+
+    device->getDefaultEngine().commandStreamReceiver = oldCsr;
+}
+
+TEST_F(AubMemoryOperationsHandlerTests, givenGfxAllocationAndNoDeviceWhenSetAubWritableToTrueThenAllocationWritablePropertyIsNotSet) {
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+
+    MemoryAllocation allocation(0, AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000u,
+                                MemoryConstants::pageSize, 0, MemoryPool::LocalMemory, false, false, MemoryManager::maxOsContextCount);
+
+    memoryOperationsInterface->setAubWritable(true, allocation, device.get());
+    ASSERT_TRUE(memoryOperationsInterface->isAubWritable(allocation, device.get()));
+
+    memoryOperationsInterface->setAubWritable(false, allocation, nullptr);
+    EXPECT_TRUE(memoryOperationsInterface->isAubWritable(allocation, device.get()));
+}
+
+TEST_F(AubMemoryOperationsHandlerTests, givenGfxAllocationAndNoDeviceWhenIsAubWritableWithoutDeviceIsCalledThenReturnFalse) {
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+
+    MemoryAllocation allocation(0, AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000u,
+                                MemoryConstants::pageSize, 0, MemoryPool::LocalMemory, false, false, MemoryManager::maxOsContextCount);
+
+    memoryOperationsInterface->setAubWritable(true, allocation, device.get());
+    ASSERT_TRUE(memoryOperationsInterface->isAubWritable(allocation, device.get()));
+    EXPECT_FALSE(memoryOperationsInterface->isAubWritable(allocation, nullptr));
+}
+
+TEST_F(AubMemoryOperationsHandlerTests, givenGfxAllocationWhenSetAubWritableToTrueThenAllocationIsAubWritable) {
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+
+    MemoryAllocation allocation(0, AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000u,
+                                MemoryConstants::pageSize, 0, MemoryPool::LocalMemory, false, false, MemoryManager::maxOsContextCount);
+
+    memoryOperationsInterface->setAubWritable(true, allocation, device.get());
+    EXPECT_TRUE(memoryOperationsInterface->isAubWritable(allocation, device.get()));
+
+    memoryOperationsInterface->setAubWritable(false, allocation, device.get());
+    EXPECT_FALSE(memoryOperationsInterface->isAubWritable(allocation, device.get()));
+}
+
+TEST_F(AubMemoryOperationsHandlerTests, givenGfxAllocationWhenSetAubWritableForNonDefaultBankThenWritablePropertyIsSetCorrectly) {
+    MockAubManager aubManager;
+    getMemoryOperationsHandler()->setAubManager(&aubManager);
+    auto memoryOperationsInterface = getMemoryOperationsHandler();
+
+    MemoryAllocation allocation(0, AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000u,
+                                MemoryConstants::pageSize, 0, MemoryPool::LocalMemory, false, false, MemoryManager::maxOsContextCount);
+
+    allocation.storageInfo.cloningOfPageTables = false;
+
+    device->deviceBitfield = 0b1111;
+
+    memoryOperationsInterface->setAubWritable(true, allocation, device.get());
+    EXPECT_TRUE(memoryOperationsInterface->isAubWritable(allocation, device.get()));
+
+    device->deviceBitfield = 0b0000;
+    memoryOperationsInterface->setAubWritable(false, allocation, device.get());
+    EXPECT_FALSE(memoryOperationsInterface->isAubWritable(allocation, device.get()));
+
+    device->deviceBitfield = 0b1000;
+    EXPECT_TRUE(memoryOperationsInterface->isAubWritable(allocation, device.get()));
+}

shared/test/unit_test/os_interface/aub_memory_operations_handler_tests.h

@@ -11,23 +11,24 @@
 #include "shared/source/os_interface/aub_memory_operations_handler.h"
 #include "shared/source/os_interface/device_factory.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
+#include "shared/test/common/mocks/mock_aub_memory_operations_handler.h"
+#include "shared/test/common/mocks/mock_device.h"
 #include "shared/test/common/mocks/mock_graphics_allocation.h"
 
 #include "gtest/gtest.h"
 
 using namespace NEO;
 
-struct MockAubMemoryOperationsHandler : public AubMemoryOperationsHandler {
-    using AubMemoryOperationsHandler::AubMemoryOperationsHandler;
-    using AubMemoryOperationsHandler::residentAllocations;
-};
-
 class AubMemoryOperationsHandlerTests : public ::testing::Test {
   public:
     void SetUp() override {
         DebugManager.flags.SetCommandStreamReceiver.set(2);
         residencyHandler = std::unique_ptr<MockAubMemoryOperationsHandler>(new MockAubMemoryOperationsHandler(nullptr));
 
+        hardwareInfo = *defaultHwInfo;
+
+        device.reset(NEO::MockDevice::createWithNewExecutionEnvironment<NEO::MockDevice>(&hardwareInfo, 0));
+
         allocPtr = &allocation;
     }
 
@@ -38,5 +39,7 @@ class AubMemoryOperationsHandlerTests : public ::testing::Test {
     MockGraphicsAllocation allocation;
     GraphicsAllocation *allocPtr;
     DebugManagerStateRestore dbgRestore;
+    HardwareInfo hardwareInfo = {};
     std::unique_ptr<MockAubMemoryOperationsHandler> residencyHandler;
+    std::unique_ptr<NEO::MockDevice> device;
 };

shared/test/unit_test/os_interface/linux/CMakeLists.txt

@@ -21,6 +21,7 @@ set(NEO_CORE_OS_INTERFACE_TESTS_LINUX
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_mapper_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_bindless_heap_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_tests.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_operations_handler_with_aub_dump_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_mock_impl.h
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_os_memory_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/drm_pci_speed_info_tests.cpp

shared/test/unit_test/os_interface/linux/drm_buffer_object_tests.cpp

@@ -231,7 +231,7 @@ TEST(DrmBufferObjectSimpleTest, givenInvalidBoWhenValidateHostptrIsCalledThenErr
     std::unique_ptr<uint32_t[]> buff(new uint32_t[256]);
     MockExecutionEnvironment executionEnvironment;
     std::unique_ptr<DrmMockCustom> mock(new DrmMockCustom(*executionEnvironment.rootDeviceEnvironments[0]));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u, false);
     OsContextLinux osContext(*mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
     ASSERT_NE(nullptr, mock.get());
     std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(0u, mock.get()));
@@ -256,7 +256,7 @@ TEST(DrmBufferObjectSimpleTest, givenInvalidBoWhenPinIsCalledThenErrorIsReturned
     std::unique_ptr<uint32_t[]> buff(new uint32_t[256]);
     MockExecutionEnvironment executionEnvironment;
     std::unique_ptr<DrmMockCustom> mock(new DrmMockCustom(*executionEnvironment.rootDeviceEnvironments[0]));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u, false);
     OsContextLinux osContext(*mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
     ASSERT_NE(nullptr, mock.get());
     std::unique_ptr<TestedBufferObject> bo(new TestedBufferObject(0u, mock.get()));
@@ -407,7 +407,7 @@ TEST(DrmBufferObject, givenDrmIoctlReturnsErrorNotSupportedThenBufferObjectRetur
     drm->errnoRetVal = EOPNOTSUPP;
     drm->baseErrno = false;
     executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
 
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
@@ -436,7 +436,7 @@ TEST(DrmBufferObject, givenPerContextVmRequiredWhenBoBoundAndUnboundThenCorrectB
     EXPECT_TRUE(drm->isPerContextVMRequired());
 
     executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
 
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
@@ -480,7 +480,7 @@ TEST(DrmBufferObject, givenPrintBOBindingResultWhenBOBindAndUnbindSucceedsThenPr
     drm->latestCreatedVmId = 1;
 
     executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
 
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
@@ -538,7 +538,7 @@ TEST(DrmBufferObject, givenPrintBOBindingResultWhenBOBindAndUnbindFailsThenPrint
     drm->latestCreatedVmId = 1;
 
     executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
 
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
 
@@ -590,7 +590,7 @@ TEST(DrmBufferObject, givenDrmWhenBindOperationFailsThenFenceValueNotGrow) {
     drm->ioctlHelper.reset(ioctlHelper.release());
 
     executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     uint64_t initFenceValue = 10u;
     drm->fenceVal[0] = initFenceValue;
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
@@ -622,7 +622,7 @@ TEST(DrmBufferObject, givenDrmWhenBindOperationSucceedsThenFenceValueGrow) {
     drm->ioctlHelper.reset(ioctlHelper.release());
 
     executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     uint64_t initFenceValue = 10u;
     drm->fenceVal[0] = initFenceValue;
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
@@ -654,7 +654,7 @@ TEST(DrmBufferObject, givenDrmWhenUnBindOperationFailsThenFenceValueNotGrow) {
     drm->ioctlHelper.reset(ioctlHelper.release());
 
     executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     uint64_t initFenceValue = 10u;
     drm->fenceVal[0] = initFenceValue;
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));
@@ -686,7 +686,7 @@ TEST(DrmBufferObject, givenDrmWhenUnBindOperationSucceedsThenFenceValueGrow) {
     drm->ioctlHelper.reset(ioctlHelper.release());
 
     executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     uint64_t initFenceValue = 10u;
     drm->fenceVal[0] = initFenceValue;
     std::unique_ptr<Device> device(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, 0));

shared/test/unit_test/os_interface/linux/drm_command_stream_mm_tests.cpp

@@ -46,7 +46,7 @@ HWTEST_F(DrmCommandStreamMMTest, GivenForcePinThenMemoryManagerCreatesPinBb) {
 
     executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     executionEnvironment.rootDeviceEnvironments[0]->initGmm();
 
     DrmCommandStreamReceiver<FamilyType> csr(executionEnvironment, 0, 1, gemCloseWorkerMode::gemCloseWorkerInactive);
@@ -87,7 +87,7 @@ HWTEST_F(DrmCommandStreamMMTest, givenExecutionEnvironmentWithMoreThanOneRootDev
         executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->initGmm();
         auto drm = new DrmMockCustom(*executionEnvironment.rootDeviceEnvironments[0]);
         executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-        executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+        executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     }
 
     auto memoryManager = new TestedDrmMemoryManager(false, true, false, executionEnvironment);

shared/test/unit_test/os_interface/linux/drm_command_stream_xehp_and_later_prelim_tests.cpp

@@ -312,7 +312,7 @@ class DrmCommandStreamForceTileTest : public ::testing::Test {
         mock->setupIoctlHelper(hwInfo->platform.eProductFamily);
         executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
         executionEnvironment.rootDeviceEnvironments[0]->initGmm();
 
         mock->createVirtualMemoryAddressSpace(GfxCoreHelper::getSubDevicesCount(hwInfo));
@@ -425,7 +425,7 @@ struct DrmImplicitScalingCommandStreamTest : ::testing::Test {
         executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(hwInfo.get());
         executionEnvironment->rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-        executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+        executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
         executionEnvironment->rootDeviceEnvironments[0]->initGmm();
 
         auto &gfxCoreHelper = executionEnvironment->rootDeviceEnvironments[0]->getHelper<GfxCoreHelper>();

shared/test/unit_test/os_interface/linux/drm_command_stream_xehp_and_later_tests.cpp

@@ -43,7 +43,7 @@ struct DrmCommandStreamMultiTileMemExecFixture {
         mock = new DrmMockCustom(*executionEnvironment->rootDeviceEnvironments[0]);
         executionEnvironment->rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0);
+        executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0, false);
 
         memoryManager = new DrmMemoryManager(gemCloseWorkerMode::gemCloseWorkerInactive,
                                              DebugManager.flags.EnableForcePin.get(),

shared/test/unit_test/os_interface/linux/drm_gem_close_worker_tests.cpp

@@ -67,7 +67,7 @@ class DrmGemCloseWorkerFixture {
 
         executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drmMock));
-        executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drmMock, 0u);
+        executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drmMock, 0u, false);
 
         this->mm = new DrmMemoryManager(gemCloseWorkerMode::gemCloseWorkerInactive,
                                         false,

shared/test/unit_test/os_interface/linux/drm_memory_manager_localmem_prelim_tests.cpp

@@ -105,7 +105,7 @@ TEST_F(DrmMemoryManagerLocalMemoryPrelimTest, givenMultiRootDeviceEnvironmentAnd
         mock->ioctlCallsCount = 0;
         executionEnvironment->rootDeviceEnvironments[i]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[i]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
         executionEnvironment->rootDeviceEnvironments[i]->initGmm();
 
         rootDeviceIndices.pushUnique(i);
@@ -155,7 +155,7 @@ TEST_F(DrmMemoryManagerLocalMemoryPrelimTest, givenMultiRootDeviceEnvironmentAnd
         mock->ioctlCallsCount = 0;
         executionEnvironment->rootDeviceEnvironments[i]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[i]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
         executionEnvironment->rootDeviceEnvironments[i]->initGmm();
 
         rootDeviceIndices.pushUnique(i);
@@ -196,7 +196,7 @@ TEST_F(DrmMemoryManagerLocalMemoryPrelimTest, givenMultiRootDeviceEnvironmentAnd
         mock->fdToHandleRetVal = -1;
         executionEnvironment->rootDeviceEnvironments[i]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[i]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
         executionEnvironment->rootDeviceEnvironments[i]->initGmm();
 
         rootDeviceIndices.pushUnique(i);
@@ -236,7 +236,7 @@ TEST_F(DrmMemoryManagerUsmSharedHandlePrelimTest, givenMultiRootDeviceEnvironmen
     mock->fdToHandleRetVal = -1;
     executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
     executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-    executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+    executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
     auto memoryManager = std::make_unique<TestedDrmMemoryManager>(true, false, false, *executionEnvironment);
 
@@ -267,7 +267,7 @@ TEST_F(DrmMemoryManagerLocalMemoryPrelimTest, givenMultiRootDeviceEnvironmentAnd
         mock->fdToHandleRetVal = -1;
         executionEnvironment->rootDeviceEnvironments[i]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[i]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
         executionEnvironment->rootDeviceEnvironments[i]->initGmm();
 
         rootDeviceIndices.pushUnique(i);

shared/test/unit_test/os_interface/linux/drm_memory_manager_localmem_upstream_tests.cpp

@@ -69,7 +69,7 @@ class DrmMemoryManagerLocalMemoryTest : public ::testing::Test {
         mock->memoryInfo.reset(new MockMemoryInfo(*mock));
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
         device.reset(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(), executionEnvironment, rootDeviceIndex));
         memoryManager = std::make_unique<TestedDrmMemoryManager>(localMemoryEnabled, false, false, *executionEnvironment);
@@ -161,7 +161,7 @@ HWTEST2_F(DrmMemoryManagerLocalMemoryTest, givenMultiRootDeviceEnvironmentAndMem
         mock->ioctlCallsCount = 0;
         executionEnvironment->rootDeviceEnvironments[i]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[i]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
         executionEnvironment->rootDeviceEnvironments[i]->initGmm();
 
         rootDeviceIndices.pushUnique(i);
@@ -208,7 +208,7 @@ TEST_F(DrmMemoryManagerLocalMemoryTest, givenMultiRootDeviceEnvironmentAndMemory
         mock->fdToHandleRetVal = -1;
         executionEnvironment->rootDeviceEnvironments[i]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[i]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
         executionEnvironment->rootDeviceEnvironments[i]->initGmm();
 
         rootDeviceIndices.pushUnique(i);
@@ -268,7 +268,7 @@ TEST_F(DrmMemoryManagerUsmSharedHandleTest, givenMultiRootDeviceEnvironmentAndMe
     mock->fdToHandleRetVal = -1;
     executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
     executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-    executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+    executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
     rootDeviceIndices.pushUnique(rootDeviceIndex);
 
@@ -300,7 +300,7 @@ TEST_F(DrmMemoryManagerLocalMemoryTest, givenMultiRootDeviceEnvironmentAndNoMemo
         mock->fdToHandleRetVal = -1;
         executionEnvironment->rootDeviceEnvironments[i]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[i]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[i]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
         executionEnvironment->rootDeviceEnvironments[i]->initGmm();
 
         rootDeviceIndices.pushUnique(i);
@@ -470,7 +470,7 @@ class DrmMemoryManagerLocalMemoryMemoryBankTest : public ::testing::Test {
         mock->memoryInfo.reset(new MockMemoryInfo(*mock));
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
-        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u);
+        executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
         device.reset(MockDevice::createWithExecutionEnvironment<MockDevice>(defaultHwInfo.get(),
                                                                             executionEnvironment,

shared/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp

@@ -3779,7 +3779,7 @@ TEST(DrmMemoryManagerFreeGraphicsMemoryCallSequenceTest, givenDrmMemoryManagerAn
     executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[0]);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     executionEnvironment.rootDeviceEnvironments[0]->initGmm();
     TestedDrmMemoryManager memoryManger(executionEnvironment);
 
@@ -3804,7 +3804,7 @@ TEST(DrmMemoryManagerFreeGraphicsMemoryUnreferenceTest,
     executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]);
     executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     executionEnvironment.rootDeviceEnvironments[0]->initGmm();
     TestedDrmMemoryManager memoryManger(executionEnvironment);
 
@@ -3831,7 +3831,7 @@ TEST(DrmMemoryManagerFreeGraphicsMemoryUnreferenceTest,
     executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]);
     executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     executionEnvironment.rootDeviceEnvironments[0]->initGmm();
     TestedDrmMemoryManager memoryManger(executionEnvironment);
 
@@ -3864,7 +3864,7 @@ TEST(DrmMemoryManagerFreeGraphicsMemoryUnreferenceTest,
     executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]);
     executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     executionEnvironment.rootDeviceEnvironments[0]->initGmm();
     TestedDrmMemoryManager memoryManger(executionEnvironment);
 
@@ -3888,7 +3888,7 @@ TEST(DrmMemoryManagerFreeGraphicsMemoryUnreferenceTest, givenDrmMemoryManagerAnd
     executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]);
     executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     executionEnvironment.rootDeviceEnvironments[0]->initGmm();
     TestedDrmMemoryManager memoryManger(executionEnvironment);
 
@@ -4114,7 +4114,7 @@ TEST(DrmMemoryManager, givenEnabledResourceRegistrationWhenSshIsAllocatedThenItI
     auto mockDrm = new DrmMockResources(*executionEnvironment->rootDeviceEnvironments[0]);
     executionEnvironment->rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
     executionEnvironment->rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mockDrm));
-    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mockDrm, 0u);
+    executionEnvironment->rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mockDrm, 0u, false);
     executionEnvironment->memoryManager = std::make_unique<TestedDrmMemoryManager>(false, false, false, *executionEnvironment);
 
     for (uint32_t i = 3; i < 3 + static_cast<uint32_t>(DrmResourceClass::MaxSize); i++) {
@@ -4820,7 +4820,7 @@ TEST(DrmMemoryManagerSimpleTest, givenDrmMemoryManagerWhenAllocateInDevicePoolIs
     executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[0]);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     TestedDrmMemoryManager memoryManager(executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
     AllocationData allocData;
@@ -4841,7 +4841,7 @@ TEST(DrmMemoryManagerSimpleTest, givenDrmMemoryManagerWhenDeviceHeapIsDepletedTh
     executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[0]);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
 
     auto mockGfxPartition = std::make_unique<MockGfxPartition>();
     mockGfxPartition->init(hwInfo->capabilityTable.gpuAddressSpace, reservedCpuAddressRangeSize, 0, 1, false, 0u);
@@ -4866,7 +4866,7 @@ TEST(DrmMemoryManagerSimpleTest, givenDrmMemoryManagerWhenAllocateInLocalDeviceM
     executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[0]);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     TestedDrmMemoryManager memoryManager(executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
     AllocationData allocData;
@@ -4884,7 +4884,7 @@ TEST(DrmMemoryManagerSimpleTest, givenDrmMemoryManagerWhenLockResourceIsCalledOn
     executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[0]);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     TestedDrmMemoryManager memoryManager(executionEnvironment);
     DrmAllocation drmAllocation(0, AllocationType::UNKNOWN, nullptr, nullptr, 0u, 0u, MemoryPool::LocalMemory);
 
@@ -4899,7 +4899,7 @@ TEST(DrmMemoryManagerSimpleTest, givenDrmMemoryManagerWhenFreeGraphicsMemoryIsCa
     executionEnvironment.rootDeviceEnvironments[0]->osInterface = std::make_unique<OSInterface>();
     auto drm = Drm::create(nullptr, *executionEnvironment.rootDeviceEnvironments[0]);
     executionEnvironment.rootDeviceEnvironments[0]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(drm));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*drm, 0u, false);
     TestedDrmMemoryManager memoryManager(executionEnvironment);
 
     auto drmAllocation = new DrmAllocation(0, AllocationType::UNKNOWN, nullptr, nullptr, 0u, 0u, MemoryPool::LocalMemory);

shared/test/unit_test/os_interface/linux/drm_memory_operations_handler_with_aub_dump_tests.cpp

@@ -0,0 +1,81 @@
+/*
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/execution_environment/execution_environment.h"
+#include "shared/source/execution_environment/root_device_environment.h"
+#include "shared/source/os_interface/linux/drm_memory_operations_handler_default.h"
+#include "shared/source/os_interface/linux/drm_memory_operations_handler_with_aub_dump.h"
+#include "shared/test/common/mocks/mock_aub_memory_operations_handler.h"
+#include "shared/test/common/mocks/mock_device.h"
+#include "shared/test/common/mocks/mock_graphics_allocation.h"
+#include "shared/test/common/test_macros/test.h"
+
+#include <memory>
+
+using namespace NEO;
+
+template <typename BaseMemoryOperationsHandler>
+class DrmMemoryOperationsHandlerWithAubDumpMock : public DrmMemoryOperationsHandlerWithAubDump<BaseMemoryOperationsHandler> {
+  public:
+    using DrmMemoryOperationsHandlerWithAubDump<BaseMemoryOperationsHandler>::aubMemoryOperationsHandler;
+    using DrmMemoryOperationsHandlerWithAubDump<BaseMemoryOperationsHandler>::residency;
+
+    DrmMemoryOperationsHandlerWithAubDumpMock(RootDeviceEnvironment &rootDeviceEnvironment, uint32_t rootDeviceIndex)
+        : DrmMemoryOperationsHandlerWithAubDump<BaseMemoryOperationsHandler>(rootDeviceEnvironment, rootDeviceIndex) {
+        aubMemoryOperationsHandler = std::make_unique<MockAubMemoryOperationsHandler>(nullptr);
+    }
+};
+
+struct DrmMemoryOperationsHandlerWithAubDumpTest : public ::testing::Test {
+    void SetUp() override {
+        device.reset(NEO::MockDevice::createWithNewExecutionEnvironment<NEO::MockDevice>(defaultHwInfo.get(), 0));
+        device->executionEnvironment->prepareRootDeviceEnvironments(1);
+        device->executionEnvironment->rootDeviceEnvironments[0]->setHwInfoAndInitHelpers(defaultHwInfo.get());
+        auto rootDeviceEnv = device->executionEnvironment->rootDeviceEnvironments[0].get();
+        drmMemoryOperationsHandlerWithAubDumpMock = std::make_unique<DrmMemoryOperationsHandlerWithAubDumpMock<DrmMemoryOperationsHandlerDefault>>(*rootDeviceEnv, 0);
+        mockAubMemoryOperationsHandler = static_cast<MockAubMemoryOperationsHandler *>(drmMemoryOperationsHandlerWithAubDumpMock->aubMemoryOperationsHandler.get());
+
+        allocationPtr = &graphicsAllocation;
+    }
+
+    MockGraphicsAllocation graphicsAllocation;
+    GraphicsAllocation *allocationPtr;
+    std::unique_ptr<DrmMemoryOperationsHandlerWithAubDumpMock<DrmMemoryOperationsHandlerDefault>> drmMemoryOperationsHandlerWithAubDumpMock;
+    MockAubMemoryOperationsHandler *mockAubMemoryOperationsHandler;
+    std::unique_ptr<NEO::MockDevice> device;
+};
+
+TEST_F(DrmMemoryOperationsHandlerWithAubDumpTest, whenMakingAllocationResidentThenAllocationIsResident) {
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 1u);
+    EXPECT_TRUE(drmMemoryOperationsHandlerWithAubDumpMock->residency.find(allocationPtr) != drmMemoryOperationsHandlerWithAubDumpMock->residency.end());
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, graphicsAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
+}
+
+TEST_F(DrmMemoryOperationsHandlerWithAubDumpTest, whenEvictingResidentAllocationThenAllocationIsNotResident) {
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 0u);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, graphicsAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 1u);
+    EXPECT_TRUE(drmMemoryOperationsHandlerWithAubDumpMock->residency.find(allocationPtr) != drmMemoryOperationsHandlerWithAubDumpMock->residency.end());
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->evict(nullptr, graphicsAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, graphicsAllocation), MemoryOperationsStatus::MEMORY_NOT_FOUND);
+    EXPECT_EQ(drmMemoryOperationsHandlerWithAubDumpMock->residency.size(), 0u);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->evictCalled);
+}
+
+TEST_F(DrmMemoryOperationsHandlerWithAubDumpTest, whenConstructingDrmMemoryOperationsHandlerWithAubDumpWithoutAubCenterThenAubCenterIsInitialized) {
+    auto rootDeviceEnv = device->executionEnvironment->rootDeviceEnvironments[0].get();
+    rootDeviceEnv->aubCenter.reset();
+    ASSERT_EQ(nullptr, rootDeviceEnv->aubCenter.get());
+    auto drmMemoryOperationsHandlerWithAubDump = std::make_unique<DrmMemoryOperationsHandlerWithAubDumpMock<DrmMemoryOperationsHandlerDefault>>(*rootDeviceEnv, 0);
+    EXPECT_NE(nullptr, rootDeviceEnv->aubCenter.get());
+}

shared/test/unit_test/os_interface/linux/drm_residency_handler_prelim_tests.cpp

@@ -707,7 +707,7 @@ TEST_F(DrmMemoryOperationsHandlerBindTest, WhenVmBindAvaialableThenMemoryManager
 
 TEST_F(DrmMemoryOperationsHandlerBindTest, givenNoVmBindSupportInDrmWhenCheckForSupportThenDefaultResidencyHandlerIsReturned) {
     mock->bindAvailable = false;
-    auto handler = DrmMemoryOperationsHandler::create(*mock, 0u);
+    auto handler = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
     mock->context.vmBindCalled = 0u;
     auto allocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{device->getRootDeviceIndex(), MemoryConstants::pageSize});
@@ -722,7 +722,7 @@ TEST_F(DrmMemoryOperationsHandlerBindTest, givenVmBindSupportAndNoMultiTileWhenC
     DebugManager.flags.CreateMultipleSubDevices.set(1u);
     mock->bindAvailable = false;
 
-    auto handler = DrmMemoryOperationsHandler::create(*mock, 0u);
+    auto handler = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
     mock->context.vmBindCalled = 0u;
     auto allocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{device->getRootDeviceIndex(), MemoryConstants::pageSize});
@@ -736,7 +736,7 @@ TEST_F(DrmMemoryOperationsHandlerBindTest, givenVmBindSupportAndNoMultiTileWhenC
 TEST_F(DrmMemoryOperationsHandlerBindTest, givenDisabledVmBindWhenCreateDrmHandlerThenVmBindIsNotUsed) {
     mock->context.vmBindReturn = 0;
     mock->bindAvailable = false;
-    auto handler = DrmMemoryOperationsHandler::create(*mock, 0u);
+    auto handler = DrmMemoryOperationsHandler::create(*mock, 0u, false);
 
     mock->context.vmBindCalled = false;
     auto allocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{device->getRootDeviceIndex(), MemoryConstants::pageSize});

shared/test/unit_test/os_interface/linux/os_context_linux_tests.cpp

@@ -23,7 +23,7 @@ using namespace NEO;
 TEST(OSContextLinux, givenReinitializeContextWhenContextIsInitThenContextIsStillIinitializedAfter) {
     MockExecutionEnvironment executionEnvironment;
     std::unique_ptr<DrmMockCustom> mock(new DrmMockCustom(*executionEnvironment.rootDeviceEnvironments[0]));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u, false);
     OsContextLinux osContext(*mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
     EXPECT_NO_THROW(osContext.reInitializeContext());
     EXPECT_NO_THROW(osContext.ensureContextInitialized());
@@ -43,7 +43,7 @@ TEST(OSContextLinux, givenInitializeContextWhenContextCreateIoctlFailsThenContex
 TEST(OSContextLinux, givenOsContextLinuxWhenQueryingForOfflineDumpContextIdThenCorrectValueIsReturned) {
     MockExecutionEnvironment executionEnvironment;
     std::unique_ptr<DrmMockCustom> mock(new DrmMockCustom(*executionEnvironment.rootDeviceEnvironments[0]));
-    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u);
+    executionEnvironment.rootDeviceEnvironments[0]->memoryOperationsInterface = DrmMemoryOperationsHandler::create(*mock.get(), 0u, false);
     MockOsContextLinux osContext(*mock, 0, 0u, EngineDescriptorHelper::getDefaultDescriptor());
 
     osContext.drmContextIds.clear();

shared/test/unit_test/os_interface/windows/CMakeLists.txt

@@ -49,6 +49,7 @@ if(WIN32 OR(UNIX AND NOT DISABLE_WDDM_LINUX))
                  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_preemption_tests.cpp
                  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_shared_allocations_test.cpp
                  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_memory_manager_tests.cpp
+                 ${CMAKE_CURRENT_SOURCE_DIR}/wddm_memory_operations_handler_with_aub_dump_tests.cpp
   )
 endif()
 

shared/test/unit_test/os_interface/windows/wddm_memory_operations_handler_with_aub_dump_tests.cpp

@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/os_interface/windows/wddm_memory_operations_handler_with_aub_dump.h"
+#include "shared/source/utilities/stackvec.h"
+#include "shared/test/common/helpers/ult_hw_config.h"
+#include "shared/test/common/mocks/mock_allocation_properties.h"
+#include "shared/test/common/mocks/mock_aub_memory_operations_handler.h"
+#include "shared/test/common/mocks/mock_device.h"
+#include "shared/test/common/mocks/mock_wddm.h"
+#include "shared/test/common/mocks/windows/mock_wddm_allocation.h"
+#include "shared/test/common/os_interface/windows/wddm_fixture.h"
+#include "shared/test/common/test_macros/hw_test.h"
+
+using namespace NEO;
+
+template <typename BaseMemoryOperationsHandler>
+class WddmMemoryOperationsHandlerWithAubDumpMock : public WddmMemoryOperationsHandlerWithAubDump<BaseMemoryOperationsHandler> {
+  public:
+    using WddmMemoryOperationsHandlerWithAubDump<BaseMemoryOperationsHandler>::aubMemoryOperationsHandler;
+
+    WddmMemoryOperationsHandlerWithAubDumpMock(Wddm *wddm, RootDeviceEnvironment &rootDeviceEnvironment)
+        : WddmMemoryOperationsHandlerWithAubDump<BaseMemoryOperationsHandler>(wddm, rootDeviceEnvironment) {
+        aubMemoryOperationsHandler = std::make_unique<MockAubMemoryOperationsHandler>(nullptr);
+    }
+};
+
+struct WddmMemoryOperationsHandlerWithAubDumpTest : public WddmTest {
+    void SetUp() override {
+        WddmTest::SetUp();
+        wddmMemoryOperationsHandlerWithAubDumpMock = std::make_unique<WddmMemoryOperationsHandlerWithAubDumpMock<WddmMemoryOperationsHandler>>(wddm, *rootDeviceEnvironment);
+        mockAubMemoryOperationsHandler = static_cast<MockAubMemoryOperationsHandler *>(wddmMemoryOperationsHandlerWithAubDumpMock->aubMemoryOperationsHandler.get());
+
+        wddmAllocation = std::make_unique<MockWddmAllocation>(rootDeviceEnvironment->getGmmHelper());
+        wddmFragmentedAllocation = std::make_unique<MockWddmAllocation>(rootDeviceEnvironment->getGmmHelper());
+        wddmAllocation->handle = 0x2u;
+
+        osHandleStorageFirst = std::make_unique<OsHandleWin>();
+        osHandleStorageSecond = std::make_unique<OsHandleWin>();
+
+        wddmFragmentedAllocation->fragmentsStorage.fragmentCount = 2;
+        wddmFragmentedAllocation->fragmentsStorage.fragmentStorageData[0].osHandleStorage = osHandleStorageFirst.get();
+        static_cast<OsHandleWin *>(wddmFragmentedAllocation->fragmentsStorage.fragmentStorageData[0].osHandleStorage)->handle = 0x3u;
+        wddmFragmentedAllocation->fragmentsStorage.fragmentStorageData[1].osHandleStorage = osHandleStorageSecond.get();
+        static_cast<OsHandleWin *>(wddmFragmentedAllocation->fragmentsStorage.fragmentStorageData[1].osHandleStorage)->handle = 0x4u;
+
+        allocationPtr = wddmAllocation.get();
+
+        allocationData.push_back(wddmAllocation.get());
+        allocationData.push_back(wddmFragmentedAllocation.get());
+    }
+
+    std::unique_ptr<WddmMemoryOperationsHandlerWithAubDumpMock<WddmMemoryOperationsHandler>> wddmMemoryOperationsHandlerWithAubDumpMock;
+    std::unique_ptr<MockWddmAllocation> wddmAllocation;
+    std::unique_ptr<MockWddmAllocation> wddmFragmentedAllocation;
+    std::unique_ptr<OsHandleWin> osHandleStorageFirst;
+    std::unique_ptr<OsHandleWin> osHandleStorageSecond;
+    GraphicsAllocation *allocationPtr;
+    StackVec<GraphicsAllocation *, 2> allocationData;
+    MockAubMemoryOperationsHandler *mockAubMemoryOperationsHandler;
+};
+
+TEST_F(WddmMemoryOperationsHandlerWithAubDumpTest, givenRegularAllocationWhenMakingResidentAllocationThenMakeResidentCalled) {
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *wddmAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
+}
+
+TEST_F(WddmMemoryOperationsHandlerWithAubDumpTest, givenFragmentedAllocationWhenMakingResidentAllocationThenMakeResidentCalled) {
+    allocationPtr = wddmFragmentedAllocation.get();
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *wddmFragmentedAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
+}
+
+TEST_F(WddmMemoryOperationsHandlerWithAubDumpTest, givenVariousAllocationsWhenMakingResidentAllocationThenMakeResidentCalled) {
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(allocationData)), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *wddmAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *wddmFragmentedAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
+}
+
+TEST_F(WddmMemoryOperationsHandlerWithAubDumpTest, givenRegularAllocationWhenEvictingResidentAllocationThenEvictCalled) {
+    wddm->callBaseEvict = true;
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->evict(nullptr, *wddmAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(0u, gdi->getEvictArg().Flags.EvictOnlyIfNecessary);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *wddmAllocation), MemoryOperationsStatus::MEMORY_NOT_FOUND);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
+}
+
+TEST_F(WddmMemoryOperationsHandlerWithAubDumpTest, givenFragmentedAllocationWhenEvictingResidentAllocationThenEvictCalled) {
+    allocationPtr = wddmFragmentedAllocation.get();
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(&allocationPtr, 1)), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->evict(nullptr, *wddmFragmentedAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *wddmFragmentedAllocation), MemoryOperationsStatus::MEMORY_NOT_FOUND);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
+}
+
+TEST_F(WddmMemoryOperationsHandlerWithAubDumpTest, givenVariousAllocationsWhenEvictingResidentAllocationThenEvictCalled) {
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->makeResident(nullptr, ArrayRef<GraphicsAllocation *>(allocationData)), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->evict(nullptr, *wddmAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *wddmAllocation), MemoryOperationsStatus::MEMORY_NOT_FOUND);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->evict(nullptr, *wddmFragmentedAllocation), MemoryOperationsStatus::SUCCESS);
+    EXPECT_EQ(wddmMemoryOperationsHandlerWithAubDumpMock->isResident(nullptr, *wddmFragmentedAllocation), MemoryOperationsStatus::MEMORY_NOT_FOUND);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->makeResidentCalled);
+    EXPECT_TRUE(mockAubMemoryOperationsHandler->isResidentCalled);
+}
+
+TEST_F(WddmMemoryOperationsHandlerWithAubDumpTest, whenConstructingWddmMemoryOperationsHandlerWithAubDumpWithoutAubCenterThenAubCenterIsInitialized) {
+    rootDeviceEnvironment->aubCenter.reset();
+    ASSERT_EQ(nullptr, rootDeviceEnvironment->aubCenter.get());
+    auto wddmMemoryOperationsHandlerWithAubDump = std::make_unique<WddmMemoryOperationsHandlerWithAubDump<WddmMemoryOperationsHandler>>(wddm, *rootDeviceEnvironment);
+    EXPECT_NE(nullptr, rootDeviceEnvironment->aubCenter.get());
+}