Merge pull request #57 from smoogipoo/reduce-mtl-cb-allocs

Use constant storage space for MTL command buffers

src/Veldrid.MetalBindings/MTLCommandBuffer.cs

@@ -5,7 +5,7 @@
 namespace Veldrid.MetalBindings
 {
     [StructLayout(LayoutKind.Sequential)]
-    public struct MTLCommandBuffer
+    public struct MTLCommandBuffer : IEquatable<MTLCommandBuffer>
     {
         public readonly IntPtr NativePtr;
 
@@ -29,6 +29,7 @@ public MTLComputeCommandEncoder computeCommandEncoder()
 
         public void addCompletedHandler(MTLCommandBufferHandler block)
             => objc_msgSend(NativePtr, sel_addCompletedHandler, block);
+
         public void addCompletedHandler(IntPtr block)
             => objc_msgSend(NativePtr, sel_addCompletedHandler, block);
 
@@ -42,5 +43,20 @@ public void addCompletedHandler(IntPtr block)
         private static readonly Selector sel_waitUntilCompleted = "waitUntilCompleted";
         private static readonly Selector sel_addCompletedHandler = "addCompletedHandler:";
         private static readonly Selector sel_status = "status";
+
+        public bool Equals(MTLCommandBuffer other)
+        {
+            return NativePtr == other.NativePtr;
+        }
+
+        public override bool Equals(object obj)
+        {
+            return obj is MTLCommandBuffer other && Equals(other);
+        }
+
+        public override int GetHashCode()
+        {
+            return NativePtr.GetHashCode();
+        }
     }
 }

src/Veldrid/MTL/CommandBufferUsageList.cs

@@ -0,0 +1,151 @@
+// Copyright (c) ppy Pty Ltd <[email protected]>. Licensed under the MIT Licence.
+// See the LICENCE file in the repository root for full licence text.
+
+using System.Collections;
+using System.Collections.Generic;
+using System.Runtime.InteropServices;
+using Veldrid.MetalBindings;
+
+namespace Veldrid.MTL
+{
+    internal class CommandBufferUsageList<T>
+    {
+        private readonly List<(MTLCommandBuffer buffer, T value)> items = new List<(MTLCommandBuffer buffer, T item)>();
+
+        public void Add(MTLCommandBuffer cb, T value)
+            => items.Add((cb, value));
+
+        public ItemsEnumerator EnumerateItems()
+            => new ItemsEnumerator(items);
+
+        public RemovalEnumerator EnumerateAndRemove(MTLCommandBuffer cb)
+            => new RemovalEnumerator(items, cb);
+
+        public bool Contains(MTLCommandBuffer cb)
+        {
+            foreach (var (buffer, _) in items)
+            {
+                if (buffer.Equals(cb))
+                    return true;
+            }
+
+            return false;
+        }
+
+        public void Clear()
+            => items.Clear();
+
+        /// <summary>
+        /// This is a basic enumerator for the list.
+        /// </summary>
+        public struct ItemsEnumerator : IEnumerator<T>, IEnumerable
+        {
+            private readonly List<(MTLCommandBuffer buffer, T value)> list;
+            private int index;
+
+            public ItemsEnumerator(List<(MTLCommandBuffer buffer, T value)> list)
+            {
+                this.list = list;
+            }
+
+            public bool MoveNext()
+            {
+                if (index == list.Count)
+                    return false;
+
+                Current = list[index].value;
+                index++;
+
+                return true;
+            }
+
+            public void Reset()
+            {
+                index = 0;
+            }
+
+            public T Current { get; private set; }
+
+            object IEnumerator.Current => Current;
+
+            public void Dispose()
+            {
+            }
+
+            public ItemsEnumerator GetEnumerator() => this;
+
+            IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+        }
+
+        /// <summary>
+        /// This is a combined enumerate + remove enumerator for the list.
+        ///
+        /// It works by duplicating the items that shall be retained to the end of the list
+        /// and then moving them in-place to the front of the list upon disposal.
+        ///
+        /// The combined operation has therefore O(n) time complexity.
+        /// </summary>
+        public struct RemovalEnumerator : IEnumerator<T>, IEnumerable
+        {
+            private readonly List<(MTLCommandBuffer buffer, T value)> list;
+            private readonly MTLCommandBuffer cb;
+            private readonly int count;
+            private int index;
+
+            public RemovalEnumerator(List<(MTLCommandBuffer buffer, T value)> list, MTLCommandBuffer cb)
+            {
+                this.list = list;
+                this.cb = cb;
+
+                count = list.Count;
+                list.EnsureCapacity(count * 2);
+            }
+
+            public bool MoveNext()
+            {
+                while (true)
+                {
+                    if (index == count)
+                        return false;
+
+                    if (list[index].buffer.Equals(cb))
+                        break;
+
+                    // Track the item to be kept.
+                    list.Add(list[index]);
+                    index++;
+                }
+
+                Current = list[index].value;
+                index++;
+
+                return true;
+            }
+
+            public void Reset()
+            {
+                index = 0;
+            }
+
+            public T Current { get; private set; }
+
+            object IEnumerator.Current => Current;
+
+            public void Dispose()
+            {
+                if (list.Count == 0)
+                    return;
+
+                int toKeepItemCount = list.Count - count;
+                var listSpan = CollectionsMarshal.AsSpan(list);
+
+                listSpan.Slice(count, toKeepItemCount).CopyTo(listSpan);
+                list.RemoveRange(toKeepItemCount, list.Count - toKeepItemCount);
+            }
+
+            public RemovalEnumerator GetEnumerator() => this;
+
+            IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+        }
+    }
+}

src/Veldrid/MTL/MTLCommandList.cs

@@ -15,9 +15,9 @@ internal unsafe class MtlCommandList : CommandList
         private readonly MtlGraphicsDevice gd;
 
         private readonly List<MtlBuffer> availableStagingBuffers = new List<MtlBuffer>();
-        private readonly Dictionary<MTLCommandBuffer, List<MtlBuffer>> submittedStagingBuffers = new Dictionary<MTLCommandBuffer, List<MtlBuffer>>();
+        private readonly CommandBufferUsageList<MtlBuffer> submittedStagingBuffers = new CommandBufferUsageList<MtlBuffer>();
         private readonly object submittedCommandsLock = new object();
-        private readonly Dictionary<MTLCommandBuffer, MtlFence> completionFences = new Dictionary<MTLCommandBuffer, MtlFence>();
+        private readonly CommandBufferUsageList<MtlFence> completionFences = new CommandBufferUsageList<MtlFence>();
 
         private readonly Dictionary<UIntPtr, DeviceBufferRange> boundVertexBuffers = new Dictionary<UIntPtr, DeviceBufferRange>();
         private readonly Dictionary<UIntPtr, DeviceBufferRange> boundFragmentBuffers = new Dictionary<UIntPtr, DeviceBufferRange>();
@@ -85,12 +85,11 @@ public override void Dispose()
 
                 lock (submittedStagingBuffers)
                 {
-                    foreach (var buffer in availableStagingBuffers) buffer.Dispose();
+                    foreach (var buffer in availableStagingBuffers)
+                        buffer.Dispose();
 
-                    foreach (var kvp in submittedStagingBuffers)
-                    {
-                        foreach (var buffer in kvp.Value) buffer.Dispose();
-                    }
+                    foreach (var buffer in submittedStagingBuffers.EnumerateItems())
+                        buffer.Dispose();
 
                     submittedStagingBuffers.Clear();
                 }
@@ -160,26 +159,20 @@ public void SetCompletionFence(MTLCommandBuffer cb, MtlFence fence)
         {
             lock (submittedCommandsLock)
             {
-                Debug.Assert(!completionFences.ContainsKey(cb));
-                completionFences[cb] = fence;
+                Debug.Assert(!completionFences.Contains(cb));
+                completionFences.Add(cb, fence);
             }
         }
 
         public void OnCompleted(MTLCommandBuffer cb)
         {
             lock (submittedCommandsLock)
             {
-                if (completionFences.TryGetValue(cb, out var completionFence))
-                {
-                    completionFence.Set();
-                    completionFences.Remove(cb);
-                }
+                foreach (var fence in completionFences.EnumerateAndRemove(cb))
+                    fence.Set();
 
-                if (submittedStagingBuffers.TryGetValue(cb, out var bufferList))
-                {
-                    availableStagingBuffers.AddRange(bufferList);
-                    submittedStagingBuffers.Remove(cb);
-                }
+                foreach (var buffer in submittedStagingBuffers.EnumerateAndRemove(cb))
+                    availableStagingBuffers.Add(buffer);
             }
         }
 
@@ -1267,11 +1260,7 @@ private protected override void UpdateBufferCore(DeviceBuffer buffer, uint buffe
             }
 
             lock (submittedCommandsLock)
-            {
-                if (!submittedStagingBuffers.TryGetValue(cb, out var bufferList)) submittedStagingBuffers[cb] = bufferList = new List<MtlBuffer>();
-
-                bufferList.Add(staging);
-            }
+                submittedStagingBuffers.Add(cb, staging);
         }
 
         private protected override void GenerateMipmapsCore(Texture texture)

src/Veldrid/MTL/MTLGraphicsDevice.cs

@@ -49,7 +49,7 @@ private static readonly Dictionary<IntPtr, MtlGraphicsDevice> s_aot_registered_b
         private readonly bool[] supportedSampleCounts;
 
         private readonly object submittedCommandsLock = new object();
-        private readonly Dictionary<MTLCommandBuffer, MtlCommandList> submittedCLs = new Dictionary<MTLCommandBuffer, MtlCommandList>();
+        private readonly CommandBufferUsageList<MtlCommandList> submittedCLs = new CommandBufferUsageList<MtlCommandList>();
 
         private readonly object resetEventsLock = new object();
         private readonly List<ManualResetEvent[]> resetEvents = new List<ManualResetEvent[]>();
@@ -385,11 +385,11 @@ private void OnCommandBufferCompleted(IntPtr block, MTLCommandBuffer cb)
         {
             lock (submittedCommandsLock)
             {
-                var cl = submittedCLs[cb];
-                submittedCLs.Remove(cb);
-                cl.OnCompleted(cb);
+                foreach (var cl in submittedCLs.EnumerateAndRemove(cb))
+                    cl.OnCompleted(cb);
 
-                if (latestSubmittedCb.NativePtr == cb.NativePtr) latestSubmittedCb = default;
+                if (latestSubmittedCb.NativePtr == cb.NativePtr)
+                    latestSubmittedCb = default;
             }
 
             ObjectiveCRuntime.release(cb.NativePtr);
@@ -459,7 +459,8 @@ private protected override void SubmitCommandsCore(CommandList commandList, Fenc
 
             lock (submittedCommandsLock)
             {
-                if (fence != null) mtlCl.SetCompletionFence(mtlCl.CommandBuffer, Util.AssertSubtype<Fence, MtlFence>(fence));
+                if (fence != null)
+                    mtlCl.SetCompletionFence(mtlCl.CommandBuffer, Util.AssertSubtype<Fence, MtlFence>(fence));
 
                 submittedCLs.Add(mtlCl.CommandBuffer, mtlCl);
                 latestSubmittedCb = mtlCl.Commit();