Start adding a Metal readback pool to support GPU downloads

Source/modules/render/metal/MetalDevice.cpp

@@ -11,7 +11,7 @@
 // #include "MetalDescriptorSetPools.h"
 // #include "MetalDisplay.h"
 #include "MetalResources.h"
-// #include "MetalReadbackBufferPool.h"
+#include "MetalReadbackBufferPool.h"
 #include "MetalWorkBundle.h"
 #include "MetalQueues.h"
 // #include "MetalEventPool.h"
@@ -249,7 +249,7 @@ MetalDevice::MetalDevice(const DeviceConfig& config)
     // m_gc = new MetalGc(125, *this);
     m_resources = new MetalResources(*this, m_workDb);
     // m_descriptorSetPools = new MetalDescriptorSetPools(*this);
-    // m_readbackPool = new MetalReadbackBufferPool(*this);
+    m_readbackPool = new MetalReadbackBufferPool(*this);
     // m_counterPool = new MetalCounterPool(*this);
     // m_markerCollector = new MetalMarkerCollector(*this);
 
@@ -290,8 +290,8 @@ MetalDevice::~MetalDevice()
 
     // delete m_markerCollector;
     // m_markerCollector = nullptr;
-    // delete m_readbackPool;
-    // m_readbackPool = nullptr;
+    delete m_readbackPool;
+    m_readbackPool = nullptr;
     delete m_resources;
     m_resources = nullptr;
     // delete m_descriptorSetPools;

Source/modules/render/metal/MetalDevice.h

@@ -20,7 +20,7 @@ namespace render
 {
 
 // class MetalDescriptorSetPools;
-// class MetalReadbackBufferPool;
+class MetalReadbackBufferPool;
 class MetalResources;
 class MetalQueues;
 // class MetalCounterPool;
@@ -67,7 +67,7 @@ class MetalDevice : public TDevice<MetalDevice>
     // MetalDescriptorSetPools& descriptorSetPools() { return *m_descriptorSetPools; }
     // int graphicsFamilyQueueIndex() const { return m_queueFamIndex; }
 
-    // MetalReadbackBufferPool& readbackPool() { return *m_readbackPool; }
+    MetalReadbackBufferPool& readbackPool() { return *m_readbackPool; }
 
     // bool findMemoryType(uint32_t typeFilter, MtlMemoryPropertyFlags properties, uint32_t& outMemType);
 
@@ -100,7 +100,7 @@ class MetalDevice : public TDevice<MetalDevice>
     // MetalDescriptorSetPools* m_descriptorSetPools;
     MetalQueues* m_queues;
     MetalResources* m_resources;
-    // MetalReadbackBufferPool* m_readbackPool;
+    MetalReadbackBufferPool* m_readbackPool;
     // MetalEventPool* m_eventPool;
     // MetalFencePool* m_fencePool;
     // MetalCounterPool* m_counterPool;

Source/modules/render/metal/MetalReadbackBufferPool.cpp

@@ -0,0 +1,67 @@
+#include <Config.h>
+#if ENABLE_METAL
+
+#include <Metal/Metal.h>
+#include "MetalReadbackBufferPool.h"
+#include "MetalDevice.h"
+#include "MetalResources.h"
+
+namespace coalpy
+{
+namespace render
+{
+
+enum 
+{
+    InitialBufferPoolSize = 1024 * 1024 * 5 //5 megabytes of initial size
+};
+
+bool MetalReadbackBufferPool::createNewHeap(size_t size)
+{
+    short heapIndex = (short)m_heaps.size();
+    m_heaps.emplace_back();
+    HeapState& heap = m_heaps.back();
+
+    BufferDesc bufferDesc;
+    bufferDesc.type = BufferType::Standard;
+    bufferDesc.format = Format::RGBA_8_UINT;
+    bufferDesc.elementCount = size;
+    bufferDesc.memFlags = (MemFlags)0;
+    ResourceSpecialFlags readbackFlags = ResourceSpecialFlag_CpuReadback;
+    BufferResult result = m_device.resources().createBuffer(bufferDesc, readbackFlags);
+    if (!result.success())
+        return false;
+
+    heap.buffer = result;
+    heap.largestSize = size;
+    MetalReadbackMemBlock memBlock;
+    memBlock.size = size;
+    memBlock.offset = 0;
+    memBlock.buffer = result;
+    memBlock.allocationId = m_nextAllocId++;
+    memBlock.heapIndex = heapIndex;
+
+    return true;
+}
+
+MetalReadbackBufferPool::MetalReadbackBufferPool(MetalDevice& device)
+: m_device(device)
+{
+    bool success = createNewHeap(InitialBufferPoolSize);
+    CPY_ASSERT_MSG(success, "Could not allocate heap.");
+    m_nextHeapSize = 2 * InitialBufferPoolSize;
+}
+
+MetalReadbackBufferPool::~MetalReadbackBufferPool()
+{
+    for (auto& h : m_heaps)
+    {
+        m_device.resources().release(h.buffer);
+        CPY_ASSERT_MSG(h.freeBlocks.size() == 1, "Warning: non freed heap elements found. Buffer pool has some memory blocks that have not been freed.");
+    }
+    m_heaps.clear();
+}
+
+}
+}
+#endif // ENABLE_METAL

Source/modules/render/metal/MetalReadbackBufferPool.h

@@ -0,0 +1,54 @@
+#pragma once
+
+#include <vector>
+#include <coalpy.render/Resources.h>
+
+@protocol MTLBuffer;
+
+namespace coalpy
+{
+namespace render
+{
+
+class MetalDevice;
+class MetalResource;
+
+struct MetalReadbackMemBlock
+{
+    size_t size = 0ull;
+    size_t offset = 0ull;
+    Buffer buffer = {};
+    short allocationId = 0;
+    short heapIndex = 0;
+
+    bool valid() const { return buffer.valid(); }
+};
+
+class MetalReadbackBufferPool
+{
+public:
+    MetalReadbackBufferPool(MetalDevice& device);
+    ~MetalReadbackBufferPool();
+
+    MetalReadbackMemBlock allocate(size_t size);
+    void free(const MetalReadbackMemBlock& block);
+
+private:
+    MetalDevice& m_device;
+
+    bool createNewHeap(size_t size);
+
+    struct HeapState
+    {
+        Buffer buffer;
+        std::vector<MetalReadbackMemBlock> freeBlocks;
+        size_t largestSize = 0ull;
+    };
+
+    std::vector<HeapState> m_heaps;
+    size_t m_nextHeapSize = 0;
+    short m_nextAllocId = 0;
+    bool m_isReadback;
+};
+}
+}

Source/modules/render/metal/MetalResources.h

@@ -22,7 +22,7 @@ enum ResourceSpecialFlags : int
     ResourceSpecialFlag_NoDeferDelete = 1 << 0,
     // ResourceSpecialFlag_CanDenyShaderResources = 1 << 1,
     ResourceSpecialFlag_TrackTables = 1 << 2,
-    // ResourceSpecialFlag_CpuReadback = 1 << 3,
+    ResourceSpecialFlag_CpuReadback = 1 << 3,
     // ResourceSpecialFlag_CpuUpload = 1 << 4,
     // ResourceSpecialFlag_EnableColorAttachment = 1 << 5,
 };

Source/modules/render/metal/MetalWorkBundle.cpp

@@ -57,7 +57,7 @@ static void buildDownloadCmd(
     MetalResource& metalResource = resources.unsafeGetResource(downloadCmd->source);
     if (metalResource.type == MetalResource::Type::Buffer)
     {
-        // downloadState->memoryBlock = m_device.readbackPool().allocate(metalResource.byteSize());
+        // downloadState->memoryBlock = device->readbackPool().allocate(metalResource.byteSize());
 
     }
     else

Source/modules/render/metal/MetalWorkBundle.h

@@ -2,6 +2,7 @@
 #include <stdint.h>
 #include <vector>
 #include "WorkBundleDb.h"
+#include "MetalReadbackBufferPool.h"
 
 namespace coalpy
 {
@@ -14,6 +15,7 @@ class MetalDevice;
 struct MetalResourceDownloadState
 {
     ResourceDownloadKey downloadKey;
+    MetalReadbackMemBlock memoryBlock;
 };
 
 using MetalDownloadResourceMap = std::unordered_map<ResourceDownloadKey, MetalResourceDownloadState>;

Source/tests/RenderTests.cpp

@@ -180,7 +180,7 @@ void vulkanBufferPool(TestContext& ctx)
 
 #if ENABLE_METAL
 // TODO (Apoorva)
-void metalBufferPool(TestContext& ctx)
+void MetalReadbackBufferPool(TestContext& ctx)
 {
     auto& renderTestCtx = (RenderTestContext&)ctx;
     renderTestCtx.begin();
@@ -2006,7 +2006,7 @@ static const TestCase* createCases(int& caseCounts)
         { "vulkanBufferPool", vulkanBufferPool },
 #endif
 #if ENABLE_METAL
-        // { "metalBufferPool", metalBufferPool },
+        // { "MetalReadbackBufferPool", MetalReadbackBufferPool },
 #endif
         { "createBuffer",  testCreateBuffer },
         { "createTexture", testCreateTexture },
@@ -2044,7 +2044,7 @@ static const TestCaseFilter* createCasesFilters(int& caseCounts)
         { "vulkanBufferPool", TestPlatformVulkan },
 #endif
 #if  ENABLE_METAL
-        { "metalBufferPool", TestPlatformMetal },
+        { "MetalReadbackBufferPool", TestPlatformMetal },
 #endif
     };