1 files changed, 98 insertions, 5 deletions
diff --git a/src/Ryujinx.Graphics.Vulkan/ShaderCollection.cs b/src/Ryujinx.Graphics.Vulkan/ShaderCollection.cs
index 7c25c6d1..3c35a6f0 100644
--- a/src/Ryujinx.Graphics.Vulkan/ShaderCollection.cs
+++ b/src/Ryujinx.Graphics.Vulkan/ShaderCollection.cs
@@ -108,18 +108,25 @@ namespace Ryujinx.Graphics.Vulkan
 
             _shaders = internalShaders;
 
-            bool usePushDescriptors = !isMinimal && VulkanConfiguration.UsePushDescriptors && _gd.Capabilities.SupportsPushDescriptors;
+            bool usePushDescriptors = !isMinimal &&
+                VulkanConfiguration.UsePushDescriptors &&
+                _gd.Capabilities.SupportsPushDescriptors &&
+                !IsCompute &&
+                CanUsePushDescriptors(gd, resourceLayout, IsCompute);
 
-            _plce = gd.PipelineLayoutCache.GetOrCreate(gd, device, resourceLayout.Sets, usePushDescriptors);
+            ReadOnlyCollection<ResourceDescriptorCollection> sets = usePushDescriptors ?
+                BuildPushDescriptorSets(gd, resourceLayout.Sets) : resourceLayout.Sets;
+
+            _plce = gd.PipelineLayoutCache.GetOrCreate(gd, device, sets, usePushDescriptors);
 
             HasMinimalLayout = isMinimal;
             UsePushDescriptors = usePushDescriptors;
 
             Stages = stages;
 
-            ClearSegments = BuildClearSegments(resourceLayout.Sets);
+            ClearSegments = BuildClearSegments(sets);
             BindingSegments = BuildBindingSegments(resourceLayout.SetUsages);
-            Templates = BuildTemplates();
+            Templates = BuildTemplates(usePushDescriptors);
 
             _compileTask = Task.CompletedTask;
             _firstBackgroundUse = false;
@@ -139,6 +146,76 @@ namespace Ryujinx.Graphics.Vulkan
             _firstBackgroundUse = !fromCache;
         }
 
+        private static bool CanUsePushDescriptors(VulkanRenderer gd, ResourceLayout layout, bool isCompute)
+        {
+            // If binding 3 is immediately used, use an alternate set of reserved bindings.
+            ReadOnlyCollection<ResourceUsage> uniformUsage = layout.SetUsages[0].Usages;
+            bool hasBinding3 = uniformUsage.Any(x => x.Binding == 3);
+            int[] reserved = isCompute ? Array.Empty<int>() : gd.GetPushDescriptorReservedBindings(hasBinding3);
+
+            // Can't use any of the reserved usages.
+            for (int i = 0; i < uniformUsage.Count; i++)
+            {
+                var binding = uniformUsage[i].Binding;
+
+                if (reserved.Contains(binding) ||
+                    binding >= Constants.MaxPushDescriptorBinding ||
+                    binding >= gd.Capabilities.MaxPushDescriptors + reserved.Count(id => id < binding))
+                {
+                    return false;
+                }
+            }
+
+            return true;
+        }
+
+        private static ReadOnlyCollection<ResourceDescriptorCollection> BuildPushDescriptorSets(
+            VulkanRenderer gd,
+            ReadOnlyCollection<ResourceDescriptorCollection> sets)
+        {
+            // The reserved bindings were selected when determining if push descriptors could be used.
+            int[] reserved = gd.GetPushDescriptorReservedBindings(false);
+
+            var result = new ResourceDescriptorCollection[sets.Count];
+
+            for (int i = 0; i < sets.Count; i++)
+            {
+                if (i == 0)
+                {
+                    // Push descriptors apply here. Remove reserved bindings.
+                    ResourceDescriptorCollection original = sets[i];
+
+                    var pdUniforms = new ResourceDescriptor[original.Descriptors.Count];
+                    int j = 0;
+
+                    foreach (ResourceDescriptor descriptor in original.Descriptors)
+                    {
+                        if (reserved.Contains(descriptor.Binding))
+                        {
+                            // If the binding is reserved, set its descriptor count to 0.
+                            pdUniforms[j++] = new ResourceDescriptor(
+                                descriptor.Binding,
+                                0,
+                                descriptor.Type,
+                                descriptor.Stages);
+                        }
+                        else
+                        {
+                            pdUniforms[j++] = descriptor;
+                        }
+                    }
+
+                    result[i] = new ResourceDescriptorCollection(new(pdUniforms));
+                }
+                else
+                {
+                    result[i] = sets[i];
+                }
+            }
+
+            return new(result);
+        }
+
         private static ResourceBindingSegment[][] BuildClearSegments(ReadOnlyCollection<ResourceDescriptorCollection> sets)
         {
             ResourceBindingSegment[][] segments = new ResourceBindingSegment[sets.Count][];
@@ -243,12 +320,18 @@ namespace Ryujinx.Graphics.Vulkan
             return segments;
         }
 
-        private DescriptorSetTemplate[] BuildTemplates()
+        private DescriptorSetTemplate[] BuildTemplates(bool usePushDescriptors)
         {
             var templates = new DescriptorSetTemplate[BindingSegments.Length];
 
             for (int setIndex = 0; setIndex < BindingSegments.Length; setIndex++)
             {
+                if (usePushDescriptors && setIndex == 0)
+                {
+                    // Push descriptors get updated using templates owned by the pipeline layout.
+                    continue;
+                }
+
                 ResourceBindingSegment[] segments = BindingSegments[setIndex];
 
                 if (segments != null && segments.Length > 0)
@@ -433,6 +516,11 @@ namespace Ryujinx.Graphics.Vulkan
             return null;
         }
 
+        public DescriptorSetTemplate GetPushDescriptorTemplate(long updateMask)
+        {
+            return _plce.GetPushDescriptorTemplate(IsCompute ? PipelineBindPoint.Compute : PipelineBindPoint.Graphics, updateMask);
+        }
+
         public void AddComputePipeline(ref SpecData key, Auto<DisposablePipeline> pipeline)
         {
             (_computePipelineCache ??= new()).Add(ref key, pipeline);
@@ -493,6 +581,11 @@ namespace Ryujinx.Graphics.Vulkan
             return _plce.GetNewDescriptorSetCollection(setIndex, out isNew);
         }
 
+        public bool HasSameLayout(ShaderCollection other)
+        {
+            return other != null && _plce == other._plce;
+        }
+
         protected virtual void Dispose(bool disposing)
         {
             if (disposing)