Device pool allocation in OsAgnosticMemoryManager

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>

opencl/test/unit_test/memory_manager/CMakeLists.txt

@@ -11,10 +11,8 @@ set(IGDRCL_SRCS_tests_memory_manager
     ${CMAKE_CURRENT_SOURCE_DIR}/internal_allocation_storage_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/memory_manager_multi_device_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/memory_manager_tests.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}/memory_manager_allocate_in_device_pool_tests.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/memory_manager_allocate_in_device_pool_tests.inl
-    ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}/memory_manager_allocate_in_preferred_pool_tests.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/memory_manager_allocate_in_preferred_pool_tests.inl
+    ${CMAKE_CURRENT_SOURCE_DIR}/memory_manager_allocate_in_device_pool_tests.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/memory_manager_allocate_in_preferred_pool_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/surface_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/storage_info_tests.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/unified_memory_manager_tests.cpp

opencl/test/unit_test/memory_manager/memory_manager_allocate_in_device_pool_tests.cpp

@@ -1,77 +1,401 @@
 /*
- * Copyright (C) 2018-2020 Intel Corporation
+ * Copyright (C) 2018-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
  */
 
-#include "opencl/test/unit_test/memory_manager/memory_manager_allocate_in_device_pool_tests.inl"
+#include "shared/source/aub/aub_helper.h"
+#include "shared/source/execution_environment/execution_environment.h"
+#include "shared/source/memory_manager/os_agnostic_memory_manager.h"
+#include "shared/test/common/helpers/debug_manager_state_restore.h"
+#include "shared/test/common/mocks/mock_device.h"
+#include "shared/test/common/mocks/mock_execution_environment.h"
+#include "shared/test/common/mocks/ult_device_factory.h"
 
-#include "shared/source/helpers/array_count.h"
+#include "opencl/source/helpers/memory_properties_helpers.h"
+#include "opencl/source/mem_obj/mem_obj_helper.h"
+#include "opencl/test/unit_test/mocks/mock_gmm.h"
+#include "opencl/test/unit_test/mocks/mock_memory_manager.h"
+#include "test.h"
 
-TEST(MemoryManagerTest, givenNotSetUseSystemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenAllocationIsReturned) {
+using namespace NEO;
+
+TEST(MemoryManagerTest, givenSetUseSytemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
     MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
     MockMemoryManager memoryManager(false, false, executionEnvironment);
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
+    AllocationData allocData;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.flags.useSystemMemory = true;
+    allocData.flags.allocateMemory = true;
 
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_EQ(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::RetryInNonDevicePool, status);
+}
+
+TEST(MemoryManagerTest, givenAllowed32BitAndFroce32BitWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, false, executionEnvironment);
+    memoryManager.setForce32BitAllocations(true);
+
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
+    AllocationData allocData;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.flags.allow32Bit = true;
+    allocData.flags.allocateMemory = true;
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_EQ(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::RetryInNonDevicePool, status);
+}
+
+TEST(AllocationFlagsTest, givenAllocateMemoryFlagWhenGetAllocationFlagsIsCalledThenAllocateFlagIsCorrectlySet) {
+    HardwareInfo hwInfo(*defaultHwInfo);
+    UltDeviceFactory deviceFactory{1, 0};
+    auto pDevice = deviceFactory.rootDevices[0];
+    MemoryProperties memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, 0, 0, pDevice);
+    auto allocationProperties = MemoryPropertiesHelper::getAllocationProperties(0, memoryProperties, true, 0, GraphicsAllocation::AllocationType::BUFFER, false, hwInfo, {});
+    EXPECT_TRUE(allocationProperties.flags.allocateMemory);
+
+    auto allocationProperties2 = MemoryPropertiesHelper::getAllocationProperties(0, memoryProperties, false, 0, GraphicsAllocation::AllocationType::BUFFER, false, hwInfo, {});
+    EXPECT_FALSE(allocationProperties2.flags.allocateMemory);
+}
+
+TEST(UncacheableFlagsTest, givenUncachedResourceFlagWhenGetAllocationFlagsIsCalledThenUncacheableFlagIsCorrectlySet) {
+    cl_mem_flags_intel flagsIntel = CL_MEM_LOCALLY_UNCACHED_RESOURCE;
+    UltDeviceFactory deviceFactory{1, 0};
+    auto pDevice = deviceFactory.rootDevices[0];
+    MemoryProperties memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, flagsIntel, 0, pDevice);
+    auto allocationFlags = MemoryPropertiesHelper::getAllocationProperties(
+        0, memoryProperties, false, 0, GraphicsAllocation::AllocationType::BUFFER, false, pDevice->getHardwareInfo(), {});
+    EXPECT_TRUE(allocationFlags.flags.uncacheable);
+
+    flagsIntel = 0;
+    memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, flagsIntel, 0, pDevice);
+    auto allocationFlags2 = MemoryPropertiesHelper::getAllocationProperties(
+        0, memoryProperties, false, 0, GraphicsAllocation::AllocationType::BUFFER, false, pDevice->getHardwareInfo(), {});
+    EXPECT_FALSE(allocationFlags2.flags.uncacheable);
+}
+
+TEST(AllocationFlagsTest, givenReadOnlyResourceFlagWhenGetAllocationFlagsIsCalledThenFlushL3FlagsAreCorrectlySet) {
+    cl_mem_flags flags = CL_MEM_READ_ONLY;
+    UltDeviceFactory deviceFactory{1, 0};
+    auto pDevice = deviceFactory.rootDevices[0];
+    MemoryProperties memoryProperties = MemoryPropertiesHelper::createMemoryProperties(flags, 0, 0, pDevice);
+
+    auto allocationFlags =
+        MemoryPropertiesHelper::getAllocationProperties(
+            0, memoryProperties, true, 0, GraphicsAllocation::AllocationType::BUFFER, false, pDevice->getHardwareInfo(), {});
+    EXPECT_FALSE(allocationFlags.flags.flushL3RequiredForRead);
+    EXPECT_FALSE(allocationFlags.flags.flushL3RequiredForWrite);
+
+    memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, 0, 0, pDevice);
+    auto allocationFlags2 = MemoryPropertiesHelper::getAllocationProperties(
+        0, memoryProperties, true, 0, GraphicsAllocation::AllocationType::BUFFER, false, pDevice->getHardwareInfo(), {});
+    EXPECT_TRUE(allocationFlags2.flags.flushL3RequiredForRead);
+    EXPECT_TRUE(allocationFlags2.flags.flushL3RequiredForWrite);
+}
+
+TEST(StorageInfoTest, whenStorageInfoIsCreatedWithDefaultConstructorThenReturnsOneHandle) {
+    StorageInfo storageInfo;
+    EXPECT_EQ(1u, storageInfo.getNumBanks());
+}
+
+class MockMemoryManagerBaseImplementationOfDevicePool : public MemoryManagerCreate<OsAgnosticMemoryManager> {
+  public:
+    using MemoryManagerCreate<OsAgnosticMemoryManager>::MemoryManagerCreate;
+
+    GraphicsAllocation *allocateGraphicsMemoryInDevicePool(const AllocationData &allocationData, AllocationStatus &status) override {
+        if (failInAllocate) {
+            return nullptr;
+        }
+        return OsAgnosticMemoryManager::allocateGraphicsMemoryInDevicePool(allocationData, status);
+    }
+
+    GraphicsAllocation *allocateGraphicsMemoryWithAlignment(const AllocationData &allocationData) override {
+        if (failInAllocate) {
+            return nullptr;
+        }
+        return OsAgnosticMemoryManager::allocateGraphicsMemoryWithAlignment(allocationData);
+    }
+
+    void *allocateSystemMemory(size_t size, size_t alignment) override {
+        if (failInAllocate) {
+            return nullptr;
+        }
+        return OsAgnosticMemoryManager::allocateSystemMemory(size, alignment);
+    }
+    bool failInAllocate = false;
+};
+
+using MemoryManagerTests = ::testing::Test;
+
+TEST(MemoryManagerTest, givenNotSetUseSystemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenLocalMemoryAllocationIsReturned) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
+    allocData.allFlags = 0;
     allocData.size = MemoryConstants::pageSize;
     allocData.flags.allocateMemory = true;
 
     auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
     EXPECT_NE(nullptr, allocation);
     EXPECT_EQ(MemoryManager::AllocationStatus::Success, status);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation->getMemoryPool());
 
     memoryManager.freeGraphicsMemory(allocation);
 }
 
-TEST(MemoryManagerTest, givenImageOrSharedResourceCopyWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
+TEST(MemoryManagerTest, givenNotSetUseSystemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenLocalMemoryAllocationHasCorrectStorageInfoAndFlushL3IsSet) {
     MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
-    MockMemoryManager memoryManager(false, false, executionEnvironment);
-
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
+    allocData.allFlags = 0;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.flags.allocateMemory = true;
+    allocData.storageInfo.memoryBanks = 0x1;
+    allocData.storageInfo.pageTablesVisibility = 0x2;
+    allocData.storageInfo.cloningOfPageTables = false;
+    allocData.flags.flushL3 = true;
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_NE(nullptr, allocation);
+
+    EXPECT_EQ(allocData.storageInfo.memoryBanks, allocation->storageInfo.memoryBanks);
+    EXPECT_EQ(allocData.storageInfo.pageTablesVisibility, allocation->storageInfo.pageTablesVisibility);
+    EXPECT_FALSE(allocation->storageInfo.cloningOfPageTables);
+    EXPECT_TRUE(allocation->isFlushL3Required());
+
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
+TEST(MemoryManagerTest, givenEnabledLocalMemoryAndUseSytemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
+    AllocationData allocData;
+    allocData.allFlags = 0;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.flags.useSystemMemory = true;
+    allocData.flags.allocateMemory = true;
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_EQ(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::RetryInNonDevicePool, status);
+}
+
+TEST(MemoryManagerTest, givenEnabledLocalMemoryAndAllowed32BitAndForce32BitWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+    memoryManager.setForce32BitAllocations(true);
+
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
+    AllocationData allocData;
+    allocData.allFlags = 0;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.flags.allow32Bit = true;
+    allocData.flags.allocateMemory = true;
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_EQ(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::RetryInNonDevicePool, status);
+}
+
+TEST(MemoryManagerTest, givenEnabledLocalMemoryAndAllowed32BitWhen32BitIsNotForcedThenGraphicsAllocationInDevicePoolReturnsLocalMemoryAllocation) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+    memoryManager.setForce32BitAllocations(false);
+
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
+    AllocationData allocData;
+    allocData.allFlags = 0;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.flags.allow32Bit = true;
+    allocData.flags.allocateMemory = true;
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    ASSERT_NE(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::Success, status);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation->getMemoryPool());
+
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
+TEST(MemoryManagerTest, givenEnabledLocalMemoryWhenAllocate32BitFailsThenGraphicsAllocationInDevicePoolReturnsError) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
+    AllocationData allocData;
+    allocData.type = GraphicsAllocation::AllocationType::KERNEL_ISA; // HEAP_INTERNAL will be used
+    allocData.allFlags = 0;
     allocData.size = MemoryConstants::pageSize;
     allocData.flags.allocateMemory = true;
 
-    GraphicsAllocation::AllocationType types[] = {GraphicsAllocation::AllocationType::IMAGE,
-                                                  GraphicsAllocation::AllocationType::SHARED_RESOURCE_COPY};
+    memoryManager.failAllocate32Bit = true;
 
-    for (auto type : types) {
-        allocData.type = type;
-        auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
-        EXPECT_EQ(nullptr, allocation);
-        EXPECT_EQ(MemoryManager::AllocationStatus::RetryInNonDevicePool, status);
-    }
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_EQ(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::Error, status);
+
+    memoryManager.freeGraphicsMemory(allocation);
 }
 
-TEST(MemoryManagerTest, givenSvmGpuAllocationTypeWhenAllocationSystemMemoryFailsThenReturnNull) {
+HWTEST_F(MemoryManagerTests, givenEnabledLocalMemoryWhenAllocatingDebugAreaThenHeapInternalDeviceFrontWindowIsUsed) {
     MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
-    MockMemoryManager memoryManager(false, false, executionEnvironment);
+    MemoryManagerCreate<OsAgnosticMemoryManager> osAgnosticMemoryManager(false, true, executionEnvironment);
 
+    NEO::AllocationProperties properties{0, true, MemoryConstants::pageSize64k,
+                                         NEO::GraphicsAllocation::AllocationType::DEBUG_MODULE_AREA,
+                                         false,
+                                         mockDeviceBitfield};
+    properties.flags.use32BitFrontWindow = true;
+
+    auto &hwHelper = HwHelper::get(defaultHwInfo->platform.eRenderCoreFamily);
+    auto systemMemoryPlacement = hwHelper.useSystemMemoryPlacementForISA(*defaultHwInfo);
+
+    HeapIndex expectedHeap = HeapIndex::TOTAL_HEAPS;
+    HeapIndex baseHeap = HeapIndex::TOTAL_HEAPS;
+    if (systemMemoryPlacement) {
+        expectedHeap = HeapIndex::HEAP_INTERNAL_FRONT_WINDOW;
+        baseHeap = HeapIndex::HEAP_INTERNAL;
+    } else {
+        expectedHeap = HeapIndex::HEAP_INTERNAL_DEVICE_FRONT_WINDOW;
+        baseHeap = HeapIndex::HEAP_INTERNAL_DEVICE_MEMORY;
+    }
+    auto moduleDebugArea = osAgnosticMemoryManager.allocateGraphicsMemoryWithProperties(properties);
+    auto gpuAddress = moduleDebugArea->getGpuAddress();
+    EXPECT_LE(GmmHelper::canonize(osAgnosticMemoryManager.getGfxPartition(0)->getHeapBase(expectedHeap)), gpuAddress);
+    EXPECT_GT(GmmHelper::canonize(osAgnosticMemoryManager.getGfxPartition(0)->getHeapLimit(expectedHeap)), gpuAddress);
+    EXPECT_EQ(GmmHelper::canonize(osAgnosticMemoryManager.getGfxPartition(0)->getHeapBase(expectedHeap)), moduleDebugArea->getGpuBaseAddress());
+    EXPECT_EQ(GmmHelper::canonize(osAgnosticMemoryManager.getGfxPartition(0)->getHeapBase(baseHeap)), moduleDebugArea->getGpuBaseAddress());
+
+    osAgnosticMemoryManager.freeGraphicsMemory(moduleDebugArea);
+}
+
+TEST(BaseMemoryManagerTest, givenMemoryManagerWithForced32BitsWhenSystemMemoryIsNotSetAnd32BitNotAllowedThenAllocateInDevicePoolReturnsLocalMemory) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManagerBaseImplementationOfDevicePool memoryManager(false, true, executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
+    allocData.allFlags = 0;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.flags.allocateMemory = true;
+
+    memoryManager.setForce32BitAllocations(true);
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_NE(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::Success, status);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation->getMemoryPool());
+
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
+TEST(BaseMemoryManagerTest, givenDebugVariableSetWhenCompressedBufferIsCreatedThenCreateCompressedGmm) {
+    DebugManagerStateRestore restore;
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    executionEnvironment.initGmm();
+    MemoryManagerCreate<OsAgnosticMemoryManager> memoryManager(false, true, executionEnvironment);
+
+    AllocationProperties allocPropertiesBuffer(mockRootDeviceIndex, 1, GraphicsAllocation::AllocationType::BUFFER, mockDeviceBitfield);
+    AllocationProperties allocPropertiesBufferCompressed(mockRootDeviceIndex, 1, GraphicsAllocation::AllocationType::BUFFER_COMPRESSED, mockDeviceBitfield);
+
+    DebugManager.flags.RenderCompressedBuffersEnabled.set(1);
+    auto allocationBuffer = memoryManager.allocateGraphicsMemoryInPreferredPool(allocPropertiesBuffer, nullptr);
+    auto allocationBufferCompressed = memoryManager.allocateGraphicsMemoryInPreferredPool(allocPropertiesBufferCompressed, nullptr);
+    EXPECT_EQ(nullptr, allocationBuffer->getDefaultGmm());
+    EXPECT_NE(nullptr, allocationBufferCompressed->getDefaultGmm());
+    EXPECT_TRUE(allocationBufferCompressed->getDefaultGmm()->isCompressionEnabled);
+    memoryManager.freeGraphicsMemory(allocationBuffer);
+    memoryManager.freeGraphicsMemory(allocationBufferCompressed);
+
+    DebugManager.flags.RenderCompressedBuffersEnabled.set(0);
+    allocationBuffer = memoryManager.allocateGraphicsMemoryInPreferredPool(allocPropertiesBuffer, nullptr);
+    allocationBufferCompressed = memoryManager.allocateGraphicsMemoryInPreferredPool(allocPropertiesBufferCompressed, nullptr);
+    EXPECT_EQ(nullptr, allocationBuffer->getDefaultGmm());
+    EXPECT_EQ(nullptr, allocationBufferCompressed->getDefaultGmm());
+    memoryManager.freeGraphicsMemory(allocationBuffer);
+    memoryManager.freeGraphicsMemory(allocationBufferCompressed);
+}
+
+TEST(BaseMemoryManagerTest, givenMemoryManagerWithForced32BitsWhenSystemMemoryIsNotSetAnd32BitAllowedThenAllocateInDevicePoolReturnsRetryInNonDevicePool) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManagerBaseImplementationOfDevicePool memoryManager(false, true, executionEnvironment);
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
+    AllocationData allocData;
+    allocData.allFlags = 0;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.flags.allocateMemory = true;
+    allocData.flags.allow32Bit = true;
+
+    memoryManager.setForce32BitAllocations(true);
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_EQ(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::RetryInNonDevicePool, status);
+
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
+TEST(BaseMemoryManagerTest, givenMemoryManagerWhenAllocateFailsThenAllocateInDevicePoolReturnsError) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManagerBaseImplementationOfDevicePool memoryManager(false, true, executionEnvironment);
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
+    AllocationData allocData;
+    allocData.allFlags = 0;
+    allocData.size = MemoryConstants::pageSize;
+    allocData.flags.allocateMemory = true;
+    allocData.flags.allow32Bit = true;
+
+    memoryManager.failInAllocate = true;
+    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
+    EXPECT_EQ(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::Error, status);
+
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
+TEST(BaseMemoryManagerTest, givenSvmGpuAllocationTypeWhenAllocateSystemMemoryFailsThenReturnNull) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+
+    if (!executionEnvironment.rootDeviceEnvironments[0]->isFullRangeSvm()) {
+        return;
+    }
+
+    MockMemoryManagerBaseImplementationOfDevicePool memoryManager(false, true, executionEnvironment);
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
+    AllocationData allocData;
+
     allocData.allFlags = 0;
     allocData.size = MemoryConstants::pageSize;
     allocData.type = GraphicsAllocation::AllocationType::SVM_GPU;
     allocData.hostPtr = reinterpret_cast<void *>(0x1000);
 
-    memoryManager.failAllocateSystemMemory = true;
+    memoryManager.failInAllocate = true;
     auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
     EXPECT_EQ(nullptr, allocation);
     EXPECT_EQ(MemoryManager::AllocationStatus::Error, status);
+
+    memoryManager.freeGraphicsMemory(allocation);
 }
 
-TEST(MemoryManagerTest, givenSvmGpuAllocationTypeWhenAllocationSucceedThenReturnGpuAddressAsHostPtr) {
-    if (defaultHwInfo->capabilityTable.gpuAddressSpace != maxNBitValue(48) && defaultHwInfo->capabilityTable.gpuAddressSpace != maxNBitValue(47)) {
+TEST(BaseMemoryManagerTest, givenSvmGpuAllocationTypeWhenAllocationSucceedThenReturnGpuAddressAsHostPtrAndCpuAllocation) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+
+    if (!executionEnvironment.rootDeviceEnvironments[0]->isFullRangeSvm()) {
         return;
     }
 
-    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
-    MockMemoryManager memoryManager(false, false, executionEnvironment);
-
+    MockMemoryManagerBaseImplementationOfDevicePool memoryManager(false, true, executionEnvironment);
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Error;
     AllocationData allocData;
+
+    allocData.allFlags = 0;
     allocData.size = MemoryConstants::pageSize;
     allocData.type = GraphicsAllocation::AllocationType::SVM_GPU;
     allocData.hostPtr = reinterpret_cast<void *>(0x1000);
@@ -85,9 +409,170 @@ TEST(MemoryManagerTest, givenSvmGpuAllocationTypeWhenAllocationSucceedThenReturn
     memoryManager.freeGraphicsMemory(allocation);
 }
 
+TEST(MemoryAllocationTest, givenMultiTileVisiblityWhenAskedForFlagsThenL3NeedsToBeFlushed) {
+    HardwareInfo hwInfo(*defaultHwInfo);
+    UltDeviceFactory deviceFactory{1, 0};
+    auto memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, 0, 0, deviceFactory.rootDevices[0]);
+    auto allocationFlags = MemoryPropertiesHelper::getAllocationProperties(0, memoryProperties, true, 0, GraphicsAllocation::AllocationType::BUFFER, false, hwInfo, {});
+    EXPECT_TRUE(allocationFlags.flags.flushL3RequiredForRead);
+    EXPECT_TRUE(allocationFlags.flags.flushL3RequiredForWrite);
+}
+
+TEST(MemoryAllocationTest, givenMultiTileVisiblityAndUncachedWhenAskedForFlagsThenL3DoesNotNeedToBeFlushed) {
+    UltDeviceFactory deviceFactory{1, 0};
+    auto pDevice = deviceFactory.rootDevices[0];
+    cl_mem_flags_intel flagsIntel = CL_MEM_LOCALLY_UNCACHED_RESOURCE;
+    MemoryProperties memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, flagsIntel, 0, pDevice);
+    HardwareInfo hwInfo(*defaultHwInfo);
+
+    auto allocationFlags = MemoryPropertiesHelper::getAllocationProperties(
+        0, memoryProperties, true, 0, GraphicsAllocation::AllocationType::BUFFER, false, hwInfo, {});
+    EXPECT_FALSE(allocationFlags.flags.flushL3RequiredForRead);
+    EXPECT_FALSE(allocationFlags.flags.flushL3RequiredForWrite);
+}
+
+TEST(MemoryAllocationTest, givenAubDumpForceAllToLocalMemoryWhenMemoryAllocationIsCreatedThenItHasLocalMemoryPool) {
+    DebugManagerStateRestore debugRestorer;
+    DebugManager.flags.AUBDumpForceAllToLocalMemory.set(true);
+
+    MemoryAllocation allocation(mockRootDeviceIndex, GraphicsAllocation::AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000,
+                                0x1000, 0, MemoryPool::System4KBPages, false, false, mockMaxOsContextCount);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation.getMemoryPool());
+}
+
+TEST(MemoryAllocationTest, givenAubDumpForceAllToLocalMemoryWhenMemoryAllocationIsOverridenThenItHasLocalMemoryPool) {
+    DebugManagerStateRestore debugRestorer;
+    DebugManager.flags.AUBDumpForceAllToLocalMemory.set(true);
+
+    MemoryAllocation allocation(mockRootDeviceIndex, GraphicsAllocation::AllocationType::UNKNOWN, nullptr, reinterpret_cast<void *>(0x1000), 0x1000,
+                                0x1000, 0, MemoryPool::System4KBPages, false, false, mockMaxOsContextCount);
+    allocation.overrideMemoryPool(MemoryPool::System64KBPages);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation.getMemoryPool());
+}
+
+HWTEST_F(MemoryManagerTests, givenEnabledLocalMemoryWhenAllocateInternalAllocationInDevicePoolThen32BitAllocationIsCreated) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInPreferredPool({mockRootDeviceIndex, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::INTERNAL_HEAP, mockDeviceBitfield}, nullptr);
+    EXPECT_NE(nullptr, allocation);
+    EXPECT_TRUE(allocation->is32BitAllocation());
+    EXPECT_TRUE(memoryManager.allocationInDevicePoolCreated);
+
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
+TEST(MemoryManagerTest, givenDisabledLocalMemoryWhenAllocateInternalAllocationInDevicePoolThen32BitAllocationIsCreatedInNonDevicePool) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, false, executionEnvironment);
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInPreferredPool({mockRootDeviceIndex, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::INTERNAL_HEAP, mockDeviceBitfield}, nullptr);
+    EXPECT_NE(nullptr, allocation);
+    EXPECT_TRUE(allocation->is32BitAllocation());
+    EXPECT_NE(MemoryPool::LocalMemory, allocation->getMemoryPool());
+    EXPECT_FALSE(memoryManager.allocationInDevicePoolCreated);
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
 TEST(MemoryManagerTest, givenOsAgnosticMemoryManagerWhenGetLocalMemoryIsCalledThenSizeOfLocalMemoryIsReturned) {
     MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
     MockMemoryManager memoryManager(false, false, executionEnvironment);
 
-    EXPECT_EQ(0 * GB, memoryManager.getLocalMemorySize(0u, 0xF));
+    EXPECT_EQ(AubHelper::getMemBankSize(executionEnvironment.rootDeviceEnvironments[0]->getHardwareInfo()), memoryManager.getLocalMemorySize(0u, 0xF));
+}
+
+HWTEST_F(MemoryManagerTests, givenEnabledLocalMemoryWhenAllocatingKernelIsaThenLocalMemoryPoolIsUsed) {
+    auto hwInfo = *defaultHwInfo;
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.EnableLocalMemory.set(true);
+    hwInfo.featureTable.ftrLocalMemory = true;
+
+    MockExecutionEnvironment executionEnvironment(&hwInfo);
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInPreferredPool({mockRootDeviceIndex, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::KERNEL_ISA, mockDeviceBitfield}, nullptr);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation->getMemoryPool());
+    EXPECT_TRUE(memoryManager.allocationInDevicePoolCreated);
+
+    memoryManager.freeGraphicsMemory(allocation);
+
+    allocation = memoryManager.allocateGraphicsMemoryInPreferredPool({mockRootDeviceIndex, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::KERNEL_ISA_INTERNAL, mockDeviceBitfield}, nullptr);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation->getMemoryPool());
+    EXPECT_TRUE(memoryManager.allocationInDevicePoolCreated);
+
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
+HWTEST_F(MemoryManagerTests, givenEnabledLocalMemoryWhenAllocateKernelIsaInDevicePoolThenLocalMemoryPoolIsUsed) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInPreferredPool({mockRootDeviceIndex, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::INTERNAL_HEAP, mockDeviceBitfield}, nullptr);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation->getMemoryPool());
+    EXPECT_TRUE(memoryManager.allocationInDevicePoolCreated);
+
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
+HWTEST_F(MemoryManagerTests, givenEnabledLocalMemoryWhenLinearStreamIsAllocatedInDevicePoolThenLocalMemoryPoolIsUsed) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+
+    auto allocation = memoryManager.allocateGraphicsMemoryInPreferredPool({mockRootDeviceIndex, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::LINEAR_STREAM, mockDeviceBitfield}, nullptr);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation->getMemoryPool());
+    EXPECT_TRUE(memoryManager.allocationInDevicePoolCreated);
+
+    memoryManager.freeGraphicsMemory(allocation);
+}
+
+TEST(StorageInfoTest, whenMemoryBanksAreSetToZeroThenOneHandleIsReturned) {
+    StorageInfo storageInfo;
+    storageInfo.memoryBanks = 0u;
+    EXPECT_EQ(1u, storageInfo.getNumBanks());
+}
+
+TEST(StorageInfoTest, whenMemoryBanksAreNotSetToZeroThenNumberOfEnabledBitsIsReturnedAsNumberOfHandles) {
+    StorageInfo storageInfo;
+    storageInfo.memoryBanks = 0b001;
+    EXPECT_EQ(1u, storageInfo.getNumBanks());
+    storageInfo.memoryBanks = 0b010;
+    EXPECT_EQ(1u, storageInfo.getNumBanks());
+    storageInfo.memoryBanks = 0b100;
+    EXPECT_EQ(1u, storageInfo.getNumBanks());
+    storageInfo.memoryBanks = 0b011;
+    EXPECT_EQ(2u, storageInfo.getNumBanks());
+    storageInfo.memoryBanks = 0b101;
+    EXPECT_EQ(2u, storageInfo.getNumBanks());
+    storageInfo.memoryBanks = 0b110;
+    EXPECT_EQ(2u, storageInfo.getNumBanks());
+    storageInfo.memoryBanks = 0b111;
+    EXPECT_EQ(3u, storageInfo.getNumBanks());
+}
+
+TEST(MemoryManagerTest, givenNotSetUseSystemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenlocalMemoryUsageIsUpdated) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+
+    AllocationProperties allocProperties(mockRootDeviceIndex, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::BUFFER, mockDeviceBitfield);
+    auto allocation = memoryManager.allocateGraphicsMemoryInPreferredPool(allocProperties, nullptr);
+    EXPECT_NE(nullptr, allocation);
+    EXPECT_EQ(MemoryPool::LocalMemory, allocation->getMemoryPool());
+    EXPECT_EQ(allocation->getUnderlyingBufferSize(), memoryManager.localMemoryUsageBankSelector[allocProperties.rootDeviceIndex]->getOccupiedMemorySizeForBank(0));
+
+    memoryManager.freeGraphicsMemory(allocation);
+    EXPECT_EQ(0u, memoryManager.localMemoryUsageBankSelector[allocProperties.rootDeviceIndex]->getOccupiedMemorySizeForBank(0));
+}
+
+TEST(MemoryManagerTest, givenSetUseSystemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenlocalMemoryUsageIsNotUpdated) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    MockMemoryManager memoryManager(false, true, executionEnvironment);
+
+    AllocationProperties allocProperties(mockRootDeviceIndex, MemoryConstants::pageSize, GraphicsAllocation::AllocationType::EXTERNAL_HOST_PTR, mockDeviceBitfield);
+    auto allocation = memoryManager.allocateGraphicsMemoryInPreferredPool(allocProperties, nullptr);
+    EXPECT_NE(nullptr, allocation);
+    EXPECT_EQ(0u, memoryManager.localMemoryUsageBankSelector[allocProperties.rootDeviceIndex]->getOccupiedMemorySizeForBank(0));
+
+    memoryManager.freeGraphicsMemory(allocation);
+    EXPECT_EQ(0u, memoryManager.localMemoryUsageBankSelector[allocProperties.rootDeviceIndex]->getOccupiedMemorySizeForBank(0));
 }

opencl/test/unit_test/memory_manager/memory_manager_allocate_in_device_pool_tests.inl

@@ -1,103 +0,0 @@
-/*
- * Copyright (C) 2018-2021 Intel Corporation
- *
- * SPDX-License-Identifier: MIT
- *
- */
-
-#include "shared/source/execution_environment/execution_environment.h"
-#include "shared/source/memory_manager/os_agnostic_memory_manager.h"
-#include "shared/test/common/mocks/mock_device.h"
-#include "shared/test/common/mocks/mock_execution_environment.h"
-#include "shared/test/common/mocks/ult_device_factory.h"
-
-#include "opencl/source/helpers/memory_properties_helpers.h"
-#include "opencl/source/mem_obj/mem_obj_helper.h"
-#include "opencl/test/unit_test/mocks/mock_memory_manager.h"
-#include "test.h"
-
-#include "gtest/gtest.h"
-
-using namespace NEO;
-
-TEST(MemoryManagerTest, givenSetUseSytemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
-    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
-    MockMemoryManager memoryManager(false, false, executionEnvironment);
-    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
-    AllocationData allocData;
-    allocData.size = MemoryConstants::pageSize;
-    allocData.flags.useSystemMemory = true;
-    allocData.flags.allocateMemory = true;
-
-    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
-    EXPECT_EQ(nullptr, allocation);
-    EXPECT_EQ(MemoryManager::AllocationStatus::RetryInNonDevicePool, status);
-}
-
-TEST(MemoryManagerTest, givenAllowed32BitAndFroce32BitWhenGraphicsAllocationInDevicePoolIsAllocatedThenNullptrIsReturned) {
-    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
-    MockMemoryManager memoryManager(false, false, executionEnvironment);
-    memoryManager.setForce32BitAllocations(true);
-
-    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
-    AllocationData allocData;
-    allocData.size = MemoryConstants::pageSize;
-    allocData.flags.allow32Bit = true;
-    allocData.flags.allocateMemory = true;
-
-    auto allocation = memoryManager.allocateGraphicsMemoryInDevicePool(allocData, status);
-    EXPECT_EQ(nullptr, allocation);
-    EXPECT_EQ(MemoryManager::AllocationStatus::RetryInNonDevicePool, status);
-}
-
-TEST(AllocationFlagsTest, givenAllocateMemoryFlagWhenGetAllocationFlagsIsCalledThenAllocateFlagIsCorrectlySet) {
-    HardwareInfo hwInfo(*defaultHwInfo);
-    UltDeviceFactory deviceFactory{1, 0};
-    auto pDevice = deviceFactory.rootDevices[0];
-    MemoryProperties memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, 0, 0, pDevice);
-    auto allocationProperties = MemoryPropertiesHelper::getAllocationProperties(0, memoryProperties, true, 0, GraphicsAllocation::AllocationType::BUFFER, false, hwInfo, {});
-    EXPECT_TRUE(allocationProperties.flags.allocateMemory);
-
-    auto allocationProperties2 = MemoryPropertiesHelper::getAllocationProperties(0, memoryProperties, false, 0, GraphicsAllocation::AllocationType::BUFFER, false, hwInfo, {});
-    EXPECT_FALSE(allocationProperties2.flags.allocateMemory);
-}
-
-TEST(UncacheableFlagsTest, givenUncachedResourceFlagWhenGetAllocationFlagsIsCalledThenUncacheableFlagIsCorrectlySet) {
-    cl_mem_flags_intel flagsIntel = CL_MEM_LOCALLY_UNCACHED_RESOURCE;
-    UltDeviceFactory deviceFactory{1, 0};
-    auto pDevice = deviceFactory.rootDevices[0];
-    MemoryProperties memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, flagsIntel, 0, pDevice);
-    auto allocationFlags = MemoryPropertiesHelper::getAllocationProperties(
-        0, memoryProperties, false, 0, GraphicsAllocation::AllocationType::BUFFER, false, pDevice->getHardwareInfo(), {});
-    EXPECT_TRUE(allocationFlags.flags.uncacheable);
-
-    flagsIntel = 0;
-    memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, flagsIntel, 0, pDevice);
-    auto allocationFlags2 = MemoryPropertiesHelper::getAllocationProperties(
-        0, memoryProperties, false, 0, GraphicsAllocation::AllocationType::BUFFER, false, pDevice->getHardwareInfo(), {});
-    EXPECT_FALSE(allocationFlags2.flags.uncacheable);
-}
-
-TEST(AllocationFlagsTest, givenReadOnlyResourceFlagWhenGetAllocationFlagsIsCalledThenFlushL3FlagsAreCorrectlySet) {
-    cl_mem_flags flags = CL_MEM_READ_ONLY;
-    UltDeviceFactory deviceFactory{1, 0};
-    auto pDevice = deviceFactory.rootDevices[0];
-    MemoryProperties memoryProperties = MemoryPropertiesHelper::createMemoryProperties(flags, 0, 0, pDevice);
-
-    auto allocationFlags =
-        MemoryPropertiesHelper::getAllocationProperties(
-            0, memoryProperties, true, 0, GraphicsAllocation::AllocationType::BUFFER, false, pDevice->getHardwareInfo(), {});
-    EXPECT_FALSE(allocationFlags.flags.flushL3RequiredForRead);
-    EXPECT_FALSE(allocationFlags.flags.flushL3RequiredForWrite);
-
-    memoryProperties = MemoryPropertiesHelper::createMemoryProperties(0, 0, 0, pDevice);
-    auto allocationFlags2 = MemoryPropertiesHelper::getAllocationProperties(
-        0, memoryProperties, true, 0, GraphicsAllocation::AllocationType::BUFFER, false, pDevice->getHardwareInfo(), {});
-    EXPECT_TRUE(allocationFlags2.flags.flushL3RequiredForRead);
-    EXPECT_TRUE(allocationFlags2.flags.flushL3RequiredForWrite);
-}
-
-TEST(StorageInfoTest, whenStorageInfoIsCreatedWithDefaultConstructorThenReturnsOneHandle) {
-    StorageInfo storageInfo;
-    EXPECT_EQ(1u, storageInfo.getNumBanks());
-}

shared/source/memory_manager/CMakeLists.txt

@@ -44,7 +44,6 @@ set(NEO_CORE_MEMORY_MANAGER
     ${CMAKE_CURRENT_SOURCE_DIR}/multi_graphics_allocation.h
     ${CMAKE_CURRENT_SOURCE_DIR}/os_agnostic_memory_manager.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/os_agnostic_memory_manager.h
-    ${CMAKE_CURRENT_SOURCE_DIR}${BRANCH_DIR_SUFFIX}/os_agnostic_memory_manager_allocate_in_device_pool.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/physical_address_allocator.h
     ${CMAKE_CURRENT_SOURCE_DIR}/residency.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/residency.h

shared/source/memory_manager/os_agnostic_memory_manager.cpp

@@ -8,6 +8,7 @@
 #include "shared/source/memory_manager/os_agnostic_memory_manager.h"
 
 #include "shared/source/aub/aub_center.h"
+#include "shared/source/aub/aub_helper.h"
 #include "shared/source/execution_environment/execution_environment.h"
 #include "shared/source/execution_environment/root_device_environment.h"
 #include "shared/source/gmm_helper/gmm.h"
@@ -416,4 +417,94 @@ void OsAgnosticMemoryManager::freeGpuAddress(AddressRange addressRange, uint32_t
     gfxPartition->freeGpuAddressRange(graphicsAddress, addressRange.size);
 }
 
+GraphicsAllocation *OsAgnosticMemoryManager::allocateGraphicsMemoryInDevicePool(const AllocationData &allocationData, AllocationStatus &status) {
+    MemoryAllocation *allocation = nullptr;
+    status = AllocationStatus::RetryInNonDevicePool;
+    auto numHandles = allocationData.storageInfo.getNumBanks();
+
+    if (!this->localMemorySupported[allocationData.rootDeviceIndex]) {
+        return nullptr;
+    }
+
+    if (allocationData.flags.useSystemMemory || (allocationData.flags.allow32Bit && this->force32bitAllocations)) {
+        return nullptr;
+    }
+    bool use32Allocator = heapAssigner.use32BitHeap(allocationData.type);
+    if (use32Allocator) {
+        auto adjustedAllocationData(allocationData);
+        adjustedAllocationData.size = alignUp(allocationData.size, MemoryConstants::pageSize64k);
+        adjustedAllocationData.alignment = MemoryConstants::pageSize64k;
+        allocation = static_cast<MemoryAllocation *>(allocate32BitGraphicsMemoryImpl(adjustedAllocationData, true));
+    } else if (allocationData.type == GraphicsAllocation::AllocationType::SVM_GPU) {
+        auto storage = allocateSystemMemory(allocationData.size, allocationData.alignment);
+        allocation = new MemoryAllocation(allocationData.rootDeviceIndex, numHandles, allocationData.type, storage, storage, reinterpret_cast<uint64_t>(allocationData.hostPtr),
+                                          allocationData.size, counter, MemoryPool::LocalMemory, false, allocationData.flags.flushL3, maxOsContextCount);
+        counter++;
+    } else {
+        std::unique_ptr<Gmm> gmm;
+        size_t sizeAligned64k = 0;
+        if (allocationData.type == GraphicsAllocation::AllocationType::IMAGE ||
+            allocationData.type == GraphicsAllocation::AllocationType::SHARED_RESOURCE_COPY) {
+            allocationData.imgInfo->useLocalMemory = true;
+            gmm = std::make_unique<Gmm>(executionEnvironment.rootDeviceEnvironments[allocationData.rootDeviceIndex]->getGmmClientContext(), *allocationData.imgInfo, allocationData.storageInfo);
+            sizeAligned64k = alignUp(allocationData.imgInfo->size, MemoryConstants::pageSize64k);
+        } else {
+            sizeAligned64k = alignUp(allocationData.size, MemoryConstants::pageSize64k);
+            if (DebugManager.flags.RenderCompressedBuffersEnabled.get() &&
+                allocationData.flags.preferRenderCompressed) {
+                gmm = std::make_unique<Gmm>(executionEnvironment.rootDeviceEnvironments[allocationData.rootDeviceIndex]->getGmmClientContext(),
+                                            allocationData.hostPtr,
+                                            sizeAligned64k,
+                                            MemoryConstants::pageSize64k,
+                                            allocationData.flags.uncacheable,
+                                            true,
+                                            allocationData.flags.useSystemMemory,
+                                            allocationData.storageInfo);
+            }
+        }
+
+        auto gfxPartition = getGfxPartition(allocationData.rootDeviceIndex);
+        auto heapIndex = HeapIndex::HEAP_STANDARD64KB;
+
+        if ((gfxPartition->getHeapLimit(HeapIndex::HEAP_EXTENDED) > 0) && !allocationData.flags.resource48Bit) {
+            heapIndex = HeapIndex::HEAP_EXTENDED;
+        }
+
+        auto systemMemory = allocateSystemMemory(sizeAligned64k, MemoryConstants::pageSize64k);
+        if (allocationData.type == GraphicsAllocation::AllocationType::PREEMPTION) {
+            memset(systemMemory, 0, sizeAligned64k);
+        }
+        auto sizeOfHeapChunk = sizeAligned64k;
+        auto gpuAddress = GmmHelper::canonize(gfxPartition->heapAllocate(heapIndex, sizeOfHeapChunk));
+        allocation = new MemoryAllocation(allocationData.rootDeviceIndex, numHandles, allocationData.type, systemMemory, systemMemory,
+                                          gpuAddress, sizeAligned64k, counter,
+                                          MemoryPool::LocalMemory, false, allocationData.flags.flushL3, maxOsContextCount);
+        counter++;
+        allocation->setDefaultGmm(gmm.release());
+        allocation->sizeToFree = sizeOfHeapChunk;
+    }
+
+    if (allocation) {
+        allocation->overrideMemoryPool(MemoryPool::LocalMemory);
+        allocation->storageInfo = allocationData.storageInfo;
+        status = AllocationStatus::Success;
+    } else {
+        status = AllocationStatus::Error;
+    }
+
+    return allocation;
+}
+
+uint64_t OsAgnosticMemoryManager::getLocalMemorySize(uint32_t rootDeviceIndex, uint32_t deviceBitfield) {
+    return AubHelper::getMemBankSize(executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->getHardwareInfo());
+}
+
+void MemoryAllocation::overrideMemoryPool(MemoryPool::Type pool) {
+    if (DebugManager.flags.AUBDumpForceAllToLocalMemory.get()) {
+        this->memoryPool = MemoryPool::LocalMemory;
+        return;
+    }
+    this->memoryPool = pool;
+}
+
 } // namespace NEO

shared/source/memory_manager/os_agnostic_memory_manager_allocate_in_device_pool.cpp

@@ -1,51 +0,0 @@
-/*
- * Copyright (C) 2018-2020 Intel Corporation
- *
- * SPDX-License-Identifier: MIT
- *
- */
-
-#include "shared/source/memory_manager/os_agnostic_memory_manager.h"
-
-namespace NEO {
-GraphicsAllocation *OsAgnosticMemoryManager::allocateGraphicsMemoryInDevicePool(const AllocationData &allocationData, AllocationStatus &status) {
-    status = AllocationStatus::Error;
-    switch (allocationData.type) {
-    case GraphicsAllocation::AllocationType::IMAGE:
-    case GraphicsAllocation::AllocationType::SHARED_RESOURCE_COPY:
-        break;
-    default:
-        if (!allocationData.flags.useSystemMemory && !(allocationData.flags.allow32Bit && this->force32bitAllocations)) {
-            GraphicsAllocation *allocation = nullptr;
-            if (allocationData.type == GraphicsAllocation::AllocationType::SVM_GPU) {
-                void *cpuAllocation = allocateSystemMemory(allocationData.size, allocationData.alignment);
-                if (!cpuAllocation) {
-                    return nullptr;
-                }
-                uint64_t gpuAddress = reinterpret_cast<uint64_t>(allocationData.hostPtr);
-                allocation = new MemoryAllocation(allocationData.rootDeviceIndex, allocationData.type, cpuAllocation,
-                                                  cpuAllocation, gpuAddress, allocationData.size, counter++, MemoryPool::LocalMemory, false, false, maxOsContextCount);
-            } else {
-                allocation = allocateGraphicsMemory(allocationData);
-            }
-
-            if (allocation) {
-                allocation->storageInfo = allocationData.storageInfo;
-                allocation->setFlushL3Required(allocationData.flags.flushL3);
-                status = AllocationStatus::Success;
-            }
-            return allocation;
-        }
-    }
-    status = AllocationStatus::RetryInNonDevicePool;
-    return nullptr;
-}
-
-uint64_t OsAgnosticMemoryManager::getLocalMemorySize(uint32_t rootDeviceIndex, uint32_t deviceBitfield) {
-    return 0 * GB;
-}
-
-void MemoryAllocation::overrideMemoryPool(MemoryPool::Type pool) {
-    this->memoryPool = pool;
-}
-} // namespace NEO