fix: configure ISA Pool params based on productHelper

When is2MBLocalMemAlignmentEnabled returns true,
increase pool size for builtins from 64k to 2MB.

Additionally, set appropriate alignment for kernel ISA heap allocations.
Additionally, configure isaAllocationPageSize based on productHelper.

Related-To: NEO-12287
Signed-off-by: Fabian Zwoliński <[email protected]>

Author: fzwolinski

level_zero/core/source/module/module_imp.cpp

@@ -554,11 +554,10 @@ NEO::GraphicsAllocation *ModuleTranslationUnit::getGlobalVarBufferGA() const {
 ModuleImp::ModuleImp(Device *device, ModuleBuildLog *moduleBuildLog, ModuleType type)
     : device(device), translationUnit(std::make_unique<ModuleTranslationUnit>(device)),
       moduleBuildLog(moduleBuildLog), type(type) {
-    auto &gfxCoreHelper = device->getGfxCoreHelper();
     auto &hwInfo = device->getHwInfo();
-    this->isaAllocationPageSize = gfxCoreHelper.useSystemMemoryPlacementForISA(hwInfo) ? MemoryConstants::pageSize : MemoryConstants::pageSize64k;
     this->productFamily = hwInfo.platform.eProductFamily;
     this->metadataGeneration = std::make_unique<NEO::MetadataGeneration>();
+    this->isaAllocationPageSize = getIsaAllocationPageSize();
 }
 
 ModuleImp::~ModuleImp() {
@@ -1790,4 +1789,19 @@ NEO::GraphicsAllocation *ModuleImp::getKernelsIsaParentAllocation() const {
     return sharedIsaAllocation->getGraphicsAllocation();
 }
 
+size_t ModuleImp::getIsaAllocationPageSize() const {
+    auto &gfxCoreHelper = device->getGfxCoreHelper();
+    auto &hwInfo = device->getHwInfo();
+
+    if (gfxCoreHelper.useSystemMemoryPlacementForISA(hwInfo)) {
+        return MemoryConstants::pageSize;
+    }
+
+    if (device->getProductHelper().is2MBLocalMemAlignmentEnabled()) {
+        return MemoryConstants::pageSize2M;
+    } else {
+        return MemoryConstants::pageSize64k;
+    }
+}
+
 } // namespace L0

level_zero/core/source/module/module_imp.h

@@ -201,6 +201,7 @@ struct ModuleImp : public Module {
     MOCKABLE_VIRTUAL size_t computeKernelIsaAllocationAlignedSizeWithPadding(size_t isaSize, bool lastKernel);
     MOCKABLE_VIRTUAL NEO::GraphicsAllocation *allocateKernelsIsaMemory(size_t size);
     StackVec<NEO::GraphicsAllocation *, 32> getModuleAllocations();
+    size_t getIsaAllocationPageSize() const;
 
     Device *device = nullptr;
     PRODUCT_FAMILY productFamily{};

level_zero/core/test/unit_tests/mocks/mock_module.h

@@ -73,6 +73,7 @@ struct WhiteBox<::L0::Module> : public ::L0::ModuleImp {
     using BaseClass::copyPatchedSegments;
     using BaseClass::device;
     using BaseClass::exportedFunctionsSurface;
+    using BaseClass::getIsaAllocationPageSize;
     using BaseClass::importedSymbolAllocations;
     using BaseClass::isaSegmentsForPatching;
     using BaseClass::isFullyLinked;
@@ -126,6 +127,7 @@ struct MockModule : public L0::ModuleImp {
     using ModuleImp::allocateKernelsIsaMemory;
     using ModuleImp::computeKernelIsaAllocationAlignedSizeWithPadding;
     using ModuleImp::debugModuleHandle;
+    using ModuleImp::getIsaAllocationPageSize;
     using ModuleImp::getModuleAllocations;
     using ModuleImp::initializeKernelImmutableData;
     using ModuleImp::isaAllocationPageSize;

level_zero/core/test/unit_tests/sources/module/test_module.cpp

@@ -182,7 +182,6 @@ HWTEST_F(ModuleTest, givenUserModuleWhenCreatedThenCorrectAllocationTypeIsUsedFo
 
 template <bool localMemEnabled>
 struct ModuleKernelIsaAllocationsFixture : public ModuleFixture {
-    static constexpr size_t isaAllocationPageSize = (localMemEnabled ? MemoryConstants::pageSize64k : MemoryConstants::pageSize);
     using Module = WhiteBox<::L0::Module>;
 
     void setUp() {
@@ -217,7 +216,7 @@ struct ModuleKernelIsaAllocationsFixture : public ModuleFixture {
     void givenSeparateIsaMemoryRegionPerKernelWhenGraphicsAllocationFailsThenProperErrorReturned() {
         mockModule->allocateKernelsIsaMemoryCallBase = false;
         mockModule->computeKernelIsaAllocationAlignedSizeWithPaddingCallBase = false;
-        mockModule->computeKernelIsaAllocationAlignedSizeWithPaddingResult = isaAllocationPageSize;
+        mockModule->computeKernelIsaAllocationAlignedSizeWithPaddingResult = this->mockModule->getIsaAllocationPageSize();
 
         auto result = module->initialize(&this->moduleDesc, device->getNEODevice());
         EXPECT_EQ(result, ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY);
@@ -806,6 +805,35 @@ HWTEST_F(ModuleTest, whenMultipleModulesCreatedThenModulesShareIsaAllocation) {
     }
 };
 
+TEST_F(ModuleTest, GivenLocalMemoryEnabledOrDisabledAnd2MBAlignmentEnabledOrDisabledWhenGettingIsaAllocationPageSizeThenCorrectValueIsReturned) {
+    DebugManagerStateRestore restorer;
+
+    auto mockProductHelper = new MockProductHelper;
+    device->getNEODevice()->getRootDeviceEnvironmentRef().productHelper.reset(mockProductHelper);
+
+    MockModule mockModule{device, nullptr, ModuleType::user};
+    EXPECT_EQ(mockModule.getIsaAllocationPageSize(), mockModule.isaAllocationPageSize);
+
+    {
+        debugManager.flags.EnableLocalMemory.set(0);
+        mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;
+
+        EXPECT_EQ(MemoryConstants::pageSize, mockModule.getIsaAllocationPageSize());
+    }
+    {
+        debugManager.flags.EnableLocalMemory.set(1);
+        mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;
+
+        EXPECT_EQ(MemoryConstants::pageSize2M, mockModule.getIsaAllocationPageSize());
+    }
+    {
+        debugManager.flags.EnableLocalMemory.set(1);
+        mockProductHelper->is2MBLocalMemAlignmentEnabledResult = false;
+
+        EXPECT_EQ(MemoryConstants::pageSize64k, mockModule.getIsaAllocationPageSize());
+    }
+}
+
 template <typename T1, typename T2>
 struct ModuleSpecConstantsFixture : public DeviceFixture {
     void setUp() {
@@ -4347,7 +4375,6 @@ TEST_F(ModuleTest, whenContainsStatefulAccessIsCalledThenResultIsCorrect) {
 
 template <bool localMemEnabled>
 struct ModuleIsaAllocationsFixture : public DeviceFixture {
-    static constexpr size_t isaAllocationPageSize = (localMemEnabled ? MemoryConstants::pageSize64k : MemoryConstants::pageSize);
     static constexpr NEO::MemoryPool isaAllocationMemoryPool = (localMemEnabled ? NEO::MemoryPool::localMemory : NEO::MemoryPool::system4KBPagesWith32BitGpuAddressing);
 
     void setUp() {
@@ -4363,6 +4390,7 @@ struct ModuleIsaAllocationsFixture : public DeviceFixture {
         this->mockMemoryManager->localMemorySupported[this->neoDevice->getRootDeviceIndex()] = true;
         this->mockModule.reset(new MockModule{this->device, nullptr, ModuleType::user});
         this->mockModule->translationUnit.reset(new MockModuleTranslationUnit{this->device});
+        this->isaAllocationPageSize = this->mockModule->getIsaAllocationPageSize();
     }
 
     void tearDown() {
@@ -4418,8 +4446,8 @@ struct ModuleIsaAllocationsFixture : public DeviceFixture {
         EXPECT_EQ(kernelImmData[1]->getIsaOffsetInParentAllocation(), 0lu);
         EXPECT_EQ(kernelImmData[1]->getIsaSubAllocationSize(), 0lu);
         if constexpr (localMemEnabled) {
-            EXPECT_EQ(isaAllocationPageSize, kernelImmData[0]->getIsaSize());
-            EXPECT_EQ(isaAllocationPageSize, kernelImmData[1]->getIsaSize());
+            EXPECT_EQ(alignUp<size_t>(maxAllocationSizeInPage, MemoryConstants::pageSize64k), kernelImmData[0]->getIsaSize());
+            EXPECT_EQ(alignUp<size_t>(tinyAllocationSize, MemoryConstants::pageSize64k), kernelImmData[1]->getIsaSize());
         } else {
             EXPECT_EQ(this->computeKernelIsaAllocationSizeWithPadding(maxAllocationSizeInPage), kernelImmData[0]->getIsaSize());
             EXPECT_EQ(this->computeKernelIsaAllocationSizeWithPadding(tinyAllocationSize), kernelImmData[1]->getIsaSize());
@@ -4466,6 +4494,7 @@ struct ModuleIsaAllocationsFixture : public DeviceFixture {
 
     size_t isaPadding;
     size_t kernelStartPointerAlignment;
+    size_t isaAllocationPageSize;
     NEO::Device *neoDevice = nullptr;
     MockMemoryManager *mockMemoryManager = nullptr;
     std::unique_ptr<MockModule> mockModule = nullptr;

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -2046,11 +2046,10 @@ void fillGmmsInAllocation(GmmHelper *gmmHelper, DrmAllocation *allocation) {
     }
 }
 
-inline uint64_t getCanonizedHeapAllocationAddress(HeapIndex heap, GmmHelper *gmmHelper, GfxPartition *gfxPartition, size_t &sizeAllocated, bool packed) {
-    size_t alignment = 0;
-
-    if (debugManager.flags.ExperimentalEnableCustomLocalMemoryAlignment.get() != -1) {
-        alignment = static_cast<size_t>(debugManager.flags.ExperimentalEnableCustomLocalMemoryAlignment.get());
+inline uint64_t getCanonizedHeapAllocationAddress(HeapIndex heap, GmmHelper *gmmHelper, GfxPartition *gfxPartition, size_t &sizeAllocated, size_t alignment, bool packed) {
+    if (const size_t customAlignment = static_cast<size_t>(debugManager.flags.ExperimentalEnableCustomLocalMemoryAlignment.get());
+        customAlignment > 0) {
+        alignment = customAlignment;
     }
     auto address = gfxPartition->heapAllocateWithCustomAlignment(heap, sizeAllocated, alignment);
     return gmmHelper->canonize(address);
@@ -2070,10 +2069,15 @@ AllocationStatus getGpuAddress(const AlignmentSelector &alignmentSelector, HeapA
     case AllocationType::kernelIsa:
     case AllocationType::kernelIsaInternal:
     case AllocationType::internalHeap:
-    case AllocationType::debugModuleArea:
+    case AllocationType::debugModuleArea: {
+        size_t alignment = 0;
+        if (gmmHelper->getRootDeviceEnvironment().getHelper<ProductHelper>().is2MBLocalMemAlignmentEnabled()) {
+            alignment = MemoryConstants::pageSize2M;
+        }
         gpuAddress = getCanonizedHeapAllocationAddress(heapAssigner.get32BitHeapIndex(allocType, true, hwInfo, allocationData.flags.use32BitFrontWindow),
-                                                       gmmHelper, gfxPartition, sizeAllocated, false);
+                                                       gmmHelper, gfxPartition, sizeAllocated, alignment, false);
         break;
+    }
     case AllocationType::writeCombined:
         sizeAllocated = 0;
         break;

shared/source/utilities/isa_pool_allocator.cpp

@@ -8,6 +8,7 @@
 #include "shared/source/utilities/isa_pool_allocator.h"
 
 #include "shared/source/device/device.h"
+#include "shared/source/helpers/aligned_memory.h"
 #include "shared/source/memory_manager/allocation_properties.h"
 #include "shared/source/memory_manager/memory_manager.h"
 #include "shared/source/utilities/buffer_pool_allocator.inl"
@@ -16,6 +17,9 @@ namespace NEO {
 
 ISAPool::ISAPool(Device *device, bool isBuiltin, size_t storageSize)
     : BaseType(device->getMemoryManager(), nullptr), device(device), isBuiltin(isBuiltin) {
+    DEBUG_BREAK_IF(device->getProductHelper().is2MBLocalMemAlignmentEnabled() &&
+                   !isAligned(storageSize, MemoryConstants::pageSize2M));
+
     auto allocationType = isBuiltin ? NEO::AllocationType::kernelIsaInternal : NEO::AllocationType::kernelIsa;
     AllocationProperties allocProperties = {device->getRootDeviceIndex(),
                                             storageSize,
@@ -58,6 +62,7 @@ const StackVec<NEO::GraphicsAllocation *, 1> &ISAPool::getAllocationsVector() {
 }
 
 ISAPoolAllocator::ISAPoolAllocator(Device *device) : device(device) {
+    initAllocParams();
 }
 
 /**
@@ -79,7 +84,7 @@ SharedIsaAllocation *ISAPoolAllocator::requestGraphicsAllocationForIsa(bool isBu
     auto maxAllocationSize = getAllocationSize(isBuiltin);
 
     if (sizeWithPadding > maxAllocationSize) {
-        addNewBufferPool(ISAPool(device, isBuiltin, sizeWithPadding));
+        addNewBufferPool(ISAPool(device, isBuiltin, alignToPoolSize(sizeWithPadding)));
     }
 
     auto sharedIsaAllocation = tryAllocateISA(isBuiltin, sizeWithPadding);
@@ -94,7 +99,7 @@ SharedIsaAllocation *ISAPoolAllocator::requestGraphicsAllocationForIsa(bool isBu
         return sharedIsaAllocation;
     }
 
-    addNewBufferPool(ISAPool(device, isBuiltin, getAllocationSize(isBuiltin)));
+    addNewBufferPool(ISAPool(device, isBuiltin, alignToPoolSize(getAllocationSize(isBuiltin))));
     return tryAllocateISA(isBuiltin, sizeWithPadding);
 }
 
@@ -133,4 +138,19 @@ SharedIsaAllocation *ISAPoolAllocator::tryAllocateISA(bool isBuiltin, size_t siz
     return nullptr;
 }
 
+void ISAPoolAllocator::initAllocParams() {
+    if (device->getProductHelper().is2MBLocalMemAlignmentEnabled()) {
+        userAllocationSize = MemoryConstants::pageSize2M * 2;
+        builtinAllocationSize = MemoryConstants::pageSize2M;
+        poolAlignment = MemoryConstants::pageSize2M;
+    } else {
+        userAllocationSize = MemoryConstants::pageSize2M * 2;
+        builtinAllocationSize = MemoryConstants::pageSize64k;
+    }
+}
+
+size_t ISAPoolAllocator::alignToPoolSize(size_t size) const {
+    return alignUp(size, poolAlignment);
+}
+
 } // namespace NEO

shared/source/utilities/isa_pool_allocator.h

@@ -48,15 +48,18 @@ class ISAPoolAllocator : public AbstractBuffersAllocator<ISAPool, GraphicsAlloca
     void freeSharedIsaAllocation(SharedIsaAllocation *sharedIsaAllocation);
 
   private:
+    void initAllocParams();
     SharedIsaAllocation *tryAllocateISA(bool isBuiltin, size_t size);
 
     size_t getAllocationSize(bool isBuiltin) const {
-        return isBuiltin ? buitinAllocationSize : userAllocationSize;
+        return isBuiltin ? builtinAllocationSize : userAllocationSize;
     }
+    size_t alignToPoolSize(size_t size) const;
 
     Device *device;
     size_t userAllocationSize = MemoryConstants::pageSize2M * 2;
-    size_t buitinAllocationSize = MemoryConstants::pageSize64k;
+    size_t builtinAllocationSize = MemoryConstants::pageSize64k;
+    size_t poolAlignment = 1u;
     std::mutex allocatorMtx;
 };
 

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

@@ -7961,6 +7961,39 @@ TEST_F(DrmMemoryManagerLocalMemoryAlignmentTest, givenEnabled2MBSizeAlignmentWhe
     memoryManager->freeGraphicsMemory(allocation);
 }
 
+TEST_F(DrmMemoryManagerLocalMemoryAlignmentTest, givenEnabled2MBSizeAlignmentWhenAllocatingLargeKernelIsaThenAllocationIsAlignedTo2MB) {
+    auto mockProductHelper = new MockProductHelper;
+    executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->productHelper.reset(mockProductHelper);
+    mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;
+
+    ASSERT_TRUE(executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->productHelper->is2MBLocalMemAlignmentEnabled());
+
+    MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
+    AllocationData allocData;
+    allocData.allFlags = 0;
+    allocData.size = MemoryConstants::pageSize2M + 1;
+    allocData.flags.allocateMemory = true;
+    allocData.type = AllocationType::kernelIsa;
+    allocData.rootDeviceIndex = rootDeviceIndex;
+
+    auto memoryManager = createMemoryManager();
+    auto allocation = memoryManager->allocateGraphicsMemoryInDevicePool(allocData, status);
+    ASSERT_NE(nullptr, allocation);
+    EXPECT_EQ(MemoryManager::AllocationStatus::Success, status);
+    EXPECT_EQ(MemoryPool::localMemory, allocation->getMemoryPool());
+    EXPECT_TRUE(isAligned(allocation->getGpuAddress(), MemoryConstants::pageSize2M));
+    EXPECT_TRUE(isAligned(allocation->getUnderlyingBufferSize(), MemoryConstants::pageSize2M));
+
+    auto drmAllocation = static_cast<DrmAllocation *>(allocation);
+    auto bo = drmAllocation->getBO();
+    EXPECT_NE(nullptr, bo);
+    EXPECT_EQ(allocation->getGpuAddress(), bo->peekAddress());
+    EXPECT_TRUE(isAligned(bo->peekAddress(), MemoryConstants::pageSize2M));
+    EXPECT_TRUE(isAligned(bo->peekSize(), MemoryConstants::pageSize2M));
+
+    memoryManager->freeGraphicsMemory(allocation);
+}
+
 TEST_F(DrmMemoryManagerLocalMemoryAlignmentTest, Given2MBLocalMemAlignmentEnabledWhenAllocating2MBPageTypeInDevicePoolThenAllocationIs2MBAligned) {
     auto mockProductHelper = new MockProductHelper;
     executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->productHelper.reset(mockProductHelper);

shared/test/unit_test/utilities/isa_pool_allocator_tests.cpp

@@ -5,14 +5,29 @@
  *
  */
 
+#include "shared/source/helpers/aligned_memory.h"
+#include "shared/source/utilities/heap_allocator.h"
 #include "shared/test/common/fixtures/device_fixture.h"
 #include "shared/test/common/mocks/mock_device.h"
+#include "shared/test/common/mocks/mock_product_helper.h"
 #include "shared/test/common/test_macros/test.h"
 
 #include "gtest/gtest.h"
 
 using namespace NEO;
 
+class MockISAPool : public ISAPool {
+  public:
+    using ISAPool::chunkAllocator;
+};
+
+class MockISAPoolAllocator : public ISAPoolAllocator {
+  public:
+    using ISAPoolAllocator::bufferPools;
+
+    MockISAPoolAllocator(Device *device) : ISAPoolAllocator(device) {}
+};
+
 using IsaPoolAllocatorTest = Test<DeviceFixture>;
 
 void verifySharedIsaAllocation(const SharedIsaAllocation *sharedAllocation, size_t expectedOffset, size_t expectedSize) {
@@ -99,3 +114,24 @@ TEST_F(IsaPoolAllocatorTest, givenIsaPoolAllocatorWhenRequestForLargeAllocationT
     verifySharedIsaAllocation(allocation, 0, requestAllocationSize);
     isaAllocator.freeSharedIsaAllocation(allocation);
 }
+
+TEST_F(IsaPoolAllocatorTest, Given2MBLocalMemAlignmentEnabledWhenAllocatingIsaThenISAPoolSizeIsAlignedTo2MB) {
+    auto mockProductHelper = new MockProductHelper;
+    pDevice->getRootDeviceEnvironmentRef().productHelper.reset(mockProductHelper);
+    mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;
+
+    MockISAPoolAllocator mockIsaAllocator(pDevice);
+
+    constexpr size_t requestAllocationSize = MemoryConstants::pageSize + 1;
+    auto allocation = mockIsaAllocator.requestGraphicsAllocationForIsa(true, requestAllocationSize);
+
+    EXPECT_NE(nullptr, allocation);
+    ASSERT_GE(mockIsaAllocator.bufferPools.size(), 1u);
+
+    auto *bufferPool = static_cast<MockISAPool *>(&mockIsaAllocator.bufferPools[0]);
+    auto chunkAllocatorSize = bufferPool->chunkAllocator->getLeftSize() + bufferPool->chunkAllocator->getUsedSize();
+
+    EXPECT_TRUE(isAligned(chunkAllocatorSize, MemoryConstants::pageSize2M));
+
+    mockIsaAllocator.freeSharedIsaAllocation(allocation);
+}