1 files changed, 163 insertions, 46 deletions
diff --git a/src/Ryujinx.Graphics.Vulkan/DescriptorSetUpdater.cs b/src/Ryujinx.Graphics.Vulkan/DescriptorSetUpdater.cs
index 45392b64..14e4c02f 100644
--- a/src/Ryujinx.Graphics.Vulkan/DescriptorSetUpdater.cs
+++ b/src/Ryujinx.Graphics.Vulkan/DescriptorSetUpdater.cs
@@ -1,9 +1,9 @@
-using Ryujinx.Graphics.GAL;
+using Ryujinx.Common.Memory;
+using Ryujinx.Graphics.GAL;
 using Ryujinx.Graphics.Shader;
 using Silk.NET.Vulkan;
 using System;
 using System.Runtime.CompilerServices;
-using Buffer = Silk.NET.Vulkan.Buffer;
 using CompareOp = Ryujinx.Graphics.GAL.CompareOp;
 using Format = Ryujinx.Graphics.GAL.Format;
 using SamplerCreateInfo = Ryujinx.Graphics.GAL.SamplerCreateInfo;
@@ -12,13 +12,34 @@ namespace Ryujinx.Graphics.Vulkan
 {
     class DescriptorSetUpdater
     {
+        private const ulong StorageBufferMaxMirrorable = 0x2000;
+        private record struct BufferRef
+        {
+            public Auto<DisposableBuffer> Buffer;
+            public int Offset;
+            public bool Write;
+
+            public BufferRef(Auto<DisposableBuffer> buffer)
+            {
+                Buffer = buffer;
+                Offset = 0;
+                Write = true;
+            }
+
+            public BufferRef(Auto<DisposableBuffer> buffer, ref BufferRange range)
+            {
+                Buffer = buffer;
+                Offset = range.Offset;
+                Write = range.Write;
+            }
+        }
+
         private readonly VulkanRenderer _gd;
         private readonly PipelineBase _pipeline;
-
         private ShaderCollection _program;
 
-        private readonly Auto<DisposableBuffer>[] _uniformBufferRefs;
-        private readonly Auto<DisposableBuffer>[] _storageBufferRefs;
+        private readonly BufferRef[] _uniformBufferRefs;
+        private readonly BufferRef[] _storageBufferRefs;
         private readonly Auto<DisposableImageView>[] _textureRefs;
         private readonly Auto<DisposableSampler>[] _samplerRefs;
         private readonly Auto<DisposableImageView>[] _imageRefs;
@@ -33,8 +54,10 @@ namespace Ryujinx.Graphics.Vulkan
         private readonly BufferView[] _bufferTextures;
         private readonly BufferView[] _bufferImages;
 
-        private readonly bool[] _uniformSet;
-        private readonly bool[] _storageSet;
+        private BitMapStruct<Array2<long>> _uniformSet;
+        private BitMapStruct<Array2<long>> _storageSet;
+        private BitMapStruct<Array2<long>> _uniformMirrored;
+        private BitMapStruct<Array2<long>> _storageMirrored;
 
         [Flags]
         private enum DirtyFlags
@@ -61,8 +84,8 @@ namespace Ryujinx.Graphics.Vulkan
             // Some of the bindings counts needs to be multiplied by 2 because we have buffer and
             // regular textures/images interleaved on the same descriptor set.
 
-            _uniformBufferRefs = new Auto<DisposableBuffer>[Constants.MaxUniformBufferBindings];
-            _storageBufferRefs = new Auto<DisposableBuffer>[Constants.MaxStorageBufferBindings];
+            _uniformBufferRefs = new BufferRef[Constants.MaxUniformBufferBindings];
+            _storageBufferRefs = new BufferRef[Constants.MaxStorageBufferBindings];
             _textureRefs = new Auto<DisposableImageView>[Constants.MaxTextureBindings * 2];
             _samplerRefs = new Auto<DisposableSampler>[Constants.MaxTextureBindings * 2];
             _imageRefs = new Auto<DisposableImageView>[Constants.MaxImageBindings * 2];
@@ -85,9 +108,6 @@ namespace Ryujinx.Graphics.Vulkan
             _textures.AsSpan().Fill(initialImageInfo);
             _images.AsSpan().Fill(initialImageInfo);
 
-            _uniformSet = new bool[Constants.MaxUniformBufferBindings];
-            _storageSet = new bool[Constants.MaxStorageBufferBindings];
-
             if (gd.Capabilities.SupportsNullDescriptors)
             {
                 // If null descriptors are supported, we can pass null as the handle.
@@ -138,6 +158,63 @@ namespace Ryujinx.Graphics.Vulkan
             _dummyTexture.SetData(dummyTextureData);
         }
 
+        private static bool BindingOverlaps(ref DescriptorBufferInfo info, int bindingOffset, int offset, int size)
+        {
+            return offset < bindingOffset + (int)info.Range && (offset + size) > bindingOffset;
+        }
+
+        internal void Rebind(Auto<DisposableBuffer> buffer, int offset, int size)
+        {
+            if (_program == null)
+            {
+                return;
+            }
+
+            // Check stage bindings
+
+            _uniformMirrored.Union(_uniformSet).SignalSet((int binding, int count) =>
+            {
+                for (int i = 0; i < count; i++)
+                {
+                    ref BufferRef bufferRef = ref _uniformBufferRefs[binding];
+                    if (bufferRef.Buffer == buffer)
+                    {
+                        ref DescriptorBufferInfo info = ref _uniformBuffers[binding];
+                        int bindingOffset = bufferRef.Offset;
+
+                        if (BindingOverlaps(ref info, bindingOffset, offset, size))
+                        {
+                            _uniformSet.Clear(binding);
+                            SignalDirty(DirtyFlags.Uniform);
+                        }
+                    }
+
+                    binding++;
+                }
+            });
+
+            _storageMirrored.Union(_storageSet).SignalSet((int binding, int count) =>
+            {
+                for (int i = 0; i < count; i++)
+                {
+                    ref BufferRef bufferRef = ref _storageBufferRefs[binding];
+                    if (bufferRef.Buffer == buffer)
+                    {
+                        ref DescriptorBufferInfo info = ref _storageBuffers[binding];
+                        int bindingOffset = bufferRef.Offset;
+
+                        if (BindingOverlaps(ref info, bindingOffset, offset, size))
+                        {
+                            _storageSet.Clear(binding);
+                            SignalDirty(DirtyFlags.Storage);
+                        }
+                    }
+
+                    binding++;
+                }
+            });
+        }
+
         public void SetProgram(ShaderCollection program)
         {
             _program = program;
@@ -180,22 +257,28 @@ namespace Ryujinx.Graphics.Vulkan
                 var buffer = assignment.Range;
                 int index = assignment.Binding;
 
-                Auto<DisposableBuffer> vkBuffer = _gd.BufferManager.GetBuffer(commandBuffer, buffer.Handle, false, isSSBO: true);
-                ref Auto<DisposableBuffer> currentVkBuffer = ref _storageBufferRefs[index];
+                Auto<DisposableBuffer> vkBuffer = buffer.Handle == BufferHandle.Null
+                    ? null
+                    : _gd.BufferManager.GetBuffer(commandBuffer, buffer.Handle, buffer.Write, isSSBO: true);
+
+                ref BufferRef currentBufferRef = ref _storageBufferRefs[index];
 
                 DescriptorBufferInfo info = new()
                 {
                     Offset = (ulong)buffer.Offset,
                     Range = (ulong)buffer.Size,
                 };
+
+                var newRef = new BufferRef(vkBuffer, ref buffer);
+
                 ref DescriptorBufferInfo currentInfo = ref _storageBuffers[index];
 
-                if (vkBuffer != currentVkBuffer || currentInfo.Offset != info.Offset || currentInfo.Range != info.Range)
+                if (!currentBufferRef.Equals(newRef) || currentInfo.Range != info.Range)
                 {
-                    _storageSet[index] = false;
+                    _storageSet.Clear(index);
 
                     currentInfo = info;
-                    currentVkBuffer = vkBuffer;
+                    currentBufferRef = newRef;
                 }
             }
 
@@ -209,21 +292,24 @@ namespace Ryujinx.Graphics.Vulkan
                 var vkBuffer = buffers[i];
                 int index = first + i;
 
-                ref Auto<DisposableBuffer> currentVkBuffer = ref _storageBufferRefs[index];
+                ref BufferRef currentBufferRef = ref _storageBufferRefs[index];
 
                 DescriptorBufferInfo info = new()
                 {
                     Offset = 0,
                     Range = Vk.WholeSize,
                 };
+
+                BufferRef newRef = new(vkBuffer);
+
                 ref DescriptorBufferInfo currentInfo = ref _storageBuffers[index];
 
-                if (vkBuffer != currentVkBuffer || currentInfo.Offset != info.Offset || currentInfo.Range != info.Range)
+                if (!currentBufferRef.Equals(newRef) || currentInfo.Range != info.Range)
                 {
-                    _storageSet[index] = false;
+                    _storageSet.Clear(index);
 
                     currentInfo = info;
-                    currentVkBuffer = vkBuffer;
+                    currentBufferRef = newRef;
                 }
             }
 
@@ -288,22 +374,28 @@ namespace Ryujinx.Graphics.Vulkan
                 var buffer = assignment.Range;
                 int index = assignment.Binding;
 
-                Auto<DisposableBuffer> vkBuffer = _gd.BufferManager.GetBuffer(commandBuffer, buffer.Handle, false);
-                ref Auto<DisposableBuffer> currentVkBuffer = ref _uniformBufferRefs[index];
+                Auto<DisposableBuffer> vkBuffer = buffer.Handle == BufferHandle.Null
+                    ? null
+                    : _gd.BufferManager.GetBuffer(commandBuffer, buffer.Handle, false);
+
+                ref BufferRef currentBufferRef = ref _uniformBufferRefs[index];
 
                 DescriptorBufferInfo info = new()
                 {
                     Offset = (ulong)buffer.Offset,
                     Range = (ulong)buffer.Size,
                 };
+
+                BufferRef newRef = new(vkBuffer, ref buffer);
+
                 ref DescriptorBufferInfo currentInfo = ref _uniformBuffers[index];
 
-                if (vkBuffer != currentVkBuffer || currentInfo.Offset != info.Offset || currentInfo.Range != info.Range)
+                if (!currentBufferRef.Equals(newRef) || currentInfo.Range != info.Range)
                 {
-                    _uniformSet[index] = false;
+                    _uniformSet.Clear(index);
 
                     currentInfo = info;
-                    currentVkBuffer = vkBuffer;
+                    currentBufferRef = newRef;
                 }
             }
 
@@ -353,13 +445,26 @@ namespace Ryujinx.Graphics.Vulkan
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static void UpdateBuffer(
+        private static bool UpdateBuffer(
             CommandBufferScoped cbs,
             ref DescriptorBufferInfo info,
-            Auto<DisposableBuffer> buffer,
-            Auto<DisposableBuffer> dummyBuffer)
+            ref BufferRef buffer,
+            Auto<DisposableBuffer> dummyBuffer,
+            bool mirrorable)
         {
-            info.Buffer = buffer?.Get(cbs, (int)info.Offset, (int)info.Range).Value ?? default;
+            int offset = buffer.Offset;
+            bool mirrored = false;
+
+            if (mirrorable)
+            {
+                info.Buffer = buffer.Buffer?.GetMirrorable(cbs, ref offset, (int)info.Range, out mirrored).Value ?? default;
+            }
+            else
+            {
+                info.Buffer = buffer.Buffer?.Get(cbs, offset, (int)info.Range, buffer.Write).Value ?? default;
+            }
+
+            info.Offset = (ulong)offset;
 
             // The spec requires that buffers with null handle have offset as 0 and range as VK_WHOLE_SIZE.
             if (info.Buffer.Handle == 0)
@@ -368,6 +473,8 @@ namespace Ryujinx.Graphics.Vulkan
                 info.Offset = 0;
                 info.Range = Vk.WholeSize;
             }
+
+            return mirrored;
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -404,11 +511,13 @@ namespace Ryujinx.Graphics.Vulkan
                     {
                         int index = binding + i;
 
-                        if (!_uniformSet[index])
+                        if (_uniformSet.Set(index))
                         {
-                            UpdateBuffer(cbs, ref _uniformBuffers[index], _uniformBufferRefs[index], dummyBuffer);
+                            ref BufferRef buffer = ref _uniformBufferRefs[index];
 
-                            _uniformSet[index] = true;
+                            bool mirrored = UpdateBuffer(cbs, ref _uniformBuffers[index], ref buffer, dummyBuffer, true);
+
+                            _uniformMirrored.Set(index, mirrored);
                         }
                     }
 
@@ -421,11 +530,19 @@ namespace Ryujinx.Graphics.Vulkan
                     {
                         int index = binding + i;
 
-                        if (!_storageSet[index])
+                        ref BufferRef buffer = ref _storageBufferRefs[index];
+
+                        if (_storageSet.Set(index))
                         {
-                            UpdateBuffer(cbs, ref _storageBuffers[index], _storageBufferRefs[index], dummyBuffer);
+                            ref var info = ref _storageBuffers[index];
+
+                            bool mirrored = UpdateBuffer(cbs,
+                                ref info,
+                                ref _storageBufferRefs[index],
+                                dummyBuffer,
+                                !buffer.Write && info.Range <= StorageBufferMaxMirrorable);
 
-                            _storageSet[index] = true;
+                            _storageMirrored.Set(index, mirrored);
                         }
                     }
 
@@ -464,7 +581,7 @@ namespace Ryujinx.Graphics.Vulkan
 
                         for (int i = 0; i < count; i++)
                         {
-                            bufferTextures[i] = _bufferTextureRefs[binding + i]?.GetBufferView(cbs) ?? default;
+                            bufferTextures[i] = _bufferTextureRefs[binding + i]?.GetBufferView(cbs, false) ?? default;
                         }
 
                         dsc.UpdateBufferImages(0, binding, bufferTextures[..count], DescriptorType.UniformTexelBuffer);
@@ -489,7 +606,7 @@ namespace Ryujinx.Graphics.Vulkan
 
                         for (int i = 0; i < count; i++)
                         {
-                            bufferImages[i] = _bufferImageRefs[binding + i]?.GetBufferView(cbs, _bufferImageFormats[binding + i]) ?? default;
+                            bufferImages[i] = _bufferImageRefs[binding + i]?.GetBufferView(cbs, _bufferImageFormats[binding + i], true) ?? default;
                         }
 
                         dsc.UpdateBufferImages(0, binding, bufferImages[..count], DescriptorType.StorageTexelBuffer);
@@ -546,10 +663,10 @@ namespace Ryujinx.Graphics.Vulkan
                 {
                     int index = binding + i;
 
-                    if (!_uniformSet[index])
+                    if (_uniformSet.Set(index))
                     {
-                        UpdateBuffer(cbs, ref _uniformBuffers[index], _uniformBufferRefs[index], dummyBuffer);
-                        _uniformSet[index] = true;
+                        ref BufferRef buffer = ref _uniformBufferRefs[index];
+                        UpdateBuffer(cbs, ref _uniformBuffers[index], ref buffer, dummyBuffer, true);
                         doUpdate = true;
                     }
                 }
@@ -582,17 +699,17 @@ namespace Ryujinx.Graphics.Vulkan
         {
             _dirty = DirtyFlags.All;
 
-            Array.Clear(_uniformSet);
-            Array.Clear(_storageSet);
+            _uniformSet.Clear();
+            _storageSet.Clear();
         }
 
-        private static void SwapBuffer(Auto<DisposableBuffer>[] list, Auto<DisposableBuffer> from, Auto<DisposableBuffer> to)
+        private static void SwapBuffer(BufferRef[] list, Auto<DisposableBuffer> from, Auto<DisposableBuffer> to)
         {
             for (int i = 0; i < list.Length; i++)
             {
-                if (list[i] == from)
+                if (list[i].Buffer == from)
                 {
-                    list[i] = to;
+                    list[i].Buffer = to;
                 }
             }
         }