Allow Heap to be COM-wrapped.

README.md

@@ -7,7 +7,7 @@
 
 # GPGMM
 
-GPGMM is a General-Purpose GPU Memory Management C++ library used by GPU applications or middleware libraries that rely on low-level graphics and compute APIs (D3D12 or Vulkan) for "explicit" memory management. GPGMM is a fast, multi-threaded Video Memory Manager (VidMM) implementation that replaces what older "implicit" graphics and compute APIs (D3D11 or OpenGL) had accomplished through the GPU driver.
+GPGMM is a General-Purpose GPU Memory Management C++ library used by GPU applications or middleware runtimes that rely on low-level graphics and compute APIs (D3D12 or Vulkan) for "explicit" memory management. GPGMM is a fast, multi-threaded, full-fledged GPU Memory Manager (GMM) implementation that replaces what older "implicit" graphics and compute APIs (D3D11 or OpenGL) had accomplished through the GPU driver.
 
 * [API Documentation](https://intel.github.io/GPGMM/)
 * [Changelog](https://github.com/intel/GPGMM/releases)
@@ -64,12 +64,21 @@ set.Reset();
 Residency also works for non-resources too:
 
 ```cpp
-// Wraps a ID3D12DescriptorHeap, ID3D12QueryHeap, etc.
-class HeapWrapper : public gpgmm::d3d12::Heap {};
-
-HeapWrapper heap;
-residency->InsertHeap(&heap);
-residency->LockHeap(heap) // Pin it (eg. shader-visible)
+gpgmm::d3d12::HEAP_DESC shaderVisibleHeap = {};
+shaderVisibleHeap.SizeInBytes = kHeapSize;
+shaderVisibleHeap.MemorySegment = gpgmm::d3d12::RESIDENCY_SEGMENT_LOCAL;
+
+ComPtr<gpgmm::d3d12::Heap> descriptorHeap;
+gpgmm::d3d12::Heap::CreateHeap(shaderVisibleHeap, residencyManager,
+  [&](ID3D12Pageable** ppPageableOut) -> HRESULT {
+      ComPtr<ID3D12DescriptorHeap> heap;
+      mDevice->CreateDescriptorHeap(&heapDesc, IID_PPV_ARGS(&heap));
+      *ppPageableOut = heap.Detach();
+      return S_OK;
+}, &descriptorHeap);
+
+// Shader-visible heaps should be locked or always resident.
+residency->LockHeap(descriptorHeap.get());
 ```
 
 To clean-up, simply call `Release()` once the is GPU is finished.

src/gpgmm/common/BuddyMemoryAllocator.cpp

@@ -114,7 +114,7 @@ namespace gpgmm {
         MemoryBase* memory = memoryAllocation->GetMemory();
         ASSERT(memory != nullptr);
 
-        if (memory->Unref()) {
+        if (memory->RemoveSubAllocationRef()) {
             GetNextInChain()->DeallocateMemory(std::move(memoryAllocation));
         } else {
             mUsedPool.ReturnToPool(std::move(memoryAllocation), memoryIndex);

src/gpgmm/common/Memory.cpp

@@ -20,12 +20,12 @@
 namespace gpgmm {
 
     MemoryBase::MemoryBase(uint64_t size, uint64_t alignment)
-        : RefCounted(0), mSize(size), mAlignment(alignment) {
+        : mSubAllocationRefs(0), mSize(size), mAlignment(alignment) {
         ASSERT(mSize != kInvalidSize);
     }
 
     MemoryBase::~MemoryBase() {
-        ASSERT(GetRefCount() == 0);
+        ASSERT(mSubAllocationRefs.GetRefCount() == 0);
     }
 
     uint64_t MemoryBase::GetSize() const {
@@ -45,4 +45,12 @@ namespace gpgmm {
         mPool = pool;
     }
 
+    void MemoryBase::AddSubAllocationRef() {
+        mSubAllocationRefs.Ref();
+    }
+
+    bool MemoryBase::RemoveSubAllocationRef() {
+        return mSubAllocationRefs.Unref();
+    }
+
 }  // namespace gpgmm

src/gpgmm/common/Memory.h

@@ -26,7 +26,7 @@ namespace gpgmm {
 
     When memory is sub-allocated, it will have a non-zero refcount.
     */
-    class MemoryBase : public RefCounted, public NonCopyable {
+    class MemoryBase : public NonCopyable {
       public:
         /** \brief Constructs a memory object of the specified size and alignment.
 
@@ -60,7 +60,17 @@ namespace gpgmm {
         */
         void SetPool(MemoryPool* pool);
 
+        /** \brief Increments the sub-allocation reference count on the heap.
+         */
+        void AddSubAllocationRef();
+
+        /** \brief Decrements the sub-allocation reference count on the heap.
+         */
+        bool RemoveSubAllocationRef();
+
       private:
+        RefCounted mSubAllocationRefs;
+
         const uint64_t mSize;
         const uint64_t mAlignment;
 

src/gpgmm/common/MemoryAllocator.h

@@ -300,7 +300,7 @@ namespace gpgmm {
             }
 
             ASSERT(memory != nullptr);
-            memory->Ref();
+            memory->AddSubAllocationRef();
 
             // Caller is be responsible in fully initializing the memory allocation.
             // This is because TrySubAllocateMemory() does not necessarily know how to map the

src/gpgmm/common/SlabMemoryAllocator.cpp

@@ -389,7 +389,7 @@ namespace gpgmm {
         MemoryBase* slabMemory = subAllocation->GetMemory();
         ASSERT(slabMemory != nullptr);
 
-        slabMemory->Unref();
+        slabMemory->RemoveSubAllocationRef();
 
         if (slab->IsEmpty()) {
             mMemoryAllocator->DeallocateMemory(std::move(slab->Allocation));

src/gpgmm/d3d12/HeapD3D12.h

@@ -18,6 +18,7 @@
 
 #include "gpgmm/common/Memory.h"
 #include "gpgmm/d3d12/DebugObjectD3D12.h"
+#include "gpgmm/d3d12/IUnknownImplD3D12.h"
 #include "gpgmm/utils/Limits.h"
 #include "gpgmm/utils/LinkedList.h"
 #include "gpgmm/utils/RefCount.h"
@@ -130,7 +131,10 @@ namespace gpgmm::d3d12 {
     node is removed from the cache when it is evicted from video memory due to budget constraints,
     or when the memory is released.
     */
-    class GPGMM_EXPORT Heap : public MemoryBase, public DebugObject, public LinkNode<Heap> {
+    class GPGMM_EXPORT Heap : public MemoryBase,
+                              public DebugObject,
+                              public IUnknownImpl,
+                              public LinkNode<Heap> {
       public:
         /** \brief  Create a heap managed by GPGMM.
 

src/tests/end2end/D3D12ResidencyManagerTests.cpp

@@ -99,14 +99,14 @@ TEST_F(D3D12ResidencyManagerTests, CreateResourceHeap) {
 
     constexpr uint64_t kHeapSize = GPGMM_MB_TO_BYTES(10);
 
-    auto createResourceHeapFn = [&](ID3D12Pageable** ppPageableOut) -> HRESULT {
-        D3D12_HEAP_PROPERTIES heapProperties = {};
-        heapProperties.Type = D3D12_HEAP_TYPE_DEFAULT;
+    D3D12_HEAP_PROPERTIES heapProperties = {};
+    heapProperties.Type = D3D12_HEAP_TYPE_DEFAULT;
 
-        D3D12_HEAP_DESC heapDesc = {};
-        heapDesc.Properties = heapProperties;
-        heapDesc.SizeInBytes = kHeapSize;
+    D3D12_HEAP_DESC heapDesc = {};
+    heapDesc.Properties = heapProperties;
+    heapDesc.SizeInBytes = kHeapSize;
 
+    auto createHeapFn = [&](ID3D12Pageable** ppPageableOut) -> HRESULT {
         ComPtr<ID3D12Heap> heap;
         if (FAILED(mDevice->CreateHeap(&heapDesc, IID_PPV_ARGS(&heap)))) {
             return E_FAIL;
@@ -117,12 +117,20 @@ TEST_F(D3D12ResidencyManagerTests, CreateResourceHeap) {
 
     HEAP_DESC resourceHeapDesc = {};
     resourceHeapDesc.SizeInBytes = kHeapSize;
-    resourceHeapDesc.DebugName = "Resource heap";
-    resourceHeapDesc.AlwaysInBudget = true;
     resourceHeapDesc.HeapType = D3D12_HEAP_TYPE_DEFAULT;
 
     ASSERT_SUCCEEDED(
-        Heap::CreateHeap(resourceHeapDesc, residencyManager.Get(), createResourceHeapFn, nullptr));
+        Heap::CreateHeap(resourceHeapDesc, residencyManager.Get(), createHeapFn, nullptr));
+
+    ComPtr<Heap> resourceHeap;
+    ASSERT_SUCCEEDED(
+        Heap::CreateHeap(resourceHeapDesc, residencyManager.Get(), createHeapFn, &resourceHeap));
+    ASSERT_NE(resourceHeap, nullptr);
+
+    ComPtr<ID3D12Heap> heap;
+    resourceHeap->As(&heap);
+
+    EXPECT_NE(heap, nullptr);
 }
 
 TEST_F(D3D12ResidencyManagerTests, CreateResidencySet) {