Vulkan separate descriptor set fixes (#6895)

* Ensure descriptor sets are only re-used when all command buffers using it have completed

* Fix some SPIR-V capabilities

* Set update after bind flag if we exceed limits

* Simpler fix for Intel

* Format whitespace

* Make struct readonly

* Add barriers for extra set arrays too

1 files changed, 102 insertions, 22 deletions

diff --git a/src/Ryujinx.Graphics.Vulkan/PipelineLayoutCacheEntry.cs b/src/Ryujinx.Graphics.Vulkan/PipelineLayoutCacheEntry.cs
index 7d0948d6..ae296b03 100644
--- a/src/Ryujinx.Graphics.Vulkan/PipelineLayoutCacheEntry.cs
+++ b/src/Ryujinx.Graphics.Vulkan/PipelineLayoutCacheEntry.cs
@@ -3,6 +3,7 @@ using Silk.NET.Vulkan;
 using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
+using System.Diagnostics;
 using System.Runtime.InteropServices;
 
 namespace Ryujinx.Graphics.Vulkan
@@ -15,6 +16,7 @@ namespace Ryujinx.Graphics.Vulkan
         private readonly Device _device;
 
         public DescriptorSetLayout[] DescriptorSetLayouts { get; }
+        public bool[] DescriptorSetLayoutsUpdateAfterBind { get; }
         public PipelineLayout PipelineLayout { get; }
 
         private readonly int[] _consumedDescriptorsPerSet;
@@ -31,20 +33,37 @@ namespace Ryujinx.Graphics.Vulkan
         private struct ManualDescriptorSetEntry
         {
             public Auto<DescriptorSetCollection> DescriptorSet;
-            public int CbIndex;
-            public int CbSubmissionCount;
+            public uint CbRefMask;
             public bool InUse;
 
-            public ManualDescriptorSetEntry(Auto<DescriptorSetCollection> descriptorSet, int cbIndex, int cbSubmissionCount, bool inUse)
+            public ManualDescriptorSetEntry(Auto<DescriptorSetCollection> descriptorSet, int cbIndex)
             {
                 DescriptorSet = descriptorSet;
-                CbIndex = cbIndex;
-                CbSubmissionCount = cbSubmissionCount;
-                InUse = inUse;
+                CbRefMask = 1u << cbIndex;
+                InUse = true;
+            }
+        }
+
+        private readonly struct PendingManualDsConsumption
+        {
+            public FenceHolder Fence { get; }
+            public int CommandBufferIndex { get; }
+            public int SetIndex { get; }
+            public int CacheIndex { get; }
+
+            public PendingManualDsConsumption(FenceHolder fence, int commandBufferIndex, int setIndex, int cacheIndex)
+            {
+                Fence = fence;
+                CommandBufferIndex = commandBufferIndex;
+                SetIndex = setIndex;
+                CacheIndex = cacheIndex;
+                fence.Get();
             }
         }
 
         private readonly List<ManualDescriptorSetEntry>[] _manualDsCache;
+        private readonly Queue<PendingManualDsConsumption> _pendingManualDsConsumptions;
+        private readonly Queue<int>[] _freeManualDsCacheEntries;
 
         private readonly Dictionary<long, DescriptorSetTemplate> _pdTemplates;
         private readonly ResourceDescriptorCollection _pdDescriptors;
@@ -70,6 +89,8 @@ namespace Ryujinx.Graphics.Vulkan
 
             _dsCacheCursor = new int[setsCount];
             _manualDsCache = new List<ManualDescriptorSetEntry>[setsCount];
+            _pendingManualDsConsumptions = new Queue<PendingManualDsConsumption>();
+            _freeManualDsCacheEntries = new Queue<int>[setsCount];
         }
 
         public PipelineLayoutCacheEntry(
@@ -78,7 +99,11 @@ namespace Ryujinx.Graphics.Vulkan
             ReadOnlyCollection<ResourceDescriptorCollection> setDescriptors,
             bool usePushDescriptors) : this(gd, device, setDescriptors.Count)
         {
-            (DescriptorSetLayouts, PipelineLayout) = PipelineLayoutFactory.Create(gd, device, setDescriptors, usePushDescriptors);
+            ResourceLayouts layouts = PipelineLayoutFactory.Create(gd, device, setDescriptors, usePushDescriptors);
+
+            DescriptorSetLayouts = layouts.DescriptorSetLayouts;
+            DescriptorSetLayoutsUpdateAfterBind = layouts.DescriptorSetLayoutsUpdateAfterBind;
+            PipelineLayout = layouts.PipelineLayout;
 
             _consumedDescriptorsPerSet = new int[setDescriptors.Count];
             _poolSizes = new DescriptorPoolSize[setDescriptors.Count][];
@@ -133,7 +158,7 @@ namespace Ryujinx.Graphics.Vulkan
                     _poolSizes[setIndex],
                     setIndex,
                     _consumedDescriptorsPerSet[setIndex],
-                    false);
+                    DescriptorSetLayoutsUpdateAfterBind[setIndex]);
 
                 list.Add(dsc);
                 isNew = true;
@@ -144,49 +169,99 @@ namespace Ryujinx.Graphics.Vulkan
             return list[index];
         }
 
-        public Auto<DescriptorSetCollection> GetNewManualDescriptorSetCollection(int commandBufferIndex, int setIndex, out int cacheIndex)
+        public Auto<DescriptorSetCollection> GetNewManualDescriptorSetCollection(CommandBufferScoped cbs, int setIndex, out int cacheIndex)
         {
-            int submissionCount = _gd.CommandBufferPool.GetSubmissionCount(commandBufferIndex);
+            FreeCompletedManualDescriptorSets();
 
             var list = _manualDsCache[setIndex] ??= new();
             var span = CollectionsMarshal.AsSpan(list);
 
-            for (int index = 0; index < span.Length; index++)
+            Queue<int> freeQueue = _freeManualDsCacheEntries[setIndex];
+
+            // Do we have at least one freed descriptor set? If so, just use that.
+            if (freeQueue != null && freeQueue.TryDequeue(out int freeIndex))
             {
-                ref ManualDescriptorSetEntry entry = ref span[index];
+                ref ManualDescriptorSetEntry entry = ref span[freeIndex];
 
-                if (!entry.InUse && (entry.CbIndex != commandBufferIndex || entry.CbSubmissionCount != submissionCount))
-                {
-                    entry.InUse = true;
-                    entry.CbIndex = commandBufferIndex;
-                    entry.CbSubmissionCount = submissionCount;
+                Debug.Assert(!entry.InUse && entry.CbRefMask == 0);
 
-                    cacheIndex = index;
+                entry.InUse = true;
+                entry.CbRefMask = 1u << cbs.CommandBufferIndex;
+                cacheIndex = freeIndex;
 
-                    return entry.DescriptorSet;
-                }
+                _pendingManualDsConsumptions.Enqueue(new PendingManualDsConsumption(cbs.GetFence(), cbs.CommandBufferIndex, setIndex, freeIndex));
+
+                return entry.DescriptorSet;
             }
 
+            // Otherwise create a new descriptor set, and add to our pending queue for command buffer consumption tracking.
             var dsc = _descriptorSetManager.AllocateDescriptorSet(
                 _gd.Api,
                 DescriptorSetLayouts[setIndex],
                 _poolSizes[setIndex],
                 setIndex,
                 _consumedDescriptorsPerSet[setIndex],
-                false);
+                DescriptorSetLayoutsUpdateAfterBind[setIndex]);
 
             cacheIndex = list.Count;
-            list.Add(new ManualDescriptorSetEntry(dsc, commandBufferIndex, submissionCount, inUse: true));
+            list.Add(new ManualDescriptorSetEntry(dsc, cbs.CommandBufferIndex));
+            _pendingManualDsConsumptions.Enqueue(new PendingManualDsConsumption(cbs.GetFence(), cbs.CommandBufferIndex, setIndex, cacheIndex));
 
             return dsc;
         }
 
+        public void UpdateManualDescriptorSetCollectionOwnership(CommandBufferScoped cbs, int setIndex, int cacheIndex)
+        {
+            FreeCompletedManualDescriptorSets();
+
+            var list = _manualDsCache[setIndex];
+            var span = CollectionsMarshal.AsSpan(list);
+            ref var entry = ref span[cacheIndex];
+
+            uint cbMask = 1u << cbs.CommandBufferIndex;
+
+            if ((entry.CbRefMask & cbMask) == 0)
+            {
+                entry.CbRefMask |= cbMask;
+
+                _pendingManualDsConsumptions.Enqueue(new PendingManualDsConsumption(cbs.GetFence(), cbs.CommandBufferIndex, setIndex, cacheIndex));
+            }
+        }
+
+        private void FreeCompletedManualDescriptorSets()
+        {
+            FenceHolder signalledFence = null;
+            while (_pendingManualDsConsumptions.TryPeek(out var pds) && (pds.Fence == signalledFence || pds.Fence.IsSignaled()))
+            {
+                signalledFence = pds.Fence; // Already checked - don't need to do it again.
+                var dequeued = _pendingManualDsConsumptions.Dequeue();
+                Debug.Assert(dequeued.Fence == pds.Fence);
+                pds.Fence.Put();
+
+                var span = CollectionsMarshal.AsSpan(_manualDsCache[dequeued.SetIndex]);
+                ref var entry = ref span[dequeued.CacheIndex];
+                entry.CbRefMask &= ~(1u << dequeued.CommandBufferIndex);
+
+                if (!entry.InUse && entry.CbRefMask == 0)
+                {
+                    // If not in use by any array, and not bound to any command buffer, the descriptor set can be re-used immediately.
+                    (_freeManualDsCacheEntries[dequeued.SetIndex] ??= new()).Enqueue(dequeued.CacheIndex);
+                }
+            }
+        }
+
         public void ReleaseManualDescriptorSetCollection(int setIndex, int cacheIndex)
         {
             var list = _manualDsCache[setIndex];
             var span = CollectionsMarshal.AsSpan(list);
 
             span[cacheIndex].InUse = false;
+
+            if (span[cacheIndex].CbRefMask == 0)
+            {
+                // This is no longer in use by any array, so if not bound to any command buffer, the descriptor set can be re-used immediately.
+                (_freeManualDsCacheEntries[setIndex] ??= new()).Enqueue(cacheIndex);
+            }
         }
 
         private static Span<DescriptorPoolSize> GetDescriptorPoolSizes(Span<DescriptorPoolSize> output, ResourceDescriptorCollection setDescriptor, uint multiplier)
@@ -291,6 +366,11 @@ namespace Ryujinx.Graphics.Vulkan
                     _gd.Api.DestroyDescriptorSetLayout(_device, DescriptorSetLayouts[i], null);
                 }
 
+                while (_pendingManualDsConsumptions.TryDequeue(out var pds))
+                {
+                    pds.Fence.Put();
+                }
+
                 _descriptorSetManager.Dispose();
             }
         }