Virtual address space partitioning on Linux [2/n]

PiotrFusik · Compute-Runtime-Automation · commit 75edea81bbc9 · 2019-02-20T14:43:08.000+01:00
Move selectHeap from Wddm to MemoryManager.
Set DrmAllocation::origin.

Change-Id: I5d412e35d524d1f31174893b9ce1d3b1e98eee96
diff --git a/runtime/memory_manager/graphics_allocation.h b/runtime/memory_manager/graphics_allocation.h
@@ -115,8 +115,8 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation> {
         return gpuAddress + allocationOffset - gpuBaseAddress;
     }
 
-    bool isCoherent() { return coherent; };
-    void setCoherent(bool coherentIn) { this->coherent = coherentIn; };
+    bool isCoherent() const { return coherent; }
+    void setCoherent(bool coherentIn) { this->coherent = coherentIn; }
     void setSize(size_t size) { this->size = size; }
     osHandle peekSharedHandle() { return sharedHandle; }
 
diff --git a/runtime/memory_manager/memory_manager.cpp b/runtime/memory_manager/memory_manager.cpp
@@ -15,6 +15,7 @@
 #include "runtime/gmm_helper/resource_info.h"
 #include "runtime/helpers/aligned_memory.h"
 #include "runtime/helpers/basic_math.h"
+#include "runtime/helpers/hw_info.h"
 #include "runtime/helpers/kernel_commands.h"
 #include "runtime/helpers/options.h"
 #include "runtime/memory_manager/deferrable_allocation_deletion.h"
@@ -374,4 +375,24 @@ void MemoryManager::unlockResource(GraphicsAllocation *graphicsAllocation) {
     unlockResourceImpl(*graphicsAllocation);
     graphicsAllocation->unlock();
 }
+
+HeapIndex MemoryManager::selectHeap(const GraphicsAllocation *allocation, const void *ptr, const HardwareInfo &hwInfo) {
+    if (allocation) {
+        if (allocation->origin == AllocationOrigin::INTERNAL_ALLOCATION) {
+            return internalHeapIndex;
+        } else if (allocation->is32BitAllocation) {
+            return HeapIndex::HEAP_EXTERNAL;
+        }
+    }
+    if (hwInfo.capabilityTable.gpuAddressSpace == MemoryConstants::max48BitAddress) {
+        if (ptr) {
+            return HeapIndex::HEAP_SVM;
+        }
+        if (allocation && GraphicsAllocation::isCpuAccessRequired(allocation->getAllocationType())) {
+            return HeapIndex::HEAP_STANDARD64Kb;
+        }
+        return HeapIndex::HEAP_STANDARD;
+    }
+    return HeapIndex::HEAP_LIMITED;
+}
 } // namespace OCLRT
diff --git a/runtime/memory_manager/memory_manager.h b/runtime/memory_manager/memory_manager.h
@@ -27,6 +27,7 @@ class Gmm;
 class GraphicsAllocation;
 class HostPtrManager;
 class OsContext;
+struct HardwareInfo;
 struct ImageInfo;
 
 using CsrContainer = std::vector<std::vector<std::unique_ptr<CommandStreamReceiver>>>;
@@ -187,6 +188,8 @@ class MemoryManager {
     HostPtrManager *getHostPtrManager() const { return hostPtrManager.get(); }
     void setDefaultEngineIndex(uint32_t index) { defaultEngineIndex = index; }
 
+    static HeapIndex selectHeap(const GraphicsAllocation *allocation, const void *ptr, const HardwareInfo &hwInfo);
+
   protected:
     struct AllocationData {
         union {
diff --git a/runtime/memory_manager/os_agnostic_memory_manager.h b/runtime/memory_manager/os_agnostic_memory_manager.h
@@ -8,7 +8,6 @@
 #pragma once
 #include "runtime/helpers/basic_math.h"
 #include "runtime/memory_manager/memory_manager.h"
-#include <map>
 
 namespace OCLRT {
 constexpr size_t bigAllocation = 1 * MB;
@@ -31,18 +30,15 @@ class MemoryAllocation : public GraphicsAllocation {
     void overrideMemoryPool(MemoryPool::Type pool);
 };
 
-typedef std::map<void *, MemoryAllocation *> PointerMap;
-
 class OsAgnosticMemoryManager : public MemoryManager {
   public:
     using MemoryManager::allocateGraphicsMemory;
-    using MemoryManager::createGraphicsAllocationFromSharedHandle;
 
     OsAgnosticMemoryManager(bool enable64kbPages, bool enableLocalMemory, ExecutionEnvironment &executionEnvironment) : OsAgnosticMemoryManager(enable64kbPages, enableLocalMemory, false, executionEnvironment){};
 
     OsAgnosticMemoryManager(bool enable64kbPages, bool enableLocalMemory, bool aubUsage, ExecutionEnvironment &executionEnvironment) : MemoryManager(enable64kbPages, enableLocalMemory, executionEnvironment) {
         allocator32Bit = std::unique_ptr<Allocator32bit>(create32BitAllocator(aubUsage));
-    };
+    }
 
     ~OsAgnosticMemoryManager() override;
     GraphicsAllocation *allocateGraphicsMemoryForNonSvmHostPtr(size_t size, void *cpuPtr) override;
@@ -71,8 +67,8 @@ class OsAgnosticMemoryManager : public MemoryManager {
     GraphicsAllocation *allocateGraphicsMemory64kb(AllocationData allocationData) override;
     GraphicsAllocation *allocateGraphicsMemoryForImageImpl(const AllocationData &allocationData, std::unique_ptr<Gmm> gmm) override;
 
-    void *lockResourceImpl(GraphicsAllocation &graphicsAllocation) override { return ptrOffset(graphicsAllocation.getUnderlyingBuffer(), static_cast<size_t>(graphicsAllocation.allocationOffset)); };
-    void unlockResourceImpl(GraphicsAllocation &graphicsAllocation) override{};
+    void *lockResourceImpl(GraphicsAllocation &graphicsAllocation) override { return ptrOffset(graphicsAllocation.getUnderlyingBuffer(), static_cast<size_t>(graphicsAllocation.allocationOffset)); }
+    void unlockResourceImpl(GraphicsAllocation &graphicsAllocation) override {}
     GraphicsAllocation *allocate32BitGraphicsMemoryImpl(const AllocationData &allocationData) override;
     GraphicsAllocation *allocateGraphicsMemoryInDevicePool(const AllocationData &allocationData, AllocationStatus &status) override;
 
diff --git a/runtime/os_interface/linux/drm_memory_manager.cpp b/runtime/os_interface/linux/drm_memory_manager.cpp
@@ -241,7 +241,9 @@ DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryWithAlignment(const Alloc
     if (forcePinEnabled && pinBB != nullptr && allocationData.flags.forcePin && allocationData.size >= this->pinThreshold) {
         pinBB->pin(&bo, 1, getDefaultCommandStreamReceiver(0)->getOsContext().get()->getDrmContextId());
     }
-    return new DrmAllocation(bo, res, cSize, MemoryPool::System4KBPages, allocationData.flags.multiOsContextCapable);
+    auto allocation = new DrmAllocation(bo, res, cSize, MemoryPool::System4KBPages, allocationData.flags.multiOsContextCapable);
+    allocation->origin = allocationData.allocationOrigin;
+    return allocation;
 }
 
 DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryWithHostPtr(const AllocationData &allocationData) {
@@ -256,7 +258,7 @@ DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryWithHostPtr(const Allocat
 }
 
 DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryForNonSvmHostPtr(size_t size, void *cpuPtr) {
-    if ((size == 0) || !cpuPtr)
+    if (size == 0 || !cpuPtr)
         return nullptr;
 
     auto alignedPtr = alignDown(reinterpret_cast<char *>(cpuPtr), MemoryConstants::pageSize);
@@ -329,6 +331,7 @@ GraphicsAllocation *DrmMemoryManager::allocateGraphicsMemoryForImageImpl(const A
     auto allocation = new DrmAllocation(bo, nullptr, (uint64_t)gpuRange, allocationData.imgInfo->size, MemoryPool::SystemCpuInaccessible, false);
     bo->setAllocationType(allocatorType);
     allocation->gmm = gmm.release();
+    allocation->origin = allocationData.allocationOrigin;
     return allocation;
 }
 
@@ -363,6 +366,7 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio
                                                allocationSize, MemoryPool::System4KBPagesWith32BitGpuAddressing, false);
         drmAllocation->is32BitAllocation = true;
         drmAllocation->gpuBaseAddress = allocatorToUse->getBase();
+        drmAllocation->origin = allocationData.allocationOrigin;
         return drmAllocation;
     }
 
@@ -417,6 +421,7 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio
     drmAllocation->is32BitAllocation = true;
     drmAllocation->gpuBaseAddress = allocatorToUse->getBase();
     drmAllocation->driverAllocatedCpuPointer = ptrAlloc;
+    drmAllocation->origin = allocationData.allocationOrigin;
     return drmAllocation;
 }
 
diff --git a/runtime/os_interface/linux/drm_memory_manager.h b/runtime/os_interface/linux/drm_memory_manager.h
@@ -21,9 +21,6 @@ class Drm;
 
 class DrmMemoryManager : public MemoryManager {
   public:
-    using MemoryManager::allocateGraphicsMemory;
-    using MemoryManager::createGraphicsAllocationFromSharedHandle;
-
     DrmMemoryManager(Drm *drm, gemCloseWorkerMode mode, bool forcePinAllowed, bool validateHostPtrMemory, ExecutionEnvironment &executionEnvironment);
     ~DrmMemoryManager() override;
 
diff --git a/runtime/os_interface/windows/wddm/wddm.cpp b/runtime/os_interface/windows/wddm/wddm.cpp
@@ -11,6 +11,7 @@
 #include "runtime/gmm_helper/gmm_helper.h"
 #include "runtime/gmm_helper/resource_info.h"
 #include "runtime/gmm_helper/page_table_mngr.h"
+#include "runtime/memory_manager/memory_manager.h"
 #include "runtime/os_interface/windows/wddm/wddm.h"
 #include "runtime/os_interface/hw_info_config.h"
 #include "runtime/os_interface/windows/gdi_interface.h"
@@ -288,15 +289,15 @@ bool Wddm::makeResident(D3DKMT_HANDLE *handles, uint32_t count, bool cantTrimFur
 bool Wddm::mapGpuVirtualAddress(WddmAllocation *allocation, void *cpuPtr) {
     void *mapPtr = allocation->getReservedAddress() != nullptr ? allocation->getReservedAddress() : cpuPtr;
     return mapGpuVirtualAddressImpl(allocation->gmm, allocation->handle, mapPtr, allocation->gpuPtr,
-                                    selectHeap(allocation, mapPtr));
+                                    MemoryManager::selectHeap(allocation, mapPtr, *hardwareInfoTable[gfxPlatform->eProductFamily]));
 }
 
 bool Wddm::mapGpuVirtualAddress(AllocationStorageData *allocationStorageData) {
     return mapGpuVirtualAddressImpl(allocationStorageData->osHandleStorage->gmm,
                                     allocationStorageData->osHandleStorage->handle,
                                     const_cast<void *>(allocationStorageData->cpuPtr),
                                     allocationStorageData->osHandleStorage->gpuPtr,
-                                    selectHeap(nullptr, allocationStorageData->cpuPtr));
+                                    MemoryManager::selectHeap(nullptr, allocationStorageData->cpuPtr, *hardwareInfoTable[gfxPlatform->eProductFamily]));
 }
 
 bool Wddm::mapGpuVirtualAddressImpl(Gmm *gmm, D3DKMT_HANDLE handle, void *cpuPtr, D3DGPU_VIRTUAL_ADDRESS &gpuPtr, HeapIndex heapIndex) {
@@ -994,24 +995,4 @@ std::unique_lock<SpinLock> Wddm::acquireLock(SpinLock &lock) {
     return std::unique_lock<SpinLock>{lock};
 }
 
-HeapIndex Wddm::selectHeap(const WddmAllocation *allocation, const void *ptr) const {
-    if (allocation) {
-        if (allocation->origin == AllocationOrigin::INTERNAL_ALLOCATION) {
-            return internalHeapIndex;
-        } else if (allocation->is32BitAllocation) {
-            return HeapIndex::HEAP_EXTERNAL;
-        }
-    }
-    if (hardwareInfoTable[gfxPlatform->eProductFamily]->capabilityTable.gpuAddressSpace == MemoryConstants::max48BitAddress) {
-        if (ptr) {
-            return HeapIndex::HEAP_SVM;
-        }
-        if (allocation && GraphicsAllocation::isCpuAccessRequired(allocation->getAllocationType())) {
-            return HeapIndex::HEAP_STANDARD64Kb;
-        }
-        return HeapIndex::HEAP_STANDARD;
-    }
-    return HeapIndex::HEAP_LIMITED;
-}
-
 } // namespace OCLRT
diff --git a/runtime/os_interface/windows/wddm/wddm.h b/runtime/os_interface/windows/wddm/wddm.h
@@ -148,7 +148,6 @@ class Wddm {
     MOCKABLE_VIRTUAL EvictionStatus evictTemporaryResource(WddmAllocation &allocation);
     MOCKABLE_VIRTUAL void applyBlockingMakeResident(WddmAllocation &allocation);
     MOCKABLE_VIRTUAL std::unique_lock<SpinLock> acquireLock(SpinLock &lock);
-    HeapIndex selectHeap(const WddmAllocation *allocation, const void *cpuPtr) const;
 
   protected:
     bool initialized = false;
diff --git a/runtime/os_interface/windows/wddm_memory_manager.cpp b/runtime/os_interface/windows/wddm_memory_manager.cpp
@@ -48,7 +48,7 @@ GraphicsAllocation *WddmMemoryManager::allocateGraphicsMemoryForImageImpl(const
     auto allocation = std::make_unique<WddmAllocation>(nullptr, allocationData.imgInfo->size, nullptr, MemoryPool::SystemCpuInaccessible, false);
     allocation->gmm = gmm.get();
     allocation->origin = allocationData.allocationOrigin;
-    if (!WddmMemoryManager::createWddmAllocation(allocation.get())) {
+    if (!createWddmAllocation(allocation.get())) {
         return nullptr;
     }
 
@@ -60,11 +60,10 @@ GraphicsAllocation *WddmMemoryManager::allocateGraphicsMemoryForImageImpl(const
 
 GraphicsAllocation *WddmMemoryManager::allocateGraphicsMemory64kb(AllocationData allocationData) {
     size_t sizeAligned = alignUp(allocationData.size, MemoryConstants::pageSize64k);
-    Gmm *gmm = nullptr;
 
     auto wddmAllocation = std::make_unique<WddmAllocation>(nullptr, sizeAligned, nullptr, MemoryPool::System64KBPages, !!allocationData.flags.multiOsContextCapable);
 
-    gmm = new Gmm(nullptr, sizeAligned, false, allocationData.flags.preferRenderCompressed, true, {});
+    auto gmm = new Gmm(nullptr, sizeAligned, false, allocationData.flags.preferRenderCompressed, true, {});
     wddmAllocation->gmm = gmm;
 
     if (!wddm->createAllocation64k(wddmAllocation.get())) {
@@ -77,7 +76,7 @@ GraphicsAllocation *WddmMemoryManager::allocateGraphicsMemory64kb(AllocationData
     // 64kb map is not needed
     auto status = wddm->mapGpuVirtualAddress(wddmAllocation.get(), cpuPtr);
     DEBUG_BREAK_IF(!status);
-    wddmAllocation->setCpuPtrAndGpuAddress(cpuPtr, (uint64_t)wddmAllocation->gpuPtr);
+    wddmAllocation->setCpuPtrAndGpuAddress(cpuPtr, wddmAllocation->gpuPtr);
 
     DebugManager.logAllocation(wddmAllocation.get());
     return wddmAllocation.release();
@@ -280,15 +279,15 @@ void WddmMemoryManager::removeAllocationFromHostPtrManager(GraphicsAllocation *g
 
 void *WddmMemoryManager::lockResourceImpl(GraphicsAllocation &graphicsAllocation) {
     return wddm->lockResource(static_cast<WddmAllocation &>(graphicsAllocation));
-};
+}
 void WddmMemoryManager::unlockResourceImpl(GraphicsAllocation &graphicsAllocation) {
     auto &wddmAllocation = static_cast<WddmAllocation &>(graphicsAllocation);
     wddm->unlockResource(wddmAllocation);
     if (wddmAllocation.needsMakeResidentBeforeLock) {
         auto evictionStatus = wddm->evictTemporaryResource(wddmAllocation);
         DEBUG_BREAK_IF(evictionStatus == EvictionStatus::FAILED);
     }
-};
+}
 
 void WddmMemoryManager::freeGraphicsMemoryImpl(GraphicsAllocation *gfxAllocation) {
     WddmAllocation *input = static_cast<WddmAllocation *>(gfxAllocation);
@@ -450,7 +449,7 @@ uint64_t WddmMemoryManager::getMaxApplicationAddress() {
 }
 
 uint64_t WddmMemoryManager::getInternalHeapBaseAddress() {
-    return this->wddm->getGfxPartition().Heap32[static_cast<uint32_t>(internalHeapIndex)].Base;
+    return wddm->getGfxPartition().Heap32[static_cast<uint32_t>(internalHeapIndex)].Base;
 }
 
 bool WddmMemoryManager::mapAuxGpuVA(GraphicsAllocation *graphicsAllocation) {
diff --git a/runtime/os_interface/windows/wddm_memory_manager.h b/runtime/os_interface/windows/wddm_memory_manager.h
@@ -24,10 +24,7 @@ using OsContextWin = OsContext::OsContextImpl;
 
 class WddmMemoryManager : public MemoryManager {
   public:
-    using MemoryManager::allocateGraphicsMemory;
-    using MemoryManager::createGraphicsAllocationFromSharedHandle;
-
-    ~WddmMemoryManager();
+    ~WddmMemoryManager() override;
     WddmMemoryManager(bool enable64kbPages, bool enableLocalMemory, Wddm *wddm, ExecutionEnvironment &executionEnvironment);
 
     WddmMemoryManager(const WddmMemoryManager &) = delete;
diff --git a/unit_tests/gen10/windows/os_interface_tests_gen10.cpp b/unit_tests/gen10/windows/os_interface_tests_gen10.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2017-2018 Intel Corporation
+ * Copyright (C) 2017-2019 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -13,7 +13,6 @@ typedef OsInterfaceTest OsInterfaceTestCnl;
 
 GEN10TEST_F(OsInterfaceTestCnl, askKmdIfPreemptionRegisterWhitelisted) {
     HardwareInfo *hwInfo = nullptr;
-    const HardwareInfo *refHwinfo = *platformDevices;
     size_t numDevices = 0;
 
     ExecutionEnvironment executionEnvironment;
diff --git a/unit_tests/memory_manager/memory_manager_tests.cpp b/unit_tests/memory_manager/memory_manager_tests.cpp
@@ -1617,3 +1617,79 @@ TEST(OsContextTest, givenOsContextWithNumberOfSupportedDevicesWhenConstructingTh
     OsContext osContext(nullptr, 5, 7, {EngineType::ENGINE_RCS, 0}, PreemptionMode::Disabled);
     EXPECT_EQ(7u, osContext.getNumDevicesSupported());
 }
+
+TEST(MemoryManagerTest, given32bitInternalAllocationWhenSelectingHeapThenInternalHeapIsUsed) {
+    GraphicsAllocation allocation{nullptr, 0, 0, 0, false};
+    allocation.is32BitAllocation = true;
+    allocation.origin = AllocationOrigin::INTERNAL_ALLOCATION;
+    EXPECT_EQ(internalHeapIndex, MemoryManager::selectHeap(&allocation, nullptr, *platformDevices[0]));
+}
+
+TEST(MemoryManagerTest, givenNon32bitInternalAllocationWhenSelectingHeapThenInternalHeapIsUsed) {
+    GraphicsAllocation allocation{nullptr, 0, 0, 0, false};
+    allocation.is32BitAllocation = false;
+    allocation.origin = AllocationOrigin::INTERNAL_ALLOCATION;
+    EXPECT_EQ(internalHeapIndex, MemoryManager::selectHeap(&allocation, nullptr, *platformDevices[0]));
+}
+
+TEST(MemoryManagerTest, given32bitExternalAllocationWhenSelectingHeapThenExternalHeapIsUsed) {
+    GraphicsAllocation allocation{nullptr, 0, 0, 0, false};
+    allocation.is32BitAllocation = true;
+    allocation.origin = AllocationOrigin::EXTERNAL_ALLOCATION;
+    EXPECT_EQ(HeapIndex::HEAP_EXTERNAL, MemoryManager::selectHeap(&allocation, nullptr, *platformDevices[0]));
+}
+
+TEST(MemoryManagerTest, givenLimitedAddressSpaceWhenSelectingHeapForExternalAllocationThenLimitedHeapIsUsed) {
+    GraphicsAllocation allocation{nullptr, 0, 0, 0, false};
+    EXPECT_EQ(AllocationOrigin::EXTERNAL_ALLOCATION, allocation.origin);
+    if (platformDevices[0]->capabilityTable.gpuAddressSpace == MemoryConstants::max48BitAddress) {
+        return;
+    }
+    EXPECT_EQ(HeapIndex::HEAP_LIMITED, MemoryManager::selectHeap(&allocation, nullptr, *platformDevices[0]));
+}
+
+TEST(MemoryManagerTest, givenFullAddressSpaceWhenSelectingHeapForExternalAllocationWithPtrThenSvmHeapIsUsed) {
+    GraphicsAllocation allocation{nullptr, 0, 0, 0, false};
+    EXPECT_EQ(AllocationOrigin::EXTERNAL_ALLOCATION, allocation.origin);
+    if (platformDevices[0]->capabilityTable.gpuAddressSpace != MemoryConstants::max48BitAddress) {
+        return;
+    }
+    EXPECT_EQ(HeapIndex::HEAP_SVM, MemoryManager::selectHeap(&allocation, &allocation, *platformDevices[0]));
+}
+
+TEST(MemoryManagerTest, givenFullAddressSpaceWhenSelectingHeapForExternalAllocationWithoutPtrAndCpuAccessIsRequiredThenStandard64kHeapIsUsed) {
+    GraphicsAllocation allocation{nullptr, 0, 0, 0, false};
+    EXPECT_EQ(AllocationOrigin::EXTERNAL_ALLOCATION, allocation.origin);
+    auto allocationType = GraphicsAllocation::AllocationType::LINEAR_STREAM;
+    allocation.setAllocationType(allocationType);
+    EXPECT_TRUE(GraphicsAllocation::isCpuAccessRequired(allocationType));
+    if (platformDevices[0]->capabilityTable.gpuAddressSpace != MemoryConstants::max48BitAddress) {
+        return;
+    }
+    EXPECT_EQ(HeapIndex::HEAP_STANDARD64Kb, MemoryManager::selectHeap(&allocation, nullptr, *platformDevices[0]));
+}
+
+TEST(MemoryManagerTest, givenFullAddressSpaceWhenSelectingHeapForExternalAllocationWithoutPtrAndCpuAccessIsNotRequiredThenStandardHeapIsUsed) {
+    GraphicsAllocation allocation{nullptr, 0, 0, 0, false};
+    EXPECT_EQ(AllocationOrigin::EXTERNAL_ALLOCATION, allocation.origin);
+    auto allocationType = GraphicsAllocation::AllocationType::UNDECIDED;
+    allocation.setAllocationType(allocationType);
+    EXPECT_FALSE(GraphicsAllocation::isCpuAccessRequired(allocationType));
+    if (platformDevices[0]->capabilityTable.gpuAddressSpace != MemoryConstants::max48BitAddress) {
+        return;
+    }
+    EXPECT_EQ(HeapIndex::HEAP_STANDARD, MemoryManager::selectHeap(&allocation, nullptr, *platformDevices[0]));
+}
+
+TEST(MemoryManagerTest, givenFullAddressSpaceWhenSelectingHeapForNullAllocationWithoutPtrThenStandardHeapIsUsed) {
+    if (platformDevices[0]->capabilityTable.gpuAddressSpace != MemoryConstants::max48BitAddress) {
+        return;
+    }
+    EXPECT_EQ(HeapIndex::HEAP_STANDARD, MemoryManager::selectHeap(nullptr, nullptr, *platformDevices[0]));
+}
+
+TEST(MemoryManagerTest, givenLimitedAddressSpaceWhenSelectingHeapForNullAllocationWithoutPtrThenLimitedHeapIsUsed) {
+    auto hwInfo = *platformDevices[0];
+    hwInfo.capabilityTable.gpuAddressSpace = MemoryConstants::max32BitAddress;
+    EXPECT_EQ(HeapIndex::HEAP_LIMITED, MemoryManager::selectHeap(nullptr, nullptr, hwInfo));
+}
diff --git a/unit_tests/os_interface/windows/wddm20_tests.cpp b/unit_tests/os_interface/windows/wddm20_tests.cpp
diff --git a/unit_tests/os_interface/windows/wddm_kmdaf_listener_tests.cpp b/unit_tests/os_interface/windows/wddm_kmdaf_listener_tests.cpp
diff --git a/unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp b/unit_tests/os_interface/windows/wddm_memory_manager_tests.cpp

Original file line number	Diff line number	Diff line change
`@@ -115,8 +115,8 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation> {`
`115`	`115`	`return gpuAddress + allocationOffset - gpuBaseAddress;`
`116`	`116`	`}`
`117`	`117`
`118`		`- bool isCoherent() { return coherent; };`
`119`		`- void setCoherent(bool coherentIn) { this->coherent = coherentIn; };`
	`118`	`+ bool isCoherent() const { return coherent; }`
	`119`	`+ void setCoherent(bool coherentIn) { this->coherent = coherentIn; }`
`120`	`120`	`void setSize(size_t size) { this->size = size; }`
`121`	`121`	`osHandle peekSharedHandle() { return sharedHandle; }`
`122`	`122`
Original file line number	Diff line number	Diff line change
`@@ -241,7 +241,9 @@ DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryWithAlignment(const Alloc`
`241`	`241`	`if (forcePinEnabled && pinBB != nullptr && allocationData.flags.forcePin && allocationData.size >= this->pinThreshold) {`
`242`	`242`	`pinBB->pin(&bo, 1, getDefaultCommandStreamReceiver(0)->getOsContext().get()->getDrmContextId());`
`243`	`243`	`}`
`244`		`- return new DrmAllocation(bo, res, cSize, MemoryPool::System4KBPages, allocationData.flags.multiOsContextCapable);`
	`244`	`+ auto allocation = new DrmAllocation(bo, res, cSize, MemoryPool::System4KBPages, allocationData.flags.multiOsContextCapable);`
	`245`	`+ allocation->origin = allocationData.allocationOrigin;`
	`246`	`+ return allocation;`
`245`	`247`	`}`
`246`	`248`
`247`	`249`	`DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryWithHostPtr(const AllocationData &allocationData) {`
`@@ -256,7 +258,7 @@ DrmAllocation *DrmMemoryManager::allocateGraphicsMemoryWithHostPtr(const Allocat`
`256`	`258`	`}`
`257`	`259`
`258`	`260`	`DrmAllocation DrmMemoryManager::allocateGraphicsMemoryForNonSvmHostPtr(size_t size, void cpuPtr) {`
`259`		`- if ((size == 0) \|\| !cpuPtr)`
	`261`	`+ if (size == 0 \|\| !cpuPtr)`
`260`	`262`	`return nullptr;`
`261`	`263`
`262`	`264`	`auto alignedPtr = alignDown(reinterpret_cast<char *>(cpuPtr), MemoryConstants::pageSize);`
`@@ -329,6 +331,7 @@ GraphicsAllocation *DrmMemoryManager::allocateGraphicsMemoryForImageImpl(const A`
`329`	`331`	`auto allocation = new DrmAllocation(bo, nullptr, (uint64_t)gpuRange, allocationData.imgInfo->size, MemoryPool::SystemCpuInaccessible, false);`
`330`	`332`	`bo->setAllocationType(allocatorType);`
`331`	`333`	`allocation->gmm = gmm.release();`
	`334`	`+ allocation->origin = allocationData.allocationOrigin;`
`332`	`335`	`return allocation;`
`333`	`336`	`}`
`334`	`337`
`@@ -363,6 +366,7 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio`
`363`	`366`	`allocationSize, MemoryPool::System4KBPagesWith32BitGpuAddressing, false);`
`364`	`367`	`drmAllocation->is32BitAllocation = true;`
`365`	`368`	`drmAllocation->gpuBaseAddress = allocatorToUse->getBase();`
	`369`	`+ drmAllocation->origin = allocationData.allocationOrigin;`
`366`	`370`	`return drmAllocation;`
`367`	`371`	`}`
`368`	`372`
`@@ -417,6 +421,7 @@ DrmAllocation *DrmMemoryManager::allocate32BitGraphicsMemoryImpl(const Allocatio`
`417`	`421`	`drmAllocation->is32BitAllocation = true;`
`418`	`422`	`drmAllocation->gpuBaseAddress = allocatorToUse->getBase();`
`419`	`423`	`drmAllocation->driverAllocatedCpuPointer = ptrAlloc;`
	`424`	`+ drmAllocation->origin = allocationData.allocationOrigin;`
`420`	`425`	`return drmAllocation;`
`421`	`426`	`}`
`422`	`427`