Perf/push descriptor

src/Ryujinx.Graphics.Vulkan/DescriptorSetTemplate.cs

@@ -1,12 +1,19 @@
 using Ryujinx.Graphics.GAL;
 using Silk.NET.Vulkan;
 using System;
+using System.Numerics;
 using System.Runtime.CompilerServices;
 
 namespace Ryujinx.Graphics.Vulkan
 {
     class DescriptorSetTemplate : IDisposable
     {
+        /// <summary>
+        /// Renderdoc seems to crash when doing a templated uniform update with count > 1 on a push descriptor.
+        /// When this is true, consecutive buffers are always updated individually.
+        /// </summary>
+        private const bool RenderdocPushCountBug = true;
+
         private readonly VulkanRenderer _gd;
         private readonly Device _device;
 
@@ -137,6 +144,86 @@ public unsafe DescriptorSetTemplate(VulkanRenderer gd, Device device, ResourceBi
             Template = result;
         }
 
+        public unsafe DescriptorSetTemplate(VulkanRenderer gd, Device device, ResourceDescriptorCollection descriptors, long updateMask, PipelineLayoutCacheEntry plce, PipelineBindPoint pbp, int setIndex)
+        {
+            _gd = gd;
+            _device = device;
+
+            // Create a template from the set usages. Assumes the descriptor set is updated in segment order then binding order.
+            int segmentCount = BitOperations.PopCount((ulong)updateMask);
+
+            DescriptorUpdateTemplateEntry* entries = stackalloc DescriptorUpdateTemplateEntry[segmentCount];
+            int entry = 0;
+            nuint structureOffset = 0;
+
+            void addBinding(int binding, int count)
+            {
+                entries[entry++] = new DescriptorUpdateTemplateEntry()
+                {
+                    DescriptorType = DescriptorType.UniformBuffer,
+                    DstBinding = (uint)binding,
+                    DescriptorCount = (uint)count,
+                    Offset = structureOffset,
+                    Stride = (nuint)Unsafe.SizeOf<DescriptorBufferInfo>()
+                };
+
+                structureOffset += (nuint)(Unsafe.SizeOf<DescriptorBufferInfo>() * count);
+            }
+
+            int startBinding = 0;
+            int bindingCount = 0;
+
+            foreach (ResourceDescriptor descriptor in descriptors.Descriptors)
+            {
+                for (int i = 0; i < descriptor.Count; i++)
+                {
+                    int binding = descriptor.Binding + i;
+
+                    if ((updateMask & (1L << binding)) != 0)
+                    {
+                        if (bindingCount > 0 && (RenderdocPushCountBug || startBinding + bindingCount != binding))
+                        {
+                            addBinding(startBinding, bindingCount);
+
+                            bindingCount = 0;
+                        }
+
+                        if (bindingCount == 0)
+                        {
+                            startBinding = binding;
+                        }
+
+                        bindingCount++;
+                    }
+                }
+            }
+
+            if (bindingCount > 0)
+            {
+                addBinding(startBinding, bindingCount);
+            }
+
+            Size = (int)structureOffset;
+
+            var info = new DescriptorUpdateTemplateCreateInfo()
+            {
+                SType = StructureType.DescriptorUpdateTemplateCreateInfo,
+                DescriptorUpdateEntryCount = (uint)entry,
+                PDescriptorUpdateEntries = entries,
+
+                TemplateType = DescriptorUpdateTemplateType.PushDescriptorsKhr,
+                DescriptorSetLayout = plce.DescriptorSetLayouts[setIndex],
+                PipelineBindPoint = pbp,
+                PipelineLayout = plce.PipelineLayout,
+                Set = (uint)setIndex,
+            };
+
+            DescriptorUpdateTemplate result;
+            gd.Api.CreateDescriptorUpdateTemplate(device, &info, null, &result).ThrowOnError();
+
+            Template = result;
+        }
+
         public unsafe void Dispose()
         {
             _gd.Api.DestroyDescriptorUpdateTemplate(_device, Template, null);

src/Ryujinx.Graphics.Vulkan/DescriptorSetTemplateUpdater.cs

@@ -52,11 +52,23 @@ public DescriptorSetTemplateWriter Begin(DescriptorSetTemplate template)
             return new DescriptorSetTemplateWriter(new Span<byte>(_data.Pointer, template.Size));
         }
 
+        public DescriptorSetTemplateWriter Begin(int maxSize)
+        {
+            EnsureSize(maxSize);
+
+            return new DescriptorSetTemplateWriter(new Span<byte>(_data.Pointer, maxSize));
+        }
+
         public void Commit(VulkanRenderer gd, Device device, DescriptorSet set)
         {
             gd.Api.UpdateDescriptorSetWithTemplate(device, set, _activeTemplate.Template, _data.Pointer);
         }
 
+        public void CommitPushDescriptor(VulkanRenderer gd, CommandBufferScoped cbs, DescriptorSetTemplate template, PipelineLayout layout)
+        {
+            gd.PushDescriptorApi.CmdPushDescriptorSetWithTemplate(cbs.CommandBuffer, template.Template, layout, 0, _data.Pointer);
+        }
+
         public void Dispose()
         {
             _data?.Dispose();

src/Ryujinx.Graphics.Vulkan/DescriptorSetUpdater.cs

@@ -4,6 +4,7 @@
 using Silk.NET.Vulkan;
 using System;
 using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
 using CompareOp = Ryujinx.Graphics.GAL.CompareOp;
 using Format = Ryujinx.Graphics.GAL.Format;
 using SamplerCreateInfo = Ryujinx.Graphics.GAL.SamplerCreateInfo;
@@ -61,6 +62,8 @@ public BufferRef(Auto<DisposableBuffer> buffer, ref BufferRange range)
         private BitMapStruct<Array2<long>> _storageSet;
         private BitMapStruct<Array2<long>> _uniformMirrored;
         private BitMapStruct<Array2<long>> _storageMirrored;
+        private readonly int[] _uniformSetPd;
+        private int _pdSequence = 1;
 
         private bool _updateDescriptorCacheCbIndex;
 
@@ -106,6 +109,8 @@ public DescriptorSetUpdater(VulkanRenderer gd, Device device, PipelineBase pipel
             _bufferTextures = new BufferView[Constants.MaxTexturesPerStage];
             _bufferImages = new BufferView[Constants.MaxImagesPerStage];
 
+            _uniformSetPd = new int[Constants.MaxUniformBufferBindings];
+
             var initialImageInfo = new DescriptorImageInfo
             {
                 ImageLayout = ImageLayout.General,
@@ -193,6 +198,7 @@ internal void Rebind(Auto<DisposableBuffer> buffer, int offset, int size)
                         if (BindingOverlaps(ref info, bindingOffset, offset, size))
                         {
                             _uniformSet.Clear(binding);
+                            _uniformSetPd[binding] = 0;
                             SignalDirty(DirtyFlags.Uniform);
                         }
                     }
@@ -223,8 +229,23 @@ internal void Rebind(Auto<DisposableBuffer> buffer, int offset, int size)
             });
         }
 
+        public void AdvancePdSequence()
+        {
+            if (++_pdSequence == 0)
+            {
+                _pdSequence = 1;
+            }
+        }
+
         public void SetProgram(ShaderCollection program)
         {
+            if (!program.HasSameLayout(_program))
+            {
+                // When the pipeline layout changes, push descriptor bindings are invalidated.
+
+                AdvancePdSequence();
+            }
+
             _program = program;
             _updateDescriptorCacheCbIndex = true;
             _dirty = DirtyFlags.All;
@@ -402,6 +423,7 @@ public void SetUniformBuffers(CommandBuffer commandBuffer, ReadOnlySpan<BufferAs
                 if (!currentBufferRef.Equals(newRef) || currentInfo.Range != info.Range)
                 {
                     _uniformSet.Clear(index);
+                    _uniformSetPd[index] = 0;
 
                     currentInfo = info;
                     currentBufferRef = newRef;
@@ -671,15 +693,19 @@ private unsafe void UpdateBuffers(
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private void UpdateAndBindUniformBufferPd(CommandBufferScoped cbs, PipelineBindPoint pbp)
         {
+            int sequence = _pdSequence;
             var bindingSegments = _program.BindingSegments[PipelineBase.UniformSetIndex];
             var dummyBuffer = _dummyBuffer?.GetBuffer();
 
+            long updatedBindings = 0;
+            DescriptorSetTemplateWriter writer = _templateUpdater.Begin(32 * Unsafe.SizeOf<DescriptorBufferInfo>());
+
             foreach (ResourceBindingSegment segment in bindingSegments)
             {
                 int binding = segment.Binding;
                 int count = segment.Count;
 
-                bool doUpdate = false;
+                ReadOnlySpan<DescriptorBufferInfo> uniformBuffers = _uniformBuffers;
 
                 for (int i = 0; i < count; i++)
                 {
@@ -688,17 +714,29 @@ private void UpdateAndBindUniformBufferPd(CommandBufferScoped cbs, PipelineBindP
                     if (_uniformSet.Set(index))
                     {
                         ref BufferRef buffer = ref _uniformBufferRefs[index];
-                        UpdateBuffer(cbs, ref _uniformBuffers[index], ref buffer, dummyBuffer, true);
-                        doUpdate = true;
+
+                        bool mirrored = UpdateBuffer(cbs, ref _uniformBuffers[index], ref buffer, dummyBuffer, true);
+
+                        _uniformMirrored.Set(index, mirrored);
                     }
-                }
 
-                if (doUpdate)
-                {
-                    ReadOnlySpan<DescriptorBufferInfo> uniformBuffers = _uniformBuffers;
-                    UpdateBuffers(cbs, pbp, binding, uniformBuffers.Slice(binding, count), DescriptorType.UniformBuffer);
+                    if (_uniformSetPd[index] != sequence)
+                    {
+                        // Need to set this push descriptor (even if the buffer binding has not changed)
+
+                        _uniformSetPd[index] = sequence;
+                        updatedBindings |= 1L << index;
+
+                        writer.Push(MemoryMarshal.CreateReadOnlySpan(ref _uniformBuffers[index], 1));
+                    }
                 }
             }
+
+            if (updatedBindings > 0)
+            {
+                DescriptorSetTemplate template = _program.GetPushDescriptorTemplate(updatedBindings);
+                _templateUpdater.CommitPushDescriptor(_gd, cbs, template, _program.PipelineLayout);
+            }
         }
 
         private void Initialize(CommandBufferScoped cbs, int setIndex, DescriptorSetCollection dsc)
@@ -724,6 +762,7 @@ public void SignalCommandBufferChange()
 
             _uniformSet.Clear();
             _storageSet.Clear();
+            AdvancePdSequence();
         }
 
         private static void SwapBuffer(BufferRef[] list, Auto<DisposableBuffer> from, Auto<DisposableBuffer> to)

src/Ryujinx.Graphics.Vulkan/HardwareCapabilities.cs

@@ -34,6 +34,7 @@ readonly struct HardwareCapabilities
         public readonly bool SupportsMultiView;
         public readonly bool SupportsNullDescriptors;
         public readonly bool SupportsPushDescriptors;
+        public readonly uint MaxPushDescriptors;
         public readonly bool SupportsPrimitiveTopologyListRestart;
         public readonly bool SupportsPrimitiveTopologyPatchListRestart;
         public readonly bool SupportsTransformFeedback;
@@ -71,6 +72,7 @@ public HardwareCapabilities(
             bool supportsMultiView,
             bool supportsNullDescriptors,
             bool supportsPushDescriptors,
+            uint maxPushDescriptors,
             bool supportsPrimitiveTopologyListRestart,
             bool supportsPrimitiveTopologyPatchListRestart,
             bool supportsTransformFeedback,
@@ -107,6 +109,7 @@ public HardwareCapabilities(
             SupportsMultiView = supportsMultiView;
             SupportsNullDescriptors = supportsNullDescriptors;
             SupportsPushDescriptors = supportsPushDescriptors;
+            MaxPushDescriptors = maxPushDescriptors;
             SupportsPrimitiveTopologyListRestart = supportsPrimitiveTopologyListRestart;
             SupportsPrimitiveTopologyPatchListRestart = supportsPrimitiveTopologyPatchListRestart;
             SupportsTransformFeedback = supportsTransformFeedback;

src/Ryujinx.Graphics.Vulkan/PipelineLayoutCacheEntry.cs

@@ -31,6 +31,11 @@ class PipelineLayoutCacheEntry
         private int _dsLastCbIndex;
         private int _dsLastSubmissionCount;
 
+        private readonly Dictionary<long, DescriptorSetTemplate> _pdTemplates;
+        private readonly ResourceDescriptorCollection _pdDescriptors;
+        private long _lastPdUsage;
+        private DescriptorSetTemplate _lastPdTemplate;
+
         private PipelineLayoutCacheEntry(VulkanRenderer gd, Device device, int setsCount)
         {
             _gd = gd;
@@ -72,6 +77,12 @@ public PipelineLayoutCacheEntry(
 
                 _consumedDescriptorsPerSet[setIndex] = count;
             }
+
+            if (usePushDescriptors)
+            {
+                _pdDescriptors = setDescriptors[0];
+                _pdTemplates = new();
+            }
         }
 
         public void UpdateCommandBufferIndex(int commandBufferIndex)
@@ -143,10 +154,39 @@ private static Span<DescriptorPoolSize> GetDescriptorPoolSizes(Span<DescriptorPo
             return output[..count];
         }
 
+        public DescriptorSetTemplate GetPushDescriptorTemplate(PipelineBindPoint pbp, long updateMask)
+        {
+            if (_lastPdUsage == updateMask && _lastPdTemplate != null)
+            {
+                // Most likely result is that it asks to update the same buffers.
+                return _lastPdTemplate;
+            }
+
+            if (!_pdTemplates.TryGetValue(updateMask, out DescriptorSetTemplate template))
+            {
+                template = new DescriptorSetTemplate(_gd, _device, _pdDescriptors, updateMask, this, pbp, 0);
+
+                _pdTemplates.Add(updateMask, template);
+            }
+
+            _lastPdUsage = updateMask;
+            _lastPdTemplate = template;
+
+            return template;
+        }
+
         protected virtual unsafe void Dispose(bool disposing)
         {
             if (disposing)
             {
+                if (_pdTemplates != null)
+                {
+                    foreach (DescriptorSetTemplate template in _pdTemplates.Values)
+                    {
+                        template.Dispose();
+                    }
+                }
+
                 for (int i = 0; i < _dsCache.Length; i++)
                 {
                     for (int j = 0; j < _dsCache[i].Length; j++)