Allow texture arrays to use separate descriptor sets on Vulkan (#6870)

* Report base and extra sets from the backend

* Pass texture set index everywhere

* Key textures using set and binding (rather than just binding)

* Start using extra sets for array textures

* Shader cache version bump

* Separate new commands, some PR feedback

* Introduce new manual descriptor set reservation method that prevents it from being used by something else while owned by an array

* Move bind extra sets logic to new method

* Should only use separate array is MaximumExtraSets is not zero

* Format whitespace

5 files changed, 130 insertions, 49 deletions

diff --git a/src/Ryujinx.Graphics.Gpu/Shader/CachedShaderBindings.cs b/src/Ryujinx.Graphics.Gpu/Shader/CachedShaderBindings.cs
index a80dcbc8..51be00b6 100644
--- a/src/Ryujinx.Graphics.Gpu/Shader/CachedShaderBindings.cs
+++ b/src/Ryujinx.Graphics.Gpu/Shader/CachedShaderBindings.cs
@@ -62,6 +62,7 @@ namespace Ryujinx.Graphics.Gpu.Shader
 
                     var result = new TextureBindingInfo(
                         target,
+                        descriptor.Set,
                         descriptor.Binding,
                         descriptor.ArrayLength,
                         descriptor.CbufSlot,
@@ -90,6 +91,7 @@ namespace Ryujinx.Graphics.Gpu.Shader
                     var result = new TextureBindingInfo(
                         target,
                         format,
+                        descriptor.Set,
                         descriptor.Binding,
                         descriptor.ArrayLength,
                         descriptor.CbufSlot,
diff --git a/src/Ryujinx.Graphics.Gpu/Shader/DiskCache/DiskCacheHostStorage.cs b/src/Ryujinx.Graphics.Gpu/Shader/DiskCache/DiskCacheHostStorage.cs
index ea54049c..990c6ba3 100644
--- a/src/Ryujinx.Graphics.Gpu/Shader/DiskCache/DiskCacheHostStorage.cs
+++ b/src/Ryujinx.Graphics.Gpu/Shader/DiskCache/DiskCacheHostStorage.cs
@@ -22,7 +22,7 @@ namespace Ryujinx.Graphics.Gpu.Shader.DiskCache
         private const ushort FileFormatVersionMajor = 1;
         private const ushort FileFormatVersionMinor = 2;
         private const uint FileFormatVersionPacked = ((uint)FileFormatVersionMajor << 16) | FileFormatVersionMinor;
-        private const uint CodeGenVersion = 5936;
+        private const uint CodeGenVersion = 6870;
 
         private const string SharedTocFileName = "shared.toc";
         private const string SharedDataFileName = "shared.data";
diff --git a/src/Ryujinx.Graphics.Gpu/Shader/GpuAccessorBase.cs b/src/Ryujinx.Graphics.Gpu/Shader/GpuAccessorBase.cs
index 0d562b0d..d89eebab 100644
--- a/src/Ryujinx.Graphics.Gpu/Shader/GpuAccessorBase.cs
+++ b/src/Ryujinx.Graphics.Gpu/Shader/GpuAccessorBase.cs
@@ -51,7 +51,7 @@ namespace Ryujinx.Graphics.Gpu.Shader
             _reservedImages = rrc.ReservedImages;
         }
 
-        public int CreateConstantBufferBinding(int index)
+        public SetBindingPair CreateConstantBufferBinding(int index)
         {
             int binding;
 
@@ -64,10 +64,10 @@ namespace Ryujinx.Graphics.Gpu.Shader
                 binding = _resourceCounts.UniformBuffersCount++;
             }
 
-            return binding + _reservedConstantBuffers;
+            return new SetBindingPair(_context.Capabilities.UniformBufferSetIndex, binding + _reservedConstantBuffers);
         }
 
-        public int CreateImageBinding(int count, bool isBuffer)
+        public SetBindingPair CreateImageBinding(int count, bool isBuffer)
         {
             int binding;
 
@@ -96,10 +96,10 @@ namespace Ryujinx.Graphics.Gpu.Shader
                 _resourceCounts.ImagesCount += count;
             }
 
-            return binding + _reservedImages;
+            return new SetBindingPair(_context.Capabilities.ImageSetIndex, binding + _reservedImages);
         }
 
-        public int CreateStorageBufferBinding(int index)
+        public SetBindingPair CreateStorageBufferBinding(int index)
         {
             int binding;
 
@@ -112,10 +112,10 @@ namespace Ryujinx.Graphics.Gpu.Shader
                 binding = _resourceCounts.StorageBuffersCount++;
             }
 
-            return binding + _reservedStorageBuffers;
+            return new SetBindingPair(_context.Capabilities.StorageBufferSetIndex, binding + _reservedStorageBuffers);
         }
 
-        public int CreateTextureBinding(int count, bool isBuffer)
+        public SetBindingPair CreateTextureBinding(int count, bool isBuffer)
         {
             int binding;
 
@@ -144,7 +144,7 @@ namespace Ryujinx.Graphics.Gpu.Shader
                 _resourceCounts.TexturesCount += count;
             }
 
-            return binding + _reservedTextures;
+            return new SetBindingPair(_context.Capabilities.TextureSetIndex, binding + _reservedTextures);
         }
 
         private int GetBindingFromIndex(int index, uint maxPerStage, string resourceName)
@@ -183,6 +183,16 @@ namespace Ryujinx.Graphics.Gpu.Shader
             return maxPerStage * Constants.ShaderStages;
         }
 
+        public int CreateExtraSet()
+        {
+            if (_resourceCounts.SetsCount >= _context.Capabilities.MaximumExtraSets)
+            {
+                return -1;
+            }
+
+            return _context.Capabilities.ExtraSetBaseIndex + _resourceCounts.SetsCount++;
+        }
+
         public int QueryHostGatherBiasPrecision() => _context.Capabilities.GatherBiasPrecision;
 
         public bool QueryHostReducedPrecision() => _context.Capabilities.ReduceShaderPrecision;
diff --git a/src/Ryujinx.Graphics.Gpu/Shader/ResourceCounts.cs b/src/Ryujinx.Graphics.Gpu/Shader/ResourceCounts.cs
index 126e3249..59ab378c 100644
--- a/src/Ryujinx.Graphics.Gpu/Shader/ResourceCounts.cs
+++ b/src/Ryujinx.Graphics.Gpu/Shader/ResourceCounts.cs
@@ -24,5 +24,10 @@ namespace Ryujinx.Graphics.Gpu.Shader
         /// Total of images used by the shaders.
         /// </summary>
         public int ImagesCount;
+
+        /// <summary>
+        /// Total of extra sets used by the shaders.
+        /// </summary>
+        public int SetsCount;
     }
 }
diff --git a/src/Ryujinx.Graphics.Gpu/Shader/ShaderInfoBuilder.cs b/src/Ryujinx.Graphics.Gpu/Shader/ShaderInfoBuilder.cs
index ed56db3b..42b2cbb5 100644
--- a/src/Ryujinx.Graphics.Gpu/Shader/ShaderInfoBuilder.cs
+++ b/src/Ryujinx.Graphics.Gpu/Shader/ShaderInfoBuilder.cs
@@ -1,5 +1,6 @@
 using Ryujinx.Graphics.GAL;
 using Ryujinx.Graphics.Shader;
+using System;
 using System.Collections.Generic;
 
 namespace Ryujinx.Graphics.Gpu.Shader
@@ -9,13 +10,6 @@ namespace Ryujinx.Graphics.Gpu.Shader
     /// </summary>
     class ShaderInfoBuilder
     {
-        private const int TotalSets = 4;
-
-        private const int UniformSetIndex = 0;
-        private const int StorageSetIndex = 1;
-        private const int TextureSetIndex = 2;
-        private const int ImageSetIndex = 3;
-
         private const ResourceStages SupportBufferStages =
             ResourceStages.Compute |
             ResourceStages.Vertex |
@@ -36,8 +30,8 @@ namespace Ryujinx.Graphics.Gpu.Shader
         private readonly int _reservedTextures;
         private readonly int _reservedImages;
 
-        private readonly List<ResourceDescriptor>[] _resourceDescriptors;
-        private readonly List<ResourceUsage>[] _resourceUsages;
+        private List<ResourceDescriptor>[] _resourceDescriptors;
+        private List<ResourceUsage>[] _resourceUsages;
 
         /// <summary>
         /// Creates a new shader info builder.
@@ -51,17 +45,27 @@ namespace Ryujinx.Graphics.Gpu.Shader
 
             _fragmentOutputMap = -1;
 
-            _resourceDescriptors = new List<ResourceDescriptor>[TotalSets];
-            _resourceUsages = new List<ResourceUsage>[TotalSets];
+            int uniformSetIndex = context.Capabilities.UniformBufferSetIndex;
+            int storageSetIndex = context.Capabilities.StorageBufferSetIndex;
+            int textureSetIndex = context.Capabilities.TextureSetIndex;
+            int imageSetIndex = context.Capabilities.ImageSetIndex;
+
+            int totalSets = Math.Max(uniformSetIndex, storageSetIndex);
+            totalSets = Math.Max(totalSets, textureSetIndex);
+            totalSets = Math.Max(totalSets, imageSetIndex);
+            totalSets++;
+
+            _resourceDescriptors = new List<ResourceDescriptor>[totalSets];
+            _resourceUsages = new List<ResourceUsage>[totalSets];
 
-            for (int index = 0; index < TotalSets; index++)
+            for (int index = 0; index < totalSets; index++)
             {
                 _resourceDescriptors[index] = new();
                 _resourceUsages[index] = new();
             }
 
-            AddDescriptor(SupportBufferStages, ResourceType.UniformBuffer, UniformSetIndex, 0, 1);
-            AddUsage(SupportBufferStages, ResourceType.UniformBuffer, UniformSetIndex, 0, 1);
+            AddDescriptor(SupportBufferStages, ResourceType.UniformBuffer, uniformSetIndex, 0, 1);
+            AddUsage(SupportBufferStages, ResourceType.UniformBuffer, uniformSetIndex, 0, 1);
 
             ResourceReservationCounts rrc = new(!context.Capabilities.SupportsTransformFeedback && tfEnabled, vertexAsCompute);
 
@@ -73,12 +77,20 @@ namespace Ryujinx.Graphics.Gpu.Shader
             // TODO: Handle that better? Maybe we should only set the binding that are really needed on each shader.
             ResourceStages stages = vertexAsCompute ? ResourceStages.Compute | ResourceStages.Vertex : VtgStages;
 
-            PopulateDescriptorAndUsages(stages, ResourceType.UniformBuffer, UniformSetIndex, 1, rrc.ReservedConstantBuffers - 1);
-            PopulateDescriptorAndUsages(stages, ResourceType.StorageBuffer, StorageSetIndex, 0, rrc.ReservedStorageBuffers);
-            PopulateDescriptorAndUsages(stages, ResourceType.BufferTexture, TextureSetIndex, 0, rrc.ReservedTextures);
-            PopulateDescriptorAndUsages(stages, ResourceType.BufferImage, ImageSetIndex, 0, rrc.ReservedImages);
+            PopulateDescriptorAndUsages(stages, ResourceType.UniformBuffer, uniformSetIndex, 1, rrc.ReservedConstantBuffers - 1);
+            PopulateDescriptorAndUsages(stages, ResourceType.StorageBuffer, storageSetIndex, 0, rrc.ReservedStorageBuffers);
+            PopulateDescriptorAndUsages(stages, ResourceType.BufferTexture, textureSetIndex, 0, rrc.ReservedTextures);
+            PopulateDescriptorAndUsages(stages, ResourceType.BufferImage, imageSetIndex, 0, rrc.ReservedImages);
         }
 
+        /// <summary>
+        /// Populates descriptors and usages for vertex as compute and transform feedback emulation reserved resources.
+        /// </summary>
+        /// <param name="stages">Shader stages where the resources are used</param>
+        /// <param name="type">Resource type</param>
+        /// <param name="setIndex">Resource set index where the resources are used</param>
+        /// <param name="start">First binding number</param>
+        /// <param name="count">Amount of bindings</param>
         private void PopulateDescriptorAndUsages(ResourceStages stages, ResourceType type, int setIndex, int start, int count)
         {
             AddDescriptor(stages, type, setIndex, start, count);
@@ -127,18 +139,23 @@ namespace Ryujinx.Graphics.Gpu.Shader
             int textureBinding = _reservedTextures + stageIndex * texturesPerStage * 2;
             int imageBinding = _reservedImages + stageIndex * imagesPerStage * 2;
 
-            AddDescriptor(stages, ResourceType.UniformBuffer, UniformSetIndex, uniformBinding, uniformsPerStage);
-            AddDescriptor(stages, ResourceType.StorageBuffer, StorageSetIndex, storageBinding, storagesPerStage);
-            AddDualDescriptor(stages, ResourceType.TextureAndSampler, ResourceType.BufferTexture, TextureSetIndex, textureBinding, texturesPerStage);
-            AddDualDescriptor(stages, ResourceType.Image, ResourceType.BufferImage, ImageSetIndex, imageBinding, imagesPerStage);
+            int uniformSetIndex = _context.Capabilities.UniformBufferSetIndex;
+            int storageSetIndex = _context.Capabilities.StorageBufferSetIndex;
+            int textureSetIndex = _context.Capabilities.TextureSetIndex;
+            int imageSetIndex = _context.Capabilities.ImageSetIndex;
 
-            AddArrayDescriptors(info.Textures, stages, TextureSetIndex, isImage: false);
-            AddArrayDescriptors(info.Images, stages, TextureSetIndex, isImage: true);
+            AddDescriptor(stages, ResourceType.UniformBuffer, uniformSetIndex, uniformBinding, uniformsPerStage);
+            AddDescriptor(stages, ResourceType.StorageBuffer, storageSetIndex, storageBinding, storagesPerStage);
+            AddDualDescriptor(stages, ResourceType.TextureAndSampler, ResourceType.BufferTexture, textureSetIndex, textureBinding, texturesPerStage);
+            AddDualDescriptor(stages, ResourceType.Image, ResourceType.BufferImage, imageSetIndex, imageBinding, imagesPerStage);
 
-            AddUsage(info.CBuffers, stages, UniformSetIndex, isStorage: false);
-            AddUsage(info.SBuffers, stages, StorageSetIndex, isStorage: true);
-            AddUsage(info.Textures, stages, TextureSetIndex, isImage: false);
-            AddUsage(info.Images, stages, ImageSetIndex, isImage: true);
+            AddArrayDescriptors(info.Textures, stages, isImage: false);
+            AddArrayDescriptors(info.Images, stages, isImage: true);
+
+            AddUsage(info.CBuffers, stages, isStorage: false);
+            AddUsage(info.SBuffers, stages, isStorage: true);
+            AddUsage(info.Textures, stages, isImage: false);
+            AddUsage(info.Images, stages, isImage: true);
         }
 
         /// <summary>
@@ -177,9 +194,8 @@ namespace Ryujinx.Graphics.Gpu.Shader
         /// </summary>
         /// <param name="textures">Textures to be added</param>
         /// <param name="stages">Stages where the textures are used</param>
-        /// <param name="setIndex">Descriptor set index where the textures will be bound</param>
         /// <param name="isImage">True for images, false for textures</param>
-        private void AddArrayDescriptors(IEnumerable<TextureDescriptor> textures, ResourceStages stages, int setIndex, bool isImage)
+        private void AddArrayDescriptors(IEnumerable<TextureDescriptor> textures, ResourceStages stages, bool isImage)
         {
             foreach (TextureDescriptor texture in textures)
             {
@@ -187,7 +203,7 @@ namespace Ryujinx.Graphics.Gpu.Shader
                 {
                     ResourceType type = GetTextureResourceType(texture, isImage);
 
-                    _resourceDescriptors[setIndex].Add(new ResourceDescriptor(texture.Binding, texture.ArrayLength, type, stages));
+                    GetDescriptors(texture.Set).Add(new ResourceDescriptor(texture.Binding, texture.ArrayLength, type, stages));
                 }
             }
         }
@@ -213,13 +229,12 @@ namespace Ryujinx.Graphics.Gpu.Shader
         /// </summary>
         /// <param name="buffers">Buffers to be added</param>
         /// <param name="stages">Stages where the buffers are used</param>
-        /// <param name="setIndex">Descriptor set index where the buffers will be bound</param>
         /// <param name="isStorage">True for storage buffers, false for uniform buffers</param>
-        private void AddUsage(IEnumerable<BufferDescriptor> buffers, ResourceStages stages, int setIndex, bool isStorage)
+        private void AddUsage(IEnumerable<BufferDescriptor> buffers, ResourceStages stages, bool isStorage)
         {
             foreach (BufferDescriptor buffer in buffers)
             {
-                _resourceUsages[setIndex].Add(new ResourceUsage(
+                GetUsages(buffer.Set).Add(new ResourceUsage(
                     buffer.Binding,
                     1,
                     isStorage ? ResourceType.StorageBuffer : ResourceType.UniformBuffer,
@@ -232,18 +247,65 @@ namespace Ryujinx.Graphics.Gpu.Shader
         /// </summary>
         /// <param name="textures">Textures to be added</param>
         /// <param name="stages">Stages where the textures are used</param>
-        /// <param name="setIndex">Descriptor set index where the textures will be bound</param>
         /// <param name="isImage">True for images, false for textures</param>
-        private void AddUsage(IEnumerable<TextureDescriptor> textures, ResourceStages stages, int setIndex, bool isImage)
+        private void AddUsage(IEnumerable<TextureDescriptor> textures, ResourceStages stages, bool isImage)
         {
             foreach (TextureDescriptor texture in textures)
             {
                 ResourceType type = GetTextureResourceType(texture, isImage);
 
-                _resourceUsages[setIndex].Add(new ResourceUsage(texture.Binding, texture.ArrayLength, type, stages));
+                GetUsages(texture.Set).Add(new ResourceUsage(texture.Binding, texture.ArrayLength, type, stages));
+            }
+        }
+
+        /// <summary>
+        /// Gets the list of resource descriptors for a given set index. A new list will be created if needed.
+        /// </summary>
+        /// <param name="setIndex">Resource set index</param>
+        /// <returns>List of resource descriptors</returns>
+        private List<ResourceDescriptor> GetDescriptors(int setIndex)
+        {
+            if (_resourceDescriptors.Length <= setIndex)
+            {
+                int oldLength = _resourceDescriptors.Length;
+                Array.Resize(ref _resourceDescriptors, setIndex + 1);
+
+                for (int index = oldLength; index <= setIndex; index++)
+                {
+                    _resourceDescriptors[index] = new();
+                }
+            }
+
+            return _resourceDescriptors[setIndex];
+        }
+
+        /// <summary>
+        /// Gets the list of resource usages for a given set index. A new list will be created if needed.
+        /// </summary>
+        /// <param name="setIndex">Resource set index</param>
+        /// <returns>List of resource usages</returns>
+        private List<ResourceUsage> GetUsages(int setIndex)
+        {
+            if (_resourceUsages.Length <= setIndex)
+            {
+                int oldLength = _resourceUsages.Length;
+                Array.Resize(ref _resourceUsages, setIndex + 1);
+
+                for (int index = oldLength; index <= setIndex; index++)
+                {
+                    _resourceUsages[index] = new();
+                }
             }
+
+            return _resourceUsages[setIndex];
         }
 
+        /// <summary>
+        /// Gets a resource type from a texture descriptor.
+        /// </summary>
+        /// <param name="texture">Texture descriptor</param>
+        /// <param name="isImage">Whether the texture is a image texture (writable) or not (sampled)</param>
+        /// <returns>Resource type</returns>
         private static ResourceType GetTextureResourceType(TextureDescriptor texture, bool isImage)
         {
             bool isBuffer = (texture.Type & SamplerType.Mask) == SamplerType.TextureBuffer;
@@ -278,10 +340,12 @@ namespace Ryujinx.Graphics.Gpu.Shader
         /// <returns>Shader information</returns>
         public ShaderInfo Build(ProgramPipelineState? pipeline, bool fromCache = false)
         {
-            var descriptors = new ResourceDescriptorCollection[TotalSets];
-            var usages = new ResourceUsageCollection[TotalSets];
+            int totalSets = _resourceDescriptors.Length;
+
+            var descriptors = new ResourceDescriptorCollection[totalSets];
+            var usages = new ResourceUsageCollection[totalSets];
 
-            for (int index = 0; index < TotalSets; index++)
+            for (int index = 0; index < totalSets; index++)
             {
                 descriptors[index] = new ResourceDescriptorCollection(_resourceDescriptors[index].ToArray().AsReadOnly());
                 usages[index] = new ResourceUsageCollection(_resourceUsages[index].ToArray().AsReadOnly());