1 files changed, 67 insertions, 36 deletions
diff --git a/src/Ryujinx.Graphics.Vulkan/DescriptorSetManager.cs b/src/Ryujinx.Graphics.Vulkan/DescriptorSetManager.cs
index 2f7b604c..7594384d 100644
--- a/src/Ryujinx.Graphics.Vulkan/DescriptorSetManager.cs
+++ b/src/Ryujinx.Graphics.Vulkan/DescriptorSetManager.cs
@@ -6,7 +6,7 @@ namespace Ryujinx.Graphics.Vulkan
 {
     class DescriptorSetManager : IDisposable
     {
-        private const uint DescriptorPoolMultiplier = 16;
+        public const uint MaxSets = 16;
 
         public class DescriptorPoolHolder : IDisposable
         {
@@ -14,36 +14,28 @@ namespace Ryujinx.Graphics.Vulkan
             public Device Device { get; }
 
             private readonly DescriptorPool _pool;
-            private readonly uint _capacity;
+            private int _freeDescriptors;
             private int _totalSets;
             private int _setsInUse;
             private bool _done;
 
-            public unsafe DescriptorPoolHolder(Vk api, Device device)
+            public unsafe DescriptorPoolHolder(Vk api, Device device, ReadOnlySpan<DescriptorPoolSize> poolSizes, bool updateAfterBind)
             {
                 Api = api;
                 Device = device;
 
-                var poolSizes = new[]
+                foreach (var poolSize in poolSizes)
                 {
-                    new DescriptorPoolSize(DescriptorType.UniformBuffer, (1 + Constants.MaxUniformBufferBindings) * DescriptorPoolMultiplier),
-                    new DescriptorPoolSize(DescriptorType.StorageBuffer, Constants.MaxStorageBufferBindings * DescriptorPoolMultiplier),
-                    new DescriptorPoolSize(DescriptorType.CombinedImageSampler, Constants.MaxTextureBindings * DescriptorPoolMultiplier),
-                    new DescriptorPoolSize(DescriptorType.StorageImage, Constants.MaxImageBindings * DescriptorPoolMultiplier),
-                    new DescriptorPoolSize(DescriptorType.UniformTexelBuffer, Constants.MaxTextureBindings * DescriptorPoolMultiplier),
-                    new DescriptorPoolSize(DescriptorType.StorageTexelBuffer, Constants.MaxImageBindings * DescriptorPoolMultiplier),
-                };
-
-                uint maxSets = (uint)poolSizes.Length * DescriptorPoolMultiplier;
-
-                _capacity = maxSets;
+                    _freeDescriptors += (int)poolSize.DescriptorCount;
+                }
 
                 fixed (DescriptorPoolSize* pPoolsSize = poolSizes)
                 {
                     var descriptorPoolCreateInfo = new DescriptorPoolCreateInfo
                     {
                         SType = StructureType.DescriptorPoolCreateInfo,
-                        MaxSets = maxSets,
+                        Flags = updateAfterBind ? DescriptorPoolCreateFlags.UpdateAfterBindBit : DescriptorPoolCreateFlags.None,
+                        MaxSets = MaxSets,
                         PoolSizeCount = (uint)poolSizes.Length,
                         PPoolSizes = pPoolsSize,
                     };
@@ -52,18 +44,22 @@ namespace Ryujinx.Graphics.Vulkan
                 }
             }
 
-            public DescriptorSetCollection AllocateDescriptorSets(ReadOnlySpan<DescriptorSetLayout> layouts)
+            public unsafe DescriptorSetCollection AllocateDescriptorSets(ReadOnlySpan<DescriptorSetLayout> layouts, int consumedDescriptors)
             {
-                TryAllocateDescriptorSets(layouts, isTry: false, out var dsc);
+                TryAllocateDescriptorSets(layouts, consumedDescriptors, isTry: false, out var dsc);
                 return dsc;
             }
 
-            public bool TryAllocateDescriptorSets(ReadOnlySpan<DescriptorSetLayout> layouts, out DescriptorSetCollection dsc)
+            public bool TryAllocateDescriptorSets(ReadOnlySpan<DescriptorSetLayout> layouts, int consumedDescriptors, out DescriptorSetCollection dsc)
             {
-                return TryAllocateDescriptorSets(layouts, isTry: true, out dsc);
+                return TryAllocateDescriptorSets(layouts, consumedDescriptors, isTry: true, out dsc);
             }
 
-            private unsafe bool TryAllocateDescriptorSets(ReadOnlySpan<DescriptorSetLayout> layouts, bool isTry, out DescriptorSetCollection dsc)
+            private unsafe bool TryAllocateDescriptorSets(
+                ReadOnlySpan<DescriptorSetLayout> layouts,
+                int consumedDescriptors,
+                bool isTry,
+                out DescriptorSetCollection dsc)
             {
                 Debug.Assert(!_done);
 
@@ -84,7 +80,7 @@ namespace Ryujinx.Graphics.Vulkan
                         var result = Api.AllocateDescriptorSets(Device, &descriptorSetAllocateInfo, pDescriptorSets);
                         if (isTry && result == Result.ErrorOutOfPoolMemory)
                         {
-                            _totalSets = (int)_capacity;
+                            _totalSets = (int)MaxSets;
                             _done = true;
                             DestroyIfDone();
                             dsc = default;
@@ -95,6 +91,7 @@ namespace Ryujinx.Graphics.Vulkan
                     }
                 }
 
+                _freeDescriptors -= consumedDescriptors;
                 _totalSets += layouts.Length;
                 _setsInUse += layouts.Length;
 
@@ -109,9 +106,15 @@ namespace Ryujinx.Graphics.Vulkan
                 DestroyIfDone();
             }
 
-            public bool CanFit(int count)
+            public bool CanFit(int setsCount, int descriptorsCount)
             {
-                if (_totalSets + count <= _capacity)
+                // Try to determine if an allocation with the given parameters will succeed.
+                // An allocation may fail if the sets count or descriptors count exceeds the available counts
+                // of the pool.
+                // Not getting that right is not fatal, it will just create a new pool and try again,
+                // but it is less efficient.
+
+                if (_totalSets + setsCount <= MaxSets && _freeDescriptors >= descriptorsCount)
                 {
                     return true;
                 }
@@ -148,46 +151,74 @@ namespace Ryujinx.Graphics.Vulkan
         }
 
         private readonly Device _device;
-        private DescriptorPoolHolder _currentPool;
+        private readonly DescriptorPoolHolder[] _currentPools;
 
-        public DescriptorSetManager(Device device)
+        public DescriptorSetManager(Device device, int poolCount)
         {
             _device = device;
+            _currentPools = new DescriptorPoolHolder[poolCount];
         }
 
-        public Auto<DescriptorSetCollection> AllocateDescriptorSet(Vk api, DescriptorSetLayout layout)
+        public Auto<DescriptorSetCollection> AllocateDescriptorSet(
+            Vk api,
+            DescriptorSetLayout layout,
+            ReadOnlySpan<DescriptorPoolSize> poolSizes,
+            int poolIndex,
+            int consumedDescriptors,
+            bool updateAfterBind)
         {
             Span<DescriptorSetLayout> layouts = stackalloc DescriptorSetLayout[1];
             layouts[0] = layout;
-            return AllocateDescriptorSets(api, layouts);
+            return AllocateDescriptorSets(api, layouts, poolSizes, poolIndex, consumedDescriptors, updateAfterBind);
         }
 
-        public Auto<DescriptorSetCollection> AllocateDescriptorSets(Vk api, ReadOnlySpan<DescriptorSetLayout> layouts)
+        public Auto<DescriptorSetCollection> AllocateDescriptorSets(
+            Vk api,
+            ReadOnlySpan<DescriptorSetLayout> layouts,
+            ReadOnlySpan<DescriptorPoolSize> poolSizes,
+            int poolIndex,
+            int consumedDescriptors,
+            bool updateAfterBind)
         {
             // If we fail the first time, just create a new pool and try again.
-            if (!GetPool(api, layouts.Length).TryAllocateDescriptorSets(layouts, out var dsc))
+
+            var pool = GetPool(api, poolSizes, poolIndex, layouts.Length, consumedDescriptors, updateAfterBind);
+            if (!pool.TryAllocateDescriptorSets(layouts, consumedDescriptors, out var dsc))
             {
-                dsc = GetPool(api, layouts.Length).AllocateDescriptorSets(layouts);
+                pool = GetPool(api, poolSizes, poolIndex, layouts.Length, consumedDescriptors, updateAfterBind);
+                dsc = pool.AllocateDescriptorSets(layouts, consumedDescriptors);
             }
 
             return new Auto<DescriptorSetCollection>(dsc);
         }
 
-        private DescriptorPoolHolder GetPool(Vk api, int requiredCount)
+        private DescriptorPoolHolder GetPool(
+            Vk api,
+            ReadOnlySpan<DescriptorPoolSize> poolSizes,
+            int poolIndex,
+            int setsCount,
+            int descriptorsCount,
+            bool updateAfterBind)
         {
-            if (_currentPool == null || !_currentPool.CanFit(requiredCount))
+            ref DescriptorPoolHolder currentPool = ref _currentPools[poolIndex];
+
+            if (currentPool == null || !currentPool.CanFit(setsCount, descriptorsCount))
             {
-                _currentPool = new DescriptorPoolHolder(api, _device);
+                currentPool = new DescriptorPoolHolder(api, _device, poolSizes, updateAfterBind);
             }
 
-            return _currentPool;
+            return currentPool;
         }
 
         protected virtual void Dispose(bool disposing)
         {
             if (disposing)
             {
-                _currentPool?.Dispose();
+                for (int index = 0; index < _currentPools.Length; index++)
+                {
+                    _currentPools[index]?.Dispose();
+                    _currentPools[index] = null;
+                }
             }
         }