feature: support SVM heap in reserveVirtualMem

Related-To: NEO-11981

Signed-off-by: Wenbin Lu <wenbin.lu@intel.com>

level_zero/core/source/context/context_imp.cpp

@@ -10,12 +10,13 @@
 #include "shared/source/command_container/implicit_scaling.h"
 #include "shared/source/command_stream/command_stream_receiver.h"
 #include "shared/source/execution_environment/root_device_environment.h"
-#include "shared/source/helpers/basic_math.h"
+#include "shared/source/helpers/aligned_memory.h"
 #include "shared/source/helpers/gfx_core_helper.h"
 #include "shared/source/helpers/ptr_math.h"
 #include "shared/source/memory_manager/allocation_properties.h"
 #include "shared/source/memory_manager/memory_operations_handler.h"
 #include "shared/source/memory_manager/unified_memory_manager.h"
+#include "shared/source/utilities/cpu_info.h"
 
 #include "level_zero/api/driver_experimental/public/zex_memory.h"
 #include "level_zero/core/source/cmdlist/cmdlist.h"
@@ -1047,21 +1048,60 @@ NEO::VirtualMemoryReservation *ContextImp::findSupportedVirtualReservation(const
 ze_result_t ContextImp::reserveVirtualMem(const void *pStart,
                                           size_t size,
                                           void **pptr) {
+    uint64_t maxCpuVa = 0;
+    if (this->driverHandle->getMemoryManager()->peek32bit()) {
+        maxCpuVa = maxNBitValue(32);
+    } else {
+        maxCpuVa = NEO::CpuInfo::getInstance().getVirtualAddressSize() == 57u ? maxNBitValue(56) : maxNBitValue(47);
+    }
+    bool reserveOnSvmHeap = pStart == nullptr;
+    if (castToUint64(pStart) <= maxCpuVa) {
+        reserveOnSvmHeap = true;
+    }
+    reserveOnSvmHeap &= NEO::debugManager.flags.EnableReservingInSvmRange.get();
+
+    NEO::AddressRange addressRange{};
+    uint32_t reservedOnRootDeviceIndex = 0;
+    uint64_t reservationBase = 0;
+    size_t reservationTotalSize = 0;
+
+    if (reserveOnSvmHeap) {
+        if (alignUp(size, MemoryConstants::pageSize) != size) {
+            return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
+        }
+        reservationTotalSize = alignUp(size, MemoryConstants::pageSize2M) + MemoryConstants::pageSize2M;
+        addressRange = this->driverHandle->getMemoryManager()->reserveCpuAddressWithZeroBaseRetry(castToUint64(pStart), reservationTotalSize);
+        if (addressRange.address == 0) {
+            return ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+        }
+        DEBUG_BREAK_IF(addressRange.address + reservationTotalSize > maxCpuVa);
+        reservationBase = addressRange.address;
+        addressRange.address = alignUp(addressRange.address, MemoryConstants::pageSize2M);
+        addressRange.size = size;
+    } else {
         NEO::HeapIndex heap;
         size_t pageSize;
         if ((getPageAlignedSizeRequired(size, &heap, &pageSize) != size)) {
             return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
         }
-    NEO::VirtualMemoryReservation *virtualMemoryReservation = new NEO::VirtualMemoryReservation;
+        addressRange = this->driverHandle->getMemoryManager()->reserveGpuAddressOnHeap(castToUint64(pStart), size, this->driverHandle->rootDeviceIndices, &reservedOnRootDeviceIndex, heap, pageSize);
-    virtualMemoryReservation->virtualAddressRange = this->driverHandle->getMemoryManager()->reserveGpuAddressOnHeap(reinterpret_cast<uint64_t>(pStart), size, this->driverHandle->rootDeviceIndices, &virtualMemoryReservation->rootDeviceIndex, heap, pageSize);
+        if (addressRange.address == 0) {
-    if (virtualMemoryReservation->virtualAddressRange.address == 0) {
-        delete virtualMemoryReservation;
             return ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY;
         }
+        reservationBase = addressRange.address;
+        reservationTotalSize = addressRange.size;
+    }
+
+    NEO::VirtualMemoryReservation *virtualMemoryReservation = new NEO::VirtualMemoryReservation;
+    virtualMemoryReservation->virtualAddressRange = addressRange;
+    virtualMemoryReservation->isSvmReservation = reserveOnSvmHeap;
+    virtualMemoryReservation->rootDeviceIndex = reservedOnRootDeviceIndex;
     virtualMemoryReservation->flags.readWrite = false;
     virtualMemoryReservation->flags.readOnly = false;
     virtualMemoryReservation->flags.noAccess = true;
     virtualMemoryReservation->reservationSize = size;
+    virtualMemoryReservation->reservationBase = reservationBase;
+    virtualMemoryReservation->reservationTotalSize = reservationTotalSize;
     auto lock = this->driverHandle->getMemoryManager()->lockVirtualMemoryReservationMap();
     this->driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().insert(std::pair<void *, NEO::VirtualMemoryReservation *>(reinterpret_cast<void *>(virtualMemoryReservation->virtualAddressRange.address), virtualMemoryReservation));
     *pptr = reinterpret_cast<void *>(virtualMemoryReservation->virtualAddressRange.address);
@@ -1082,7 +1122,12 @@ ze_result_t ContextImp::freeVirtualMem(const void *ptr,
         if (virtualMemoryReservation->reservationSize != size) {
             return ZE_RESULT_ERROR_INVALID_ARGUMENT;
         }
-        this->driverHandle->getMemoryManager()->freeGpuAddress(virtualMemoryReservation->virtualAddressRange, virtualMemoryReservation->rootDeviceIndex);
+        NEO::AddressRange addressRange{virtualMemoryReservation->reservationBase, virtualMemoryReservation->reservationTotalSize};
+        if (virtualMemoryReservation->isSvmReservation) {
+            this->driverHandle->getMemoryManager()->freeCpuAddress(addressRange);
+        } else {
+            this->driverHandle->getMemoryManager()->freeGpuAddress(addressRange, virtualMemoryReservation->rootDeviceIndex);
+        }
         delete virtualMemoryReservation;
         this->driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().erase(it);
         virtualMemoryReservation = nullptr;

level_zero/core/test/unit_tests/fixtures/memory_ipc_fixture.h

@@ -212,6 +212,8 @@ class MemoryManagerIpcMock : public NEO::MemoryManager {
         return MemoryConstants::pageSize64k;
     }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
+    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override { return {}; }
+    void freeCpuAddress(AddressRange addressRange) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryWithAlignment(const NEO::AllocationData &allocationData) override { return nullptr; };
@@ -338,6 +340,8 @@ class MemoryManagerIpcImplicitScalingMock : public NEO::MemoryManager {
         return MemoryConstants::pageSize64k;
     }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
+    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override { return {}; }
+    void freeCpuAddress(AddressRange addressRange) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryWithAlignment(const NEO::AllocationData &allocationData) override { return nullptr; };

level_zero/core/test/unit_tests/sources/context/test_context.cpp

@@ -10,9 +10,9 @@
 #include "shared/source/helpers/blit_properties.h"
 #include "shared/source/memory_manager/gfx_partition.h"
 #include "shared/source/os_interface/device_factory.h"
+#include "shared/source/utilities/cpu_info.h"
 #include "shared/test/common/mocks/mock_bindless_heaps_helper.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"
 #include "shared/test/common/mocks/mock_device.h"
 #include "shared/test/common/mocks/mock_graphics_allocation.h"
@@ -28,7 +28,6 @@
 #include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
 #include "level_zero/core/test/unit_tests/fixtures/host_pointer_manager_fixture.h"
 #include "level_zero/core/test/unit_tests/mocks/mock_cmdlist.h"
-#include "level_zero/core/test/unit_tests/mocks/mock_driver_handle.h"
 
 #include "gtest/gtest.h"
 
@@ -1353,6 +1352,8 @@ TEST_F(ContextTest, whenCallingVirtualMemoryGetAttributeWithInvalidValuesThenFai
 }
 
 TEST_F(ContextTest, whenCallingVirtualMemoryFreeWithInvalidValuesThenFailuresReturned) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.EnableReservingInSvmRange.set(true);
     ze_context_handle_t hContext;
     ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
 
@@ -1381,6 +1382,18 @@ TEST_F(ContextTest, whenCallingVirtualMemoryFreeWithInvalidValuesThenFailuresRet
     res = contextImp->freeVirtualMem(ptr, pagesize);
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
 
+    const auto maxCpuVa = NEO::CpuInfo::getInstance().getVirtualAddressSize() == 57u ? maxNBitValue(56) : maxNBitValue(47);
+    pStart = reinterpret_cast<void *>(maxCpuVa + 0x1234);
+    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    EXPECT_GT(static_cast<int>(driverHandle->getMemoryManager()->getVirtualMemoryReservationMap().size()), 0);
+
+    res = contextImp->freeVirtualMem(ptr, invalidSize);
+    EXPECT_EQ(ZE_RESULT_ERROR_INVALID_ARGUMENT, res);
+
+    res = contextImp->freeVirtualMem(ptr, pagesize);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
     res = contextImp->destroy();
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
 }
@@ -1416,6 +1429,13 @@ class ReserveMemoryManagerMock : public NEO::MemoryManager {
         return MemoryConstants::pageSize64k;
     }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
+    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override {
+        if (failReserveCpuAddress) {
+            return {};
+        }
+        return AddressRange{requiredStartAddress, size};
+    }
+    void freeCpuAddress(AddressRange addressRange) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryWithAlignment(const NEO::AllocationData &allocationData) override { return nullptr; };
@@ -1448,16 +1468,23 @@ class ReserveMemoryManagerMock : public NEO::MemoryManager {
     NEO::GraphicsAllocation *allocateMemoryByKMD(const NEO::AllocationData &allocationData) override { return nullptr; };
     void *lockResourceImpl(NEO::GraphicsAllocation &graphicsAllocation) override { return nullptr; };
     void unlockResourceImpl(NEO::GraphicsAllocation &graphicsAllocation) override{};
+    bool peek32bit() override {
+        return this->is32bit;
+    }
 
     bool failReserveGpuAddress = true;
+    bool failReserveCpuAddress = true;
     bool failMapVirtualMemory = true;
     bool failAllocatePhysicalGraphicsMemory = true;
+    bool is32bit = false;
     void *buffer = nullptr;
     size_t size = 0;
     std::unique_ptr<NEO::GraphicsAllocation> mockAllocation;
 };
 
-TEST_F(ContextTest, whenCallingVirtualMemReserveWithPStartWithSuccessfulAllocationThenSuccessReturned) {
+TEST_F(ContextTest, whenCallingVirtualMemReserveWithPStartInSvmRangeWithSuccessfulAllocationThenSuccessReturned) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.EnableReservingInSvmRange.set(true);
     ze_context_handle_t hContext{};
     ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
 
@@ -1471,6 +1498,40 @@ TEST_F(ContextTest, whenCallingVirtualMemReserveWithPStartWithSuccessfulAllocati
     void *ptr = nullptr;
     size_t pagesize = 0u;
 
+    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    auto reserveMemoryManager = std::make_unique<ReserveMemoryManagerMock>(*neoDevice->executionEnvironment);
+    auto memoryManager = driverHandle->getMemoryManager();
+    reserveMemoryManager->failReserveCpuAddress = false;
+    driverHandle->setMemoryManager(reserveMemoryManager.get());
+    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    EXPECT_GT(reserveMemoryManager->getVirtualMemoryReservationMap().size(), 0u);
+    res = contextImp->freeVirtualMem(ptr, pagesize);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    driverHandle->setMemoryManager(memoryManager);
+
+    res = contextImp->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+}
+
+TEST_F(ContextTest, whenCallingVirtualMemReserveWithPStartAboveSvmRangeWithSuccessfulAllocationThenSuccessReturned) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.EnableReservingInSvmRange.set(true);
+    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));
+
+    const auto maxCpuVa = NEO::CpuInfo::getInstance().getVirtualAddressSize() == 57u ? maxNBitValue(56) : maxNBitValue(47);
+    void *pStart = reinterpret_cast<void *>(maxCpuVa + 0x1234);
+    size_t size = 4096u;
+    void *ptr = nullptr;
+    size_t pagesize = 0u;
+
     res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
     auto reserveMemoryManager = std::make_unique<ReserveMemoryManagerMock>(*neoDevice->executionEnvironment);
@@ -1630,6 +1691,8 @@ TEST_F(ContextTest, whenUsingOffsetsIntoReservedVirtualMemoryWithMultiplePhysica
 }
 
 TEST_F(ContextTest, whenCallingVirtualMemoryReservationWhenOutOfMemoryThenOutOfMemoryReturned) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.EnableReservingInSvmRange.set(true);
     ze_context_handle_t hContext;
     ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
 
@@ -1649,9 +1712,13 @@ TEST_F(ContextTest, whenCallingVirtualMemoryReservationWhenOutOfMemoryThenOutOfM
     auto memoryManager = driverHandle->getMemoryManager();
     driverHandle->setMemoryManager(failingReserveMemoryManager);
     res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
-    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);
+    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);
     pStart = reinterpret_cast<void *>(0x1234);
     res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);
+    const auto maxCpuVa = NEO::CpuInfo::getInstance().getVirtualAddressSize() == 57u ? maxNBitValue(56) : maxNBitValue(47);
+    pStart = reinterpret_cast<void *>(maxCpuVa + 0x1234);
+    res = contextImp->reserveVirtualMem(pStart, pagesize, &ptr);
     EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);
     driverHandle->setMemoryManager(memoryManager);
     delete failingReserveMemoryManager;
@@ -1661,6 +1728,8 @@ TEST_F(ContextTest, whenCallingVirtualMemoryReservationWhenOutOfMemoryThenOutOfM
 }
 
 TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithInvalidArgumentsThenUnsupportedSizeReturned) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.EnableReservingInSvmRange.set(true);
     ze_context_handle_t hContext;
     ze_context_desc_t desc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
 
@@ -1682,6 +1751,10 @@ TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithInvalidArgumentsThenU
     driverHandle->setMemoryManager(failingReserveMemoryManager);
     res = contextImp->reserveVirtualMem(pStart, size, &ptr);
     EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);
+    const auto maxCpuVa = NEO::CpuInfo::getInstance().getVirtualAddressSize() == 57u ? maxNBitValue(56) : maxNBitValue(47);
+    pStart = reinterpret_cast<void *>(maxCpuVa + 0x1234);
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_UNSUPPORTED_SIZE, res);
     driverHandle->setMemoryManager(memoryManager);
     delete failingReserveMemoryManager;
 
@@ -1753,6 +1826,207 @@ TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithValidMultiPageSizeInA
     EXPECT_EQ(ZE_RESULT_SUCCESS, res);
 }
 
+TEST_F(ContextTest, whenCallingVirtualMemoryReservationWithOverlappingReservationRangeThenSuccessReturned) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.EnableReservingInSvmRange.set(true);
+    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 = 0;
+    void *ptr = nullptr;
+    size_t pagesize = 0u;
+
+    res = contextImp->queryVirtualMemPageSize(device, size, &pagesize);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    size = 16 * pagesize;
+
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    EXPECT_NE(nullptr, ptr);
+
+    void *newPStart = addrToPtr(castToUint64(ptr) + pagesize);
+    void *newPtr = nullptr;
+    res = contextImp->reserveVirtualMem(newPStart, pagesize, &newPtr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    EXPECT_NE(nullptr, newPtr);
+
+    EXPECT_NE(ptr, newPtr);
+
+    res = contextImp->freeVirtualMem(newPtr, pagesize);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    res = contextImp->freeVirtualMem(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    res = contextImp->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+}
+
+class MockCpuInfoOverrideVirtualAddressSize {
+  public:
+    class MockCpuInfo : public CpuInfo {
+      public:
+        using CpuInfo::cpuFlags;
+        using CpuInfo::virtualAddressSize;
+    } *mockCpuInfo = reinterpret_cast<MockCpuInfo *>(const_cast<CpuInfo *>(&CpuInfo::getInstance()));
+
+    MockCpuInfoOverrideVirtualAddressSize(uint32_t newCpuVirtualAddressSize) {
+        virtualAddressSizeSave = mockCpuInfo->getVirtualAddressSize();
+        mockCpuInfo->virtualAddressSize = newCpuVirtualAddressSize;
+    }
+
+    ~MockCpuInfoOverrideVirtualAddressSize() {
+        mockCpuInfo->virtualAddressSize = virtualAddressSizeSave;
+    }
+
+    uint32_t virtualAddressSizeSave = 0;
+};
+
+TEST_F(ContextTest, Given32BitCpuAddressWidthWhenCallingVirtualMemoryReservationCorrectAllocationMethodIsSelected) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.EnableReservingInSvmRange.set(true);
+    MockCpuInfoOverrideVirtualAddressSize overrideCpuInfo(32);
+
+    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 = addrToPtr(0x1234);
+    size_t size = MemoryConstants::pageSize2M;
+    void *ptr = nullptr;
+
+    auto reserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
+    reserveMemoryManager->is32bit = true;
+    auto memoryManager = driverHandle->getMemoryManager();
+    driverHandle->setMemoryManager(reserveMemoryManager);
+
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);
+    reserveMemoryManager->failReserveCpuAddress = false;
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    res = contextImp->freeVirtualMem(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    pStart = addrToPtr(maxNBitValue(32) + 0x1234);
+
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);
+    reserveMemoryManager->failReserveGpuAddress = false;
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    res = contextImp->freeVirtualMem(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    driverHandle->setMemoryManager(memoryManager);
+    delete reserveMemoryManager;
+
+    res = contextImp->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+}
+
+TEST_F(ContextTest, Given48BitCpuAddressWidthWhenCallingVirtualMemoryReservationCorrectAllocationMethodIsSelected) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.EnableReservingInSvmRange.set(true);
+    MockCpuInfoOverrideVirtualAddressSize overrideCpuInfo(48);
+
+    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 = addrToPtr(0x1234);
+    size_t size = MemoryConstants::pageSize2M;
+    void *ptr = nullptr;
+
+    auto reserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
+    auto memoryManager = driverHandle->getMemoryManager();
+    driverHandle->setMemoryManager(reserveMemoryManager);
+
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);
+    reserveMemoryManager->failReserveCpuAddress = false;
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    res = contextImp->freeVirtualMem(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    pStart = addrToPtr(maxNBitValue(47) + 0x1234);
+
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);
+    reserveMemoryManager->failReserveGpuAddress = false;
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    res = contextImp->freeVirtualMem(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    driverHandle->setMemoryManager(memoryManager);
+    delete reserveMemoryManager;
+
+    res = contextImp->destroy();
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+}
+
+TEST_F(ContextTest, Given57BitCpuAddressWidthWhenCallingVirtualMemoryReservationCorrectAllocationMethodIsSelected) {
+    DebugManagerStateRestore restore;
+    debugManager.flags.EnableReservingInSvmRange.set(true);
+    MockCpuInfoOverrideVirtualAddressSize overrideCpuInfo(57);
+
+    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 = addrToPtr(0x1234);
+    size_t size = MemoryConstants::pageSize2M;
+    void *ptr = nullptr;
+
+    auto reserveMemoryManager = new ReserveMemoryManagerMock(*neoDevice->executionEnvironment);
+    auto memoryManager = driverHandle->getMemoryManager();
+    driverHandle->setMemoryManager(reserveMemoryManager);
+
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY, res);
+    reserveMemoryManager->failReserveCpuAddress = false;
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    res = contextImp->freeVirtualMem(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    pStart = addrToPtr(maxNBitValue(56) + 0x1234);
+
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY, res);
+    reserveMemoryManager->failReserveGpuAddress = false;
+    res = contextImp->reserveVirtualMem(pStart, size, &ptr);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+    res = contextImp->freeVirtualMem(ptr, size);
+    EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+
+    driverHandle->setMemoryManager(memoryManager);
+    delete reserveMemoryManager;
+
+    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};

level_zero/core/test/unit_tests/sources/event/test_event.cpp

@@ -12,7 +12,6 @@
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/helpers/engine_descriptor_helper.h"
 #include "shared/test/common/helpers/variable_backup.h"
-#include "shared/test/common/mocks/mock_compilers.h"
 #include "shared/test/common/mocks/mock_csr.h"
 #include "shared/test/common/mocks/mock_device.h"
 #include "shared/test/common/mocks/mock_graphics_allocation.h"
@@ -89,6 +88,8 @@ class MemoryManagerEventPoolFailMock : public NEO::MemoryManager {
         return MemoryConstants::pageSize64k;
     }
     void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
+    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override { return {}; }
+    void freeCpuAddress(AddressRange addressRange) override{};
     NEO::GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const NEO::AllocationData &allocationData) override { return nullptr; };
     NEO::GraphicsAllocation *allocateGraphicsMemoryWithAlignment(const NEO::AllocationData &allocationData) override { return nullptr; };

shared/source/debug_settings/debug_variables_base.inl

@@ -463,6 +463,7 @@ DECLARE_DEBUG_VARIABLE(bool, EnableFreeMemory, true, "Enable freeMemory in memor
 DECLARE_DEBUG_VARIABLE(bool, ForceSamplerLowFilteringPrecision, false, "Force Low Filtering Precision Sampler mode")
 DECLARE_DEBUG_VARIABLE(bool, EnablePrivateBO, false, "Enable PRELIM_I915_GEM_CREATE_EXT_VM_PRIVATE extension creating VM_PRIVATE BOs")
 DECLARE_DEBUG_VARIABLE(bool, EnableAIL, true, "Enables AIL")
+DECLARE_DEBUG_VARIABLE(bool, EnableReservingInSvmRange, false, "Enables reserving virtual memory in the SVM range")
 DECLARE_DEBUG_VARIABLE(int64_t, VmBindWaitUserFenceTimeout, -1, "-1: default, >0: time in ns for wait function timeout")
 DECLARE_DEBUG_VARIABLE(int32_t, ForceRunAloneContext, -1, "Control creation of run-alone HW context, -1:default, 0:disable, 1:enable")
 DECLARE_DEBUG_VARIABLE(int32_t, AddClGlSharing, -1, "Add cl-gl extension")

shared/source/memory_manager/memory_manager.cpp

@@ -135,6 +135,14 @@ GmmHelper *MemoryManager::getGmmHelper(uint32_t rootDeviceIndex) {
     return executionEnvironment.rootDeviceEnvironments[rootDeviceIndex]->getGmmHelper();
 }
 
+AddressRange MemoryManager::reserveCpuAddressWithZeroBaseRetry(const uint64_t requiredStartAddress, size_t size) {
+    auto addressRange = reserveCpuAddress(requiredStartAddress, size);
+    if ((addressRange.address == 0) && (requiredStartAddress != 0)) {
+        addressRange = reserveCpuAddress(0, size);
+    }
+    return addressRange;
+}
+
 HeapIndex MemoryManager::selectInternalHeap(bool useLocalMemory) {
     return useLocalMemory ? HeapIndex::heapInternalDeviceMemory : HeapIndex::heapInternal;
 }

shared/source/memory_manager/memory_manager.h

@@ -70,7 +70,10 @@ struct VirtualMemoryReservation {
     MemoryFlags flags;
     std::map<void *, MemoryMappedRange *> mappedAllocations;
     uint32_t rootDeviceIndex;
+    bool isSvmReservation;
     size_t reservationSize;
+    uint64_t reservationBase;
+    size_t reservationTotalSize;
 };
 
 struct PhysicalMemoryAllocation {
@@ -248,6 +251,9 @@ class MemoryManager {
     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;
+    virtual AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) = 0;
+    AddressRange reserveCpuAddressWithZeroBaseRetry(const uint64_t requiredStartAddress, size_t size);
+    virtual void freeCpuAddress(AddressRange addressRange) = 0;
     static HeapIndex selectInternalHeap(bool useLocalMemory);
     static HeapIndex selectExternalHeap(bool useLocalMemory);
 

shared/source/memory_manager/os_agnostic_memory_manager.cpp

@@ -555,6 +555,14 @@ void OsAgnosticMemoryManager::releaseReservedCpuAddressRange(void *reserved, siz
     alignedFreeWrapper(reserved);
 }
 
+AddressRange OsAgnosticMemoryManager::reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) {
+    return {castToUint64(alignedMallocWrapper(size, MemoryConstants::pageSize)), size};
+}
+
+void OsAgnosticMemoryManager::freeCpuAddress(AddressRange addressRange) {
+    alignedFreeWrapper(addrToPtr(addressRange.address));
+}
+
 MemoryAllocation *OsAgnosticMemoryManager::createMemoryAllocation(AllocationType allocationType, void *driverAllocatedCpuPointer,
                                                                   void *pMem, uint64_t gpuAddress, size_t memSize, uint64_t count,
                                                                   MemoryPool pool, uint32_t rootDeviceIndex, bool uncacheable,

shared/source/memory_manager/os_agnostic_memory_manager.h

@@ -48,6 +48,8 @@ class OsAgnosticMemoryManager : public MemoryManager {
     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;
+    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override;
+    void freeCpuAddress(AddressRange addressRange) override;
     bool is64kbPagesEnabled(const HardwareInfo *hwInfo);
 
   protected:

shared/source/os_interface/linux/drm_memory_manager.cpp

@@ -27,8 +27,6 @@
 #include "shared/source/memory_manager/allocation_properties.h"
 #include "shared/source/memory_manager/gfx_partition.h"
 #include "shared/source/memory_manager/host_ptr_manager.h"
-#include "shared/source/memory_manager/local_memory_usage.h"
-#include "shared/source/memory_manager/memory_banks.h"
 #include "shared/source/memory_manager/memory_pool.h"
 #include "shared/source/memory_manager/multi_graphics_allocation.h"
 #include "shared/source/memory_manager/residency.h"
@@ -47,7 +45,6 @@
 #include "shared/source/os_interface/product_helper.h"
 
 #include <cstring>
-#include <iostream>
 #include <memory>
 #include <sys/ioctl.h>
 
@@ -95,6 +92,7 @@ DrmMemoryManager::DrmMemoryManager(GemCloseWorkerMode mode,
         const auto heapIndex = customAlignment >= MemoryConstants::pageSize2M ? HeapIndex::heapStandard2MB : HeapIndex::heapStandard64KB;
         alignmentSelector.addCandidateAlignment(customAlignment, true, AlignmentSelector::anyWastage, heapIndex);
     }
+    osMemory = OSMemory::create();
 
     initialize(mode);
 }
@@ -1527,6 +1525,18 @@ void DrmMemoryManager::freeGpuAddress(AddressRange addressRange, uint32_t rootDe
     releaseGpuRange(reinterpret_cast<void *>(addressRange.address), addressRange.size, rootDeviceIndex);
 }
 
+AddressRange DrmMemoryManager::reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) {
+    void *ptr = osMemory->osReserveCpuAddressRange(addrToPtr(requiredStartAddress), size, false);
+    if (ptr == MAP_FAILED) {
+        ptr = nullptr;
+    }
+    return {castToUint64(ptr), size};
+}
+
+void DrmMemoryManager::freeCpuAddress(AddressRange addressRange) {
+    osMemory->osReleaseCpuAddressRange(addrToPtr(addressRange.address), addressRange.size);
+}
+
 std::unique_lock<std::mutex> DrmMemoryManager::acquireAllocLock() {
     return std::unique_lock<std::mutex>(this->allocMutex);
 }

shared/source/os_interface/linux/drm_memory_manager.h

@@ -9,8 +9,8 @@
 #include "shared/source/command_stream/submission_status.h"
 #include "shared/source/memory_manager/memory_manager.h"
 #include "shared/source/os_interface/linux/drm_buffer_object.h"
+#include "shared/source/os_interface/os_memory.h"
 
-#include <limits>
 #include <map>
 #include <sys/mman.h>
 #include <unistd.h>
@@ -77,6 +77,8 @@ class DrmMemoryManager : public MemoryManager {
     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;
+    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override;
+    void freeCpuAddress(AddressRange addressRange) override;
     MOCKABLE_VIRTUAL BufferObject *createBufferObjectInMemoryRegion(uint32_t rootDeviceIndex, Gmm *gmm, AllocationType allocationType, uint64_t gpuAddress, size_t size,
                                                                     DeviceBitfield memoryBanks, size_t maxOsContextCount, int32_t pairHandle, bool isSystemMemoryPool, bool isUsmHostAllocation);
 
@@ -186,6 +188,7 @@ class DrmMemoryManager : public MemoryManager {
     bool forcePinEnabled = false;
     const bool validateHostPtrMemory;
     std::unique_ptr<DrmGemCloseWorker> gemCloseWorker;
+    std::unique_ptr<OSMemory> osMemory;
     decltype(&mmap) mmapFunction = mmap;
     decltype(&munmap) munmapFunction = munmap;
     decltype(&close) closeFunction = close;

shared/source/os_interface/windows/os_memory_win.cpp

@@ -25,7 +25,7 @@ void *OSMemoryWindows::osReserveCpuAddressRange(void *baseAddress, size_t sizeTo
 }
 
 void OSMemoryWindows::osReleaseCpuAddressRange(void *reservedCpuAddressRange, size_t size) {
-    virtualFreeWrapper(reservedCpuAddressRange, size, MEM_RELEASE);
+    virtualFreeWrapper(reservedCpuAddressRange, 0, MEM_RELEASE);
 }
 
 LPVOID OSMemoryWindows::virtualAllocWrapper(LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWORD flProtect) {

shared/source/os_interface/windows/wddm_memory_manager.cpp

@@ -40,7 +40,6 @@
 #include "shared/source/os_interface/windows/wddm_allocation.h"
 #include "shared/source/os_interface/windows/wddm_residency_allocations_container.h"
 #include "shared/source/os_interface/windows/wddm_residency_controller.h"
-#include "shared/source/release_helper/release_helper.h"
 
 #include <algorithm>
 #include <emmintrin.h>
@@ -74,6 +73,7 @@ WddmMemoryManager::WddmMemoryManager(ExecutionEnvironment &executionEnvironment)
     if (customAlignment > 0) {
         alignmentSelector.addCandidateAlignment(customAlignment, false, AlignmentSelector::anyWastage);
     }
+    osMemory = OSMemory::create();
 
     initialized = true;
 }
@@ -1015,6 +1015,15 @@ void WddmMemoryManager::freeGpuAddress(AddressRange addressRange, uint32_t rootD
     getWddm(rootDeviceIndex).freeGpuVirtualAddress(addressRange.address, addressRange.size);
 }
 
+AddressRange WddmMemoryManager::reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) {
+    void *ptr = osMemory->osReserveCpuAddressRange(addrToPtr(requiredStartAddress), size, false);
+    return {castToUint64(ptr), size};
+}
+
+void WddmMemoryManager::freeCpuAddress(AddressRange addressRange) {
+    osMemory->osReleaseCpuAddressRange(addrToPtr(addressRange.address), addressRange.size);
+}
+
 bool WddmMemoryManager::mapGpuVaForOneHandleAllocation(WddmAllocation *allocation, const void *preferredGpuVirtualAddress) {
     D3DGPU_VIRTUAL_ADDRESS addressToMap = castToUint64(preferredGpuVirtualAddress);
     auto heapIndex = selectHeap(allocation, preferredGpuVirtualAddress != nullptr, is32bit || executionEnvironment.rootDeviceEnvironments[allocation->getRootDeviceIndex()]->isFullRangeSvm(), allocation->isAllocInFrontWindowPool());

shared/source/os_interface/windows/wddm_memory_manager.h

@@ -8,6 +8,7 @@
 #pragma once
 #include "shared/source/memory_manager/memory_manager.h"
 #include "shared/source/os_interface/os_context.h"
+#include "shared/source/os_interface/os_memory.h"
 #include "shared/source/os_interface/windows/wddm_allocation.h"
 
 #include <d3dkmthk.h>
@@ -67,6 +68,8 @@ class WddmMemoryManager : public MemoryManager {
     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;
+    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override;
+    void freeCpuAddress(AddressRange addressRange) override;
     bool verifyHandle(osHandle handle, uint32_t rootDeviceIndex, bool ntHandle) override;
     bool isNTHandle(osHandle handle, uint32_t rootDeviceIndex) override;
     void releaseDeviceSpecificMemResources(uint32_t rootDeviceIndex) override{};
@@ -115,6 +118,7 @@ class WddmMemoryManager : public MemoryManager {
     void adjustGpuPtrToHostAddressSpace(WddmAllocation &wddmAllocation, void *&requiredGpuVa);
     bool isStatelessAccessRequired(AllocationType type);
     AlignedMallocRestrictions mallocRestrictions;
+    std::unique_ptr<OSMemory> osMemory;
 
     Wddm &getWddm(uint32_t rootDeviceIndex) const;
 };

shared/test/common/mocks/mock_memory_manager.h

@@ -134,6 +134,17 @@ class MockMemoryManager : public MemoryManagerCreate<OsAgnosticMemoryManager> {
         return OsAgnosticMemoryManager::reserveCpuAddressRange(size, rootDeviceIndex);
     }
 
+    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override {
+        if (failReserveAddress) {
+            return {};
+        }
+        return OsAgnosticMemoryManager::reserveCpuAddress(requiredStartAddress, size);
+    }
+
+    void freeCpuAddress(AddressRange addressRange) override {
+        return OsAgnosticMemoryManager::freeCpuAddress(addressRange);
+    }
+
     void *createMultiGraphicsAllocationInSystemMemoryPool(RootDeviceIndicesContainer &rootDeviceIndices,
                                                           AllocationProperties &properties,
                                                           MultiGraphicsAllocation &multiGraphicsAllocation) override {

shared/test/common/test_files/igdrcl.config

@@ -495,6 +495,7 @@ OverridePlatformName = unk
 WddmResidencyLoggerOutputDirectory = unk
 ToggleBitIn57GpuVa = unk
 EnablePrivateBO = 0
+EnableReservingInSvmRange = 0
 ExperimentalEnableDeviceAllocationCache = -1
 OverrideL1CachePolicyInSurfaceStateAndStateless = -1
 EnableBcsSwControlWa = -1

shared/test/unit_test/device/neo_device_tests.cpp

@@ -9,7 +9,6 @@
 #include "shared/source/gmm_helper/gmm.h"
 #include "shared/source/helpers/api_specific_config.h"
 #include "shared/source/helpers/array_count.h"
-#include "shared/source/helpers/driver_model_type.h"
 #include "shared/source/helpers/gfx_core_helper.h"
 #include "shared/source/helpers/ray_tracing_helper.h"
 #include "shared/source/memory_manager/allocations_list.h"
@@ -19,7 +18,6 @@
 #include "shared/source/os_interface/driver_info.h"
 #include "shared/source/os_interface/os_context.h"
 #include "shared/source/os_interface/os_interface.h"
-#include "shared/source/os_interface/product_helper_hw.h"
 #include "shared/source/release_helper/release_helper.h"
 #include "shared/test/common/fixtures/device_fixture.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
@@ -539,6 +537,8 @@ TEST_F(DeviceGetCapsTest, givenFlagEnabled64kbPagesWhenCallConstructorMemoryMana
             return MemoryConstants::pageSize64k;
         }
         void freeGpuAddress(AddressRange addressRange, uint32_t rootDeviceIndex) override{};
+        AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override { return {}; }
+        void freeCpuAddress(AddressRange addressRange) override{};
         GraphicsAllocation *createGraphicsAllocation(OsHandleStorage &handleStorage, const AllocationData &allocationData) override { return nullptr; };
         GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(const AllocationData &allocationData) override { return nullptr; };
         GraphicsAllocation *allocateGraphicsMemoryWithAlignment(const AllocationData &allocationData) override { return nullptr; };

shared/test/unit_test/memory_manager/memory_manager_tests.cpp

@@ -1909,6 +1909,16 @@ TEST(OsAgnosticMemoryManager, givenOsAgnosticMemoryManagerWithFlagEnable64kbpage
     EXPECT_FALSE(memoryManager.is64kbPagesEnabled(&hwInfo));
 }
 
+TEST(OsAgnosticMemoryManager, givenStartAddressAndSizeWhenReservingCpuAddressThenPageAlignedAddressRangeIsReturned) {
+    MockExecutionEnvironment executionEnvironment;
+    OsAgnosticMemoryManager memoryManager(executionEnvironment);
+    auto addressRange = memoryManager.reserveCpuAddress(0, 1234);
+    EXPECT_NE(0u, addressRange.address);
+    EXPECT_EQ(1234u, addressRange.size);
+    EXPECT_EQ(0u, addressRange.address & (MemoryConstants::pageSize - 1));
+    memoryManager.freeCpuAddress(addressRange);
+}
+
 TEST(MemoryManager, givenSharedResourceCopyWhenAllocatingGraphicsMemoryThenAllocateGraphicsMemoryForImageIsCalled) {
     MockExecutionEnvironment executionEnvironment{};
     MockMemoryManager memoryManager(false, true, executionEnvironment);
@@ -2750,6 +2760,83 @@ TEST(MemoryManagerTest, whenMemoryManagerReturnsNullptrThenAllocateGlobalsSurfac
     EXPECT_EQ(deviceBitfield, memoryManager->recentlyPassedDeviceBitfield);
 }
 
+class FailFirstCpuReserveMemoryManager : public MockMemoryManager {
+  public:
+    using MockMemoryManager::MockMemoryManager;
+    AddressRange reserveCpuAddress(const uint64_t requiredStartAddress, size_t size) override {
+        cpuReservationCallCount++;
+        if (isFirstCpuReservationCall || alwaysFail) {
+            isFirstCpuReservationCall = false;
+            return {};
+        }
+        return AddressRange{0xDEADBEEF, size};
+    }
+    void freeCpuAddress(AddressRange addressRange) override{};
+
+    bool alwaysFail = false;
+    bool isFirstCpuReservationCall = true;
+    int8_t cpuReservationCallCount = 0;
+};
+
+TEST(MemoryManagerTest, givenFirstCpuReservationFailsAndRequiredStartAddressIsZeroThenReservationIsNotTriedAgain) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    auto failFirstMemoryManager = std::make_unique<FailFirstCpuReserveMemoryManager>(executionEnvironment);
+
+    EXPECT_EQ(0, failFirstMemoryManager->cpuReservationCallCount);
+    auto addressRange = failFirstMemoryManager->reserveCpuAddressWithZeroBaseRetry(0, 1234);
+    EXPECT_EQ(0u, addressRange.address);
+    EXPECT_EQ(0u, addressRange.size);
+    EXPECT_EQ(1, failFirstMemoryManager->cpuReservationCallCount);
+}
+
+TEST(MemoryManagerTest, givenFirstCpuReservationFailsAndRequiredStartAddressIsNotZeroThenReservationIsTriedAgain) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    auto failFirstMemoryManager = std::make_unique<FailFirstCpuReserveMemoryManager>(executionEnvironment);
+
+    EXPECT_EQ(0, failFirstMemoryManager->cpuReservationCallCount);
+    auto addressRange = failFirstMemoryManager->reserveCpuAddressWithZeroBaseRetry(42, 1234);
+    EXPECT_EQ(0xDEADBEEF, addressRange.address);
+    EXPECT_EQ(1234u, addressRange.size);
+    EXPECT_EQ(2, failFirstMemoryManager->cpuReservationCallCount);
+    failFirstMemoryManager->freeCpuAddress(addressRange);
+}
+
+TEST(MemoryManagerTest, givenCpuReservationFailsThenReservationIsTriedAgain) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    auto failFirstMemoryManager = std::make_unique<FailFirstCpuReserveMemoryManager>(executionEnvironment);
+
+    failFirstMemoryManager->alwaysFail = true;
+
+    EXPECT_EQ(0, failFirstMemoryManager->cpuReservationCallCount);
+    auto addressRangeZeroStart = failFirstMemoryManager->reserveCpuAddressWithZeroBaseRetry(0, 1234);
+    EXPECT_EQ(0u, addressRangeZeroStart.address);
+    EXPECT_EQ(0u, addressRangeZeroStart.size);
+    EXPECT_EQ(1, failFirstMemoryManager->cpuReservationCallCount);
+    auto addressRangeNonZeroStart = failFirstMemoryManager->reserveCpuAddressWithZeroBaseRetry(42, 1234);
+    EXPECT_EQ(0u, addressRangeNonZeroStart.address);
+    EXPECT_EQ(0u, addressRangeNonZeroStart.size);
+    EXPECT_EQ(3, failFirstMemoryManager->cpuReservationCallCount);
+}
+
+TEST(MemoryManagerTest, givenCpuReservationPassesThenReservationIsNotTriedAgain) {
+    MockExecutionEnvironment executionEnvironment(defaultHwInfo.get());
+    auto failFirstMemoryManager = std::make_unique<FailFirstCpuReserveMemoryManager>(executionEnvironment);
+
+    failFirstMemoryManager->isFirstCpuReservationCall = false;
+
+    EXPECT_EQ(0, failFirstMemoryManager->cpuReservationCallCount);
+    auto addressRangeZeroStart = failFirstMemoryManager->reserveCpuAddressWithZeroBaseRetry(0, 1234);
+    EXPECT_EQ(0xDEADBEEF, addressRangeZeroStart.address);
+    EXPECT_EQ(1234u, addressRangeZeroStart.size);
+    EXPECT_EQ(1, failFirstMemoryManager->cpuReservationCallCount);
+    failFirstMemoryManager->freeCpuAddress(addressRangeZeroStart);
+    auto addressRangeNonZeroStart = failFirstMemoryManager->reserveCpuAddressWithZeroBaseRetry(42, 1234);
+    EXPECT_EQ(0xDEADBEEF, addressRangeNonZeroStart.address);
+    EXPECT_EQ(1234u, addressRangeNonZeroStart.size);
+    EXPECT_EQ(2, failFirstMemoryManager->cpuReservationCallCount);
+    failFirstMemoryManager->freeCpuAddress(addressRangeNonZeroStart);
+}
+
 HWTEST_F(MemoryAllocatorTest, givenMemoryManagerWhenEnableHostPtrTrackingFlagIsSetTo0ThenHostPointerTrackingIsDisabled) {
     DebugManagerStateRestore dbgRestore;
     debugManager.flags.EnableHostPtrTracking.set(0);

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

@@ -17,7 +17,6 @@
 #include "shared/source/os_interface/linux/drm_memory_operations_handler.h"
 #include "shared/source/os_interface/linux/i915.h"
 #include "shared/source/os_interface/linux/os_context_linux.h"
-#include "shared/test/common/fixtures/memory_allocator_multi_device_fixture.h"
 #include "shared/test/common/helpers/engine_descriptor_helper.h"
 #include "shared/test/common/helpers/gtest_helpers.h"
 #include "shared/test/common/mocks/linux/mock_drm_allocation.h"
@@ -46,6 +45,12 @@
 #include <memory>
 #include <vector>
 
+namespace NEO {
+namespace SysCalls {
+extern bool failMmap;
+} // namespace SysCalls
+} // namespace NEO
+
 namespace {
 using DrmMemoryManagerTest = Test<DrmMemoryManagerFixture>;
 using DrmMemoryManagerWithLocalMemoryTest = Test<DrmMemoryManagerWithLocalMemoryFixture>;
@@ -7836,6 +7841,31 @@ TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenGp
     memoryManager->freeGpuAddress(addressRange, 1);
 }
 
+TEST_F(DrmMemoryManagerWithExplicitExpectationsTest, givenDrmMemoryManagerWhenCpuAddressReservationIsAttemptedThenCorrectAddressRangesAreReturned) {
+    auto memoryManager = std::make_unique<TestedDrmMemoryManager>(false, true, false, *executionEnvironment);
+    RootDeviceIndicesContainer rootDeviceIndices;
+    rootDeviceIndices.pushUnique(1);
+
+    auto addressRange = memoryManager->reserveCpuAddress(0, 0);
+    EXPECT_EQ(0u, addressRange.address);
+    EXPECT_EQ(0u, addressRange.size);
+
+    addressRange = memoryManager->reserveCpuAddress(0, MemoryConstants::pageSize);
+    EXPECT_NE(0u, addressRange.address);
+    EXPECT_EQ(MemoryConstants::pageSize, addressRange.size);
+    memoryManager->freeCpuAddress(addressRange);
+
+    addressRange = memoryManager->reserveCpuAddress(MemoryConstants::pageSize * 1234, MemoryConstants::pageSize);
+    EXPECT_NE(0u, addressRange.address);
+    EXPECT_EQ(MemoryConstants::pageSize, addressRange.size);
+    memoryManager->freeCpuAddress(addressRange);
+
+    VariableBackup<bool> backup(&SysCalls::failMmap, true);
+    addressRange = memoryManager->reserveCpuAddress(0, 0);
+    EXPECT_EQ(0u, addressRange.address);
+    EXPECT_EQ(0u, addressRange.size);
+}
+
 TEST_F(DrmMemoryManagerTest, given57bAddressSpaceCpuAndGpuWhenAllocatingHostUSMThenAddressFromExtendedHeapIsPassedAsHintAndSetAsGpuAddressAndReservedAddress) {
     if (defaultHwInfo->capabilityTable.gpuAddressSpace < maxNBitValue(57)) {
         GTEST_SKIP();

shared/test/unit_test/os_interface/windows/wddm_memory_manager_tests.cpp

@@ -14,7 +14,6 @@
 #include "shared/source/os_interface/windows/dxgi_wrapper.h"
 #include "shared/source/os_interface/windows/wddm/um_km_data_translator.h"
 #include "shared/source/os_interface/windows/windows_wrapper.h"
-#include "shared/source/release_helper/release_helper.h"
 #include "shared/source/utilities/tag_allocator.h"
 #include "shared/test/common/helpers/debug_manager_state_restore.h"
 #include "shared/test/common/helpers/execution_environment_helper.h"
@@ -29,7 +28,6 @@
 #include "shared/test/common/mocks/mock_gmm_page_table_mngr.h"
 #include "shared/test/common/mocks/mock_gmm_resource_info.h"
 #include "shared/test/common/mocks/mock_memory_manager.h"
-#include "shared/test/common/mocks/mock_os_context.h"
 #include "shared/test/common/mocks/mock_product_helper.h"
 #include "shared/test/common/mocks/mock_release_helper.h"
 #include "shared/test/common/mocks/windows/mock_wddm_allocation.h"
@@ -1318,6 +1316,13 @@ TEST_F(WddmMemoryManagerSimpleTest, givenWddmMemoryManagerWhenGpuAddressIsReserv
     memoryManager->freeGpuAddress(addressRange, 0);
 }
 
+TEST_F(WddmMemoryManagerSimpleTest, givenWddmMemoryManagerWhenCpuAddressIsReservedAndFreedThenAddressRangeIsNonZero) {
+    auto addressRange = memoryManager->reserveCpuAddress(0, 1234);
+    EXPECT_NE(0u, addressRange.address);
+    EXPECT_EQ(1234u, addressRange.size);
+    memoryManager->freeCpuAddress(addressRange);
+}
+
 TEST_F(WddmMemoryManagerSimpleTest, givenWddmMemoryManagerWhenAllocatingWithGpuVaThenNullptrIsReturned) {
     AllocationData allocationData;