Create buffer object with multiple lmem regions for kmd-migrated buffers

This commit enables cross-tile kmd migration for buffers in local memory

Related-To: NEO-6977

Signed-off-by: Milczarek, Slawomir <[email protected]>

Author: smilczar

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -1483,7 +1483,14 @@ BufferObject *DrmMemoryManager::createBufferObjectInMemoryRegion(Drm *drm, Gmm *
     }
 
     uint32_t handle = 0;
-    auto ret = memoryInfo->createGemExtWithSingleRegion(drm, memoryBanks, size, handle);
+    uint32_t ret = 0;
+
+    auto banks = std::bitset<32>(memoryBanks);
+    if (banks.count() > 1) {
+        ret = memoryInfo->createGemExtWithMultipleRegions(drm, memoryBanks, size, handle);
+    } else {
+        ret = memoryInfo->createGemExtWithSingleRegion(drm, memoryBanks, size, handle);
+    }
 
     if (ret != 0) {
         return nullptr;
@@ -1508,6 +1515,7 @@ bool DrmMemoryManager::createDrmAllocation(Drm *drm, DrmAllocation *allocation,
     auto currentBank = 0u;
     auto iterationOffset = 0u;
     auto banksCnt = storageInfo.getTotalBanksCnt();
+    auto useKmdMigrationForBuffers = (AllocationType::BUFFER == allocation->getAllocationType() && (DebugManager.flags.UseKmdMigrationForBuffers.get() > 0));
 
     auto handles = storageInfo.getNumBanks();
     if (storageInfo.colouringPolicy == ColouringPolicy::ChunkSizeBased) {
@@ -1522,14 +1530,16 @@ bool DrmMemoryManager::createDrmAllocation(Drm *drm, DrmAllocation *allocation,
             currentBank = 0;
             iterationOffset += banksCnt;
         }
-        uint32_t memoryBanks = static_cast<uint32_t>(storageInfo.memoryBanks.to_ulong());
-        if (storageInfo.getNumBanks() > 1) {
-            // check if we have this bank, if not move to next one
-            // we may have holes in memoryBanks that we need to skip i.e. memoryBanks 1101 and 3 handle allocation
-            while (!(memoryBanks & (1u << currentBank))) {
-                currentBank++;
+        auto memoryBanks = static_cast<uint32_t>(storageInfo.memoryBanks.to_ulong());
+        if (!useKmdMigrationForBuffers) {
+            if (storageInfo.getNumBanks() > 1) {
+                // check if we have this bank, if not move to next one
+                // we may have holes in memoryBanks that we need to skip i.e. memoryBanks 1101 and 3 handle allocation
+                while (!(memoryBanks & (1u << currentBank))) {
+                    currentBank++;
+                }
+                memoryBanks &= 1u << currentBank;
             }
-            memoryBanks &= 1u << currentBank;
         }
         auto gmm = allocation->getGmm(handleId);
         auto boSize = alignUp(gmm->gmmResourceInfo->getSizeAllocation(), MemoryConstants::pageSize64k);

shared/source/os_interface/linux/memory_info.cpp

@@ -122,4 +122,22 @@ uint32_t MemoryInfo::createGemExtWithSingleRegion(Drm *drm, uint32_t memoryBanks
     return ret;
 }
 
+uint32_t MemoryInfo::createGemExtWithMultipleRegions(Drm *drm, uint32_t memoryBanks, size_t allocSize, uint32_t &handle) {
+    auto pHwInfo = drm->getRootDeviceEnvironment().getHardwareInfo();
+    auto banks = std::bitset<32>(memoryBanks);
+    MemRegionsVec memRegions{};
+    size_t currentBank = 0;
+    size_t i = 0;
+    while (i < banks.count()) {
+        if (banks.test(currentBank)) {
+            auto regionClassAndInstance = getMemoryRegionClassAndInstance(1u << currentBank, *pHwInfo);
+            memRegions.push_back(regionClassAndInstance);
+            i++;
+        }
+        currentBank++;
+    }
+    auto ret = createGemExt(drm, memRegions, allocSize, handle, {});
+    return ret;
+}
+
 } // namespace NEO

shared/source/os_interface/linux/memory_info.h

@@ -37,6 +37,7 @@ class MemoryInfo {
     uint32_t getTileIndex(uint32_t memoryBank, const HardwareInfo &hwInfo);
 
     MOCKABLE_VIRTUAL uint32_t createGemExtWithSingleRegion(Drm *drm, uint32_t memoryBanks, size_t allocSize, uint32_t &handle);
+    MOCKABLE_VIRTUAL uint32_t createGemExtWithMultipleRegions(Drm *drm, uint32_t memoryBanks, size_t allocSize, uint32_t &handle);
 
     const RegionContainer &getDrmRegionInfos() const { return drmQueryRegions; }
 

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

@@ -83,6 +83,15 @@ struct MockedMemoryInfo : public NEO::MemoryInfo {
         handle = 1u;
         return 0u;
     }
+    uint32_t createGemExtWithMultipleRegions(Drm *drm, uint32_t memoryBanks, size_t allocSize, uint32_t &handle) override {
+        if (allocSize == 0) {
+            return EINVAL;
+        }
+        handle = 1u;
+        banks = memoryBanks;
+        return 0u;
+    }
+    uint32_t banks = 0;
 };
 
 class DrmMemoryManagerFixtureWithoutQuietIoctlExpectation {

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

@@ -505,3 +505,44 @@ TEST(MemoryInfo, givenMemoryInfoWithRegionsAndPrivateBOSupportedAndIsPerContextV
     ASSERT_TRUE(createExt);
     EXPECT_EQ(std::nullopt, createExt->vmPrivateExt.vmId);
 }
+
+TEST(MemoryInfo, givenMemoryInfoWithRegionsWhenCreatingGemExtWithMultipleRegionsThenReturnCorrectValues) {
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.EnableLocalMemory.set(1);
+
+    std::vector<MemoryRegion> regionInfo(5);
+    regionInfo[0].region = {I915_MEMORY_CLASS_SYSTEM, 0};
+    regionInfo[0].probedSize = 8 * GB;
+    regionInfo[1].region = {I915_MEMORY_CLASS_DEVICE, 0};
+    regionInfo[1].probedSize = 16 * GB;
+    regionInfo[2].region = {I915_MEMORY_CLASS_DEVICE, 1};
+    regionInfo[2].probedSize = 16 * GB;
+    regionInfo[3].region = {I915_MEMORY_CLASS_DEVICE, 2};
+    regionInfo[3].probedSize = 16 * GB;
+    regionInfo[4].region = {I915_MEMORY_CLASS_DEVICE, 3};
+    regionInfo[4].probedSize = 16 * GB;
+
+    auto memoryInfo = std::make_unique<MemoryInfo>(regionInfo);
+    ASSERT_NE(nullptr, memoryInfo);
+
+    auto executionEnvironment = std::make_unique<ExecutionEnvironment>();
+    executionEnvironment->prepareRootDeviceEnvironments(1);
+    auto drm = std::make_unique<DrmQueryMock>(*executionEnvironment->rootDeviceEnvironments[0]);
+    uint32_t handle = 0;
+    uint32_t memoryRegions = 0b1011;
+    auto ret = memoryInfo->createGemExtWithMultipleRegions(drm.get(), memoryRegions, 1024, handle);
+    EXPECT_EQ(1u, handle);
+    EXPECT_EQ(0u, ret);
+    EXPECT_EQ(1u, drm->ioctlCallsCount);
+
+    const auto &createExt = drm->context.receivedCreateGemExt;
+    ASSERT_TRUE(createExt);
+    ASSERT_EQ(3u, createExt->memoryRegions.size());
+    EXPECT_EQ(I915_MEMORY_CLASS_DEVICE, createExt->memoryRegions[0].memoryClass);
+    EXPECT_EQ(0u, createExt->memoryRegions[0].memoryInstance);
+    EXPECT_EQ(I915_MEMORY_CLASS_DEVICE, createExt->memoryRegions[1].memoryClass);
+    EXPECT_EQ(1u, createExt->memoryRegions[1].memoryInstance);
+    EXPECT_EQ(I915_MEMORY_CLASS_DEVICE, createExt->memoryRegions[2].memoryClass);
+    EXPECT_EQ(3u, createExt->memoryRegions[2].memoryInstance);
+    EXPECT_EQ(1024u, drm->context.receivedCreateGemExt->size);
+}

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

@@ -4140,6 +4140,26 @@ TEST_F(DrmMemoryManagerWithLocalMemoryAndExplicitExpectationsTest, givenDrmMemor
     EXPECT_EQ(nullptr, bo);
 }
 
+TEST_F(DrmMemoryManagerWithLocalMemoryAndExplicitExpectationsTest, givenUseKmdMigrationForBuffersWhenGraphicsAllocationInDevicePoolIsAllocatedForBufferWithSeveralMemoryBanksThenCreateGemObjectWithMultipleRegions) {
+    DebugManager.flags.UseKmdMigrationForBuffers.set(1);
+
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
+    AllocationData allocData;
+    allocData.allFlags = 0;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.type = AllocationType::BUFFER;
+    allocData.rootDeviceIndex = rootDeviceIndex;
+    allocData.storageInfo.memoryBanks = 0b11;
+
+    auto allocation = memoryManager->allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_NE(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::Success, status);
+
+    EXPECT_EQ(allocData.storageInfo.memoryBanks, static_cast<MockedMemoryInfo *>(mock->getMemoryInfo())->banks);
+
+    memoryManager->freeGraphicsMemory(allocation);
+}
+
 TEST_F(DrmMemoryManagerWithLocalMemoryAndExplicitExpectationsTest, givenUseSystemMemoryFlagWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
     AllocationData allocData;