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 <fabian.zwolinski@intel.com>

level_zero/core/source/module/module_imp.cpp

@@ -503,10 +503,9 @@ void ModuleTranslationUnit::processDebugData() {
 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->isaAllocationPageSize = getIsaAllocationPageSize();
 }
 
 ModuleImp::~ModuleImp() {
@@ -1714,4 +1713,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

@@ -183,6 +183,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

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2024 Intel Corporation
+ * Copyright (C) 2020-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -58,6 +58,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;
@@ -111,6 +112,7 @@ struct MockModule : public L0::ModuleImp {
     using ModuleImp::allocateKernelsIsaMemory;
     using ModuleImp::computeKernelIsaAllocationAlignedSizeWithPadding;
     using ModuleImp::debugModuleHandle;
+    using ModuleImp::getIsaAllocationPageSize;
     using ModuleImp::getModuleAllocations;
     using ModuleImp::initializeKernelImmutableDatas;
     using ModuleImp::isaAllocationPageSize;

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

@@ -34,6 +34,7 @@
 #include "shared/test/common/mocks/mock_l0_debugger.h"
 #include "shared/test/common/mocks/mock_memory_operations_handler.h"
 #include "shared/test/common/mocks/mock_modules_zebin.h"
+#include "shared/test/common/mocks/mock_product_helper.h"
 #include "shared/test/common/test_macros/hw_test.h"
 
 #include "level_zero/core/source/kernel/kernel_imp.h"
@@ -174,7 +175,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() {
@@ -209,7 +209,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);
@@ -793,6 +793,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() {
@@ -4048,7 +4077,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() {
@@ -4064,6 +4092,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() {
@@ -4119,8 +4148,8 @@ struct ModuleIsaAllocationsFixture : public DeviceFixture {
         EXPECT_EQ(kernelImmDatas[1]->getIsaOffsetInParentAllocation(), 0lu);
         EXPECT_EQ(kernelImmDatas[1]->getIsaSubAllocationSize(), 0lu);
         if constexpr (localMemEnabled) {
-            EXPECT_EQ(isaAllocationPageSize, kernelImmDatas[0]->getIsaSize());
-            EXPECT_EQ(isaAllocationPageSize, kernelImmDatas[1]->getIsaSize());
+            EXPECT_EQ(alignUp<size_t>(maxAllocationSizeInPage, MemoryConstants::pageSize64k), kernelImmDatas[0]->getIsaSize());
+            EXPECT_EQ(alignUp<size_t>(tinyAllocationSize, MemoryConstants::pageSize64k), kernelImmDatas[1]->getIsaSize());
         } else {
             EXPECT_EQ(this->computeKernelIsaAllocationSizeWithPadding(maxAllocationSizeInPage), kernelImmDatas[0]->getIsaSize());
             EXPECT_EQ(this->computeKernelIsaAllocationSizeWithPadding(tinyAllocationSize), kernelImmDatas[1]->getIsaSize());
@@ -4167,6 +4196,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;

opencl/test/unit_test/mem_obj/buffer_pool_alloc_tests.cpp

@@ -200,9 +200,9 @@ TEST_F(AggregatedSmallBuffersEnabledTest, givenAggregatedSmallBuffersEnabledWhen
     EXPECT_TRUE(poolAllocator->isAggregatedSmallBuffersEnabled(context.get()));
     EXPECT_EQ(1u, poolAllocator->bufferPools.size());
     EXPECT_NE(nullptr, poolAllocator->bufferPools[0].mainStorage.get());
-    EXPECT_NE(nullptr, mockMemoryManager->lastAllocationProperties);
-    EXPECT_TRUE(mockMemoryManager->lastAllocationProperties->makeDeviceBufferLockable);
-    EXPECT_FALSE(mockMemoryManager->lastAllocationProperties->flags.preferCompressed);
+    EXPECT_NE(nullptr, mockMemoryManager->lastAllocationPropertiesWithPtr);
+    EXPECT_TRUE(mockMemoryManager->lastAllocationPropertiesWithPtr->makeDeviceBufferLockable);
+    EXPECT_FALSE(mockMemoryManager->lastAllocationPropertiesWithPtr->flags.preferCompressed);
 }
 
 TEST_F(AggregatedSmallBuffersEnabledTest, givenAggregatedSmallBuffersEnabledAndSizeLargerThanThresholdWhenBufferCreateCalledThenDoNotUsePool) {

opencl/test/unit_test/program/kernel_info_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -10,6 +10,7 @@
 #include "shared/source/program/kernel_info.h"
 #include "shared/test/common/mocks/mock_execution_environment.h"
 #include "shared/test/common/mocks/mock_graphics_allocation.h"
+#include "shared/test/common/mocks/mock_product_helper.h"
 #include "shared/test/common/mocks/mock_sip.h"
 #include "shared/test/common/mocks/ult_device_factory.h"
 
@@ -84,6 +85,62 @@ TEST(KernelInfoTest, givenKernelInfoWhenCreatingKernelAllocationWithInternalIsaT
     device->getMemoryManager()->checkGpuUsageAndDestroyGraphicsAllocations(allocation);
 }
 
+TEST(KernelInfoTest, Given2MBAlignmentEnabledByProductHelperWhenCreatingKernelAllocationThenAllocationPropertiesAlignmentIsSetTo2M) {
+    DebugManagerStateRestore restorer;
+
+    KernelInfo kernelInfo;
+    auto factory = UltDeviceFactory{1, 0};
+    auto device = factory.rootDevices[0];
+    const size_t heapSize = 0x40;
+    char heap[heapSize];
+    kernelInfo.heapInfo.kernelHeapSize = heapSize;
+    kernelInfo.heapInfo.pKernelHeap = &heap;
+
+    auto mockProductHelper = new MockProductHelper;
+    device->getRootDeviceEnvironmentRef().productHelper.reset(mockProductHelper);
+    mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;
+    debugManager.flags.AlignLocalMemoryVaTo2MB.set(0);
+
+    auto mockMemoryManager = static_cast<MockMemoryManager *>(device->getMemoryManager());
+    mockMemoryManager->shouldStoreLastAllocationProperties = true;
+
+    auto retVal = kernelInfo.createKernelAllocation(*device, false);
+    EXPECT_TRUE(retVal);
+
+    ASSERT_NE(nullptr, mockMemoryManager->lastAllocationProperties);
+    EXPECT_EQ(MemoryConstants::pageSize2M, mockMemoryManager->lastAllocationProperties->alignment);
+
+    device->getMemoryManager()->checkGpuUsageAndDestroyGraphicsAllocations(kernelInfo.kernelAllocation);
+}
+
+TEST(KernelInfoTest, Given2MBAlignmentForcedByDebugFlagWhenCreatingKernelAllocationThenAllocationPropertiesAlignmentIsSetTo2M) {
+    DebugManagerStateRestore restorer;
+
+    KernelInfo kernelInfo;
+    auto factory = UltDeviceFactory{1, 0};
+    auto device = factory.rootDevices[0];
+    const size_t heapSize = 0x40;
+    char heap[heapSize];
+    kernelInfo.heapInfo.kernelHeapSize = heapSize;
+    kernelInfo.heapInfo.pKernelHeap = &heap;
+
+    auto mockProductHelper = new MockProductHelper;
+    device->getRootDeviceEnvironmentRef().productHelper.reset(mockProductHelper);
+    mockProductHelper->is2MBLocalMemAlignmentEnabledResult = false;
+    debugManager.flags.AlignLocalMemoryVaTo2MB.set(1);
+
+    auto mockMemoryManager = static_cast<MockMemoryManager *>(device->getMemoryManager());
+    mockMemoryManager->shouldStoreLastAllocationProperties = true;
+
+    auto retVal = kernelInfo.createKernelAllocation(*device, false);
+    EXPECT_TRUE(retVal);
+
+    ASSERT_NE(nullptr, mockMemoryManager->lastAllocationProperties);
+    EXPECT_EQ(MemoryConstants::pageSize2M, mockMemoryManager->lastAllocationProperties->alignment);
+
+    device->getMemoryManager()->checkGpuUsageAndDestroyGraphicsAllocations(kernelInfo.kernelAllocation);
+}
+
 class MyMemoryManager : public OsAgnosticMemoryManager {
   public:
     using OsAgnosticMemoryManager::OsAgnosticMemoryManager;

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -1880,11 +1880,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);
@@ -1904,10 +1903,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/program/kernel_info.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2024 Intel Corporation
+ * Copyright (C) 2018-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -62,7 +62,8 @@ bool KernelInfo::createKernelAllocation(const Device &device, bool internalIsa)
 
     AllocationProperties properties = {device.getRootDeviceIndex(), kernelIsaSize, allocType, device.getDeviceBitfield()};
 
-    if (debugManager.flags.AlignLocalMemoryVaTo2MB.get() == 1) {
+    if (device.getProductHelper().is2MBLocalMemAlignmentEnabled() ||
+        debugManager.flags.AlignLocalMemoryVaTo2MB.get() == 1) {
         properties.alignment = MemoryConstants::pageSize2M;
     }
 

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));
+
     this->chunkAllocator.reset(new NEO::HeapAllocator(params.startingOffset, storageSize, MemoryConstants::pageSize, 0u));
 
     auto allocationType = isBuiltin ? NEO::AllocationType::kernelIsaInternal : NEO::AllocationType::kernelIsa;
@@ -55,6 +59,7 @@ const StackVec<NEO::GraphicsAllocation *, 1> &ISAPool::getAllocationsVector() {
 }
 
 ISAPoolAllocator::ISAPoolAllocator(Device *device) : device(device) {
+    initAllocParams();
 }
 
 /**
@@ -76,7 +81,7 @@ SharedIsaAllocation *ISAPoolAllocator::requestGraphicsAllocationForIsa(bool isBu
     auto maxAllocationSize = getAllocationSize(isBuiltin);
 
     if (size > maxAllocationSize) {
-        addNewBufferPool(ISAPool(device, isBuiltin, size));
+        addNewBufferPool(ISAPool(device, isBuiltin, alignToPoolSize(size)));
     }
 
     auto sharedIsaAllocation = tryAllocateISA(isBuiltin, size);
@@ -91,7 +96,7 @@ SharedIsaAllocation *ISAPoolAllocator::requestGraphicsAllocationForIsa(bool isBu
         return sharedIsaAllocation;
     }
 
-    addNewBufferPool(ISAPool(device, isBuiltin, getAllocationSize(isBuiltin)));
+    addNewBufferPool(ISAPool(device, isBuiltin, alignToPoolSize(getAllocationSize(isBuiltin))));
     return tryAllocateISA(isBuiltin, size);
 }
 
@@ -130,4 +135,19 @@ SharedIsaAllocation *ISAPoolAllocator::tryAllocateISA(bool isBuiltin, size_t siz
     return nullptr;
 }
 
+void ISAPoolAllocator::initAllocParams() {
+    if (device->getProductHelper().is2MBLocalMemAlignmentEnabled()) {
+        userAllocationSize = MemoryConstants::pageSize2M * 2;
+        buitinAllocationSize = MemoryConstants::pageSize2M;
+        poolAlignment = MemoryConstants::pageSize2M;
+    } else {
+        userAllocationSize = MemoryConstants::pageSize2M * 2;
+        buitinAllocationSize = MemoryConstants::pageSize64k;
+    }
+}
+
+size_t ISAPoolAllocator::alignToPoolSize(size_t size) const {
+    return alignUp(size, poolAlignment);
+}
+
 } // namespace NEO

shared/source/utilities/isa_pool_allocator.h

@@ -73,15 +73,17 @@ 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;
     }
+    size_t alignToPoolSize(size_t size) const;
 
     Device *device;
     size_t userAllocationSize = MemoryConstants::pageSize2M * 2;
     size_t buitinAllocationSize = MemoryConstants::pageSize64k;
+    size_t poolAlignment = 1u;
     std::mutex allocatorMtx;
 };
 

shared/test/common/mocks/mock_memory_manager.cpp

@@ -72,6 +72,9 @@ void MockMemoryManager::waitForEnginesCompletion(GraphicsAllocation &graphicsAll
 
 GraphicsAllocation *MockMemoryManager::allocateGraphicsMemoryWithProperties(const AllocationProperties &properties) {
     validateAllocateProperties(properties);
+    if (shouldStoreLastAllocationProperties) {
+        lastAllocationProperties.reset(new AllocationProperties(properties));
+    }
     if (isMockHostMemoryManager) {
         allocateGraphicsMemoryWithPropertiesCount++;
         if (forceFailureInPrimaryAllocation) {
@@ -92,7 +95,7 @@ GraphicsAllocation *MockMemoryManager::allocateGraphicsMemoryWithProperties(cons
 
 GraphicsAllocation *MockMemoryManager::allocateGraphicsMemoryWithProperties(const AllocationProperties &properties, const void *ptr) {
     validateAllocateProperties(properties);
-    lastAllocationProperties.reset(new AllocationProperties(properties));
+    lastAllocationPropertiesWithPtr.reset(new AllocationProperties(properties));
     if (returnFakeAllocation) {
         return new GraphicsAllocation(properties.rootDeviceIndex, 1u /*num gmms*/, properties.allocationType, const_cast<void *>(ptr), dummyAddress, properties.size, 0, MemoryPool::system4KBPages, maxOsContextCount);
     }

shared/test/common/mocks/mock_memory_manager.h

@@ -346,6 +346,7 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
     bool failMapPhysicalToVirtualMemory = false;
     bool returnMockGAFromDevicePool = false;
     bool returnMockGAFromHostPool = false;
+    bool shouldStoreLastAllocationProperties = false;
     std::unique_ptr<MockExecutionEnvironment> mockExecutionEnvironment;
     DeviceBitfield recentlyPassedDeviceBitfield{};
     std::unique_ptr<MultiGraphicsAllocation> waitAllocations = nullptr;
@@ -354,6 +355,7 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
     MemoryManager::AllocationStatus populateOsHandlesResult = MemoryManager::AllocationStatus::Success;
     GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtrResult = nullptr;
     std::unique_ptr<AllocationProperties> lastAllocationProperties = nullptr;
+    std::unique_ptr<AllocationProperties> lastAllocationPropertiesWithPtr = nullptr;
     std::function<void(const AllocationProperties &)> validateAllocateProperties = [](const AllocationProperties &) -> void {};
     AddressRange reserveGpuAddressOnHeapResult = AddressRange{0u, 0u};
 };

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

@@ -7423,6 +7423,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(DrmMemoryManagerWithLocalMemoryAndExplicitExpectationsTest, givenNotSetUseSystemMemoryWhenGraphicsAllocationInDevicePoolIsAllocatedForBufferThenLocalMemoryAllocationIsReturnedFromStandard64KbHeap) {
     MemoryManager::AllocationStatus status = MemoryManager::AllocationStatus::Success;
     AllocationData allocData;

shared/test/unit_test/utilities/isa_pool_allocator_tests.cpp

@@ -1,18 +1,32 @@
 /*
- * Copyright (C) 2023-2024 Intel Corporation
+ * Copyright (C) 2023-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
  */
 
+#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 +113,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);
+}