diff --git a/runtime/memory_manager/gfx_partition.h b/runtime/memory_manager/gfx_partition.h
index cb40e5cd12..592120fbd9 100644
--- a/runtime/memory_manager/gfx_partition.h
+++ b/runtime/memory_manager/gfx_partition.h
@@ -31,7 +31,7 @@ constexpr auto internalHeapIndex = is32bit ? HeapIndex::HEAP_INTERNAL : HeapInde
 class GfxPartition {
   public:
     GfxPartition() {}
-    ~GfxPartition();
+    MOCKABLE_VIRTUAL ~GfxPartition();
 
     void init(uint64_t gpuAddressSpace, size_t cpuAddressRangeSizeToReserve);
 
@@ -47,7 +47,7 @@ class GfxPartition {
         getHeap(heapIndex).free(ptr, size);
     }
 
-    void freeGpuAddressRange(uint64_t ptr, size_t size);
+    MOCKABLE_VIRTUAL void freeGpuAddressRange(uint64_t ptr, size_t size);
 
     uint64_t getHeapBase(HeapIndex heapIndex) {
         return getHeap(heapIndex).getBase();
diff --git a/runtime/memory_manager/memory_manager.cpp b/runtime/memory_manager/memory_manager.cpp
index afa7b3e266..de77df01b7 100644
--- a/runtime/memory_manager/memory_manager.cpp
+++ b/runtime/memory_manager/memory_manager.cpp
@@ -42,6 +42,7 @@ MemoryManager::MemoryManager(ExecutionEnvironment &executionEnvironment) : execu
         this->enable64kbpages = DebugManager.flags.Enable64kbpages.get() != 0;
     }
     localMemoryUsageBankSelector.reset(new LocalMemoryUsageBankSelector(getBanksCount()));
+    gfxPartition = std::make_unique<GfxPartition>();
 }
 
 MemoryManager::~MemoryManager() {
diff --git a/runtime/memory_manager/memory_manager.h b/runtime/memory_manager/memory_manager.h
index a11f5e9dfd..b9abd9a34c 100644
--- a/runtime/memory_manager/memory_manager.h
+++ b/runtime/memory_manager/memory_manager.h
@@ -106,10 +106,10 @@ class MemoryManager {
     virtual uint64_t getLocalMemorySize() = 0;
 
     uint64_t getMaxApplicationAddress() { return is64bit ? MemoryConstants::max64BitAppAddress : MemoryConstants::max32BitAppAddress; };
-    uint64_t getInternalHeapBaseAddress() { return gfxPartition.getHeapBase(internalHeapIndex); }
-    uint64_t getExternalHeapBaseAddress() { return gfxPartition.getHeapBase(HeapIndex::HEAP_EXTERNAL); }
+    uint64_t getInternalHeapBaseAddress() { return gfxPartition->getHeapBase(internalHeapIndex); }
+    uint64_t getExternalHeapBaseAddress() { return gfxPartition->getHeapBase(HeapIndex::HEAP_EXTERNAL); }
 
-    bool isLimitedRange() { return gfxPartition.isLimitedRange(); }
+    bool isLimitedRange() { return gfxPartition->isLimitedRange(); }
 
     bool peek64kbPagesEnabled() const { return enable64kbpages; }
     bool peekForce32BitAllocations() const { return force32bitAllocations; }
@@ -222,7 +222,7 @@ class MemoryManager {
     uint32_t latestContextId = std::numeric_limits<uint32_t>::max();
     uint32_t defaultEngineIndex = 0;
     std::unique_ptr<DeferredDeleter> multiContextResourceDestructor;
-    GfxPartition gfxPartition;
+    std::unique_ptr<GfxPartition> gfxPartition;
     std::unique_ptr<LocalMemoryUsageBankSelector> localMemoryUsageBankSelector;
 };
 
diff --git a/runtime/memory_manager/os_agnostic_memory_manager.cpp b/runtime/memory_manager/os_agnostic_memory_manager.cpp
index b6612afb56..0eb81103eb 100644
--- a/runtime/memory_manager/os_agnostic_memory_manager.cpp
+++ b/runtime/memory_manager/os_agnostic_memory_manager.cpp
@@ -29,7 +29,7 @@ OsAgnosticMemoryManager::OsAgnosticMemoryManager(bool aubUsage, ExecutionEnviron
 
     // 4 x sizeof(Heap32) + 2 x sizeof(Standard/Standard64k)
     size_t reservedCpuAddressRangeSize = is64bit ? (4 * 4 + 2 * (aubUsage ? 32 : 4)) * GB : 0;
-    gfxPartition.init(gpuAddressSpace, reservedCpuAddressRangeSize);
+    gfxPartition->init(gpuAddressSpace, reservedCpuAddressRangeSize);
 }
 
 OsAgnosticMemoryManager::~OsAgnosticMemoryManager() {
@@ -104,7 +104,7 @@ GraphicsAllocation *OsAgnosticMemoryManager::allocate32BitGraphicsMemoryImpl(con
     auto heap = useInternal32BitAllocator(allocationData.type) ? internalHeapIndex : HeapIndex::HEAP_EXTERNAL;
     if (allocationData.hostPtr) {
         auto allocationSize = alignSizeWholePage(allocationData.hostPtr, allocationData.size);
-        auto gpuVirtualAddress = gfxPartition.heapAllocate(heap, allocationSize);
+        auto gpuVirtualAddress = gfxPartition->heapAllocate(heap, allocationSize);
         if (!gpuVirtualAddress) {
             return nullptr;
         }
@@ -114,7 +114,7 @@ GraphicsAllocation *OsAgnosticMemoryManager::allocate32BitGraphicsMemoryImpl(con
             allocationData.size, counter, MemoryPool::System4KBPagesWith32BitGpuAddressing, false, false, false);
 
         memAlloc->set32BitAllocation(true);
-        memAlloc->setGpuBaseAddress(GmmHelper::canonize(gfxPartition.getHeapBase(heap)));
+        memAlloc->setGpuBaseAddress(GmmHelper::canonize(gfxPartition->getHeapBase(heap)));
         memAlloc->sizeToFree = allocationSize;
 
         counter++;
@@ -123,7 +123,7 @@ GraphicsAllocation *OsAgnosticMemoryManager::allocate32BitGraphicsMemoryImpl(con
 
     auto allocationSize = alignUp(allocationData.size, MemoryConstants::pageSize);
     void *ptrAlloc = nullptr;
-    auto gpuAddress = gfxPartition.heapAllocate(heap, allocationSize);
+    auto gpuAddress = gfxPartition->heapAllocate(heap, allocationSize);
 
     if (allocationData.size < 0xfffff000) {
         if (fakeBigAllocations) {
@@ -140,7 +140,7 @@ GraphicsAllocation *OsAgnosticMemoryManager::allocate32BitGraphicsMemoryImpl(con
                                                 false, false);
 
         memoryAllocation->set32BitAllocation(true);
-        memoryAllocation->setGpuBaseAddress(GmmHelper::canonize(gfxPartition.getHeapBase(heap)));
+        memoryAllocation->setGpuBaseAddress(GmmHelper::canonize(gfxPartition->getHeapBase(heap)));
         memoryAllocation->sizeToFree = allocationSize;
     }
     counter++;
@@ -206,7 +206,7 @@ void OsAgnosticMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllo
 
     if (sizeToFree) {
         auto gpuAddressToFree = GmmHelper::decanonize(memoryAllocation->getGpuAddress()) & ~MemoryConstants::pageMask;
-        gfxPartition.freeGpuAddressRange(gpuAddressToFree, sizeToFree);
+        gfxPartition->freeGpuAddressRange(gpuAddressToFree, sizeToFree);
     }
 
     alignedFreeWrapper(gfxAllocation->getDriverAllocatedCpuPtr());
@@ -312,12 +312,12 @@ MemoryAllocation *OsAgnosticMemoryManager::createMemoryAllocation(GraphicsAlloca
 
     auto heap = (force32bitAllocations || requireSpecificBitness) ? HeapIndex::HEAP_EXTERNAL : HeapIndex::HEAP_STANDARD;
 
-    uint64_t limitedGpuAddress = gfxPartition.heapAllocate(heap, alignedSize);
+    uint64_t limitedGpuAddress = gfxPartition->heapAllocate(heap, alignedSize);
 
     auto memoryAllocation = new MemoryAllocation(allocationType, driverAllocatedCpuPointer, pMem, limitedGpuAddress, memSize,
                                                  count, pool, multiOsContextCapable, uncacheable, flushL3Required);
 
-    memoryAllocation->setGpuBaseAddress(GmmHelper::canonize(gfxPartition.getHeapBase(heap)));
+    memoryAllocation->setGpuBaseAddress(GmmHelper::canonize(gfxPartition->getHeapBase(heap)));
     memoryAllocation->sizeToFree = alignedSize;
 
     return memoryAllocation;
diff --git a/runtime/os_interface/linux/drm_memory_manager.cpp b/runtime/os_interface/linux/drm_memory_manager.cpp
index 1440fb632f..882d3ea4be 100644
--- a/runtime/os_interface/linux/drm_memory_manager.cpp
+++ b/runtime/os_interface/linux/drm_memory_manager.cpp
@@ -39,7 +39,7 @@ DrmMemoryManager::DrmMemoryManager(gemCloseWorkerMode mode,
                                                                                  forcePinEnabled(forcePinAllowed),
                                                                                  validateHostPtrMemory(validateHostPtrMemory) {
     supportsMultiStorageResources = false;
-    gfxPartition.init(platformDevices[0]->capabilityTable.gpuAddressSpace, getSizeToReserve());
+    gfxPartition->init(platformDevices[0]->capabilityTable.gpuAddressSpace, getSizeToReserve());
     MemoryManager::virtualPaddingAvailable = true;
     if (mode != gemCloseWorkerMode::gemCloseWorkerInactive) {
         gemCloseWorker.reset(new DrmGemCloseWorker(*this));
@@ -130,12 +130,12 @@ uint32_t DrmMemoryManager::unreference(NEO::BufferObject *bo, bool synchronousDe
 uint64_t DrmMemoryManager::acquireGpuRange(size_t &size, StorageAllocatorType &storageType, bool specificBitness) {
     if (specificBitness && this->force32bitAllocations) {
         storageType = BIT32_ALLOCATOR_EXTERNAL;
-        return GmmHelper::canonize(gfxPartition.heapAllocate(HeapIndex::HEAP_EXTERNAL, size));
+        return GmmHelper::canonize(gfxPartition->heapAllocate(HeapIndex::HEAP_EXTERNAL, size));
     }
 
     if (isLimitedRange()) {
         storageType = INTERNAL_ALLOCATOR_WITH_DYNAMIC_BITRANGE;
-        return GmmHelper::canonize(gfxPartition.heapAllocate(HeapIndex::HEAP_STANDARD, size));
+        return GmmHelper::canonize(gfxPartition->heapAllocate(HeapIndex::HEAP_STANDARD, size));
     }
 
     storageType = MMAP_ALLOCATOR;
@@ -150,7 +150,7 @@ void DrmMemoryManager::releaseGpuRange(void *address, size_t unmapSize, StorageA
 
     uint64_t graphicsAddress = static_cast<uint64_t>(reinterpret_cast<uintptr_t>(address));
     graphicsAddress = GmmHelper::decanonize(graphicsAddress);
-    gfxPartition.freeGpuAddressRange(graphicsAddress, unmapSize);
+    gfxPartition->freeGpuAddressRange(graphicsAddress, unmapSize);
 }
 
 NEO::BufferObject *DrmMemoryManager::allocUserptr(uintptr_t address, size_t size, uint64_t flags) {
@@ -228,7 +228,7 @@ DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryWithAlignment(const Alloc
     if (allocationData.type == GraphicsAllocation::AllocationType::SVM_CPU) {
         //add 2MB padding in case reserved addr is not 2MB aligned
         reserveSizeAligned = alignUp(cSize, cAlignment) + cAlignment;
-        reservedGpuAddress = GmmHelper::canonize(gfxPartition.heapAllocate(HeapIndex::HEAP_STANDARD, reserveSizeAligned));
+        reservedGpuAddress = GmmHelper::canonize(gfxPartition->heapAllocate(HeapIndex::HEAP_STANDARD, reserveSizeAligned));
         if (!reservedGpuAddress) {
             bo->close();
             delete bo;
@@ -346,7 +346,7 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio
         uintptr_t inputPtr = reinterpret_cast<uintptr_t>(allocationData.hostPtr);
         auto allocationSize = alignSizeWholePage(allocationData.hostPtr, allocationData.size);
         auto realAllocationSize = allocationSize;
-        auto gpuVirtualAddress = gfxPartition.heapAllocate(allocatorToUse, realAllocationSize);
+        auto gpuVirtualAddress = gfxPartition->heapAllocate(allocatorToUse, realAllocationSize);
         if (!gpuVirtualAddress) {
             return nullptr;
         }
@@ -355,7 +355,7 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio
 
         BufferObject *bo = allocUserptr(alignedUserPointer, allocationSize, 0);
         if (!bo) {
-            gfxPartition.heapFree(allocatorToUse, gpuVirtualAddress, realAllocationSize);
+            gfxPartition->heapFree(allocatorToUse, gpuVirtualAddress, realAllocationSize);
             return nullptr;
         }
 
@@ -366,13 +366,13 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio
         auto allocation = new DrmAllocation(allocationData.type, bo, const_cast<void *>(allocationData.hostPtr), GmmHelper::canonize(ptrOffset(gpuVirtualAddress, inputPointerOffset)),
                                             allocationSize, MemoryPool::System4KBPagesWith32BitGpuAddressing, false);
         allocation->set32BitAllocation(true);
-        allocation->setGpuBaseAddress(GmmHelper::canonize(gfxPartition.getHeapBase(allocatorToUse)));
+        allocation->setGpuBaseAddress(GmmHelper::canonize(gfxPartition->getHeapBase(allocatorToUse)));
         return allocation;
     }
 
     size_t alignedAllocationSize = alignUp(allocationData.size, MemoryConstants::pageSize);
     auto allocationSize = alignedAllocationSize;
-    auto res = gfxPartition.heapAllocate(allocatorToUse, allocationSize);
+    auto res = gfxPartition->heapAllocate(allocatorToUse, allocationSize);
 
     if (!res) {
         return nullptr;
@@ -381,7 +381,7 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio
     auto ptrAlloc = alignedMallocWrapper(alignedAllocationSize, MemoryConstants::allocationAlignment);
 
     if (!ptrAlloc) {
-        gfxPartition.heapFree(allocatorToUse, res, allocationSize);
+        gfxPartition->heapFree(allocatorToUse, res, allocationSize);
         return nullptr;
     }
 
@@ -389,7 +389,7 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio
 
     if (!bo) {
         alignedFreeWrapper(ptrAlloc);
-        gfxPartition.heapFree(allocatorToUse, res, allocationSize);
+        gfxPartition->heapFree(allocatorToUse, res, allocationSize);
         return nullptr;
     }
 
@@ -404,7 +404,7 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio
                                         MemoryPool::System4KBPagesWith32BitGpuAddressing, false);
 
     allocation->set32BitAllocation(true);
-    allocation->setGpuBaseAddress(GmmHelper::canonize(gfxPartition.getHeapBase(allocatorToUse)));
+    allocation->setGpuBaseAddress(GmmHelper::canonize(gfxPartition->getHeapBase(allocatorToUse)));
     allocation->setDriverAllocatedCpuPtr(ptrAlloc);
     return allocation;
 }
@@ -480,7 +480,7 @@ GraphicsAllocation *DrmMemoryManager::createGraphicsAllocationFromSharedHandle(o
         drmAllocation->set32BitAllocation(true);
         drmAllocation->setGpuBaseAddress(getExternalHeapBaseAddress());
     } else if (isLimitedRange()) {
-        drmAllocation->setGpuBaseAddress(gfxPartition.getHeapBase(HeapIndex::HEAP_STANDARD));
+        drmAllocation->setGpuBaseAddress(gfxPartition->getHeapBase(HeapIndex::HEAP_STANDARD));
     }
 
     if (properties.imgInfo) {
@@ -570,7 +570,7 @@ void DrmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation)
     void *reserveAddress = gfxAllocation->getReservedAddressPtr();
     if (reserveAddress) {
         auto gpuAddressToFree = GmmHelper::decanonize(reinterpret_cast<uint64_t>(reserveAddress));
-        gfxPartition.freeGpuAddressRange(gpuAddressToFree, gfxAllocation->getReservedAddressSize());
+        gfxPartition->freeGpuAddressRange(gpuAddressToFree, gfxAllocation->getReservedAddressSize());
     }
     delete gfxAllocation;
 
diff --git a/runtime/os_interface/windows/wddm_memory_manager.cpp b/runtime/os_interface/windows/wddm_memory_manager.cpp
index 333aab3763..c44b6bad4a 100644
--- a/runtime/os_interface/windows/wddm_memory_manager.cpp
+++ b/runtime/os_interface/windows/wddm_memory_manager.cpp
@@ -45,7 +45,7 @@ WddmMemoryManager::WddmMemoryManager(ExecutionEnvironment &executionEnvironment)
     if (asyncDeleterEnabled)
         deferredDeleter = createDeferredDeleter();
     mallocRestrictions.minAddress = wddm->getWddmMinAddress();
-    wddm->initGfxPartition(gfxPartition);
+    wddm->initGfxPartition(*gfxPartition);
 }
 
 GraphicsAllocation *WddmMemoryManager::allocateGraphicsMemoryForImageImpl(const AllocationData &allocationData, std::unique_ptr<Gmm> gmm) {
@@ -504,13 +504,13 @@ bool WddmMemoryManager::mapGpuVaForOneHandleAllocation(WddmAllocation *allocatio
         addressToMap = allocation->reservedGpuVirtualAddress;
     }
     auto status = wddm->mapGpuVirtualAddress(allocation->getDefaultGmm(), allocation->getDefaultHandle(),
-                                             gfxPartition.getHeapMinimalAddress(heapIndex), gfxPartition.getHeapLimit(heapIndex),
+                                             gfxPartition->getHeapMinimalAddress(heapIndex), gfxPartition->getHeapLimit(heapIndex),
                                              addressToMap, allocation->getGpuAddressToModify());
 
     if (!status && deferredDeleter) {
         deferredDeleter->drain(true);
         status = wddm->mapGpuVirtualAddress(allocation->getDefaultGmm(), allocation->getDefaultHandle(),
-                                            gfxPartition.getHeapMinimalAddress(heapIndex), gfxPartition.getHeapLimit(heapIndex),
+                                            gfxPartition->getHeapMinimalAddress(heapIndex), gfxPartition->getHeapLimit(heapIndex),
                                             addressToMap, allocation->getGpuAddressToModify());
     }
     if (!status) {
@@ -542,8 +542,8 @@ void WddmMemoryManager::obtainGpuAddressIfNeeded(WddmAllocation *allocation) {
     if (allocation->getNumHandles() > 1u) {
         auto heapIndex = selectHeap(allocation, false, executionEnvironment.isFullRangeSvm());
         allocation->reservedSizeForGpuVirtualAddress = allocation->getAlignedSize();
-        allocation->reservedGpuVirtualAddress = wddm->reserveGpuVirtualAddress(gfxPartition.getHeapMinimalAddress(heapIndex),
-                                                                               gfxPartition.getHeapLimit(heapIndex),
+        allocation->reservedGpuVirtualAddress = wddm->reserveGpuVirtualAddress(gfxPartition->getHeapMinimalAddress(heapIndex),
+                                                                               gfxPartition->getHeapLimit(heapIndex),
                                                                                allocation->reservedSizeForGpuVirtualAddress);
     }
 }
diff --git a/unit_tests/mocks/linux/mock_drm_memory_manager.h b/unit_tests/mocks/linux/mock_drm_memory_manager.h
index 174a7c89ea..b3d4ca1a97 100644
--- a/unit_tests/mocks/linux/mock_drm_memory_manager.h
+++ b/unit_tests/mocks/linux/mock_drm_memory_manager.h
@@ -53,6 +53,7 @@ class TestedDrmMemoryManager : public MemoryManagerCreate<DrmMemoryManager> {
     using DrmMemoryManager::getDefaultDrmContextId;
     using DrmMemoryManager::gfxPartition;
     using DrmMemoryManager::pinThreshold;
+    using DrmMemoryManager::releaseGpuRange;
     using DrmMemoryManager::setDomainCpu;
     using DrmMemoryManager::sharingBufferObjects;
     using DrmMemoryManager::supportsMultiStorageResources;
@@ -103,7 +104,8 @@ class TestedDrmMemoryManager : public MemoryManagerCreate<DrmMemoryManager> {
     }
 
     DrmGemCloseWorker *getgemCloseWorker() { return this->gemCloseWorker.get(); }
-    void forceLimitedRangeAllocator(uint64_t range) { gfxPartition.init(range, getSizeToReserve()); }
+    void forceLimitedRangeAllocator(uint64_t range) { gfxPartition->init(range, getSizeToReserve()); }
+    void overrideGfxPartition(GfxPartition *newGfxPartition) { gfxPartition.reset(newGfxPartition); }
 
     DrmAllocation *allocate32BitGraphicsMemory(size_t size, const void *ptr, GraphicsAllocation::AllocationType allocationType) {
         bool allocateMemory = ptr == nullptr;
diff --git a/unit_tests/mocks/mock_gfx_partition.h b/unit_tests/mocks/mock_gfx_partition.h
index 82459c78ab..43426d2324 100644
--- a/unit_tests/mocks/mock_gfx_partition.h
+++ b/unit_tests/mocks/mock_gfx_partition.h
@@ -7,6 +7,8 @@
 
 #include "runtime/memory_manager/gfx_partition.h"
 
+#include "gmock/gmock.h"
+
 using namespace NEO;
 
 class MockGfxPartition : public GfxPartition {
@@ -27,5 +29,7 @@ class MockGfxPartition : public GfxPartition {
         return reservedCpuAddressRangeSize;
     }
 
+    MOCK_METHOD2(freeGpuAddressRange, void(uint64_t gpuAddress, size_t size));
+
     static std::array<HeapIndex, static_cast<uint32_t>(HeapIndex::TOTAL_HEAPS)> allHeapNames;
 };
diff --git a/unit_tests/os_interface/linux/drm_memory_manager_tests.cpp b/unit_tests/os_interface/linux/drm_memory_manager_tests.cpp
index d86f646ad6..c9e4d727db 100644
--- a/unit_tests/os_interface/linux/drm_memory_manager_tests.cpp
+++ b/unit_tests/os_interface/linux/drm_memory_manager_tests.cpp
@@ -30,6 +30,7 @@
 #include "test.h"
 #include "unit_tests/mocks/linux/mock_drm_command_stream_receiver.h"
 #include "unit_tests/mocks/mock_context.h"
+#include "unit_tests/mocks/mock_gfx_partition.h"
 #include "unit_tests/mocks/mock_gmm.h"
 
 #include "drm/i915_drm.h"
@@ -909,15 +910,15 @@ TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhenLimitedRangeAllocatorSetThenH
     memoryManager->forceLimitedRangeAllocator(0xFFFFFFFFF);
 
     uint64_t sizeBig = 4 * MemoryConstants::megaByte + MemoryConstants::pageSize;
-    auto gpuAddressLimitedRange = memoryManager->gfxPartition.heapAllocate(HeapIndex::HEAP_STANDARD, sizeBig);
-    EXPECT_LT(memoryManager->gfxPartition.getHeapBase(HeapIndex::HEAP_STANDARD), gpuAddressLimitedRange);
-    EXPECT_GT(memoryManager->gfxPartition.getHeapLimit(HeapIndex::HEAP_STANDARD), gpuAddressLimitedRange + sizeBig);
-    EXPECT_EQ(memoryManager->gfxPartition.getHeapMinimalAddress(HeapIndex::HEAP_STANDARD), gpuAddressLimitedRange);
+    auto gpuAddressLimitedRange = memoryManager->gfxPartition->heapAllocate(HeapIndex::HEAP_STANDARD, sizeBig);
+    EXPECT_LT(memoryManager->gfxPartition->getHeapBase(HeapIndex::HEAP_STANDARD), gpuAddressLimitedRange);
+    EXPECT_GT(memoryManager->gfxPartition->getHeapLimit(HeapIndex::HEAP_STANDARD), gpuAddressLimitedRange + sizeBig);
+    EXPECT_EQ(memoryManager->gfxPartition->getHeapMinimalAddress(HeapIndex::HEAP_STANDARD), gpuAddressLimitedRange);
 
-    auto gpuInternal32BitAlloc = memoryManager->gfxPartition.heapAllocate(internalHeapIndex, sizeBig);
-    EXPECT_LT(memoryManager->gfxPartition.getHeapBase(internalHeapIndex), gpuInternal32BitAlloc);
-    EXPECT_GT(memoryManager->gfxPartition.getHeapLimit(internalHeapIndex), gpuInternal32BitAlloc + sizeBig);
-    EXPECT_EQ(memoryManager->gfxPartition.getHeapMinimalAddress(internalHeapIndex), gpuInternal32BitAlloc);
+    auto gpuInternal32BitAlloc = memoryManager->gfxPartition->heapAllocate(internalHeapIndex, sizeBig);
+    EXPECT_LT(memoryManager->gfxPartition->getHeapBase(internalHeapIndex), gpuInternal32BitAlloc);
+    EXPECT_GT(memoryManager->gfxPartition->getHeapLimit(internalHeapIndex), gpuInternal32BitAlloc + sizeBig);
+    EXPECT_EQ(memoryManager->gfxPartition->getHeapMinimalAddress(internalHeapIndex), gpuInternal32BitAlloc);
 }
 
 TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhenAskedForAllocationWithAlignmentThenCorrectAllocatorTypeSelected) {
@@ -967,7 +968,7 @@ TEST_F(DrmMemoryManagerTest, givenMemoryManagerWhenAskedForAllocationWithAlignme
 
     // emulate GPU address space exhaust
     memoryManager->forceLimitedRangeAllocator(0xFFFFFFFFF);
-    memoryManager->gfxPartition.heapInit(HeapIndex::HEAP_STANDARD, 0x0, 0x10000);
+    memoryManager->gfxPartition->heapInit(HeapIndex::HEAP_STANDARD, 0x0, 0x10000);
 
     // set size to something bigger than allowed space
     allocationData.size = 0x20000;
@@ -1136,7 +1137,7 @@ TEST_F(DrmMemoryManagerTest, Given32BitDeviceWithMemoryManagerWhenInternalHeapIs
     std::unique_ptr<Device> pDevice(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr));
 
     size_t size = MemoryConstants::pageSize64k;
-    auto alloc = memoryManager->gfxPartition.heapAllocate(internalHeapIndex, size);
+    auto alloc = memoryManager->gfxPartition->heapAllocate(internalHeapIndex, size);
     EXPECT_NE(0llu, alloc);
 
     size_t allocationSize = 4 * GB;
@@ -2135,12 +2136,12 @@ TEST_F(DrmMemoryManagerTest, given32BitAllocatorWithHeapAllocatorWhenLargerFragm
     memoryManager->setForce32BitAllocations(true);
 
     size_t allocationSize = 4 * MemoryConstants::pageSize;
-    auto ptr = memoryManager->gfxPartition.heapAllocate(HeapIndex::HEAP_EXTERNAL, allocationSize);
+    auto ptr = memoryManager->gfxPartition->heapAllocate(HeapIndex::HEAP_EXTERNAL, allocationSize);
     size_t smallAllocationSize = MemoryConstants::pageSize;
-    memoryManager->gfxPartition.heapAllocate(HeapIndex::HEAP_EXTERNAL, smallAllocationSize);
+    memoryManager->gfxPartition->heapAllocate(HeapIndex::HEAP_EXTERNAL, smallAllocationSize);
 
     //now free first allocation , this will move it to chunks
-    memoryManager->gfxPartition.heapFree(HeapIndex::HEAP_EXTERNAL, ptr, allocationSize);
+    memoryManager->gfxPartition->heapFree(HeapIndex::HEAP_EXTERNAL, ptr, allocationSize);
 
     //now ask for 3 pages, this will give ptr from chunks
     size_t pages3size = 3 * MemoryConstants::pageSize;
@@ -2460,7 +2461,7 @@ TEST_F(DrmMemoryManagerBasic, givenDefaultDrmMemoryManagerWhenItIsQueriedForInte
                                                                                                     true,
                                                                                                     true,
                                                                                                     executionEnvironment));
-    auto heapBase = memoryManager->gfxPartition.getHeapBase(internalHeapIndex);
+    auto heapBase = memoryManager->gfxPartition->getHeapBase(internalHeapIndex);
     EXPECT_EQ(heapBase, memoryManager->getInternalHeapBaseAddress());
 }
 
@@ -3022,14 +3023,14 @@ TEST_F(DrmMemoryManagerBasic, ifLimitedRangeAllocatorAvailableWhenAskedForAlloca
     memoryManager->forceLimitedRangeAllocator(0xFFFFFFFFF);
 
     size_t size = 100u;
-    auto ptr = memoryManager->gfxPartition.heapAllocate(HeapIndex::HEAP_STANDARD, size);
-    auto address64bit = ptrDiff(ptr, memoryManager->gfxPartition.getHeapBase(HeapIndex::HEAP_STANDARD));
+    auto ptr = memoryManager->gfxPartition->heapAllocate(HeapIndex::HEAP_STANDARD, size);
+    auto address64bit = ptrDiff(ptr, memoryManager->gfxPartition->getHeapBase(HeapIndex::HEAP_STANDARD));
 
     EXPECT_LT(address64bit, platformDevices[0]->capabilityTable.gpuAddressSpace);
 
     EXPECT_LT(0u, address64bit);
 
-    memoryManager->gfxPartition.heapFree(HeapIndex::HEAP_STANDARD, ptr, size);
+    memoryManager->gfxPartition->heapFree(HeapIndex::HEAP_STANDARD, ptr, size);
 }
 
 TEST_F(DrmMemoryManagerBasic, givenDisabledHostPtrTrackingWhenAllocateGraphicsMemoryForNonSvmHostPtrIsCalledWithNotAlignedPtrIsPassedThenAllocationIsCreated) {
@@ -3151,8 +3152,8 @@ TEST_F(DrmMemoryManagerTest, givenSvmCpuAllocationWhenSizeAndAlignmentProvidedTh
 
     EXPECT_NE(0llu, bo->peekAddress());
 
-    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapBase(HeapIndex::HEAP_STANDARD)), bo->peekAddress());
-    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapLimit(HeapIndex::HEAP_STANDARD)), bo->peekAddress());
+    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapBase(HeapIndex::HEAP_STANDARD)), bo->peekAddress());
+    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapLimit(HeapIndex::HEAP_STANDARD)), bo->peekAddress());
 
     EXPECT_EQ(reinterpret_cast<void *>(allocation->getGpuAddress()), alignUp(allocation->getReservedAddressPtr(), allocationData.alignment));
     EXPECT_EQ(alignUp(allocationData.size, allocationData.alignment) + allocationData.alignment, allocation->getReservedAddressSize());
@@ -3165,7 +3166,7 @@ TEST_F(DrmMemoryManagerTest, givenSvmCpuAllocationWhenSizeAndAlignmentProvidedBu
     mock->ioctl_expected.gemWait = 0;
     mock->ioctl_expected.gemClose = 1;
 
-    memoryManager->gfxPartition.heapInit(HeapIndex::HEAP_STANDARD, 0, 0);
+    memoryManager->gfxPartition->heapInit(HeapIndex::HEAP_STANDARD, 0, 0);
 
     TestedDrmMemoryManager::AllocationData allocationData;
     allocationData.size = 2 * MemoryConstants::megaByte;
@@ -3175,3 +3176,18 @@ TEST_F(DrmMemoryManagerTest, givenSvmCpuAllocationWhenSizeAndAlignmentProvidedBu
     DrmAllocation *allocation = memoryManager->allocateGraphicsMemoryWithAlignment(allocationData);
     EXPECT_EQ(nullptr, allocation);
 }
+
+TEST_F(DrmMemoryManagerTest, givenDrmMemoryManagerAndReleaseGpuRangeIsCalledThenGpuAddressIsDecanonized) {
+    auto mockGfxPartition = std::make_unique<MockGfxPartition>();
+    mockGfxPartition->init(maxNBitValue<48>, 0);
+    auto size = 2 * MemoryConstants::megaByte;
+    auto gpuAddress = mockGfxPartition->heapAllocate(HeapIndex::HEAP_STANDARD, size);
+    auto gpuAddressCanonized = GmmHelper::canonize(gpuAddress);
+    EXPECT_NE(gpuAddress, gpuAddressCanonized);
+    EXPECT_LE(gpuAddress, gpuAddressCanonized);
+
+    EXPECT_CALL(*mockGfxPartition.get(), freeGpuAddressRange(gpuAddress, size));
+
+    memoryManager->overrideGfxPartition(mockGfxPartition.release());
+    memoryManager->releaseGpuRange(reinterpret_cast<void *>(gpuAddressCanonized), size, UNKNOWN_ALLOCATOR);
+}
diff --git a/unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp b/unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp
index c2776dcfb1..ad223d16d6 100644
--- a/unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp
+++ b/unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp
@@ -845,8 +845,8 @@ TEST_F(WddmMemoryManagerTest, Allocate32BitMemoryWithNullptr) {
     auto *gpuAllocation = memoryManager->allocate32BitGraphicsMemory(3 * MemoryConstants::pageSize, nullptr, GraphicsAllocation::AllocationType::BUFFER);
 
     ASSERT_NE(nullptr, gpuAllocation);
-    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapBase(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress());
-    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapLimit(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress() + gpuAllocation->getUnderlyingBufferSize());
+    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapBase(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress());
+    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapLimit(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress() + gpuAllocation->getUnderlyingBufferSize());
 
     EXPECT_EQ(0u, gpuAllocation->fragmentsStorage.fragmentCount);
     memoryManager->freeGraphicsMemory(gpuAllocation);
@@ -857,8 +857,8 @@ TEST_F(WddmMemoryManagerTest, given32BitAllocationWhenItIsCreatedThenItHasNonZer
 
     ASSERT_NE(nullptr, gpuAllocation);
     EXPECT_NE(0llu, gpuAllocation->getGpuAddressToPatch());
-    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapBase(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress());
-    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapLimit(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress() + gpuAllocation->getUnderlyingBufferSize());
+    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapBase(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress());
+    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapLimit(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress() + gpuAllocation->getUnderlyingBufferSize());
     memoryManager->freeGraphicsMemory(gpuAllocation);
 }
 
@@ -872,8 +872,8 @@ TEST_F(WddmMemoryManagerTest, Allocate32BitMemoryWithMisalignedHostPtrDoesNotDoT
 
     EXPECT_EQ(alignSizeWholePage(misalignedPtr, misalignedSize), gpuAllocation->getUnderlyingBufferSize());
 
-    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapBase(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress());
-    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapLimit(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress() + gpuAllocation->getUnderlyingBufferSize());
+    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapBase(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress());
+    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapLimit(HeapIndex::HEAP_EXTERNAL)), gpuAllocation->getGpuAddress() + gpuAllocation->getUnderlyingBufferSize());
 
     EXPECT_EQ(0u, gpuAllocation->fragmentsStorage.fragmentCount);
 
@@ -889,7 +889,7 @@ TEST_F(WddmMemoryManagerTest, Allocate32BitMemorySetsCannonizedGpuBaseAddress) {
 
     ASSERT_NE(nullptr, gpuAllocation);
 
-    uint64_t cannonizedAddress = GmmHelper::canonize(memoryManager->gfxPartition.getHeapBase(HeapIndex::HEAP_EXTERNAL));
+    uint64_t cannonizedAddress = GmmHelper::canonize(memoryManager->gfxPartition->getHeapBase(HeapIndex::HEAP_EXTERNAL));
     EXPECT_EQ(cannonizedAddress, gpuAllocation->getGpuBaseAddress());
 
     memoryManager->freeGraphicsMemory(gpuAllocation);
@@ -1011,7 +1011,7 @@ TEST_F(WddmMemoryManagerTest, givenNullPtrAndSizePassedToCreateInternalAllocatio
     EXPECT_EQ(4096u, wddmAllocation->getUnderlyingBufferSize());
     EXPECT_NE((uint64_t)wddmAllocation->getUnderlyingBuffer(), wddmAllocation->getGpuAddress());
     auto cannonizedHeapBase = GmmHelper::canonize(memoryManager->getInternalHeapBaseAddress());
-    auto cannonizedHeapEnd = GmmHelper::canonize(memoryManager->gfxPartition.getHeapLimit(internalHeapIndex));
+    auto cannonizedHeapEnd = GmmHelper::canonize(memoryManager->gfxPartition->getHeapLimit(internalHeapIndex));
 
     EXPECT_GT(wddmAllocation->getGpuAddress(), cannonizedHeapBase);
     EXPECT_LT(wddmAllocation->getGpuAddress() + wddmAllocation->getUnderlyingBufferSize(), cannonizedHeapEnd);
@@ -1030,7 +1030,7 @@ TEST_F(WddmMemoryManagerTest, givenPtrAndSizePassedToCreateInternalAllocationWhe
     EXPECT_EQ(4096u, wddmAllocation->getUnderlyingBufferSize());
     EXPECT_NE((uint64_t)wddmAllocation->getUnderlyingBuffer(), wddmAllocation->getGpuAddress());
     auto cannonizedHeapBase = GmmHelper::canonize(memoryManager->getInternalHeapBaseAddress());
-    auto cannonizedHeapEnd = GmmHelper::canonize(memoryManager->gfxPartition.getHeapLimit(internalHeapIndex));
+    auto cannonizedHeapEnd = GmmHelper::canonize(memoryManager->gfxPartition->getHeapLimit(internalHeapIndex));
 
     EXPECT_GT(wddmAllocation->getGpuAddress(), cannonizedHeapBase);
     EXPECT_LT(wddmAllocation->getGpuAddress() + wddmAllocation->getUnderlyingBufferSize(), cannonizedHeapEnd);
@@ -1386,7 +1386,7 @@ TEST_F(WddmMemoryManagerTest, givenWddmMemoryManagerWithNoRegisteredOsContextsWh
 
 TEST_F(WddmMemoryManagerTest, givenWddmMemoryManagerAnd32bitBuildThenSvmPartitionIsAlwaysInitialized) {
     if (is32bit) {
-        EXPECT_EQ(memoryManager->gfxPartition.getHeapLimit(HeapIndex::HEAP_SVM), MemoryConstants::max32BitAddress);
+        EXPECT_EQ(memoryManager->gfxPartition->getHeapLimit(HeapIndex::HEAP_SVM), MemoryConstants::max32BitAddress);
     }
 }
 
@@ -1751,8 +1751,8 @@ TEST_F(WddmMemoryManagerSimpleTest, givenBufferHostMemoryAllocationAndLimitedRan
 
     auto heap = is32bit ? HeapIndex::HEAP_SVM : HeapIndex::HEAP_STANDARD;
 
-    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapBase(heap)), allocation->getGpuAddress());
-    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition.getHeapLimit(heap)), allocation->getGpuAddress());
+    EXPECT_LT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapBase(heap)), allocation->getGpuAddress());
+    EXPECT_GT(GmmHelper::canonize(memoryManager->gfxPartition->getHeapLimit(heap)), allocation->getGpuAddress());
 
     memoryManager->freeGraphicsMemory(allocation);
 }