diff --git a/level_zero/core/source/context/context_imp.cpp b/level_zero/core/source/context/context_imp.cpp
index 70804032e4..740f04bc28 100644
--- a/level_zero/core/source/context/context_imp.cpp
+++ b/level_zero/core/source/context/context_imp.cpp
@@ -994,11 +994,13 @@ NEO::VirtualMemoryReservation *ContextImp::findSupportedVirtualReservation(const
 ze_result_t ContextImp::reserveVirtualMem(const void *pStart,
                                           size_t size,
                                           void **pptr) {
-    if ((getPageSizeRequired(size) != size)) {
+    NEO::HeapIndex heap;
+    size_t pageSize;
+    if ((getPageAlignedSizeRequired(size, &heap, &pageSize) != size)) {
         return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
     }
     NEO::VirtualMemoryReservation *virtualMemoryReservation = new NEO::VirtualMemoryReservation;
-    virtualMemoryReservation->virtualAddressRange = this->driverHandle->getMemoryManager()->reserveGpuAddress(reinterpret_cast<uint64_t>(pStart), size, this->driverHandle->rootDeviceIndices, &virtualMemoryReservation->rootDeviceIndex);
+    virtualMemoryReservation->virtualAddressRange = this->driverHandle->getMemoryManager()->reserveGpuAddressOnHeap(reinterpret_cast<uint64_t>(pStart), size, this->driverHandle->rootDeviceIndices, &virtualMemoryReservation->rootDeviceIndex, heap, pageSize);
     if (virtualMemoryReservation->virtualAddressRange.address == 0) {
         delete virtualMemoryReservation;
         return ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY;
@@ -1037,26 +1039,42 @@ ze_result_t ContextImp::freeVirtualMem(const void *ptr,
     }
 }
 
-size_t ContextImp::getPageSizeRequired(size_t size) {
-    return std::max(Math::prevPowerOfTwo(size), MemoryConstants::pageSize64k);
+size_t ContextImp::getPageAlignedSizeRequired(size_t size, NEO::HeapIndex *heapRequired, size_t *pageSizeRequired) {
+    [[maybe_unused]] NEO::HeapIndex heap;
+    size_t pageSize;
+    pageSize = this->driverHandle->getMemoryManager()->selectAlignmentAndHeap(size, &heap);
+    if (heapRequired) {
+        *heapRequired = heap;
+    }
+    if (pageSizeRequired) {
+        *pageSizeRequired = pageSize;
+    }
+    // Given a size larger than the pageSize, then round the size up to the next pageSize alignment if unaligned
+    if (size > pageSize) {
+        return ((size + pageSize) & (~pageSize));
+    }
+    return pageSize;
 }
 
 ze_result_t ContextImp::queryVirtualMemPageSize(ze_device_handle_t hDevice,
                                                 size_t size,
                                                 size_t *pagesize) {
-    *pagesize = getPageSizeRequired(size);
+    // Retrieve the page size and heap required for this allocation size requested.
+    getPageAlignedSizeRequired(size, nullptr, pagesize);
     return ZE_RESULT_SUCCESS;
 }
 
 ze_result_t ContextImp::createPhysicalMem(ze_device_handle_t hDevice,
                                           ze_physical_mem_desc_t *desc,
                                           ze_physical_mem_handle_t *phPhysicalMemory) {
-    if ((getPageSizeRequired(desc->size) != desc->size)) {
-        return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
-    }
+
     auto device = Device::fromHandle(hDevice);
     auto neoDevice = device->getNEODevice();
 
+    if ((getPageAlignedSizeRequired(desc->size, nullptr, nullptr) != desc->size)) {
+        return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
+    }
+
     NEO::AllocationProperties physicalDeviceMemoryProperties{neoDevice->getRootDeviceIndex(),
                                                              true,
                                                              desc->size,
@@ -1101,9 +1119,6 @@ ze_result_t ContextImp::mapVirtualMem(const void *ptr,
     std::map<void *, NEO::PhysicalMemoryAllocation *>::iterator physicalIt;
     NEO::PhysicalMemoryAllocation *allocationNode = nullptr;
 
-    if ((getPageSizeRequired(size) != size)) {
-        return ZE_RESULT_ERROR_UNSUPPORTED_ALIGNMENT;
-    }
     auto lockPhysical = this->driverHandle->getMemoryManager()->lockPhysicalMemoryAllocationMap();
     physicalIt = this->driverHandle->getMemoryManager()->getPhysicalMemoryAllocationMap().find(static_cast<void *>(hPhysicalMemory));
     if (physicalIt != this->driverHandle->getMemoryManager()->getPhysicalMemoryAllocationMap().end()) {
@@ -1112,6 +1127,10 @@ ze_result_t ContextImp::mapVirtualMem(const void *ptr,
         return ZE_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
+    if ((getPageAlignedSizeRequired(size, nullptr, nullptr) != size)) {
+        return ZE_RESULT_ERROR_UNSUPPORTED_ALIGNMENT;
+    }
+
     NEO::VirtualMemoryReservation *virtualMemoryReservation = nullptr;
     auto lockVirtual = this->driverHandle->getMemoryManager()->lockVirtualMemoryReservationMap();
     virtualMemoryReservation = findSupportedVirtualReservation(ptr, size);
diff --git a/level_zero/core/source/context/context_imp.h b/level_zero/core/source/context/context_imp.h
index ff4f338020..b65428e19e 100644
--- a/level_zero/core/source/context/context_imp.h
+++ b/level_zero/core/source/context/context_imp.h
@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include "shared/source/memory_manager/gfx_partition.h"
 #include "shared/source/memory_manager/memory_manager.h"
 #include "shared/source/utilities/stackvec.h"
 
@@ -174,7 +175,7 @@ struct ContextImp : Context {
   protected:
     void setIPCHandleData(NEO::GraphicsAllocation *graphicsAllocation, uint64_t handle, IpcMemoryData &ipcData, uint64_t ptrAddress, uint8_t type);
     bool isAllocationSuitableForCompression(const StructuresLookupTable &structuresLookupTable, Device &device, size_t allocSize);
-    size_t getPageSizeRequired(size_t size);
+    size_t getPageAlignedSizeRequired(size_t size, NEO::HeapIndex *heapRequired, size_t *pageSizeRequired);
 
     std::map<uint32_t, ze_device_handle_t> devices;
     std::vector<ze_device_handle_t> deviceHandles;
diff --git a/level_zero/core/test/unit_tests/fixtures/memory_ipc_fixture.h b/level_zero/core/test/unit_tests/fixtures/memory_ipc_fixture.h
index 33fa8d2b92..5e8f6d3d9a 100644
--- a/level_zero/core/test/unit_tests/fixtures/memory_ipc_fixture.h
+++ b/level_zero/core/test/unit_tests/fixtures/memory_ipc_fixture.h
@@ -8,6 +8,7 @@
 #pragma once
 
 #include "shared/source/gmm_helper/gmm_helper.h"
+#include "shared/source/memory_manager/gfx_partition.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/mocks/mock_compilers.h"
 #include "shared/test/common/mocks/mock_device.h"
@@ -411,6 +412,13 @@ class MemoryManagerIpcMock : public NEO::MemoryManager {
     AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override {
         return {};
     }
+    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override {
+        return {};
+    }
+    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override {
+        *heap = HeapIndex::HEAP_STANDARD;
+        return MemoryConstants::pageSize64k;
+    }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
@@ -629,6 +637,13 @@ class MemoryManagerIpcImplicitScalingMock : public NEO::MemoryManager {
     AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override {
         return {};
     }
+    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override {
+        return {};
+    }
+    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override {
+        *heap = HeapIndex::HEAP_STANDARD;
+        return MemoryConstants::pageSize64k;
+    }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
diff --git a/level_zero/core/test/unit_tests/sources/context/test_context.cpp b/level_zero/core/test/unit_tests/sources/context/test_context.cpp
index 0d66fb56e5..1dd0126ce2 100644
--- a/level_zero/core/test/unit_tests/sources/context/test_context.cpp
+++ b/level_zero/core/test/unit_tests/sources/context/test_context.cpp
@@ -8,6 +8,7 @@
 #include "shared/source/built_ins/sip.h"
 #include "shared/source/gmm_helper/gmm.h"
 #include "shared/source/helpers/blit_properties.h"
+#include "shared/source/memory_manager/gfx_partition.h"
 #include "shared/test/common/mocks/mock_command_stream_receiver.h"
 #include "shared/test/common/mocks/mock_compilers.h"
 #include "shared/test/common/mocks/mock_cpu_page_fault_manager.h"
@@ -1047,13 +1048,13 @@ TEST_F(ContextTest, whenCallingQueryVirtualMemPageSizeCorrectAlignmentIsReturned
     pagesize = 0u;
     res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
-    EXPECT_EQ(pagesize, MemoryConstants::pageSize2M / 2);
+    EXPECT_EQ(pagesize, MemoryConstants::pageSize64k);
 
     size = MemoryConstants::pageSize2M + 1000;
     pagesize = 0u;
     res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
-    EXPECT_EQ(pagesize, MemoryConstants::pageSize2M);
+    EXPECT_EQ(pagesize, MemoryConstants::pageSize64k);
 
     res = contextImp->destroy();
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
@@ -1350,28 +1351,6 @@ TEST_F(ContextTest, whenCallingVirtualMemoryFreeWithInvalidValuesThenFailuresRet
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
 }
 
-TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithInvalidArgumentsThenFailureReturned) {
-    ze_context_handle_t hContext;
-    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
-
-    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
-
-    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
-
-    void *pStart = 0x0;
-    size_t size = 64u;
-    void *ptr = nullptr;
-    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
-    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);
-    size_t pagesize = 0u;
-    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
-    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
-
-    res = contextImp->destroy();
-    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
-}
-
 class ReserveMemoryManagerMock : public NEO::MemoryManager {
   public:
     ReserveMemoryManagerMock(NEO::ExecutionEnvironment &executionEnvironment) : NEO::MemoryManager(executionEnvironment) {}
@@ -1393,6 +1372,16 @@ class ReserveMemoryManagerMock : public NEO::MemoryManager {
         }
         return AddressRange{requiredStartAddress, size};
     }
+    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override {
+        if (failReserveGpuAddress) {
+            return {};
+        }
+        return AddressRange{requiredStartAddress, size};
+    }
+    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override {
+        *heap = HeapIndex::HEAP_STANDARD;
+        return MemoryConstants::pageSize64k;
+    }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
@@ -1638,6 +1627,35 @@ TEST_F(ContextTest, whenCallingVirtualMemoryReservationWhenOutOfMemoryThenOutOfM
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
 }
 
+TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithInvalidArgumentsThenUnsupportedSizeReturned) {
+    ze_context_handle_t hContext;
+    ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
+
+    ze_result_t res = driverHandle->createContext(&desc, 0u, nullptr, &hContext);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    ContextImp *contextImp = static_cast<ContextImp *>(L0::Context::fromHandle(hContext));
+
+    void *pStart = 0x0;
+    size_t size = 64u;
+    void *ptr = nullptr;
+    size_t pagesize = 0u;
+
+    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
+
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    NEO::MemoryManager *failingReserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
+    auto memoryManager = driverHandle->getMemoryManager();
+    driverHandle->setMemoryManager(failingReserveMemoryManager);
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);
+    driverHandle->setMemoryManager(memoryManager);
+    delete failingReserveMemoryManager;
+
+    res = contextImp->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+}
+
 TEST_F(ContextTest, whenCallingPhysicalMemoryAllocateWhenOutOfMemoryThenOutofMemoryReturned) {
     ze_context_handle_t hContext;
     ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
diff --git a/level_zero/core/test/unit_tests/sources/event/test_event.cpp b/level_zero/core/test/unit_tests/sources/event/test_event.cpp
index 9ba266f4db..4366227557 100644
--- a/level_zero/core/test/unit_tests/sources/event/test_event.cpp
+++ b/level_zero/core/test/unit_tests/sources/event/test_event.cpp
@@ -8,6 +8,7 @@
 #include "shared/source/built_ins/sip.h"
 #include "shared/source/gmm_helper/gmm.h"
 #include "shared/source/helpers/aligned_memory.h"
+#include "shared/source/memory_manager/gfx_partition.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/helpers/variable_backup.h"
 #include "shared/test/common/mocks/mock_compilers.h"
@@ -76,6 +77,13 @@ class MemoryManagerEventPoolFailMock : public NEO::MemoryManager {
     AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override {
         return {};
     }
+    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override {
+        return {};
+    }
+    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override {
+        *heap = HeapIndex::HEAP_STANDARD;
+        return MemoryConstants::pageSize64k;
+    }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
diff --git a/level_zero/core/test/unit_tests/sources/memory/linux/test_memory_linux.cpp b/level_zero/core/test/unit_tests/sources/memory/linux/test_memory_linux.cpp
index ab1546d90a..076bf3bfa8 100644
--- a/level_zero/core/test/unit_tests/sources/memory/linux/test_memory_linux.cpp
+++ b/level_zero/core/test/unit_tests/sources/memory/linux/test_memory_linux.cpp
@@ -6,6 +6,7 @@
  */
 
 #include "shared/source/built_ins/sip.h"
+#include "shared/source/memory_manager/gfx_partition.h"
 #include "shared/source/os_interface/linux/drm_allocation.h"
 #include "shared/source/os_interface/linux/drm_buffer_object.h"
 #include "shared/source/os_interface/linux/drm_gem_close_worker.h"
@@ -74,6 +75,13 @@ class MemoryManagerIpcImplicitScalingObtainFdMock : public NEO::DrmMemoryManager
     AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override {
         return {};
     }
+    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override {
+        return {};
+    }
+    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override {
+        *heap = HeapIndex::HEAP_STANDARD;
+        return MemoryConstants::pageSize64k;
+    }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
@@ -483,6 +491,13 @@ class MemoryManagerIpcObtainFdMock : public NEO::DrmMemoryManager {
     AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override {
         return {};
     }
+    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override {
+        return {};
+    }
+    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override {
+        *heap = HeapIndex::HEAP_STANDARD;
+        return MemoryConstants::pageSize64k;
+    }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
diff --git a/shared/source/debugger/debugger_l0.cpp b/shared/source/debugger/debugger_l0.cpp
index 4f516f5f34..57996728a2 100644
--- a/shared/source/debugger/debugger_l0.cpp
+++ b/shared/source/debugger/debugger_l0.cpp
@@ -14,6 +14,7 @@
 #include "shared/source/gmm_helper/gmm_helper.h"
 #include "shared/source/helpers/constants.h"
 #include "shared/source/memory_manager/allocation_properties.h"
+#include "shared/source/memory_manager/gfx_partition.h"
 #include "shared/source/memory_manager/memory_manager.h"
 #include "shared/source/memory_manager/memory_operations_handler.h"
 #include "shared/source/os_interface/os_context.h"
diff --git a/shared/source/memory_manager/memory_manager.h b/shared/source/memory_manager/memory_manager.h
index 6b32e14383..9d7c214690 100644
--- a/shared/source/memory_manager/memory_manager.h
+++ b/shared/source/memory_manager/memory_manager.h
@@ -233,6 +233,8 @@ class MemoryManager {
     GfxPartition *getGfxPartition(uint32_t rootDeviceIndex) { return gfxPartitions.at(rootDeviceIndex).get(); }
     GmmHelper *getGmmHelper(uint32_t rootDeviceIndex);
     virtual AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) = 0;
+    virtual AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) = 0;
+    virtual size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) = 0;
     virtual void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) = 0;
     static HeapIndex selectInternalHeap(bool useLocalMemory);
     static HeapIndex selectExternalHeap(bool useLocalMemory);
diff --git a/shared/source/memory_manager/os_agnostic_memory_manager.cpp b/shared/source/memory_manager/os_agnostic_memory_manager.cpp
index c1a03ce38c..f5ae4ebf15 100644
--- a/shared/source/memory_manager/os_agnostic_memory_manager.cpp
+++ b/shared/source/memory_manager/os_agnostic_memory_manager.cpp
@@ -543,13 +543,22 @@ MemoryAllocation *OsAgnosticMemoryManager::createMemoryAllocation(AllocationType
     return memoryAllocation;
 }
 
+size_t OsAgnosticMemoryManager::selectAlignmentAndHeap(size_t size, HeapIndex *heap) {
+    *heap = HeapIndex::HEAP_STANDARD;
+    return MemoryConstants::pageSize64k;
+}
+
 AddressRange OsAgnosticMemoryManager::reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) {
+    return reserveGpuAddressOnHeap(requiredStartAddress, size, rootDeviceIndices, reservedOnRootDeviceIndex, HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
+}
+
+AddressRange OsAgnosticMemoryManager::reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) {
     uint64_t gpuVa = 0u;
     *reservedOnRootDeviceIndex = 0;
     for (auto rootDeviceIndex : rootDeviceIndices) {
         auto gfxPartition = getGfxPartition(rootDeviceIndex);
         auto gmmHelper = getGmmHelper(rootDeviceIndex);
-        gpuVa = gmmHelper->canonize(gfxPartition->heapAllocate(HeapIndex::HEAP_STANDARD, size));
+        gpuVa = gmmHelper->canonize(gfxPartition->heapAllocate(heap, size));
         if (gpuVa != 0u) {
             *reservedOnRootDeviceIndex = rootDeviceIndex;
             break;
diff --git a/shared/source/memory_manager/os_agnostic_memory_manager.h b/shared/source/memory_manager/os_agnostic_memory_manager.h
index f9730380f6..15eab71e24 100644
--- a/shared/source/memory_manager/os_agnostic_memory_manager.h
+++ b/shared/source/memory_manager/os_agnostic_memory_manager.h
@@ -46,6 +46,8 @@ class OsAgnosticMemoryManager : public MemoryManager {
     void releaseReservedCpuAddressRange(void *reserved, size_t size, uint32_t rootDeviceIndex) override;
 
     AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override;
+    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override;
+    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override;
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override;
     bool is64kbPagesEnabled(const HardwareInfo *hwInfo);
 
diff --git a/shared/source/os_interface/linux/drm_memory_manager.cpp b/shared/source/os_interface/linux/drm_memory_manager.cpp
index e110c9e0da..b67caa6605 100644
--- a/shared/source/os_interface/linux/drm_memory_manager.cpp
+++ b/shared/source/os_interface/linux/drm_memory_manager.cpp
@@ -1357,11 +1357,47 @@ uint32_t DrmMemoryManager::getRootDeviceIndex(const Drm *drm) {
     return CommonConstants::unspecifiedDeviceIndex;
 }
 
+size_t DrmMemoryManager::selectAlignmentAndHeap(size_t size, HeapIndex *heap) {
+    AlignmentSelector::CandidateAlignment alignmentBase = alignmentSelector.selectAlignment(size);
+    size_t pageSizeAlignment = alignmentBase.alignment;
+    auto rootDeviceCount = this->executionEnvironment.rootDeviceEnvironments.size();
+
+    // If all devices can support HEAP EXTENDED, then that heap is used, otherwise the HEAP based on the size is used.
+    for (auto rootDeviceIndex = 0u; rootDeviceIndex < rootDeviceCount; rootDeviceIndex++) {
+        auto gfxPartition = getGfxPartition(rootDeviceIndex);
+        if (gfxPartition->getHeapLimit(HeapIndex::HEAP_EXTENDED) > 0) {
+            auto alignSize = size >= 8 * MemoryConstants::gigaByte && Math::isPow2(size);
+            if (DebugManager.flags.UseHighAlignmentForHeapExtended.get() != -1) {
+                alignSize = !!DebugManager.flags.UseHighAlignmentForHeapExtended.get();
+            }
+
+            if (alignSize) {
+                pageSizeAlignment = Math::prevPowerOfTwo(size);
+            }
+
+            *heap = HeapIndex::HEAP_EXTENDED;
+        } else {
+            pageSizeAlignment = alignmentBase.alignment;
+            *heap = alignmentBase.heap;
+            break;
+        }
+    }
+    return pageSizeAlignment;
+}
+
 AddressRange DrmMemoryManager::reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) {
+    return reserveGpuAddressOnHeap(requiredStartAddress, size, rootDeviceIndices, reservedOnRootDeviceIndex, HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
+}
+
+AddressRange DrmMemoryManager::reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) {
     uint64_t gpuVa = 0u;
     *reservedOnRootDeviceIndex = 0;
     for (auto rootDeviceIndex : rootDeviceIndices) {
-        gpuVa = acquireGpuRange(size, rootDeviceIndex, HeapIndex::HEAP_STANDARD);
+        if (heap == HeapIndex::HEAP_EXTENDED) {
+            gpuVa = acquireGpuRangeWithCustomAlignment(size, rootDeviceIndex, heap, alignment);
+        } else {
+            gpuVa = acquireGpuRange(size, rootDeviceIndex, heap);
+        }
         if (gpuVa != 0u) {
             *reservedOnRootDeviceIndex = rootDeviceIndex;
             break;
diff --git a/shared/source/os_interface/linux/drm_memory_manager.h b/shared/source/os_interface/linux/drm_memory_manager.h
index 0de39c7e61..32d73f9de2 100644
--- a/shared/source/os_interface/linux/drm_memory_manager.h
+++ b/shared/source/os_interface/linux/drm_memory_manager.h
@@ -69,6 +69,8 @@ class DrmMemoryManager : public MemoryManager {
 
     MOCKABLE_VIRTUAL int obtainFdFromHandle(int boHandle, uint32_t rootDeviceindex);
     AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override;
+    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override;
+    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override;
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override;
     MOCKABLE_VIRTUAL BufferObject *createBufferObjectInMemoryRegion(uint32_t rootDeviceIndex, Gmm *gmm, AllocationType allocationType, uint64_t gpuAddress, size_t size,
                                                                     uint32_t memoryBanks, size_t maxOsContextCount, int32_t pairHandle, bool isSystemMemoryPool);
diff --git a/shared/source/os_interface/windows/wddm_memory_manager.cpp b/shared/source/os_interface/windows/wddm_memory_manager.cpp
index 10446f19b9..1b46831aea 100644
--- a/shared/source/os_interface/windows/wddm_memory_manager.cpp
+++ b/shared/source/os_interface/windows/wddm_memory_manager.cpp
@@ -913,16 +913,26 @@ bool WddmMemoryManager::createWddmAllocation(WddmAllocation *allocation, void *r
     return mapGpuVirtualAddress(allocation, requiredGpuPtr);
 }
 
+size_t WddmMemoryManager::selectAlignmentAndHeap(size_t size, HeapIndex *heap) {
+    AlignmentSelector::CandidateAlignment alignment = alignmentSelector.selectAlignment(size);
+    *heap = HeapIndex::HEAP_STANDARD64KB;
+    return alignment.alignment;
+}
+
 AddressRange WddmMemoryManager::reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) {
+    return reserveGpuAddressOnHeap(requiredStartAddress, size, rootDeviceIndices, reservedOnRootDeviceIndex, HeapIndex::HEAP_STANDARD64KB, MemoryConstants::pageSize64k);
+}
+
+AddressRange WddmMemoryManager::reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) {
     uint64_t gpuVa = 0u;
     *reservedOnRootDeviceIndex = 0;
     size_t reservedSize = 0;
     NTSTATUS status = STATUS_UNSUCCESSFUL;
     for (auto rootDeviceIndex : rootDeviceIndices) {
         auto gfxPartition = getGfxPartition(rootDeviceIndex);
-        status = getWddm(rootDeviceIndex).reserveGpuVirtualAddress(requiredStartAddress, gfxPartition->getHeapMinimalAddress(HeapIndex::HEAP_STANDARD64KB), gfxPartition->getHeapLimit(HeapIndex::HEAP_STANDARD64KB), size, &gpuVa);
+        status = getWddm(rootDeviceIndex).reserveGpuVirtualAddress(requiredStartAddress, gfxPartition->getHeapMinimalAddress(heap), gfxPartition->getHeapLimit(heap), size, &gpuVa);
         if (requiredStartAddress != 0ull && status != STATUS_SUCCESS) {
-            status = getWddm(rootDeviceIndex).reserveGpuVirtualAddress(0ull, gfxPartition->getHeapMinimalAddress(HeapIndex::HEAP_STANDARD64KB), gfxPartition->getHeapLimit(HeapIndex::HEAP_STANDARD64KB), size, &gpuVa);
+            status = getWddm(rootDeviceIndex).reserveGpuVirtualAddress(0ull, gfxPartition->getHeapMinimalAddress(heap), gfxPartition->getHeapLimit(heap), size, &gpuVa);
         }
         if (status == STATUS_SUCCESS) {
             *reservedOnRootDeviceIndex = rootDeviceIndex;
diff --git a/shared/source/os_interface/windows/wddm_memory_manager.h b/shared/source/os_interface/windows/wddm_memory_manager.h
index 1e08924f5a..b10ae1fc28 100644
--- a/shared/source/os_interface/windows/wddm_memory_manager.h
+++ b/shared/source/os_interface/windows/wddm_memory_manager.h
@@ -65,6 +65,8 @@ class WddmMemoryManager : public MemoryManager {
     bool isWCMemory(const void *ptr) override;
 
     AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override;
+    AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override;
+    size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override;
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override;
     bool verifyHandle(osHandle handle, uint32_t rootDeviceIndex, bool ntHandle) override;
     bool isNTHandle(osHandle handle, uint32_t rootDeviceIndex) override;
diff --git a/shared/test/common/mocks/linux/mock_drm_memory_manager.h b/shared/test/common/mocks/linux/mock_drm_memory_manager.h
index 73c6dafa1d..8377eda457 100644
--- a/shared/test/common/mocks/linux/mock_drm_memory_manager.h
+++ b/shared/test/common/mocks/linux/mock_drm_memory_manager.h
@@ -69,6 +69,7 @@ class TestedDrmMemoryManager : public MemoryManagerCreate<DrmMemoryManager> {
     using DrmMemoryManager::registerSharedBoHandleAllocation;
     using DrmMemoryManager::releaseGpuRange;
     using DrmMemoryManager::retrieveMmapOffsetForBufferObject;
+    using DrmMemoryManager::selectAlignmentAndHeap;
     using DrmMemoryManager::setDomainCpu;
     using DrmMemoryManager::sharedBoHandles;
     using DrmMemoryManager::sharingBufferObjects;
diff --git a/shared/test/common/os_interface/windows/mock_wddm_memory_manager.h b/shared/test/common/os_interface/windows/mock_wddm_memory_manager.h
index ed29b35c9b..15e748ed75 100644
--- a/shared/test/common/os_interface/windows/mock_wddm_memory_manager.h
+++ b/shared/test/common/os_interface/windows/mock_wddm_memory_manager.h
@@ -39,6 +39,7 @@ class MockWddmMemoryManager : public MemoryManagerCreate<WddmMemoryManager> {
     using BaseClass::getHugeGfxMemoryChunkSize;
     using BaseClass::getPreferredAllocationMethod;
     using BaseClass::isStatelessAccessRequired;
+    using BaseClass::selectAlignmentAndHeap;
 
     GraphicsAllocation *allocateGraphicsMemory64kb(const AllocationData &allocationData) override {
         allocationGraphicsMemory64kbCreated = true;
diff --git a/shared/test/unit_test/device/neo_device_tests.cpp b/shared/test/unit_test/device/neo_device_tests.cpp
index ca22d303f4..5c4b2f6ff2 100644
--- a/shared/test/unit_test/device/neo_device_tests.cpp
+++ b/shared/test/unit_test/device/neo_device_tests.cpp
@@ -10,6 +10,7 @@
 #include "shared/source/helpers/compiler_product_helper.h"
 #include "shared/source/helpers/driver_model_type.h"
 #include "shared/source/helpers/gfx_core_helper.h"
+#include "shared/source/memory_manager/gfx_partition.h"
 #include "shared/source/os_interface/device_factory.h"
 #include "shared/source/os_interface/driver_info.h"
 #include "shared/source/os_interface/os_interface.h"
@@ -473,6 +474,13 @@ TEST_F(DeviceGetCapsTest, givenFlagEnabled64kbPagesWhenCallConstructorMemoryMana
         AddressRange reserveGpuAddress(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex) override {
             return {};
         }
+        AddressRange reserveGpuAddressOnHeap(const uint64_t requiredStartAddress, size_t size, RootDeviceIndicesContainer rootDeviceIndices, uint32_t *reservedOnRootDeviceIndex, HeapIndex heap, size_t alignment) override {
+            return {};
+        }
+        size_t selectAlignmentAndHeap(size_t size, HeapIndex *heap) override {
+            *heap = HeapIndex::HEAP_STANDARD;
+            return MemoryConstants::pageSize64k;
+        }
         void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
         GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const AllocationData &allocationData) override { return nullptr; };
         GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const AllocationData &allocationData) override { return nullptr; };
diff --git a/shared/test/unit_test/memory_manager/memory_manager_tests.cpp b/shared/test/unit_test/memory_manager/memory_manager_tests.cpp
index 8932b2cd04..5ca8d5267d 100644
--- a/shared/test/unit_test/memory_manager/memory_manager_tests.cpp
+++ b/shared/test/unit_test/memory_manager/memory_manager_tests.cpp
@@ -234,14 +234,40 @@ HWTEST_F(MemoryhManagerMultiContextResourceTests, givenAllocationUsedByManyOsCon
     EXPECT_TRUE(defaultCsr->getInternalAllocationStorage()->getTemporaryAllocations().peekIsEmpty());
 }
 
+TEST(OsAgnosticMemoryManager, givenOsAgnosticMemoryManagerWhenGpuAddressIsReservedOnSpecifiedHeapAndFreedThenAddressFromGfxPartitionIsUsed) {
+    MockExecutionEnvironment executionEnvironment;
+    OsAgnosticMemoryManager memoryManager(executionEnvironment);
+    RootDeviceIndicesContainer rootDeviceIndices;
+    rootDeviceIndices.pushUnique(0);
+    uint32_t rootDeviceIndexReserved = 10;
+    auto gmmHelper = memoryManager.getGmmHelper(0);
+    HeapIndex heap = HeapIndex::HEAP_STANDARD64KB;
+    auto alignment = memoryManager.selectAlignmentAndHeap(MemoryConstants::pageSize, &heap);
+    EXPECT_EQ(heap, HeapIndex::HEAP_STANDARD);
+    EXPECT_EQ(MemoryConstants::pageSize64k, alignment);
+    auto addressRange = memoryManager.reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, heap, alignment);
+    EXPECT_EQ(0u, rootDeviceIndexReserved);
+    EXPECT_LE(memoryManager.getGfxPartition(0)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memoryManager.getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+
+    memoryManager.freeGpuAddress(addressRange, 0);
+}
+
 TEST(OsAgnosticMemoryManager, givenOsAgnosticMemoryManagerWhenGpuAddressIsReservedAndFreedThenAddressFromGfxPartitionIsUsed) {
     MockExecutionEnvironment executionEnvironment;
     OsAgnosticMemoryManager memoryManager(executionEnvironment);
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(0);
     uint32_t rootDeviceIndexReserved = 10;
-    auto addressRange = memoryManager.reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
     auto gmmHelper = memoryManager.getGmmHelper(0);
+    auto addressRange = memoryManager.reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
+    EXPECT_EQ(0u, rootDeviceIndexReserved);
+    EXPECT_LE(memoryManager.getGfxPartition(0)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memoryManager.getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+
+    memoryManager.freeGpuAddress(addressRange, 0);
+
+    addressRange = memoryManager.reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
     EXPECT_EQ(0u, rootDeviceIndexReserved);
     EXPECT_LE(memoryManager.getGfxPartition(0)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
     EXPECT_GT(memoryManager.getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
@@ -255,8 +281,15 @@ TEST(OsAgnosticMemoryManager, givenOsAgnosticMemoryManagerWhenGpuAddressIsReserv
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(1);
     uint32_t rootDeviceIndexReserved = 10;
-    auto addressRange = memoryManager.reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
     auto gmmHelper = memoryManager.getGmmHelper(1);
+    auto addressRange = memoryManager.reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
+    EXPECT_EQ(1u, rootDeviceIndexReserved);
+    EXPECT_LE(memoryManager.getGfxPartition(1)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memoryManager.getGfxPartition(1)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+
+    memoryManager.freeGpuAddress(addressRange, 1);
+
+    addressRange = memoryManager.reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
     EXPECT_EQ(1u, rootDeviceIndexReserved);
     EXPECT_LE(memoryManager.getGfxPartition(1)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
     EXPECT_GT(memoryManager.getGfxPartition(1)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
@@ -272,7 +305,9 @@ TEST(OsAgnosticMemoryManager, givenOsAgnosticMemoryManagerWhenGpuAddressReservat
     uint32_t rootDeviceIndexReserved = 10;
     // emulate GPU address space exhaust
     memoryManager.getGfxPartition(0)->heapInit(HeapIndex::HEAP_STANDARD, 0x0, 0x10000);
-    auto addressRange = memoryManager.reserveGpuAddress(0ull, (size_t)(memoryManager.getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD) * 2), rootDeviceIndices, &rootDeviceIndexReserved);
+    auto addressRange = memoryManager.reserveGpuAddressOnHeap(0ull, (size_t)(memoryManager.getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD) * 2), rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
+    EXPECT_EQ(static_cast<int>(addressRange.address), 0);
+    addressRange = memoryManager.reserveGpuAddress(0ull, (size_t)(memoryManager.getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD) * 2), rootDeviceIndices, &rootDeviceIndexReserved);
     EXPECT_EQ(static_cast<int>(addressRange.address), 0);
 }
 
@@ -282,7 +317,11 @@ TEST(OsAgnosticMemoryManager, givenOsAgnosticMemoryManagerWhenGpuAddressReservat
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(0);
     uint32_t rootDeviceIndexReserved = 10;
-    auto addressRange = memoryManager.reserveGpuAddress(0x1234, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
+    auto addressRange = memoryManager.reserveGpuAddressOnHeap(0x1234, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
+    EXPECT_EQ(0u, rootDeviceIndexReserved);
+    EXPECT_NE(static_cast<int>(addressRange.address), 0x1234);
+    EXPECT_NE(static_cast<int>(addressRange.size), 0);
+    addressRange = memoryManager.reserveGpuAddress(0x1234, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
     EXPECT_EQ(0u, rootDeviceIndexReserved);
     EXPECT_NE(static_cast<int>(addressRange.address), 0x1234);
     EXPECT_NE(static_cast<int>(addressRange.size), 0);
diff --git a/shared/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp b/shared/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp
index 5814086b29..b1faaee43d 100644
--- a/shared/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp
+++ b/shared/test/unit_test/os_interface/linux/drm_memory_manager_tests.cpp
@@ -8,6 +8,7 @@
 #include "shared/source/built_ins/sip.h"
 #include "shared/source/command_stream/tag_allocation_layout.h"
 #include "shared/source/gmm_helper/gmm_helper.h"
+#include "shared/source/helpers/basic_math.h"
 #include "shared/source/helpers/common_types.h"
 #include "shared/source/helpers/surface_format_info.h"
 #include "shared/source/indirect_heap/indirect_heap.h"
@@ -300,8 +301,15 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenGp
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(1);
     uint32_t rootDeviceIndexReserved = 0;
-    auto addressRange = memoryManager->reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
     auto gmmHelper = memoryManager->getGmmHelper(1);
+    auto addressRange = memoryManager->reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
+
+    EXPECT_EQ(rootDeviceIndexReserved, 1u);
+    EXPECT_LE(memoryManager->getGfxPartition(1)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memoryManager->getGfxPartition(1)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    memoryManager->freeGpuAddress(addressRange, 1);
+
+    addressRange = memoryManager->reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
 
     EXPECT_EQ(rootDeviceIndexReserved, 1u);
     EXPECT_LE(memoryManager->getGfxPartition(1)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
@@ -314,12 +322,107 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenGp
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(0);
     uint32_t rootDeviceIndexReserved = 1;
-    auto addressRange = memoryManager->reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
     auto gmmHelper = memoryManager->getGmmHelper(0);
+    auto addressRange = memoryManager->reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
 
     EXPECT_LE(memoryManager->getGfxPartition(0)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
     EXPECT_GT(memoryManager->getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
     memoryManager->freeGpuAddress(addressRange, 0);
+
+    addressRange = memoryManager->reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
+
+    EXPECT_LE(memoryManager->getGfxPartition(0)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memoryManager->getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    memoryManager->freeGpuAddress(addressRange, 0);
+}
+
+TEST_F(DrmMemoryManagerWithLocalMemoryAndExplicitExpectationsTest, givenDebugVariableToDisableAddressAlignmentAndCallToSelectAlignmentAndHeapWithPow2MemoryThenAlignmentIs2Mb) {
+    if (!memoryManager->getGfxPartition(rootDeviceIndex)->getHeapLimit(HeapIndex::HEAP_EXTENDED)) {
+        GTEST_SKIP();
+    }
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.UseHighAlignmentForHeapExtended.set(0);
+
+    auto size = 8 * MemoryConstants::gigaByte;
+
+    HeapIndex heap = HeapIndex::HEAP_STANDARD;
+    auto alignment = memoryManager->selectAlignmentAndHeap(size, &heap);
+    EXPECT_EQ(heap, HeapIndex::HEAP_EXTENDED);
+    EXPECT_EQ(MemoryConstants::pageSize2M, alignment);
+}
+
+TEST_F(DrmMemoryManagerWithLocalMemoryAndExplicitExpectationsTest, givenCalltoSelectAlignmentAndHeapWithPow2MemoryThenAlignmentIsPreviousPow2) {
+    if (!memoryManager->getGfxPartition(rootDeviceIndex)->getHeapLimit(HeapIndex::HEAP_EXTENDED)) {
+        GTEST_SKIP();
+    }
+
+    auto size = Math::nextPowerOfTwo(8 * MemoryConstants::gigaByte);
+
+    HeapIndex heap = HeapIndex::HEAP_STANDARD;
+    auto alignment = memoryManager->selectAlignmentAndHeap(size, &heap);
+    EXPECT_EQ(heap, HeapIndex::HEAP_EXTENDED);
+    EXPECT_EQ(Math::prevPowerOfTwo(size), alignment);
+}
+
+TEST_F(DrmMemoryManagerWithLocalMemoryAndExplicitExpectationsTest, givenCalltoSelectAlignmentAndHeapWithNonPow2SizeThenAlignmentIs2MB) {
+    if (!memoryManager->getGfxPartition(rootDeviceIndex)->getHeapLimit(HeapIndex::HEAP_EXTENDED)) {
+        GTEST_SKIP();
+    }
+
+    auto size = (8 * MemoryConstants::gigaByte) + 5;
+
+    HeapIndex heap = HeapIndex::HEAP_STANDARD;
+    auto alignment = memoryManager->selectAlignmentAndHeap(size, &heap);
+    EXPECT_EQ(heap, HeapIndex::HEAP_EXTENDED);
+    EXPECT_EQ(MemoryConstants::pageSize2M, alignment);
+}
+
+TEST_F(DrmMemoryManagerWithLocalMemoryAndExplicitExpectationsTest, givenDebugVariableToDisableAddressAlignmentWhenSelectAlignmentAndHeapNewCustomAlignmentReturned) {
+    if (!memoryManager->getGfxPartition(rootDeviceIndex)->getHeapLimit(HeapIndex::HEAP_EXTENDED)) {
+        GTEST_SKIP();
+    }
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.UseHighAlignmentForHeapExtended.set(0);
+
+    auto size = 16 * MemoryConstants::megaByte;
+
+    HeapIndex heap = HeapIndex::HEAP_STANDARD;
+    auto alignment = memoryManager->selectAlignmentAndHeap(size, &heap);
+    EXPECT_EQ(heap, HeapIndex::HEAP_EXTENDED);
+    EXPECT_NE(MemoryConstants::pageSize64k, alignment);
+}
+
+TEST_F(DrmMemoryManagerWithLocalMemoryAndExplicitExpectationsTest, givenSelectAlignmentAndHeapThen64KbAlignmentReturned) {
+    if (!memoryManager->getGfxPartition(rootDeviceIndex)->getHeapLimit(HeapIndex::HEAP_EXTENDED)) {
+        GTEST_SKIP();
+    }
+
+    auto size = 16 * MemoryConstants::megaByte;
+
+    HeapIndex heap = HeapIndex::HEAP_STANDARD;
+    auto alignment = memoryManager->selectAlignmentAndHeap(size, &heap);
+    EXPECT_EQ(heap, HeapIndex::HEAP_EXTENDED);
+    EXPECT_NE(MemoryConstants::pageSize64k, alignment);
+}
+
+TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenGpuAddressReservationIsAttemptedWithTheQuieriedHeapThenSuccessReturned) {
+    auto memoryManager = std::make_unique<TestedDrmMemoryManager>(false, true, false, *executionEnvironment);
+    RootDeviceIndicesContainer rootDeviceIndices;
+    rootDeviceIndices.pushUnique(0);
+    uint32_t rootDeviceIndexReserved = 1;
+    HeapIndex heap = HeapIndex::HEAP_STANDARD;
+    auto alignment = memoryManager->selectAlignmentAndHeap(MemoryConstants::pageSize, &heap);
+    auto gfxPartition = memoryManager->getGfxPartition(0);
+    if (gfxPartition->getHeapLimit(HeapIndex::HEAP_EXTENDED) > 0) {
+        EXPECT_EQ(heap, HeapIndex::HEAP_EXTENDED);
+    } else {
+        EXPECT_EQ(heap, HeapIndex::HEAP_STANDARD64KB);
+    }
+    EXPECT_EQ(MemoryConstants::pageSize64k, alignment);
+    auto addressRange = memoryManager->reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, heap, alignment);
+    EXPECT_NE(static_cast<int>(addressRange.address), 0);
+    EXPECT_NE(static_cast<int>(addressRange.size), 0);
+    memoryManager->freeGpuAddress(addressRange, 0);
 }
 
 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenGpuAddressReservationIsAttemptedWithAnInvalidRequiredPtrThenDifferentRangeReturned) {
@@ -327,9 +430,14 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenGp
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(0);
     uint32_t rootDeviceIndexReserved = 1;
-    auto addressRange = memoryManager->reserveGpuAddress(0x1234, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
+    auto addressRange = memoryManager->reserveGpuAddressOnHeap(0x1234, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
     EXPECT_NE(static_cast<int>(addressRange.address), 0x1234);
     EXPECT_NE(static_cast<int>(addressRange.size), 0);
+    memoryManager->freeGpuAddress(addressRange, 0);
+    addressRange = memoryManager->reserveGpuAddress(0x1234, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
+    EXPECT_NE(static_cast<int>(addressRange.address), 0x1234);
+    EXPECT_NE(static_cast<int>(addressRange.size), 0);
+    memoryManager->freeGpuAddress(addressRange, 0);
 }
 
 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenGpuAddressReservationIsAttemptedWhichFailsThenNullRangeReturned) {
@@ -341,8 +449,25 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenGp
     memoryManager->forceLimitedRangeAllocator(0xFFFFFFFFF);
     memoryManager->getGfxPartition(0)->heapInit(HeapIndex::HEAP_STANDARD, 0x0, 0x10000);
     size_t invalidSize = (size_t)memoryManager->getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD) + MemoryConstants::pageSize;
-    auto addressRange = memoryManager->reserveGpuAddress(0ull, invalidSize, rootDeviceIndices, &rootDeviceIndexReserved);
+    auto addressRange = memoryManager->reserveGpuAddressOnHeap(0ull, invalidSize, rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
     EXPECT_EQ(static_cast<int>(addressRange.address), 0);
+    addressRange = memoryManager->reserveGpuAddress(0ull, invalidSize, rootDeviceIndices, &rootDeviceIndexReserved);
+    EXPECT_EQ(static_cast<int>(addressRange.address), 0);
+}
+
+TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenHeapAndAlignmentRequestedWithoutAllExtendedHeapsForRootDevicesThenHeapStandardReturned) {
+    auto memoryManager = std::make_unique<TestedDrmMemoryManager>(false, true, false, *executionEnvironment);
+    // emulate GPU address space exhaust
+    memoryManager->forceLimitedRangeAllocator(0xFFFFFFFFF);
+    memoryManager->getGfxPartition(0)->heapInit(HeapIndex::HEAP_EXTENDED, 0x11000, 0x10000);
+    memoryManager->getGfxPartition(1)->heapInit(HeapIndex::HEAP_STANDARD, 0, 0x10000);
+    memoryManager->getGfxPartition(1)->heapInit(HeapIndex::HEAP_EXTENDED, 0, 0);
+    auto size = MemoryConstants::pageSize64k;
+
+    HeapIndex heap = HeapIndex::HEAP_STANDARD;
+    auto alignment = memoryManager->selectAlignmentAndHeap(size, &heap);
+    EXPECT_EQ(heap, HeapIndex::HEAP_STANDARD64KB);
+    EXPECT_EQ(MemoryConstants::pageSize64k, alignment);
 }
 
 TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenSmallSizeAndGpuAddressSetWhenGraphicsMemoryIsAllocatedThenAllocationWithSpecifiedGpuAddressInSystemMemoryIsCreated) {
@@ -6829,14 +6954,29 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenGp
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(1);
     uint32_t rootDeviceIndexReserved = 0;
-    auto addressRange = memoryManager->reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
     auto gmmHelper = memoryManager->getGmmHelper(1);
+    auto addressRange = memoryManager->reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
 
     EXPECT_EQ(rootDeviceIndexReserved, 1u);
     EXPECT_LE(memoryManager->getGfxPartition(1)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
     EXPECT_GT(memoryManager->getGfxPartition(1)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
     uint64_t requiredAddr = addressRange.address;
     memoryManager->freeGpuAddress(addressRange, 1);
+    addressRange = memoryManager->reserveGpuAddressOnHeap(requiredAddr, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved, NEO::HeapIndex::HEAP_STANDARD, MemoryConstants::pageSize64k);
+
+    EXPECT_EQ(rootDeviceIndexReserved, 1u);
+    EXPECT_EQ(addressRange.address, requiredAddr);
+    EXPECT_LE(memoryManager->getGfxPartition(1)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memoryManager->getGfxPartition(1)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    memoryManager->freeGpuAddress(addressRange, 1);
+
+    addressRange = memoryManager->reserveGpuAddress(0ull, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
+
+    EXPECT_EQ(rootDeviceIndexReserved, 1u);
+    EXPECT_LE(memoryManager->getGfxPartition(1)->getHeapBase(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memoryManager->getGfxPartition(1)->getHeapLimit(HeapIndex::HEAP_STANDARD), gmmHelper->decanonize(addressRange.address));
+    requiredAddr = addressRange.address;
+    memoryManager->freeGpuAddress(addressRange, 1);
     addressRange = memoryManager->reserveGpuAddress(requiredAddr, MemoryConstants::pageSize, rootDeviceIndices, &rootDeviceIndexReserved);
 
     EXPECT_EQ(rootDeviceIndexReserved, 1u);
diff --git a/shared/test/unit_test/os_interface/wddm_linux/configure_device_address_space_drm_or_wddm_test.cpp b/shared/test/unit_test/os_interface/wddm_linux/configure_device_address_space_drm_or_wddm_test.cpp
index e182084395..25f30e2224 100644
--- a/shared/test/unit_test/os_interface/wddm_linux/configure_device_address_space_drm_or_wddm_test.cpp
+++ b/shared/test/unit_test/os_interface/wddm_linux/configure_device_address_space_drm_or_wddm_test.cpp
@@ -709,7 +709,7 @@ TEST_F(WddmLinuxTest, givenAllocatedMemoryAndCloseInternalHandleThenSharedHandle
     memoryManager.freeGraphicsMemoryImpl(alloc);
 }
 
-TEST_F(WddmLinuxTest, givenAllocatedMemoryAndCloseInternalHandleWihtoutAllocationThenSharedHandleStillClosed) {
+TEST_F(WddmLinuxTest, givenAllocatedMemoryAndCloseInternalHandleWithoutAllocationThenSharedHandleStillClosed) {
     osEnvironment->gdi->reserveGpuVirtualAddress = reserveDeviceAddressSpaceMock;
     osEnvironment->gdi->createAllocation2 = createAllocation2Mock;
     osEnvironment->gdi->mapGpuVirtualAddress = mapGpuVirtualAddressMock;
diff --git a/shared/test/unit_test/os_interface/windows/wddm_memory_manager_tests.cpp b/shared/test/unit_test/os_interface/windows/wddm_memory_manager_tests.cpp
index 5a7a9e5ce1..0b301ae41c 100644
--- a/shared/test/unit_test/os_interface/windows/wddm_memory_manager_tests.cpp
+++ b/shared/test/unit_test/os_interface/windows/wddm_memory_manager_tests.cpp
@@ -1113,17 +1113,35 @@ TEST_F(WddmMemoryManagerSimpleTest, givenNonZeroFenceValueOnSomeOfMultipleEngine
     memoryManager->freeGraphicsMemory(allocation);
 }
 
+TEST_F(WddmMemoryManagerSimpleTest, givenWddmMemoryManagerWhenSelectAlignmentAndHeapCalledThenCorrectHeapReturned) {
+    HeapIndex heap = HeapIndex::HEAP_STANDARD;
+    auto alignment = memoryManager->selectAlignmentAndHeap(MemoryConstants::pageSize64k, &heap);
+    EXPECT_EQ(heap, HeapIndex::HEAP_STANDARD64KB);
+    EXPECT_EQ(MemoryConstants::pageSize64k, alignment);
+}
+
 TEST_F(WddmMemoryManagerSimpleTest, givenWddmMemoryManagerWhenGpuAddressIsReservedAndFreedThenAddressRangeIsNonZero) {
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(0);
     uint32_t rootDeviceIndexReserved = 1;
-    auto addressRange = memoryManager->reserveGpuAddress(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
+    HeapIndex heap = HeapIndex::HEAP_STANDARD;
+    auto alignment = memoryManager->selectAlignmentAndHeap(MemoryConstants::pageSize64k, &heap);
+    EXPECT_EQ(heap, HeapIndex::HEAP_STANDARD64KB);
+    EXPECT_EQ(MemoryConstants::pageSize64k, alignment);
+    auto addressRange = memoryManager->reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved, heap, alignment);
     auto gmmHelper = memoryManager->getGmmHelper(0);
     EXPECT_EQ(0u, rootDeviceIndexReserved);
     EXPECT_NE(0u, gmmHelper->decanonize(addressRange.address));
     EXPECT_EQ(MemoryConstants::pageSize64k, addressRange.size);
 
     memoryManager->freeGpuAddress(addressRange, 0);
+
+    addressRange = memoryManager->reserveGpuAddress(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
+    EXPECT_EQ(0u, rootDeviceIndexReserved);
+    EXPECT_NE(0u, gmmHelper->decanonize(addressRange.address));
+    EXPECT_EQ(MemoryConstants::pageSize64k, addressRange.size);
+
+    memoryManager->freeGpuAddress(addressRange, 0);
 }
 
 TEST_F(WddmMemoryManagerSimpleTest, givenWddmMemoryManagerWhenAllocatingWithGpuVaThenNullptrIsReturned) {
diff --git a/shared/test/unit_test/os_interface/windows/wddm_memory_reservation_tests.cpp b/shared/test/unit_test/os_interface/windows/wddm_memory_reservation_tests.cpp
index 08f4aac025..bbc5f81bff 100644
--- a/shared/test/unit_test/os_interface/windows/wddm_memory_reservation_tests.cpp
+++ b/shared/test/unit_test/os_interface/windows/wddm_memory_reservation_tests.cpp
@@ -22,6 +22,7 @@ namespace NEO {
 struct MemoryReservationMock : public MockWddmMemoryManager {
     using MemoryManager::freeGpuAddress;
     using MemoryManager::reserveGpuAddress;
+    using MemoryManager::reserveGpuAddressOnHeap;
     MemoryReservationMock(NEO::ExecutionEnvironment &executionEnvironment) : MockWddmMemoryManager(executionEnvironment) {}
 };
 
@@ -50,8 +51,15 @@ TEST_F(WddmMemoryReservationTests, givenWddmMemoryManagerWhenGpuAddressIsReserve
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(0);
     uint32_t rootDeviceIndexReserved = 1;
-    auto addressRange = memManager->reserveGpuAddress(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
     auto gmmHelper = memManager->getGmmHelper(0);
+    auto addressRange = memManager->reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved, HeapIndex::HEAP_STANDARD64KB, MemoryConstants::pageSize64k);
+
+    EXPECT_EQ(rootDeviceIndexReserved, 0u);
+    EXPECT_LE(memManager->getGfxPartition(0)->getHeapBase(HeapIndex::HEAP_STANDARD64KB), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memManager->getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD64KB), gmmHelper->decanonize(addressRange.address));
+    memManager->freeGpuAddress(addressRange, 0);
+
+    addressRange = memManager->reserveGpuAddress(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
 
     EXPECT_EQ(rootDeviceIndexReserved, 0u);
     EXPECT_LE(memManager->getGfxPartition(0)->getHeapBase(HeapIndex::HEAP_STANDARD64KB), gmmHelper->decanonize(addressRange.address));
@@ -64,8 +72,16 @@ TEST_F(WddmMemoryReservationTests, givenWddmMemoryManagerWhenGpuAddressIsReserve
     rootDeviceIndices.pushUnique(0);
     uint32_t rootDeviceIndexReserved = 1;
     uint64_t invalidAddress = 0x1234;
-    auto addressRange = memManager->reserveGpuAddress(invalidAddress, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
     auto gmmHelper = memManager->getGmmHelper(0);
+    auto addressRange = memManager->reserveGpuAddressOnHeap(invalidAddress, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved, HeapIndex::HEAP_STANDARD64KB, MemoryConstants::pageSize64k);
+    EXPECT_NE(invalidAddress, addressRange.address);
+
+    EXPECT_EQ(rootDeviceIndexReserved, 0u);
+    EXPECT_LE(memManager->getGfxPartition(0)->getHeapBase(HeapIndex::HEAP_STANDARD64KB), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memManager->getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD64KB), gmmHelper->decanonize(addressRange.address));
+    memManager->freeGpuAddress(addressRange, 0);
+
+    addressRange = memManager->reserveGpuAddress(invalidAddress, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
     EXPECT_NE(invalidAddress, addressRange.address);
 
     EXPECT_EQ(rootDeviceIndexReserved, 0u);
@@ -78,11 +94,22 @@ TEST_F(WddmMemoryReservationTests, givenWddmMemoryManagerWhenGpuAddressIsReserve
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(0);
     uint32_t rootDeviceIndexReserved = 1;
-    auto addressRange = memManager->reserveGpuAddress(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
+    auto gmmHelper = memManager->getGmmHelper(0);
+    auto addressRange = memManager->reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved, HeapIndex::HEAP_STANDARD64KB, MemoryConstants::pageSize64k);
     auto previousAddress = addressRange.address;
     memManager->freeGpuAddress(addressRange, 0);
-    auto newAddressRange = memManager->reserveGpuAddress(previousAddress, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
-    auto gmmHelper = memManager->getGmmHelper(0);
+    auto newAddressRange = memManager->reserveGpuAddressOnHeap(previousAddress, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved, HeapIndex::HEAP_STANDARD64KB, MemoryConstants::pageSize64k);
+    EXPECT_EQ(previousAddress, addressRange.address);
+
+    EXPECT_EQ(rootDeviceIndexReserved, 0u);
+    EXPECT_LE(memManager->getGfxPartition(0)->getHeapBase(HeapIndex::HEAP_STANDARD64KB), gmmHelper->decanonize(addressRange.address));
+    EXPECT_GT(memManager->getGfxPartition(0)->getHeapLimit(HeapIndex::HEAP_STANDARD64KB), gmmHelper->decanonize(addressRange.address));
+    memManager->freeGpuAddress(addressRange, 0);
+
+    addressRange = memManager->reserveGpuAddress(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
+    previousAddress = addressRange.address;
+    memManager->freeGpuAddress(addressRange, 0);
+    newAddressRange = memManager->reserveGpuAddress(previousAddress, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
     EXPECT_EQ(previousAddress, addressRange.address);
 
     EXPECT_EQ(rootDeviceIndexReserved, 0u);
@@ -109,7 +136,10 @@ TEST(WddmMemoryReservationFailTest, givenWddmMemoryManagerWhenGpuAddressReservat
     RootDeviceIndicesContainer rootDeviceIndices;
     rootDeviceIndices.pushUnique(0);
     uint32_t rootDeviceIndexReserved = 1;
-    auto addressRange = memManager->reserveGpuAddress(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
+    auto addressRange = memManager->reserveGpuAddressOnHeap(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved, HeapIndex::HEAP_STANDARD64KB, MemoryConstants::pageSize64k);
+    EXPECT_EQ(addressRange.address, 0ull);
+    EXPECT_EQ(addressRange.size, 0u);
+    addressRange = memManager->reserveGpuAddress(0ull, MemoryConstants::pageSize64k, rootDeviceIndices, &rootDeviceIndexReserved);
     EXPECT_EQ(addressRange.address, 0ull);
     EXPECT_EQ(addressRange.size, 0u);
 }