Add memory prefetch modes for single and multiple subdevices

Single-subdevice prefetch for cmd list copy-only (with bcs) and acc mode. Multi-subdevice prefetch (default) for shared allocation with multiple BOs. Related-To: NEO-6740 Signed-off-by: Slawomir Milczarek <slawomir.milczarek@intel.com>
2025-12-19 06:24:51 +08:00 · 2022-11-30 00:48:56 +00:00
parent e3ede4bb92
commit 5c1b50bccf
22 changed files with 228 additions and 59 deletions
--- a/level_zero/core/source/cmdlist/cmdlist_hw_immediate.inl
+++ b/level_zero/core/source/cmdlist/cmdlist_hw_immediate.inl
@@ -204,7 +204,10 @@ ze_result_t CommandListCoreFamilyImmediate<gfxCoreFamily>::executeCommandListImm

    if (this->performMemoryPrefetch) {
        auto prefetchManager = this->device->getDriverHandle()->getMemoryManager()->getPrefetchManager();
-        prefetchManager->migrateAllocationsToGpu(this->getPrefetchContext(), *this->device->getDriverHandle()->getSvmAllocsManager(), *this->device->getNEODevice());
+        prefetchManager->migrateAllocationsToGpu(this->getPrefetchContext(),
+                                                 *this->device->getDriverHandle()->getSvmAllocsManager(),
+                                                 *this->device->getNEODevice(),
+                                                 *this->csr);
    }

    NEO::CompletionStamp completionStamp;
--- a/level_zero/core/source/cmdqueue/cmdqueue_hw.inl
+++ b/level_zero/core/source/cmdqueue/cmdqueue_hw.inl
@@ -963,7 +963,8 @@ void CommandQueueHw<gfxCoreFamily>::prefetchMemoryToDeviceAssociatedWithCmdList(
        auto prefetchManager = this->device->getDriverHandle()->getMemoryManager()->getPrefetchManager();
        prefetchManager->migrateAllocationsToGpu(commandList->getPrefetchContext(),
                                                 *this->device->getDriverHandle()->getSvmAllocsManager(),
-                                                 *commandList->device->getNEODevice());
+                                                 *commandList->device->getNEODevice(),
+                                                 *commandList->csr);
    }
 }

--- a/level_zero/core/test/unit_tests/xe_hpc_core/test_cmdlist_xe_hpc_core.cpp
+++ b/level_zero/core/test/unit_tests/xe_hpc_core/test_cmdlist_xe_hpc_core.cpp
@@ -311,7 +311,7 @@ HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigr
    EXPECT_EQ(ZE_RESULT_SUCCESS, result);

    EXPECT_TRUE(memoryManager->setMemPrefetchCalled);
-    EXPECT_EQ(1u, memoryManager->memPrefetchSubDeviceId);
+    EXPECT_EQ(1u, memoryManager->memPrefetchSubDeviceIds[0]);

    context->freeMem(ptr);
    commandList->destroy();
@@ -375,7 +375,7 @@ HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigr
    EXPECT_EQ(ZE_RESULT_SUCCESS, result);

    EXPECT_TRUE(memoryManager->setMemPrefetchCalled);
-    EXPECT_EQ(3u, memoryManager->memPrefetchSubDeviceId);
+    EXPECT_EQ(3u, memoryManager->memPrefetchSubDeviceIds[0]);

    EXPECT_TRUE(prefetchManager->migrateAllocationsToGpuCalled);
    EXPECT_EQ(1u, commandList->getPrefetchContext().allocations.size());
@@ -454,7 +454,7 @@ HWTEST2_F(CommandListStatePrefetchXeHpcCore, givenAppendMemoryPrefetchForKmdMigr
    EXPECT_EQ(ZE_RESULT_SUCCESS, result);

    EXPECT_TRUE(memoryManager->setMemPrefetchCalled);
-    EXPECT_EQ(0u, memoryManager->memPrefetchSubDeviceId);
+    EXPECT_EQ(0u, memoryManager->memPrefetchSubDeviceIds[0]);

    EXPECT_TRUE(prefetchManager->migrateAllocationsToGpuCalled);
    EXPECT_EQ(2u, commandList->getPrefetchContext().allocations.size());
--- a/opencl/source/api/api.cpp
+++ b/opencl/source/api/api.cpp
@@ -4199,7 +4199,7 @@ CL_API_ENTRY cl_int CL_API_CALL clEnqueueMigrateMemINTEL(

            auto allocData = pSvmAllocMgr->getSVMAlloc(ptr);
            if (allocData) {
-                pSvmAllocMgr->prefetchMemory(pCommandQueue->getDevice(), *allocData);
+                pSvmAllocMgr->prefetchMemory(pCommandQueue->getDevice(), pCommandQueue->getGpgpuCommandStreamReceiver(), *allocData);
            }
        }
    }
--- a/opencl/test/unit_test/api/cl_unified_shared_memory_tests.inl
+++ b/opencl/test/unit_test/api/cl_unified_shared_memory_tests.inl
@@ -918,7 +918,7 @@ TEST(clUnifiedSharedMemoryTests, givenUseKmdMigrationAndAppendMemoryPrefetchForK

    auto mockMemoryManager = static_cast<MockMemoryManager *>(device->getMemoryManager());
    EXPECT_TRUE(mockMemoryManager->setMemPrefetchCalled);
-    EXPECT_EQ(0u, mockMemoryManager->memPrefetchSubDeviceId);
+    EXPECT_EQ(0u, mockMemoryManager->memPrefetchSubDeviceIds[0]);

    clMemFreeINTEL(&mockContext, unifiedMemorySharedAllocation);
 }
@@ -947,7 +947,7 @@ TEST(clUnifiedSharedMemoryTests, givenContextWithMultipleSubdevicesWhenClEnqueue

    auto mockMemoryManager = static_cast<MockMemoryManager *>(subDevice->getMemoryManager());
    EXPECT_TRUE(mockMemoryManager->setMemPrefetchCalled);
-    EXPECT_EQ(1u, mockMemoryManager->memPrefetchSubDeviceId);
+    EXPECT_EQ(1u, mockMemoryManager->memPrefetchSubDeviceIds[0]);

    clMemFreeINTEL(&multiTileContext, unifiedMemorySharedAllocation);
 }
@@ -976,7 +976,9 @@ TEST(clUnifiedSharedMemoryTests, givenContextWithMultipleSubdevicesWhenClEnqueue

    auto mockMemoryManager = static_cast<MockMemoryManager *>(device->getMemoryManager());
    EXPECT_TRUE(mockMemoryManager->setMemPrefetchCalled);
-    EXPECT_EQ(0u, mockMemoryManager->memPrefetchSubDeviceId);
+    for (auto index = 0u; index < mockMemoryManager->memPrefetchSubDeviceIds.size(); index++) {
+        EXPECT_EQ(index, mockMemoryManager->memPrefetchSubDeviceIds[index]);
+    }

    clMemFreeINTEL(&multiTileContext, unifiedMemorySharedAllocation);
 }
--- a/shared/source/helpers/common_types.h
+++ b/shared/source/helpers/common_types.h
@@ -6,6 +6,8 @@
 */

 #pragma once
+#include "shared/source/utilities/stackvec.h"
+
 #include <bitset>
 #include <cstdint>
 #include <memory>
@@ -17,6 +19,7 @@ using EngineControlContainer = std::vector<EngineControl>;
 using DeviceBitfield = std::bitset<4>;
 class Device;
 using DeviceVector = std::vector<std::unique_ptr<Device>>;
+using SubDeviceIdsVec = StackVec<uint32_t, 4>;

 enum class DebugPauseState : uint32_t {
    disabled,
--- a/shared/source/memory_manager/memory_manager.h
+++ b/shared/source/memory_manager/memory_manager.h
@@ -225,7 +225,7 @@ class MemoryManager {
    virtual AllocationStatus registerLocalMemAlloc(GraphicsAllocation *allocation, uint32_t rootDeviceIndex) { return AllocationStatus::Success; };

    virtual bool setMemAdvise(GraphicsAllocation *gfxAllocation, MemAdviseFlags flags, uint32_t rootDeviceIndex) { return true; }
-    virtual bool setMemPrefetch(GraphicsAllocation *gfxAllocation, uint32_t subDeviceId, uint32_t rootDeviceIndex) { return true; }
+    virtual bool setMemPrefetch(GraphicsAllocation *gfxAllocation, SubDeviceIdsVec &subDeviceIds, uint32_t rootDeviceIndex) { return true; }

    bool isExternalAllocation(AllocationType allocationType);
    LocalMemoryUsageBankSelector *getLocalMemoryUsageBankSelector(AllocationType allocationType, uint32_t rootDeviceIndex);
--- a/shared/source/memory_manager/prefetch_manager.cpp
+++ b/shared/source/memory_manager/prefetch_manager.cpp
@@ -23,10 +23,10 @@ void PrefetchManager::insertAllocation(PrefetchContext &context, SvmAllocationDa
    }
 }

-void PrefetchManager::migrateAllocationsToGpu(PrefetchContext &context, SVMAllocsManager &unifiedMemoryManager, Device &device) {
+void PrefetchManager::migrateAllocationsToGpu(PrefetchContext &context, SVMAllocsManager &unifiedMemoryManager, Device &device, CommandStreamReceiver &csr) {
    std::unique_lock<SpinLock> lock{context.lock};
    for (auto allocData : context.allocations) {
-        unifiedMemoryManager.prefetchMemory(device, allocData);
+        unifiedMemoryManager.prefetchMemory(device, csr, allocData);
    }
 }

--- a/shared/source/memory_manager/prefetch_manager.h
+++ b/shared/source/memory_manager/prefetch_manager.h
@@ -16,6 +16,7 @@

 namespace NEO {

+class CommandStreamReceiver;
 class Device;
 class SVMAllocsManager;

@@ -32,7 +33,7 @@ class PrefetchManager : public NonCopyableOrMovableClass {

    void insertAllocation(PrefetchContext &context, SvmAllocationData &svmData);

-    MOCKABLE_VIRTUAL void migrateAllocationsToGpu(PrefetchContext &context, SVMAllocsManager &unifiedMemoryManager, Device &device);
+    MOCKABLE_VIRTUAL void migrateAllocationsToGpu(PrefetchContext &context, SVMAllocsManager &unifiedMemoryManager, Device &device, CommandStreamReceiver &csr);

    MOCKABLE_VIRTUAL void removeAllocations(PrefetchContext &context);
 };
--- a/shared/source/memory_manager/unified_memory_manager.cpp
+++ b/shared/source/memory_manager/unified_memory_manager.cpp
@@ -671,7 +671,7 @@ AllocationType SVMAllocsManager::getGraphicsAllocationTypeAndCompressionPreferen
    return allocationType;
 }

-void SVMAllocsManager::prefetchMemory(Device &device, SvmAllocationData &svmData) {
+void SVMAllocsManager::prefetchMemory(Device &device, CommandStreamReceiver &commandStreamReceiver, SvmAllocationData &svmData) {
    auto getSubDeviceId = [](Device &device) {
        if (!device.isSubDevice()) {
            uint32_t deviceBitField = static_cast<uint32_t>(device.getDeviceBitfield().to_ulong());
@@ -683,13 +683,24 @@ void SVMAllocsManager::prefetchMemory(Device &device, SvmAllocationData &svmData
        return static_cast<NEO::SubDevice *>(&device)->getSubDeviceIndex();
    };

+    auto getSubDeviceIds = [](CommandStreamReceiver &csr) {
+        SubDeviceIdsVec subDeviceIds;
+        for (auto subDeviceId = 0u; subDeviceId < csr.getOsContext().getDeviceBitfield().size(); subDeviceId++) {
+            if (csr.getOsContext().getDeviceBitfield().test(subDeviceId)) {
+                subDeviceIds.push_back(subDeviceId);
+            }
+        }
+        return subDeviceIds;
+    };
+
    if (memoryManager->isKmdMigrationAvailable(device.getRootDeviceIndex()) &&
        (svmData.memoryType == InternalMemoryType::SHARED_UNIFIED_MEMORY)) {
        auto gfxAllocation = svmData.gpuAllocations.getGraphicsAllocation(device.getRootDeviceIndex());
-        auto subDeviceId = getSubDeviceId(device);
-        memoryManager->setMemPrefetch(gfxAllocation, subDeviceId, device.getRootDeviceIndex());
+        auto subDeviceIds = commandStreamReceiver.getActivePartitions() > 1 ? getSubDeviceIds(commandStreamReceiver) : SubDeviceIdsVec{getSubDeviceId(device)};
+        memoryManager->setMemPrefetch(gfxAllocation, subDeviceIds, device.getRootDeviceIndex());
    }
 }
+
 std::unique_lock<std::mutex> SVMAllocsManager::obtainOwnership() {
    return std::unique_lock<std::mutex>(mtxForIndirectAccess);
 }
--- a/shared/source/memory_manager/unified_memory_manager.h
+++ b/shared/source/memory_manager/unified_memory_manager.h
@@ -186,7 +186,7 @@ class SVMAllocsManager {
    std::atomic<uint32_t> allocationsCounter = 0;
    MOCKABLE_VIRTUAL void makeIndirectAllocationsResident(CommandStreamReceiver &commandStreamReceiver, TaskCountType taskCount);
    void prepareIndirectAllocationForDestruction(SvmAllocationData *);
-    void prefetchMemory(Device &device, SvmAllocationData &svmData);
+    void prefetchMemory(Device &device, CommandStreamReceiver &commandStreamReceiver, SvmAllocationData &svmData);
    std::unique_lock<std::mutex> obtainOwnership();

    std::map<CommandStreamReceiver *, InternalAllocationsTracker> indirectAllocationsResidency;
--- a/shared/source/os_interface/linux/drm_allocation.cpp
+++ b/shared/source/os_interface/linux/drm_allocation.cpp
@@ -183,18 +183,18 @@ int DrmAllocation::bindBOs(OsContext *osContext, uint32_t vmHandleId, std::vecto
    return 0;
 }

-bool DrmAllocation::prefetchBO(BufferObject *bo, uint32_t subDeviceId) {
+bool DrmAllocation::prefetchBO(BufferObject *bo, uint32_t vmHandleId, uint32_t subDeviceId) {
    auto drm = bo->peekDrm();
    auto ioctlHelper = drm->getIoctlHelper();
    auto memoryClassDevice = ioctlHelper->getDrmParamValue(DrmParam::MemoryClassDevice);
    auto region = static_cast<uint32_t>((memoryClassDevice << 16u) | subDeviceId);
-    auto vmId = drm->getVirtualMemoryAddressSpace(subDeviceId);
+    auto vmId = drm->getVirtualMemoryAddressSpace(vmHandleId);

    auto result = ioctlHelper->setVmPrefetch(bo->peekAddress(), bo->peekSize(), region, vmId);

    PRINT_DEBUG_STRING(DebugManager.flags.PrintBOPrefetchingResult.get(), stdout,
-                       "prefetch BO=%d to VM %u, range: %llx - %llx, size: %lld, region: %x, result: %d\n",
-                       bo->peekHandle(), vmId, bo->peekAddress(), ptrOffset(bo->peekAddress(), bo->peekSize()), bo->peekSize(), region, result);
+                       "prefetch BO=%d to VM %u, drmVmId=%u, range: %llx - %llx, size: %lld, region: %x, result: %d\n",
+                       bo->peekHandle(), vmId, vmHandleId, bo->peekAddress(), ptrOffset(bo->peekAddress(), bo->peekSize()), bo->peekSize(), region, result);
    return result;
 }

@@ -335,20 +335,24 @@ bool DrmAllocation::setMemAdvise(Drm *drm, MemAdviseFlags flags) {
    return success;
 }

-bool DrmAllocation::setMemPrefetch(Drm *drm, uint32_t subDeviceId) {
+bool DrmAllocation::setMemPrefetch(Drm *drm, SubDeviceIdsVec &subDeviceIds) {
+    UNRECOVERABLE_IF(subDeviceIds.size() == 0);
+
    bool success = true;
    auto numHandles = GraphicsAllocation::getNumHandlesForKmdSharedAllocation(storageInfo.getNumBanks());

    if (numHandles > 1) {
-        for (uint8_t subDeviceId = 0u; subDeviceId < EngineLimits::maxHandleCount; subDeviceId++) {
-            if (storageInfo.memoryBanks.test(subDeviceId)) {
-                auto bo = this->getBOs()[subDeviceId];
-                success &= prefetchBO(bo, subDeviceId);
+        for (uint8_t handleId = 0u; handleId < EngineLimits::maxHandleCount; handleId++) {
+            if (storageInfo.memoryBanks.test(handleId)) {
+                auto bo = this->getBOs()[handleId];
+                auto vmHandleId = subDeviceIds[handleId % subDeviceIds.size()];
+                auto subDeviceId = DebugManager.flags.CreateContextWithAccessCounters.get() > 0 ? vmHandleId : handleId;
+                success &= prefetchBO(bo, vmHandleId, subDeviceId);
            }
        }
    } else {
        auto bo = this->getBO();
-        success = prefetchBO(bo, subDeviceId);
+        success = prefetchBO(bo, subDeviceIds[0], subDeviceIds[0]);
    }

    return success;
--- a/shared/source/os_interface/linux/drm_allocation.h
+++ b/shared/source/os_interface/linux/drm_allocation.h
@@ -99,7 +99,7 @@ class DrmAllocation : public GraphicsAllocation {
    void setCachePolicy(CachePolicy memType);

    bool setMemAdvise(Drm *drm, MemAdviseFlags flags);
-    bool setMemPrefetch(Drm *drm, uint32_t subDeviceId);
+    bool setMemPrefetch(Drm *drm, SubDeviceIdsVec &subDeviceIds);

    void *getMmapPtr() { return this->mmapPtr; }
    void setMmapPtr(void *ptr) { this->mmapPtr = ptr; }
@@ -109,7 +109,7 @@ class DrmAllocation : public GraphicsAllocation {
    MOCKABLE_VIRTUAL int makeBOsResident(OsContext *osContext, uint32_t vmHandleId, std::vector<BufferObject *> *bufferObjects, bool bind);
    MOCKABLE_VIRTUAL int bindBO(BufferObject *bo, OsContext *osContext, uint32_t vmHandleId, std::vector<BufferObject *> *bufferObjects, bool bind);
    MOCKABLE_VIRTUAL int bindBOs(OsContext *osContext, uint32_t vmHandleId, std::vector<BufferObject *> *bufferObjects, bool bind);
-    MOCKABLE_VIRTUAL bool prefetchBO(BufferObject *bo, uint32_t subDeviceId);
+    MOCKABLE_VIRTUAL bool prefetchBO(BufferObject *bo, uint32_t vmHandleId, uint32_t subDeviceId);
    MOCKABLE_VIRTUAL void registerBOBindExtHandle(Drm *drm);
    void freeRegisteredBOBindExtHandles(Drm *drm);
    void linkWithRegisteredHandle(uint32_t handle);
--- a/shared/source/os_interface/linux/drm_memory_manager.cpp
+++ b/shared/source/os_interface/linux/drm_memory_manager.cpp
@@ -237,18 +237,20 @@ bool DrmMemoryManager::setMemAdvise(GraphicsAllocation *gfxAllocation, MemAdvise
    return drmAllocation->setMemAdvise(&this->getDrm(rootDeviceIndex), flags);
 }

-bool DrmMemoryManager::setMemPrefetch(GraphicsAllocation *gfxAllocation, uint32_t subDeviceId, uint32_t rootDeviceIndex) {
+bool DrmMemoryManager::setMemPrefetch(GraphicsAllocation *gfxAllocation, SubDeviceIdsVec &subDeviceIds, uint32_t rootDeviceIndex) {
    auto drmAllocation = static_cast<DrmAllocation *>(gfxAllocation);
    auto osContextLinux = static_cast<OsContextLinux *>(registeredEngines[defaultEngineIndex[rootDeviceIndex]].osContext);
-    auto vmHandleId = subDeviceId;

-    auto retVal = drmAllocation->bindBOs(osContextLinux, vmHandleId, nullptr, true);
-    if (retVal != 0) {
-        DEBUG_BREAK_IF(true);
-        return false;
+    for (auto subDeviceId : subDeviceIds) {
+        auto vmHandleId = subDeviceId;
+        auto retVal = drmAllocation->bindBOs(osContextLinux, vmHandleId, nullptr, true);
+        if (retVal != 0) {
+            DEBUG_BREAK_IF(true);
+            return false;
+        }
    }

-    return drmAllocation->setMemPrefetch(&this->getDrm(rootDeviceIndex), subDeviceId);
+    return drmAllocation->setMemPrefetch(&this->getDrm(rootDeviceIndex), subDeviceIds);
 }

 NEO::BufferObject *DrmMemoryManager::allocUserptr(uintptr_t address, size_t size, uint32_t rootDeviceIndex) {
--- a/shared/source/os_interface/linux/drm_memory_manager.h
+++ b/shared/source/os_interface/linux/drm_memory_manager.h
@@ -70,7 +70,7 @@ class DrmMemoryManager : public MemoryManager {
    bool isKmdMigrationAvailable(uint32_t rootDeviceIndex) override;

    bool setMemAdvise(GraphicsAllocation *gfxAllocation, MemAdviseFlags flags, uint32_t rootDeviceIndex) override;
-    bool setMemPrefetch(GraphicsAllocation *gfxAllocation, uint32_t subDeviceId, uint32_t rootDeviceIndex) override;
+    bool setMemPrefetch(GraphicsAllocation *gfxAllocation, SubDeviceIdsVec &subDeviceIds, uint32_t rootDeviceIndex) override;

    [[nodiscard]] std::unique_lock<std::mutex> acquireAllocLock();
    std::vector<GraphicsAllocation *> &getSysMemAllocs();
--- a/shared/test/common/memory_manager/mock_prefetch_manager.h
+++ b/shared/test/common/memory_manager/mock_prefetch_manager.h
@@ -14,8 +14,8 @@ using namespace NEO;

 class MockPrefetchManager : public PrefetchManager {
  public:
-    void migrateAllocationsToGpu(PrefetchContext &prefetchContext, SVMAllocsManager &unifiedMemoryManager, Device &device) override {
-        PrefetchManager::migrateAllocationsToGpu(prefetchContext, unifiedMemoryManager, device);
+    void migrateAllocationsToGpu(PrefetchContext &prefetchContext, SVMAllocsManager &unifiedMemoryManager, Device &device, CommandStreamReceiver &csr) override {
+        PrefetchManager::migrateAllocationsToGpu(prefetchContext, unifiedMemoryManager, device, csr);
        migrateAllocationsToGpuCalled = true;
    }

--- a/shared/test/common/mocks/linux/mock_drm_allocation.h
+++ b/shared/test/common/mocks/linux/mock_drm_allocation.h
@@ -66,10 +66,11 @@ class MockDrmAllocation : public DrmAllocation {
        return bindBOsRetValue;
    }

-    bool prefetchBO(BufferObject *bo, uint32_t subDeviceId) override {
+    bool prefetchBO(BufferObject *bo, uint32_t vmHandleId, uint32_t subDeviceId) override {
        prefetchBOCalled = true;
+        vmHandleIdsReceived.push_back(vmHandleId);
        subDeviceIdsReceived.push_back(subDeviceId);
-        return DrmAllocation::prefetchBO(bo, subDeviceId);
+        return DrmAllocation::prefetchBO(bo, vmHandleId, subDeviceId);
    }

    ADDMETHOD_NOBASE(makeBOsResident, int, 0, (OsContext * osContext, uint32_t vmHandleId, std::vector<BufferObject *> *bufferObjects, bool bind));
@@ -79,6 +80,7 @@ class MockDrmAllocation : public DrmAllocation {
    bool bindBOsCalled = false;
    int bindBOsRetValue = 0;
    bool prefetchBOCalled = false;
+    std::vector<uint32_t> vmHandleIdsReceived;
    std::vector<uint32_t> subDeviceIdsReceived;
 };

--- a/shared/test/common/mocks/mock_memory_manager.h
+++ b/shared/test/common/mocks/mock_memory_manager.h
@@ -157,10 +157,10 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
        return MemoryManager::setMemAdvise(gfxAllocation, flags, rootDeviceIndex);
    }

-    bool setMemPrefetch(GraphicsAllocation *gfxAllocation, uint32_t subDeviceId, uint32_t rootDeviceIndex) override {
-        memPrefetchSubDeviceId = subDeviceId;
+    bool setMemPrefetch(GraphicsAllocation *gfxAllocation, SubDeviceIdsVec &subDeviceIds, uint32_t rootDeviceIndex) override {
+        memPrefetchSubDeviceIds = subDeviceIds;
        setMemPrefetchCalled = true;
-        return MemoryManager::setMemPrefetch(gfxAllocation, subDeviceId, rootDeviceIndex);
+        return MemoryManager::setMemPrefetch(gfxAllocation, subDeviceIds, rootDeviceIndex);
    }

    bool isKmdMigrationAvailable(uint32_t rootDeviceIndex) override {
@@ -258,7 +258,7 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
    std::unique_ptr<MockExecutionEnvironment> mockExecutionEnvironment;
    DeviceBitfield recentlyPassedDeviceBitfield{};
    std::unique_ptr<MultiGraphicsAllocation> waitAllocations = nullptr;
-    uint32_t memPrefetchSubDeviceId = 0;
+    SubDeviceIdsVec memPrefetchSubDeviceIds;
    MemAdviseFlags memAdviseFlags{};
    MemoryManager::AllocationStatus populateOsHandlesResult = MemoryManager::AllocationStatus::Success;
    GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtrResult = nullptr;
--- a/shared/test/unit_test/memory_manager/prefetch_manager_tests.cpp
+++ b/shared/test/unit_test/memory_manager/prefetch_manager_tests.cpp
@@ -7,6 +7,7 @@

 #include "shared/source/memory_manager/prefetch_manager.h"
 #include "shared/test/common/memory_manager/mock_prefetch_manager.h"
+#include "shared/test/common/mocks/mock_command_stream_receiver.h"
 #include "shared/test/common/mocks/mock_device.h"
 #include "shared/test/common/mocks/mock_svm_manager.h"
 #include "shared/test/common/mocks/ult_device_factory.h"
@@ -20,6 +21,7 @@ TEST(PrefetchManagerTests, givenPrefetchManagerWhenCallingInterfaceFunctionsThen
    RootDeviceIndicesContainer rootDeviceIndices = {mockRootDeviceIndex};
    std::map<uint32_t, DeviceBitfield> deviceBitfields{{mockRootDeviceIndex, mockDeviceBitfield}};
    auto device = deviceFactory->rootDevices[0];
+    auto csr = std::make_unique<MockCommandStreamReceiver>(*device->getExecutionEnvironment(), device->getRootDeviceIndex(), device->getDeviceBitfield());
    auto svmManager = std::make_unique<MockSVMAllocsManager>(device->getMemoryManager(), false);
    auto prefetchManager = std::make_unique<MockPrefetchManager>();
    PrefetchContext prefetchContext;
@@ -36,7 +38,7 @@ TEST(PrefetchManagerTests, givenPrefetchManagerWhenCallingInterfaceFunctionsThen
    prefetchManager->insertAllocation(prefetchContext, *svmData);
    EXPECT_EQ(1u, prefetchContext.allocations.size());

-    prefetchManager->migrateAllocationsToGpu(prefetchContext, *svmManager.get(), *device);
+    prefetchManager->migrateAllocationsToGpu(prefetchContext, *svmManager.get(), *device, *csr.get());
    EXPECT_EQ(1u, prefetchContext.allocations.size());

    prefetchManager->removeAllocations(prefetchContext);
--- a/shared/test/unit_test/memory_manager/unified_memory_manager_tests.cpp
+++ b/shared/test/unit_test/memory_manager/unified_memory_manager_tests.cpp
@@ -6,6 +6,7 @@
 */

 #include "shared/test/common/helpers/debug_manager_state_restore.h"
+#include "shared/test/common/mocks/mock_command_stream_receiver.h"
 #include "shared/test/common/mocks/mock_device.h"
 #include "shared/test/common/mocks/mock_memory_manager.h"
 #include "shared/test/common/mocks/mock_svm_manager.h"
@@ -58,4 +59,36 @@ TEST_F(SVMLocalMemoryAllocatorTest, whenFreeSharedAllocWithOffsetPointerThenReso
    pageFaultMemoryData = mockPageFaultManager->memoryData.find(ptr);
    EXPECT_EQ(svmData, nullptr);
    EXPECT_EQ(pageFaultMemoryData, mockPageFaultManager->memoryData.end());
-}
+}
+
+TEST_F(SVMLocalMemoryAllocatorTest, givenKmdMigratedSharedAllocationWhenPrefetchMemoryIsCalledForMultipleActivePartitionsThenPrefetchAllocationToSubDevices) {
+    DebugManagerStateRestore restore;
+    DebugManager.flags.UseKmdMigration.set(1);
+
+    std::unique_ptr<UltDeviceFactory> deviceFactory(new UltDeviceFactory(1, 2));
+    auto device = deviceFactory->rootDevices[0];
+    auto svmManager = std::make_unique<MockSVMAllocsManager>(device->getMemoryManager(), false);
+    auto csr = std::make_unique<MockCommandStreamReceiver>(*device->getExecutionEnvironment(), device->getRootDeviceIndex(), device->getDeviceBitfield());
+    csr->setupContext(*device->getDefaultEngine().osContext);
+    void *cmdQ = reinterpret_cast<void *>(0x12345);
+
+    SVMAllocsManager::UnifiedMemoryProperties unifiedMemoryProperties(InternalMemoryType::SHARED_UNIFIED_MEMORY, rootDeviceIndices, deviceBitfields);
+
+    auto ptr = svmManager->createSharedUnifiedMemoryAllocation(4096, unifiedMemoryProperties, &cmdQ);
+    EXPECT_NE(nullptr, ptr);
+
+    auto svmData = svmManager->getSVMAlloc(ptr);
+
+    csr->setActivePartitions(2);
+    svmManager->prefetchMemory(*device, *csr, *svmData);
+
+    auto mockMemoryManager = static_cast<MockMemoryManager *>(device->getMemoryManager());
+    EXPECT_TRUE(mockMemoryManager->setMemPrefetchCalled);
+    EXPECT_EQ(2u, mockMemoryManager->memPrefetchSubDeviceIds.size());
+
+    for (auto index = 0u; index < mockMemoryManager->memPrefetchSubDeviceIds.size(); index++) {
+        EXPECT_EQ(index, mockMemoryManager->memPrefetchSubDeviceIds[index]);
+    }
+
+    svmManager->freeSVMAlloc(ptr);
+}
--- a/shared/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp
+++ b/shared/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp
@@ -4298,8 +4298,9 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenSetMemAdviseIsCalledThenUp
    }
 }

-TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenSetMemPrefetchIsCalledThenPrefetchBOsToMultipleSubDevices) {
+TEST_F(DrmMemoryManagerTest, givenKmdMigratedSharedAllocationWithMultipleBOsWhenSetMemPrefetchIsCalledWithSubDevicesThenPrefetchBOsToTheseSubDevices) {
    TestedDrmMemoryManager memoryManager(false, false, false, *executionEnvironment);
+    SubDeviceIdsVec subDeviceIds{0, 1};
    BufferObject bo0(mock, 3, 1, 1024, 0);
    BufferObject bo1(mock, 3, 2, 1024, 0);
    BufferObjects bos{&bo0, &bo1};
@@ -4312,22 +4313,120 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenSetMemPrefetchIsCalledThen
        engine.osContext->incRefInternal();
    }

-    EXPECT_TRUE(memoryManager.setMemPrefetch(&drmAllocation, 0, rootDeviceIndex));
+    EXPECT_TRUE(memoryManager.setMemPrefetch(&drmAllocation, subDeviceIds, rootDeviceIndex));

    EXPECT_TRUE(drmAllocation.bindBOsCalled);
    EXPECT_TRUE(drmAllocation.prefetchBOCalled);

+    ASSERT_EQ(2u, drmAllocation.vmHandleIdsReceived.size());
    ASSERT_EQ(2u, drmAllocation.subDeviceIdsReceived.size());
-    for (uint32_t subDeviceId = 0; subDeviceId < drmAllocation.subDeviceIdsReceived.size(); subDeviceId++) {
-        EXPECT_EQ(subDeviceId, drmAllocation.subDeviceIdsReceived[subDeviceId]);
-    }
+
+    EXPECT_EQ(0u, drmAllocation.vmHandleIdsReceived[0]);
+    EXPECT_EQ(0u, drmAllocation.subDeviceIdsReceived[0]);
+
+    EXPECT_EQ(1u, drmAllocation.vmHandleIdsReceived[1]);
+    EXPECT_EQ(1u, drmAllocation.subDeviceIdsReceived[1]);
 }

-TEST_F(DrmMemoryManagerTest, givenCreateKmdMigratedSharedAllocationWithMultipleBOsSetToFalseWhenSetMemPrefetchIsCalledThenPrefetchBOToSubdevice) {
-    DebugManagerStateRestore restorer;
-    DebugManager.flags.CreateKmdMigratedSharedAllocationWithMultipleBOs.set(false);
+TEST_F(DrmMemoryManagerTest, givenKmdMigratedSharedAllocationWithMultipleBOsWhenSetMemPrefetchIsCalledWithASubDeviceThenPrefetchBOsToTheSubDevices) {
+    TestedDrmMemoryManager memoryManager(false, false, false, *executionEnvironment);
+    SubDeviceIdsVec subDeviceIds{0};
+    BufferObject bo0(mock, 3, 1, 1024, 0);
+    BufferObject bo1(mock, 3, 2, 1024, 0);
+    BufferObjects bos{&bo0, &bo1};
+
+    MockDrmAllocation drmAllocation(AllocationType::UNIFIED_SHARED_MEMORY, MemoryPool::LocalMemory, bos);
+    drmAllocation.storageInfo.memoryBanks = 0x3;
+
+    memoryManager.registeredEngines = EngineControlContainer{this->device->allEngines};
+    for (auto engine : memoryManager.registeredEngines) {
+        engine.osContext->incRefInternal();
+    }
+
+    EXPECT_TRUE(memoryManager.setMemPrefetch(&drmAllocation, subDeviceIds, rootDeviceIndex));
+
+    EXPECT_TRUE(drmAllocation.bindBOsCalled);
+    EXPECT_TRUE(drmAllocation.prefetchBOCalled);
+
+    ASSERT_EQ(2u, drmAllocation.vmHandleIdsReceived.size());
+    ASSERT_EQ(2u, drmAllocation.subDeviceIdsReceived.size());
+
+    EXPECT_EQ(0u, drmAllocation.vmHandleIdsReceived[0]);
+    EXPECT_EQ(0u, drmAllocation.subDeviceIdsReceived[0]);
+
+    EXPECT_EQ(0u, drmAllocation.vmHandleIdsReceived[1]);
+    EXPECT_EQ(1u, drmAllocation.subDeviceIdsReceived[1]);
+}
+
+TEST_F(DrmMemoryManagerTest, givenContextWithAccessCountersAndKmdMigratedSharedAllocationWithMultipleBOsWhenSetMemPrefetchIsCalledWithASubDeviceThenPrefetchBOsToThisSubDevice) {
+    DebugManagerStateRestore restore;
+    DebugManager.flags.CreateContextWithAccessCounters.set(1);

    TestedDrmMemoryManager memoryManager(false, false, false, *executionEnvironment);
+    SubDeviceIdsVec subDeviceIds{0};
+    BufferObject bo0(mock, 3, 1, 1024, 0);
+    BufferObject bo1(mock, 3, 2, 1024, 0);
+    BufferObjects bos{&bo0, &bo1};
+
+    MockDrmAllocation drmAllocation(AllocationType::UNIFIED_SHARED_MEMORY, MemoryPool::LocalMemory, bos);
+    drmAllocation.storageInfo.memoryBanks = 0x3;
+
+    memoryManager.registeredEngines = EngineControlContainer{this->device->allEngines};
+    for (auto engine : memoryManager.registeredEngines) {
+        engine.osContext->incRefInternal();
+    }
+
+    EXPECT_TRUE(memoryManager.setMemPrefetch(&drmAllocation, subDeviceIds, rootDeviceIndex));
+
+    EXPECT_TRUE(drmAllocation.bindBOsCalled);
+    EXPECT_TRUE(drmAllocation.prefetchBOCalled);
+
+    ASSERT_EQ(2u, drmAllocation.vmHandleIdsReceived.size());
+    ASSERT_EQ(2u, drmAllocation.subDeviceIdsReceived.size());
+
+    EXPECT_EQ(0u, drmAllocation.vmHandleIdsReceived[0]);
+    EXPECT_EQ(0u, drmAllocation.subDeviceIdsReceived[0]);
+
+    EXPECT_EQ(0u, drmAllocation.vmHandleIdsReceived[1]);
+    EXPECT_EQ(0u, drmAllocation.subDeviceIdsReceived[1]);
+}
+
+TEST_F(DrmMemoryManagerTest, givenContextWithAccessCountersAndKmdMigratedSharedAllocationWithMultipleBOsWhenSetMemPrefetchIsCalledWithSubDevicesThenPrefetchBOsToTheseSubDevices) {
+    DebugManagerStateRestore restore;
+    DebugManager.flags.CreateContextWithAccessCounters.set(1);
+
+    TestedDrmMemoryManager memoryManager(false, false, false, *executionEnvironment);
+    SubDeviceIdsVec subDeviceIds{0, 1};
+    BufferObject bo0(mock, 3, 1, 1024, 0);
+    BufferObject bo1(mock, 3, 2, 1024, 0);
+    BufferObjects bos{&bo0, &bo1};
+
+    MockDrmAllocation drmAllocation(AllocationType::UNIFIED_SHARED_MEMORY, MemoryPool::LocalMemory, bos);
+    drmAllocation.storageInfo.memoryBanks = 0x3;
+
+    memoryManager.registeredEngines = EngineControlContainer{this->device->allEngines};
+    for (auto engine : memoryManager.registeredEngines) {
+        engine.osContext->incRefInternal();
+    }
+
+    EXPECT_TRUE(memoryManager.setMemPrefetch(&drmAllocation, subDeviceIds, rootDeviceIndex));
+
+    EXPECT_TRUE(drmAllocation.bindBOsCalled);
+    EXPECT_TRUE(drmAllocation.prefetchBOCalled);
+
+    ASSERT_EQ(2u, drmAllocation.vmHandleIdsReceived.size());
+    ASSERT_EQ(2u, drmAllocation.subDeviceIdsReceived.size());
+
+    EXPECT_EQ(0u, drmAllocation.vmHandleIdsReceived[0]);
+    EXPECT_EQ(0u, drmAllocation.subDeviceIdsReceived[0]);
+
+    EXPECT_EQ(1u, drmAllocation.vmHandleIdsReceived[1]);
+    EXPECT_EQ(1u, drmAllocation.subDeviceIdsReceived[1]);
+}
+
+TEST_F(DrmMemoryManagerTest, givenKmdMigratedSharedAllocationWithSingleBOWhenSetMemPrefetchIsCalledWithASubDeviceThenPrefetchBOToThisSubdevice) {
+    TestedDrmMemoryManager memoryManager(false, false, false, *executionEnvironment);
+    SubDeviceIdsVec subDeviceIds{0};
    BufferObject bo(mock, 3, 1, 1024, 0);

    MockDrmAllocation drmAllocation(AllocationType::UNIFIED_SHARED_MEMORY, MemoryPool::LocalMemory);
@@ -4338,17 +4437,21 @@ TEST_F(DrmMemoryManagerTest, givenCreateKmdMigratedSharedAllocationWithMultipleB
        engine.osContext->incRefInternal();
    }

-    EXPECT_TRUE(memoryManager.setMemPrefetch(&drmAllocation, 0, rootDeviceIndex));
+    EXPECT_TRUE(memoryManager.setMemPrefetch(&drmAllocation, subDeviceIds, rootDeviceIndex));

    EXPECT_TRUE(drmAllocation.bindBOsCalled);
    EXPECT_TRUE(drmAllocation.prefetchBOCalled);

+    ASSERT_EQ(1u, drmAllocation.vmHandleIdsReceived.size());
    ASSERT_EQ(1u, drmAllocation.subDeviceIdsReceived.size());
+
+    EXPECT_EQ(0u, drmAllocation.vmHandleIdsReceived[0]);
    EXPECT_EQ(0u, drmAllocation.subDeviceIdsReceived[0]);
 }

 TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenSetMemPrefetchFailsToBindBufferObjectThenReturnFalse) {
    TestedDrmMemoryManager memoryManager(false, false, false, *executionEnvironment);
+    SubDeviceIdsVec subDeviceIds{0};
    BufferObject bo(mock, 3, 1, 1024, 0);

    MockDrmAllocation drmAllocation(AllocationType::UNIFIED_SHARED_MEMORY, MemoryPool::LocalMemory);
@@ -4361,7 +4464,7 @@ TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerWhenSetMemPrefetchFailsToBindB

    drmAllocation.bindBOsRetValue = -1;

-    EXPECT_FALSE(memoryManager.setMemPrefetch(&drmAllocation, 0, rootDeviceIndex));
+    EXPECT_FALSE(memoryManager.setMemPrefetch(&drmAllocation, subDeviceIds, rootDeviceIndex));

    EXPECT_TRUE(drmAllocation.bindBOsCalled);
    EXPECT_FALSE(drmAllocation.prefetchBOCalled);
--- a/shared/test/unit_test/os_interface/linux/drm_with_prelim_tests.cpp
+++ b/shared/test/unit_test/os_interface/linux/drm_with_prelim_tests.cpp
@@ -284,21 +284,23 @@ TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemAdviseWithDevicePre
 }

 TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemPrefetchSucceedsThenReturnTrue) {
+    SubDeviceIdsVec subDeviceIds{0};
    MockBufferObject bo(drm.get(), 3, 0, 0, 1);
    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::LocalMemory);
    allocation.bufferObjects[0] = &bo;

    drm->ioctlRetVal = 0;
-    EXPECT_TRUE(allocation.setMemPrefetch(drm.get(), 0));
+    EXPECT_TRUE(allocation.setMemPrefetch(drm.get(), subDeviceIds));
 }

 TEST_F(IoctlHelperPrelimFixture, givenDrmAllocationWhenSetMemPrefetchFailsThenReturnFalse) {
+    SubDeviceIdsVec subDeviceIds{0};
    MockBufferObject bo(drm.get(), 3, 0, 0, 1);
    MockDrmAllocation allocation(AllocationType::BUFFER, MemoryPool::LocalMemory);
    allocation.bufferObjects[0] = &bo;

    drm->ioctlRetVal = EINVAL;
-    EXPECT_FALSE(allocation.setMemPrefetch(drm.get(), 0));
+    EXPECT_FALSE(allocation.setMemPrefetch(drm.get(), subDeviceIds));
 }

 TEST_F(IoctlHelperPrelimFixture, givenVariousDirectSubmissionFlagSettingWhenCreateDrmContextIsCalledThenCorrectFlagsArePassedToIoctl) {