1 files changed, 80 insertions, 13 deletions
diff --git a/src/Ryujinx.Graphics.Vulkan/PipelineLayoutCacheEntry.cs b/src/Ryujinx.Graphics.Vulkan/PipelineLayoutCacheEntry.cs
index eeb25dc0..2840dda0 100644
--- a/src/Ryujinx.Graphics.Vulkan/PipelineLayoutCacheEntry.cs
+++ b/src/Ryujinx.Graphics.Vulkan/PipelineLayoutCacheEntry.cs
@@ -1,5 +1,6 @@
 using Ryujinx.Graphics.GAL;
 using Silk.NET.Vulkan;
+using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 
@@ -7,15 +8,28 @@ namespace Ryujinx.Graphics.Vulkan
 {
     class PipelineLayoutCacheEntry
     {
+        // Those were adjusted based on current descriptor usage and the descriptor counts usually used on pipeline layouts.
+        // It might be a good idea to tweak them again if those change, or maybe find a way to calculate an optimal value dynamically.
+        private const uint DefaultUniformBufferPoolCapacity = 19 * DescriptorSetManager.MaxSets;
+        private const uint DefaultStorageBufferPoolCapacity = 16 * DescriptorSetManager.MaxSets;
+        private const uint DefaultTexturePoolCapacity = 128 * DescriptorSetManager.MaxSets;
+        private const uint DefaultImagePoolCapacity = 8 * DescriptorSetManager.MaxSets;
+
+        private const int MaxPoolSizesPerSet = 2;
+
         private readonly VulkanRenderer _gd;
         private readonly Device _device;
 
         public DescriptorSetLayout[] DescriptorSetLayouts { get; }
         public PipelineLayout PipelineLayout { get; }
 
+        private readonly int[] _consumedDescriptorsPerSet;
+
         private readonly List<Auto<DescriptorSetCollection>>[][] _dsCache;
+        private List<Auto<DescriptorSetCollection>>[] _currentDsCache;
         private readonly int[] _dsCacheCursor;
         private int _dsLastCbIndex;
+        private int _dsLastSubmissionCount;
 
         private PipelineLayoutCacheEntry(VulkanRenderer gd, Device device, int setsCount)
         {
@@ -44,29 +58,55 @@ namespace Ryujinx.Graphics.Vulkan
             bool usePushDescriptors) : this(gd, device, setDescriptors.Count)
         {
             (DescriptorSetLayouts, PipelineLayout) = PipelineLayoutFactory.Create(gd, device, setDescriptors, usePushDescriptors);
-        }
 
-        public Auto<DescriptorSetCollection> GetNewDescriptorSetCollection(
-            VulkanRenderer gd,
-            int commandBufferIndex,
-            int setIndex,
-            out bool isNew)
-        {
-            if (_dsLastCbIndex != commandBufferIndex)
+            _consumedDescriptorsPerSet = new int[setDescriptors.Count];
+
+            for (int setIndex = 0; setIndex < setDescriptors.Count; setIndex++)
             {
-                _dsLastCbIndex = commandBufferIndex;
+                int count = 0;
 
-                for (int i = 0; i < _dsCacheCursor.Length; i++)
+                foreach (var descriptor in setDescriptors[setIndex].Descriptors)
                 {
-                    _dsCacheCursor[i] = 0;
+                    count += descriptor.Count;
                 }
+
+                _consumedDescriptorsPerSet[setIndex] = count;
             }
+        }
+
+        public void UpdateCommandBufferIndex(int commandBufferIndex)
+        {
+            int submissionCount = _gd.CommandBufferPool.GetSubmissionCount(commandBufferIndex);
 
-            var list = _dsCache[commandBufferIndex][setIndex];
+            if (_dsLastCbIndex != commandBufferIndex || _dsLastSubmissionCount != submissionCount)
+            {
+                _dsLastCbIndex = commandBufferIndex;
+                _dsLastSubmissionCount = submissionCount;
+                Array.Clear(_dsCacheCursor);
+            }
+
+            _currentDsCache = _dsCache[commandBufferIndex];
+        }
+
+        public Auto<DescriptorSetCollection> GetNewDescriptorSetCollection(int setIndex, out bool isNew)
+        {
+            var list = _currentDsCache[setIndex];
             int index = _dsCacheCursor[setIndex]++;
             if (index == list.Count)
             {
-                var dsc = gd.DescriptorSetManager.AllocateDescriptorSet(gd.Api, DescriptorSetLayouts[setIndex]);
+                Span<DescriptorPoolSize> poolSizes = stackalloc DescriptorPoolSize[MaxPoolSizesPerSet];
+                poolSizes = GetDescriptorPoolSizes(poolSizes, setIndex);
+
+                int consumedDescriptors = _consumedDescriptorsPerSet[setIndex];
+
+                var dsc = _gd.DescriptorSetManager.AllocateDescriptorSet(
+                    _gd.Api,
+                    DescriptorSetLayouts[setIndex],
+                    poolSizes,
+                    setIndex,
+                    consumedDescriptors,
+                    false);
+
                 list.Add(dsc);
                 isNew = true;
                 return dsc;
@@ -76,6 +116,33 @@ namespace Ryujinx.Graphics.Vulkan
             return list[index];
         }
 
+        private static Span<DescriptorPoolSize> GetDescriptorPoolSizes(Span<DescriptorPoolSize> output, int setIndex)
+        {
+            int count = 1;
+
+            switch (setIndex)
+            {
+                case PipelineBase.UniformSetIndex:
+                    output[0] = new(DescriptorType.UniformBuffer, DefaultUniformBufferPoolCapacity);
+                    break;
+                case PipelineBase.StorageSetIndex:
+                    output[0] = new(DescriptorType.StorageBuffer, DefaultStorageBufferPoolCapacity);
+                    break;
+                case PipelineBase.TextureSetIndex:
+                    output[0] = new(DescriptorType.CombinedImageSampler, DefaultTexturePoolCapacity);
+                    output[1] = new(DescriptorType.UniformTexelBuffer, DefaultTexturePoolCapacity);
+                    count = 2;
+                    break;
+                case PipelineBase.ImageSetIndex:
+                    output[0] = new(DescriptorType.StorageImage, DefaultImagePoolCapacity);
+                    output[1] = new(DescriptorType.StorageTexelBuffer, DefaultImagePoolCapacity);
+                    count = 2;
+                    break;
+            }
+
+            return output[..count];
+        }
+
         protected virtual unsafe void Dispose(bool disposing)
         {
             if (disposing)