Metal: Support vertex stride 0 and null buffers, and fix resource release race (#38)

* Fix race on resource release + fixes for unleashed

* Add nullBuffer and stride=0 support

* Use array instead of vector for vertex slot tracking.

* Remove extra comment

* Fix slotUsed initialization

---------

Co-authored-by: tanmaysachan <tnmysachan@gmail.com>

Author: squidbus

plume_metal.cpp

@@ -26,11 +26,6 @@ namespace plume {
     static constexpr size_t PUSH_CONSTANT_MAX_INDEX = 15;
     static constexpr size_t VERTEX_BUFFER_MAX_INDEX = 30;
 
-    static constexpr uint32_t COLOR_COUNT = DESCRIPTOR_SET_MAX_INDEX;
-    static constexpr uint32_t DEPTH_INDEX = COLOR_COUNT;
-    static constexpr uint32_t STENCIL_INDEX = DEPTH_INDEX + 1;
-    static constexpr uint32_t ATTACHMENT_COUNT = STENCIL_INDEX + 1;
-
     // MARK: - Prototypes
 
     MTL::PixelFormat mapPixelFormat(RenderFormat format);
@@ -1440,9 +1435,17 @@ namespace plume {
             // Set index right after push constants, clamp at Metal's limit of 31
             const uint32_t vertexBufferIndex = std::min(PUSH_CONSTANT_MAX_INDEX + 1 + inputSlot.index, VERTEX_BUFFER_MAX_INDEX);
             MTL::VertexBufferLayoutDescriptor *layout = vertexDescriptor->layouts()->object(vertexBufferIndex);
-            layout->setStride(inputSlot.stride);
-            layout->setStepFunction(mapVertexStepFunction(inputSlot.classification));
-            layout->setStepRate((layout->stepFunction() == MTL::VertexStepFunctionPerInstance) ? inputSlot.stride : 1);
+            if (inputSlot.stride == 0) {
+                // Metal does not support stride 0, we must provide a
+                // substitute "null" buffer to match behaviour of other robust APIs.
+                layout->setStride(1);
+                layout->setStepFunction(MTL::VertexStepFunctionConstant);
+                layout->setStepRate(0);
+            } else {
+                layout->setStride(inputSlot.stride);
+                layout->setStepFunction(mapVertexStepFunction(inputSlot.classification));
+                layout->setStepRate((layout->stepFunction() == MTL::VertexStepFunctionPerInstance) ? inputSlot.stride : 1);
+            }
         }
 
         for (uint32_t i = 0; i < desc.inputElementsCount; i++) {
@@ -1610,6 +1613,10 @@ namespace plume {
     }
 
     MetalDescriptorSet::~MetalDescriptorSet() {
+        for (const auto resource: toReleaseOnDestruction) {
+            resource->release();
+        }
+
         for (const auto &entry : resourceEntries) {
             if (entry.resource != nullptr) {
                 entry.resource->release();
@@ -1705,8 +1712,7 @@ namespace plume {
 
         if (dtype != MTL::DataTypeSampler) {
             if (resourceEntries[descriptorIndex].resource != nullptr) {
-                resourceEntries[descriptorIndex].resource->release();
-                resourceEntries[descriptorIndex].resource = nullptr;
+                toReleaseOnDestruction.push_back(resourceEntries[descriptorIndex].resource);
             }
         }
 
@@ -2216,7 +2222,7 @@ namespace plume {
         pushConstants[rangeIndex].binding = range.binding;
         pushConstants[rangeIndex].set = range.set;
         pushConstants[rangeIndex].offset = range.offset;
-        pushConstants[rangeIndex].size = range.size;
+        pushConstants[rangeIndex].size = alignUp(range.size);
         pushConstants[rangeIndex].stageFlags = range.stageFlags;
 
         dirtyComputeState.pushConstants = 1;
@@ -2267,7 +2273,7 @@ namespace plume {
         pushConstants[rangeIndex].binding = range.binding;
         pushConstants[rangeIndex].set = range.set;
         pushConstants[rangeIndex].offset = range.offset;
-        pushConstants[rangeIndex].size = range.size;
+        pushConstants[rangeIndex].size = alignUp(range.size);
         pushConstants[rangeIndex].stageFlags = range.stageFlags;
 
         dirtyGraphicsState.pushConstants = 1;
@@ -2322,9 +2328,14 @@ namespace plume {
             // Check for changes in bindings
             for (uint32_t i = 0; i < viewCount; i++) {
                 const MetalBuffer* interfaceBuffer = static_cast<const MetalBuffer*>(views[i].buffer.ref);
-                const uint64_t newOffset = views[i].buffer.offset;
+                uint64_t newOffset = views[i].buffer.offset;
                 const uint32_t newIndex = startSlot + i;
 
+                if (interfaceBuffer == nullptr) {
+                    interfaceBuffer = static_cast<const MetalBuffer*>(queue->device->nullBuffer.get());
+                    newOffset = 0;
+                }
+
                 vertexBuffers[i] = interfaceBuffer->mtl;
                 vertexBufferOffsets[i] = newOffset;
                 vertexBufferIndices[i] = newIndex;
@@ -2962,10 +2973,21 @@ namespace plume {
         }
 
         if (dirtyGraphicsState.vertexBuffers) {
+            std::array<bool, 31> slotUsed{};
+
             for (uint32_t i = 0; i < viewCount; i++) {
                 // Bind right after the push constants, up till the max vertex buffer index
                 const uint32_t bindIndex = std::min(PUSH_CONSTANT_MAX_INDEX + 1 + vertexBufferIndices[i], VERTEX_BUFFER_MAX_INDEX);
                 activeRenderEncoder->setVertexBuffer(vertexBuffers[i], vertexBufferOffsets[i], bindIndex);
+                slotUsed[bindIndex] = true;
+            }
+
+            // Bind all unbound slots to the null buffer
+            auto nullBuffer = static_cast<const MetalBuffer*>(queue->device->nullBuffer.get());
+            for (uint32_t i = PUSH_CONSTANT_MAX_INDEX + 1; i <= VERTEX_BUFFER_MAX_INDEX; i++) {
+                if (!slotUsed[i]) {
+                    activeRenderEncoder->setVertexBuffer(nullBuffer->mtl, 0, i);
+                }
             }
 
             stateCache.lastVertexBuffers = vertexBuffers;
@@ -3262,6 +3284,8 @@ namespace plume {
         capabilities.dynamicDepthBias = true;
         capabilities.uma = mtl->hasUnifiedMemory();
         capabilities.queryPools = false;
+
+        nullBuffer = createBuffer(RenderBufferDesc::DefaultBuffer(16, RenderBufferFlag::VERTEX));
     }
 
     MetalDevice::~MetalDevice() {

plume_metal.h

@@ -209,6 +209,7 @@ namespace plume {
         std::vector<Descriptor> descriptors;
         MetalArgumentBuffer argumentBuffer;
         std::vector<ResourceEntry> resourceEntries;
+        std::vector<MTL::Resource *> toReleaseOnDestruction;
 
         MetalDescriptorSet(MetalDevice *device, const RenderDescriptorSetDesc &desc);
         MetalDescriptorSet(MetalDevice *device, uint32_t entryCount);
@@ -608,6 +609,8 @@ namespace plume {
         // Blit functionality
         MTL::BlitPassDescriptor *sharedBlitDescriptor = nullptr;
 
+        std::unique_ptr<RenderBuffer> nullBuffer;
+
         explicit MetalDevice(MetalInterface *renderInterface, const std::string &preferredDeviceName);
         ~MetalDevice() override;
         std::unique_ptr<RenderDescriptorSet> createDescriptorSet(const RenderDescriptorSetDesc &desc) override;