From 274d2ff4b4d5c145ea2cb7f6c1caa7af9212006c Mon Sep 17 00:00:00 2001
From: Krzysztof Gibala <krzysztof.gibala@intel.com>
Date: Tue, 16 Mar 2021 11:38:53 +0100
Subject: [PATCH] Correct createMultiGraphicsAllocationInSystemMemoryPool logic

isUSMHostAllocation flag will be set according to limited range product setup

Related-To: NEO-5508
Signed-off-by: Krzysztof Gibala <krzysztof.gibala@intel.com>
---
 .../memory_manager_multi_device_tests.cpp     | 29 +++++++++++++++++++
 .../linux/drm_memory_manager_tests.cpp        | 13 +++++++--
 .../source/memory_manager/memory_manager.cpp  |  8 +++--
 3 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/opencl/test/unit_test/memory_manager/memory_manager_multi_device_tests.cpp b/opencl/test/unit_test/memory_manager/memory_manager_multi_device_tests.cpp
index cc3fc44d59..c92ef57e6c 100644
--- a/opencl/test/unit_test/memory_manager/memory_manager_multi_device_tests.cpp
+++ b/opencl/test/unit_test/memory_manager/memory_manager_multi_device_tests.cpp
@@ -85,3 +85,32 @@ TEST_P(MemoryManagerMultiDeviceTest, givenRootDeviceIndexSpecifiedWhenAllocateGr
     }
     delete tagsMultiAllocation;
 }
+
+TEST_P(MemoryManagerMultiDeviceTest, givenRootDeviceIndexSpecifiedWhenAllocateGraphicsMemoryIsCalledThenAllocationPropertiesUsmFlagIsSetAccordingToAddressRange) {
+    std::vector<uint32_t> rootDeviceIndices;
+
+    for (uint32_t rootDeviceIndex = 0; rootDeviceIndex < getNumRootDevices(); ++rootDeviceIndex) {
+        rootDeviceIndices.push_back(rootDeviceIndex);
+    }
+
+    auto maxRootDeviceIndex = *std::max_element(rootDeviceIndices.begin(), rootDeviceIndices.end(), std::less<uint32_t const>());
+    auto tagsMultiAllocation = new MultiGraphicsAllocation(maxRootDeviceIndex);
+
+    AllocationProperties unifiedMemoryProperties{rootDeviceIndices.at(0), MemoryConstants::pageSize, GraphicsAllocation::AllocationType::TAG_BUFFER, systemMemoryBitfield};
+
+    memoryManager->createMultiGraphicsAllocationInSystemMemoryPool(rootDeviceIndices, unifiedMemoryProperties, *tagsMultiAllocation);
+    EXPECT_NE(nullptr, tagsMultiAllocation);
+
+    for (auto rootDeviceIndex : rootDeviceIndices) {
+        if (memoryManager->isLimitedRange(rootDeviceIndex)) {
+            EXPECT_EQ(unifiedMemoryProperties.flags.isUSMHostAllocation, false);
+        } else {
+            EXPECT_EQ(unifiedMemoryProperties.flags.isUSMHostAllocation, true);
+        }
+    }
+
+    for (auto graphicsAllocation : tagsMultiAllocation->getGraphicsAllocations()) {
+        memoryManager->freeGraphicsMemory(graphicsAllocation);
+    }
+    delete tagsMultiAllocation;
+}
diff --git a/opencl/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp b/opencl/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp
index 22d05ee275..d6e225953d 100644
--- a/opencl/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp
+++ b/opencl/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp
@@ -1110,7 +1110,11 @@ TEST_F(DrmMemoryManagerTest, GivenMisalignedHostPtrAndMultiplePagesSizeWhenAsked
 
     auto hostPtrManager = static_cast<MockHostPtrManager *>(memoryManager->getHostPtrManager());
 
-    ASSERT_EQ(3u, hostPtrManager->getFragmentCount());
+    if (memoryManager->isLimitedRange(rootDeviceIndex)) {
+        ASSERT_EQ(6u, hostPtrManager->getFragmentCount());
+    } else {
+        ASSERT_EQ(3u, hostPtrManager->getFragmentCount());
+    }
 
     auto reqs = MockHostPtrManager::getAllocationRequirements(rootDeviceIndex, ptr, size);
 
@@ -1120,7 +1124,12 @@ TEST_F(DrmMemoryManagerTest, GivenMisalignedHostPtrAndMultiplePagesSizeWhenAsked
         EXPECT_EQ(reqs.allocationFragments[i].allocationPtr, reinterpret_cast<void *>(graphicsAllocation->fragmentsStorage.fragmentStorageData[i].osHandleStorage->bo->peekAddress()));
     }
     memoryManager->freeGraphicsMemory(graphicsAllocation);
-    EXPECT_EQ(0u, hostPtrManager->getFragmentCount());
+
+    if (memoryManager->isLimitedRange(rootDeviceIndex)) {
+        EXPECT_EQ(3u, hostPtrManager->getFragmentCount());
+    } else {
+        EXPECT_EQ(0u, hostPtrManager->getFragmentCount());
+    }
 }
 
 TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhenAskedFor32BitAllocationThen32BitDrmAllocationIsBeingReturned) {
diff --git a/shared/source/memory_manager/memory_manager.cpp b/shared/source/memory_manager/memory_manager.cpp
index fd934393f6..3153448a42 100644
--- a/shared/source/memory_manager/memory_manager.cpp
+++ b/shared/source/memory_manager/memory_manager.cpp
@@ -144,9 +144,14 @@ void *MemoryManager::createMultiGraphicsAllocationInSystemMemoryPool(std::vector
 
     for (auto &rootDeviceIndex : rootDeviceIndices) {
         properties.rootDeviceIndex = rootDeviceIndex;
+        properties.flags.isUSMHostAllocation = true;
+
+        if (isLimitedRange(properties.rootDeviceIndex)) {
+            properties.flags.isUSMHostAllocation = false;
+            DEBUG_BREAK_IF(rootDeviceIndices.size() > 1);
+        }
 
         if (!ptr) {
-            properties.flags.isUSMHostAllocation = true;
             auto graphicsAllocation = allocateGraphicsMemoryWithProperties(properties);
             if (!graphicsAllocation) {
                 return nullptr;
@@ -155,7 +160,6 @@ void *MemoryManager::createMultiGraphicsAllocationInSystemMemoryPool(std::vector
             ptr = reinterpret_cast<void *>(graphicsAllocation->getUnderlyingBuffer());
         } else {
             properties.flags.allocateMemory = false;
-            properties.flags.isUSMHostAllocation = true;
 
             auto graphicsAllocation = createGraphicsAllocationFromExistingStorage(properties, ptr, multiGraphicsAllocation);