zeCommandListAppendMemoryPrefetch with memory prefetch manager

Related-To: NEO-6740

Signed-off-by: Milczarek, Slawomir <slawomir.milczarek@intel.com>

level_zero/core/source/cmdlist/cmdlist.h

@@ -229,6 +229,9 @@ struct CommandList : _ze_command_list_handle_t {
     bool containsCooperativeKernels() const {
         return containsCooperativeKernelsFlag;
     }
+    bool isMemoryPrefetchRequested() const {
+        return performMemoryPrefetch;
+    }
 
     enum CommandListType : uint32_t {
         TYPE_REGULAR = 0u,
@@ -293,6 +296,7 @@ struct CommandList : _ze_command_list_handle_t {
     bool internalUsage = false;
     bool containsCooperativeKernelsFlag = false;
     bool containsStatelessUncachedResource = false;
+    bool performMemoryPrefetch = false;
 };
 
 using CommandListAllocatorFn = CommandList *(*)(uint32_t);

level_zero/core/source/cmdlist/cmdlist_hw.inl

@@ -85,6 +85,7 @@ ze_result_t CommandListCoreFamily<gfxCoreFamily>::reset() {
     removeHostPtrAllocations();
     commandContainer.reset();
     containsStatelessUncachedResource = false;
+    performMemoryPrefetch = false;
     indirectAllocationsAllowed = false;
     unifiedMemoryControls.indirectHostAllocationsAllowed = false;
     unifiedMemoryControls.indirectSharedAllocationsAllowed = false;

level_zero/core/source/cmdlist/cmdlist_hw_immediate.inl

@@ -12,6 +12,7 @@
 #include "shared/source/helpers/hw_info.h"
 #include "shared/source/helpers/logical_state_helper.h"
 #include "shared/source/memory_manager/internal_allocation_storage.h"
+#include "shared/source/memory_manager/prefetch_manager.h"
 
 #include "level_zero/core/source/cmdlist/cmdlist_hw_immediate.h"
 
@@ -108,6 +109,12 @@ ze_result_t CommandListCoreFamilyImmediate<gfxCoreFamily>::executeCommandListImm
         this->migrateSharedAllocations();
     }
 
+    if (this->performMemoryPrefetch) {
+        auto prefetchManager = this->device->getDriverHandle()->getMemoryManager()->getPrefetchManager();
+        prefetchManager->migrateAllocationsToGpu(*this->device->getDriverHandle()->getSvmAllocsManager(), *this->device->getNEODevice());
+        this->performMemoryPrefetch = false;
+    }
+
     auto ioh = (this->commandContainer.getIndirectHeap(NEO::IndirectHeap::Type::INDIRECT_OBJECT));
     NEO::IndirectHeap *dsh = nullptr;
     NEO::IndirectHeap *ssh = nullptr;

level_zero/core/source/cmdqueue/cmdqueue_hw.inl

@@ -27,6 +27,7 @@
 #include "shared/source/helpers/preamble.h"
 #include "shared/source/memory_manager/graphics_allocation.h"
 #include "shared/source/memory_manager/memory_manager.h"
+#include "shared/source/memory_manager/prefetch_manager.h"
 #include "shared/source/memory_manager/residency_container.h"
 #include "shared/source/os_interface/hw_info_config.h"
 #include "shared/source/os_interface/os_context.h"
@@ -83,6 +84,7 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
     auto anyCommandListWithoutCooperativeKernels = false;
     bool anyCommandListRequiresDisabledEUFusion = false;
     bool cachedMOCSAllowed = true;
+    bool performMemoryPrefetch = false;
 
     for (auto i = 0u; i < numCommandLists; i++) {
         auto commandList = CommandList::fromHandle(phCommandLists[i]);
@@ -108,6 +110,10 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
         if (commandList->requiresQueueUncachedMocs && cachedMOCSAllowed == true) {
             cachedMOCSAllowed = false;
         }
+
+        if (commandList->isMemoryPrefetchRequested()) {
+            performMemoryPrefetch = true;
+        }
     }
 
     bool isMixingRegularAndCooperativeKernelsAllowed = NEO::DebugManager.flags.AllowMixingRegularAndCooperativeKernels.get();
@@ -446,6 +452,12 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
         commandList->migrateSharedAllocations();
     }
 
+    if (performMemoryPrefetch) {
+        auto prefetchManager = device->getDriverHandle()->getMemoryManager()->getPrefetchManager();
+        prefetchManager->migrateAllocationsToGpu(*this->device->getDriverHandle()->getSvmAllocsManager(), *this->device->getNEODevice());
+        performMemoryPrefetch = false;
+    }
+
     if (!isCopyOnlyCommandQueue && stateSipRequired) {
         NEO::PreemptionHelper::programStateSipEndWa<GfxFamily>(child, *neoDevice);
     }

level_zero/core/source/xe_hpc_core/cmdlist_xe_hpc_core.cpp

@@ -5,6 +5,8 @@
  *
  */
 
+#include "shared/source/memory_manager/prefetch_manager.h"
+
 #include "level_zero/core/source/cmdlist/cmdlist_hw.h"
 #include "level_zero/core/source/cmdlist/cmdlist_hw.inl"
 #include "level_zero/core/source/cmdlist/cmdlist_hw_immediate.h"
@@ -33,7 +35,11 @@ ze_result_t CommandListCoreFamily<IGFX_XE_HPC_CORE>::appendMemoryPrefetch(const
     }
 
     if (NEO::DebugManager.flags.AppendMemoryPrefetchForKmdMigratedSharedAllocations.get() > 0) {
-        svmAllocMgr->prefetchMemory(*device->getNEODevice(), *allocData);
+        this->performMemoryPrefetch = true;
+        auto prefetchManager = device->getDriverHandle()->getMemoryManager()->getPrefetchManager();
+        if (prefetchManager) {
+            prefetchManager->insertAllocation(*allocData);
+        }
     }
 
     if (NEO::DebugManager.flags.AddStatePrefetchCmdToMemoryPrefetchAPI.get() != 1) {

level_zero/core/test/unit_tests/xe_hpc_core/test_cmdlist_xe_hpc_core.cpp

@@ -9,6 +9,7 @@
 #include "shared/test/common/cmd_parse/gen_cmd_parse.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/test_macros/hw_test.h"
+#include "shared/test/unit_test/memory_manager/mock_prefetch_manager.h"
 
 #include "level_zero/core/source/event/event.h"
 #include "level_zero/core/test/unit_tests/fixtures/module_fixture.h"
@@ -64,7 +65,7 @@ HWTEST2_F(CommandListAppendLaunchKernelXeHpcCore, givenKernelUsingSyncBufferWhen
 
 using CommandListStatePrefetchXeHpcCore = Test<ModuleFixture>;
 
-HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenUnifiedSharedMemoryWhenPrefetchApiCalledThenDontSetMemPrefetch, IsXeHpcCore) {
+HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenUnifiedSharedMemoryWhenPrefetchApiIsCalledThenDontRequestMemoryPrefetchByDefault, IsXeHpcCore) {
     auto pCommandList = std::make_unique<WhiteBox<::L0::CommandListCoreFamily<gfxCoreFamily>>>();
     auto result = pCommandList->initialize(device, NEO::EngineGroupType::Compute, 0u);
     ASSERT_EQ(ZE_RESULT_SUCCESS, result);
@@ -75,20 +76,19 @@ HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenUnifiedSharedMemoryWhenPrefetc
 
     ze_device_mem_alloc_desc_t deviceDesc = {};
     ze_host_mem_alloc_desc_t hostDesc = {};
-    auto res = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, size, alignment, &ptr);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    result = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, size, alignment, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
     EXPECT_NE(nullptr, ptr);
 
-    auto ret = pCommandList->appendMemoryPrefetch(ptr, size);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, ret);
+    result = pCommandList->appendMemoryPrefetch(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
 
-    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
-    EXPECT_FALSE(memoryManager->setMemPrefetchCalled);
+    EXPECT_FALSE(pCommandList->isMemoryPrefetchRequested());
 
     context->freeMem(ptr);
 }
 
-HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsWhenPrefetchApiCalledThenDontCallSetMemPrefetchByDefault, IsXeHpcCore) {
+HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsWhenPrefetchApiIsCalledThenRequestMemoryPrefetch, IsXeHpcCore) {
     DebugManagerStateRestore restore;
     DebugManager.flags.AppendMemoryPrefetchForKmdMigratedSharedAllocations.set(1);
 
@@ -102,60 +102,195 @@ HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigr
 
     ze_device_mem_alloc_desc_t deviceDesc = {};
     ze_host_mem_alloc_desc_t hostDesc = {};
-    auto res = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, size, alignment, &ptr);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    result = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, size, alignment, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
     EXPECT_NE(nullptr, ptr);
 
-    auto ret = pCommandList->appendMemoryPrefetch(ptr, size);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, ret);
+    result = pCommandList->appendMemoryPrefetch(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
 
-    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
-    EXPECT_FALSE(memoryManager->setMemPrefetchCalled);
+    EXPECT_TRUE(pCommandList->isMemoryPrefetchRequested());
 
     context->freeMem(ptr);
 }
 
-HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsSetWhenPrefetchApiCalledOnUnifiedSharedMemoryThenCallSetMemPrefetch, IsXeHpcCore) {
+HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsSetWhenPrefetchApiIsCalledOnUnifiedSharedMemoryThenAppendAllocationForPrefetch, IsXeHpcCore) {
+    DebugManagerStateRestore restore;
+    DebugManager.flags.AppendMemoryPrefetchForKmdMigratedSharedAllocations.set(1);
+    DebugManager.flags.UseKmdMigration.set(1);
+
+    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
+    memoryManager->prefetchManager.reset(new MockPrefetchManager());
+
+    auto pCommandList = std::make_unique<WhiteBox<::L0::CommandListCoreFamily<gfxCoreFamily>>>();
+    auto result = pCommandList->initialize(device, NEO::EngineGroupType::Compute, 0u);
+    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+
+    size_t size = 10;
+    size_t alignment = 1u;
+    void *ptr = nullptr;
+
+    ze_device_mem_alloc_desc_t deviceDesc = {};
+    ze_host_mem_alloc_desc_t hostDesc = {};
+    result = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, size, alignment, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_NE(nullptr, ptr);
+
+    result = pCommandList->appendMemoryPrefetch(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    EXPECT_TRUE(pCommandList->isMemoryPrefetchRequested());
+
+    auto prefetchManager = static_cast<MockPrefetchManager *>(memoryManager->prefetchManager.get());
+    EXPECT_EQ(1u, prefetchManager->allocations.size());
+
+    context->freeMem(ptr);
+}
+
+HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsSetWhenPrefetchApiIsCalledOnUnifiedDeviceMemoryThenDontAppendAllocationForPrefetch, IsXeHpcCore) {
+    DebugManagerStateRestore restore;
+    DebugManager.flags.AppendMemoryPrefetchForKmdMigratedSharedAllocations.set(1);
+    DebugManager.flags.UseKmdMigration.set(1);
+
+    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
+    memoryManager->prefetchManager.reset(new MockPrefetchManager());
+
+    auto pCommandList = std::make_unique<WhiteBox<::L0::CommandListCoreFamily<gfxCoreFamily>>>();
+    auto result = pCommandList->initialize(device, NEO::EngineGroupType::Compute, 0u);
+    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+
+    size_t size = 10;
+    size_t alignment = 1u;
+    void *ptr = nullptr;
+
+    ze_device_mem_alloc_desc_t deviceDesc = {};
+    result = context->allocDeviceMem(device->toHandle(), &deviceDesc, size, alignment, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_NE(nullptr, ptr);
+
+    result = pCommandList->appendMemoryPrefetch(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    EXPECT_TRUE(pCommandList->isMemoryPrefetchRequested());
+
+    auto prefetchManager = static_cast<MockPrefetchManager *>(memoryManager->prefetchManager.get());
+    EXPECT_EQ(0u, prefetchManager->allocations.size());
+
+    context->freeMem(ptr);
+}
+
+HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsSetWhenPrefetchApiIsCalledOnUnifiedHostMemoryThenDontAppendAllocationForPrefetch, IsXeHpcCore) {
+    DebugManagerStateRestore restore;
+    DebugManager.flags.AppendMemoryPrefetchForKmdMigratedSharedAllocations.set(1);
+    DebugManager.flags.UseKmdMigration.set(1);
+
+    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
+    memoryManager->prefetchManager.reset(new MockPrefetchManager());
+
+    auto pCommandList = std::make_unique<WhiteBox<::L0::CommandListCoreFamily<gfxCoreFamily>>>();
+    auto result = pCommandList->initialize(device, NEO::EngineGroupType::Compute, 0u);
+    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+
+    size_t size = 10;
+    size_t alignment = 1u;
+    void *ptr = nullptr;
+
+    ze_host_mem_alloc_desc_t hostDesc = {};
+    result = context->allocHostMem(&hostDesc, size, alignment, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+    EXPECT_NE(nullptr, ptr);
+
+    result = pCommandList->appendMemoryPrefetch(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    EXPECT_TRUE(pCommandList->isMemoryPrefetchRequested());
+
+    auto prefetchManager = static_cast<MockPrefetchManager *>(memoryManager->prefetchManager.get());
+    EXPECT_EQ(0u, prefetchManager->allocations.size());
+
+    context->freeMem(ptr);
+}
+
+HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsSetWhenPrefetchApiIsCalledOnUnifiedDeviceMemoryThenDontCallSetMemPrefetchOnTheAssociatedDevice, IsXeHpcCore) {
+    using GfxFamily = typename NEO::GfxFamilyMapper<gfxCoreFamily>::GfxFamily;
+    using POSTSYNC_DATA = typename FamilyType::POSTSYNC_DATA;
+    using WALKER_TYPE = typename FamilyType::WALKER_TYPE;
+
     DebugManagerStateRestore restore;
     DebugManager.flags.AppendMemoryPrefetchForKmdMigratedSharedAllocations.set(1);
     DebugManager.flags.UseKmdMigration.set(1);
 
     EXPECT_EQ(0b0001u, neoDevice->deviceBitfield.to_ulong());
 
-    auto pCommandList = std::make_unique<WhiteBox<::L0::CommandListCoreFamily<gfxCoreFamily>>>();
-    auto result = pCommandList->initialize(device, NEO::EngineGroupType::Compute, 0u);
-    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
+    memoryManager->prefetchManager.reset(new MockPrefetchManager());
+
+    createKernel();
+    ze_result_t returnValue;
+    ze_command_queue_desc_t queueDesc = {};
+    auto commandList = CommandList::createImmediate(productFamily, device, &queueDesc, false, NEO::EngineGroupType::RenderCompute, returnValue);
+    ze_event_pool_desc_t eventPoolDesc = {};
+    eventPoolDesc.count = 1;
+    eventPoolDesc.flags = ZE_EVENT_POOL_FLAG_KERNEL_TIMESTAMP;
+
+    ze_event_desc_t eventDesc = {};
+    eventDesc.index = 0;
+
+    auto eventPool = std::unique_ptr<EventPool>(EventPool::create(driverHandle.get(), context, 0, nullptr, &eventPoolDesc, returnValue));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, returnValue);
+    auto event = std::unique_ptr<Event>(Event::create<uint32_t>(eventPool.get(), &eventDesc, device));
 
     size_t size = 10;
     size_t alignment = 1u;
     void *ptr = nullptr;
 
     ze_device_mem_alloc_desc_t deviceDesc = {};
-    ze_host_mem_alloc_desc_t hostDesc = {};
-    auto res = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, size, alignment, &ptr);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    auto result = context->allocDeviceMem(device->toHandle(), &deviceDesc, size, alignment, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
     EXPECT_NE(nullptr, ptr);
 
-    auto ret = pCommandList->appendMemoryPrefetch(ptr, size);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, ret);
+    result = commandList->appendMemoryPrefetch(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
 
-    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
-    EXPECT_TRUE(memoryManager->setMemPrefetchCalled);
-    EXPECT_EQ(0u, memoryManager->memPrefetchSubDeviceId);
+    ze_group_count_t groupCount{1, 1, 1};
+    CmdListKernelLaunchParams launchParams = {};
+    result = commandList->appendLaunchKernel(kernel->toHandle(), &groupCount, event->toHandle(), 0, nullptr, launchParams);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    EXPECT_FALSE(memoryManager->setMemPrefetchCalled);
 
     context->freeMem(ptr);
+    commandList->destroy();
 }
 
-HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsSetWhenPrefetchApiCalledOnUnifiedSharedMemoryThenCallSetMemPrefetchOnTheAssociatedDevice, IsXeHpcCore) {
+HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsSetWhenPrefetchApiIsCalledOnUnifiedSharedMemoryThenCallSetMemPrefetchOnTheAssociatedDevice, IsXeHpcCore) {
+    using GfxFamily = typename NEO::GfxFamilyMapper<gfxCoreFamily>::GfxFamily;
+    using POSTSYNC_DATA = typename FamilyType::POSTSYNC_DATA;
+    using WALKER_TYPE = typename FamilyType::WALKER_TYPE;
+
     DebugManagerStateRestore restore;
     DebugManager.flags.AppendMemoryPrefetchForKmdMigratedSharedAllocations.set(1);
     DebugManager.flags.UseKmdMigration.set(1);
 
     neoDevice->deviceBitfield = 0b0010;
 
-    auto pCommandList = std::make_unique<WhiteBox<::L0::CommandListCoreFamily<gfxCoreFamily>>>();
-    auto result = pCommandList->initialize(device, NEO::EngineGroupType::Compute, 0u);
-    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
+    memoryManager->prefetchManager.reset(new MockPrefetchManager());
+
+    createKernel();
+    ze_result_t returnValue;
+    ze_command_queue_desc_t queueDesc = {};
+    auto commandList = CommandList::createImmediate(productFamily, device, &queueDesc, false, NEO::EngineGroupType::RenderCompute, returnValue);
+    ze_event_pool_desc_t eventPoolDesc = {};
+    eventPoolDesc.count = 1;
+    eventPoolDesc.flags = ZE_EVENT_POOL_FLAG_KERNEL_TIMESTAMP;
+
+    ze_event_desc_t eventDesc = {};
+    eventDesc.index = 0;
+
+    auto eventPool = std::unique_ptr<EventPool>(EventPool::create(driverHandle.get(), context, 0, nullptr, &eventPoolDesc, returnValue));
+    EXPECT_EQ(ZE_RESULT_SUCCESS, returnValue);
+    auto event = std::unique_ptr<Event>(Event::create<uint32_t>(eventPool.get(), &eventDesc, device));
 
     size_t size = 10;
     size_t alignment = 1u;
@@ -163,70 +298,23 @@ HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigr
 
     ze_device_mem_alloc_desc_t deviceDesc = {};
     ze_host_mem_alloc_desc_t hostDesc = {};
-    auto res = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, size, alignment, &ptr);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    auto result = context->allocSharedMem(device->toHandle(), &deviceDesc, &hostDesc, size, alignment, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
     EXPECT_NE(nullptr, ptr);
 
-    auto ret = pCommandList->appendMemoryPrefetch(ptr, size);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, ret);
+    result = commandList->appendMemoryPrefetch(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
+
+    ze_group_count_t groupCount{1, 1, 1};
+    CmdListKernelLaunchParams launchParams = {};
+    result = commandList->appendLaunchKernel(kernel->toHandle(), &groupCount, event->toHandle(), 0, nullptr, launchParams);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, result);
 
-    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
     EXPECT_TRUE(memoryManager->setMemPrefetchCalled);
     EXPECT_EQ(1u, memoryManager->memPrefetchSubDeviceId);
 
     context->freeMem(ptr);
-}
-
-HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsSetWhenPrefetchApiCalledOnUnifiedDeviceMemoryThenDontCallSetMemPrefetch, IsXeHpcCore) {
-    DebugManagerStateRestore restore;
-    DebugManager.flags.AppendMemoryPrefetchForKmdMigratedSharedAllocations.set(1);
-    DebugManager.flags.UseKmdMigration.set(1);
-
-    auto pCommandList = std::make_unique<WhiteBox<::L0::CommandListCoreFamily<gfxCoreFamily>>>();
-    auto result = pCommandList->initialize(device, NEO::EngineGroupType::Compute, 0u);
-    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
-
-    size_t size = 10;
-    size_t alignment = 1u;
-    void *ptr = nullptr;
-
-    ze_device_mem_alloc_desc_t deviceDesc = {};
-    context->allocDeviceMem(device->toHandle(), &deviceDesc, size, alignment, &ptr);
-    EXPECT_NE(nullptr, ptr);
-
-    auto ret = pCommandList->appendMemoryPrefetch(ptr, size);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, ret);
-
-    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
-    EXPECT_FALSE(memoryManager->setMemPrefetchCalled);
-
-    context->freeMem(ptr);
-}
-
-HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigratedSharedAllocationsSetWhenPrefetchApiCalledOnUnifiedHostMemoryThenDontCallSetMemPrefetch, IsXeHpcCore) {
-    DebugManagerStateRestore restore;
-    DebugManager.flags.AppendMemoryPrefetchForKmdMigratedSharedAllocations.set(1);
-    DebugManager.flags.UseKmdMigration.set(1);
-
-    auto pCommandList = std::make_unique<WhiteBox<::L0::CommandListCoreFamily<gfxCoreFamily>>>();
-    auto result = pCommandList->initialize(device, NEO::EngineGroupType::Compute, 0u);
-    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
-
-    size_t size = 10;
-    size_t alignment = 1u;
-    void *ptr = nullptr;
-
-    ze_host_mem_alloc_desc_t hostDesc = {};
-    context->allocHostMem(&hostDesc, size, alignment, &ptr);
-    EXPECT_NE(nullptr, ptr);
-
-    auto ret = pCommandList->appendMemoryPrefetch(ptr, size);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, ret);
-
-    auto memoryManager = static_cast<MockMemoryManager *>(device->getDriverHandle()->getMemoryManager());
-    EXPECT_FALSE(memoryManager->setMemPrefetchCalled);
-
-    context->freeMem(ptr);
+    commandList->destroy();
 }
 
 using CommandListEventFenceTestsXeHpcCore = Test<ModuleFixture>;

shared/source/memory_manager/CMakeLists.txt

@@ -57,6 +57,8 @@ set(NEO_CORE_MEMORY_MANAGER
     ${CMAKE_CURRENT_SOURCE_DIR}/page_table.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/page_table.h
     ${CMAKE_CURRENT_SOURCE_DIR}/page_table.inl
+    ${CMAKE_CURRENT_SOURCE_DIR}/prefetch_manager.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/prefetch_manager.h
 )
 
 if(ENABLE_DYNAMIC_MEMORY_TRACKING)

shared/source/memory_manager/memory_manager.cpp

@@ -27,6 +27,7 @@
 #include "shared/source/memory_manager/deferred_deleter.h"
 #include "shared/source/memory_manager/host_ptr_manager.h"
 #include "shared/source/memory_manager/internal_allocation_storage.h"
+#include "shared/source/memory_manager/prefetch_manager.h"
 #include "shared/source/os_interface/hw_info_config.h"
 #include "shared/source/os_interface/os_context.h"
 #include "shared/source/os_interface/os_interface.h"
@@ -62,6 +63,7 @@ MemoryManager::MemoryManager(ExecutionEnvironment &executionEnvironment) : execu
 
     if (anyLocalMemorySupported) {
         pageFaultManager = PageFaultManager::create();
+        prefetchManager = PrefetchManager::create();
     }
 
     if (DebugManager.flags.EnableMultiStorageResources.get() != -1) {

shared/source/memory_manager/memory_manager.h

@@ -34,6 +34,7 @@ class ExecutionEnvironment;
 class Gmm;
 class HostPtrManager;
 class OsContext;
+class PrefetchManager;
 
 enum AllocationUsage {
     TEMPORARY_ALLOCATION,
@@ -154,6 +155,10 @@ class MemoryManager {
         return pageFaultManager.get();
     }
 
+    PrefetchManager *getPrefetchManager() const {
+        return prefetchManager.get();
+    }
+
     void waitForDeletions();
     MOCKABLE_VIRTUAL void waitForEnginesCompletion(GraphicsAllocation &graphicsAllocation);
     void cleanTemporaryAllocationListOnAllEngines(bool waitForCompletion);
@@ -313,6 +318,7 @@ class MemoryManager {
     std::vector<std::unique_ptr<LocalMemoryUsageBankSelector>> externalLocalMemoryUsageBankSelector;
     void *reservedMemory = nullptr;
     std::unique_ptr<PageFaultManager> pageFaultManager;
+    std::unique_ptr<PrefetchManager> prefetchManager;
     OSMemory::ReservedCpuAddressRange reservedCpuAddressRange;
     HeapAssigner heapAssigner;
     AlignmentSelector alignmentSelector = {};

shared/source/memory_manager/prefetch_manager.cpp

@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2022 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/memory_manager/prefetch_manager.h"
+
+#include "shared/source/device/device.h"
+#include "shared/source/memory_manager/unified_memory_manager.h"
+
+namespace NEO {
+
+std::unique_ptr<PrefetchManager> PrefetchManager::create() {
+    return std::make_unique<PrefetchManager>();
+}
+
+void PrefetchManager::insertAllocation(SvmAllocationData &svmData) {
+    std::unique_lock<SpinLock> lock{mtx};
+    if (svmData.memoryType == InternalMemoryType::SHARED_UNIFIED_MEMORY) {
+        allocations.push_back(svmData);
+    }
+}
+
+void PrefetchManager::migrateAllocationsToGpu(SVMAllocsManager &unifiedMemoryManager, Device &device) {
+    std::unique_lock<SpinLock> lock{mtx};
+    for (auto allocData : allocations) {
+        unifiedMemoryManager.prefetchMemory(device, allocData);
+    }
+    allocations.clear();
+}
+
+} // namespace NEO

shared/source/memory_manager/prefetch_manager.h

@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2022 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+
+#include "shared/source/helpers/non_copyable_or_moveable.h"
+#include "shared/source/memory_manager/unified_memory_manager.h"
+#include "shared/source/utilities/spinlock.h"
+
+#include <memory>
+#include <vector>
+
+namespace NEO {
+
+class Device;
+class SVMAllocsManager;
+
+class PrefetchManager : public NonCopyableOrMovableClass {
+  public:
+    static std::unique_ptr<PrefetchManager> create();
+
+    virtual ~PrefetchManager() = default;
+
+    void insertAllocation(SvmAllocationData &svmData);
+
+    void migrateAllocationsToGpu(SVMAllocsManager &unifiedMemoryManager, Device &device);
+
+  protected:
+    std::vector<SvmAllocationData> allocations;
+    SpinLock mtx;
+};
+
+} // namespace NEO

shared/test/common/mocks/mock_memory_manager.h

@@ -43,6 +43,7 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
     using MemoryManager::multiContextResourceDestructor;
     using MemoryManager::overrideAllocationData;
     using MemoryManager::pageFaultManager;
+    using MemoryManager::prefetchManager;
     using MemoryManager::registeredEngines;
     using MemoryManager::supportsMultiStorageResources;
     using MemoryManager::useNonSvmHostPtrAlloc;

shared/test/unit_test/memory_manager/mock_prefetch_manager.h

@@ -0,0 +1,17 @@
+/*
+ * Copyright (C) 2022 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#pragma once
+
+#include "shared/source/memory_manager/prefetch_manager.h"
+
+using namespace NEO;
+
+class MockPrefetchManager : public PrefetchManager {
+  public:
+    using PrefetchManager::allocations;
+};