Vulkan: Fix some issues with CacheByRange (#3743)

* Fix some issues with CacheByRange

- Cache now clears under more circumstances, the most important being the fast path write.
- Cache supports partial clear which should help when more buffers join.
- Fixed an issue with I8->I16 conversion where it wouldn't register the buffer for use on dispose.

Should hopefully fix issues with https://github.com/Ryujinx/Ryujinx-Games-List/issues/4010 and maybe others.

* Fix collection modified exception

* Fix accidental use of parameterless constructor

* Replay DynamicState when restoring from helper shader

Ryujinx.Graphics.Vulkan/BufferHolder.cs

@@ -109,12 +109,34 @@ namespace Ryujinx.Graphics.Vulkan
         {
             if (isWrite)
             {
-                _cachedConvertedBuffers.Clear();
+                SignalWrite(0, Size);
             }
 
             return _buffer;
         }
 
+        public Auto<DisposableBuffer> GetBuffer(CommandBuffer commandBuffer, int offset, int size, bool isWrite = false)
+        {
+            if (isWrite)
+            {
+                SignalWrite(offset, size);
+            }
+
+            return _buffer;
+        }
+
+        public void SignalWrite(int offset, int size)
+        {
+            if (offset == 0 && size == Size)
+            {
+                _cachedConvertedBuffers.Clear();
+            }
+            else
+            {
+                _cachedConvertedBuffers.ClearRange(offset, size);
+            }
+        }
+
         public BufferHandle GetHandle()
         {
             var handle = _bufferHandle;
@@ -183,6 +205,8 @@ namespace Ryujinx.Graphics.Vulkan
 
                     data.Slice(0, dataSize).CopyTo(new Span<byte>((void*)(_map + offset), dataSize));
 
+                    SignalWrite(offset, dataSize);
+
                     return;
                 }
             }
@@ -240,7 +264,7 @@ namespace Ryujinx.Graphics.Vulkan
 
             endRenderPass?.Invoke();
 
-            var dstBuffer = GetBuffer(cbs.CommandBuffer, true).Get(cbs, dstOffset, data.Length).Value;
+            var dstBuffer = GetBuffer(cbs.CommandBuffer, dstOffset, data.Length, true).Get(cbs, dstOffset, data.Length).Value;
 
             InsertBufferBarrier(
                 _gd,
@@ -364,7 +388,7 @@ namespace Ryujinx.Graphics.Vulkan
 
         public Auto<DisposableBuffer> GetBufferI8ToI16(CommandBufferScoped cbs, int offset, int size)
         {
-            var key = new I8ToI16CacheKey();
+            var key = new I8ToI16CacheKey(_gd);
 
             if (!_cachedConvertedBuffers.TryGetValue(offset, size, key, out var holder))
             {
@@ -373,6 +397,8 @@ namespace Ryujinx.Graphics.Vulkan
                 _gd.PipelineInternal.EndRenderPass();
                 _gd.HelperShader.ConvertI8ToI16(_gd, cbs, this, holder, offset, size);
 
+                key.SetBuffer(holder.GetBuffer());
+
                 _cachedConvertedBuffers.Add(offset, size, key, holder);
             }
 
@@ -417,6 +443,8 @@ namespace Ryujinx.Graphics.Vulkan
                 _gd.PipelineInternal.EndRenderPass();
                 _gd.HelperShader.ConvertIndexBuffer(_gd, cbs, this, holder, pattern, indexSize, offset, indexCount);
 
+                key.SetBuffer(holder.GetBuffer());
+
                 _cachedConvertedBuffers.Add(offset, size, key, holder);
             }
 

Ryujinx.Graphics.Vulkan/BufferManager.cs

@@ -124,6 +124,16 @@ namespace Ryujinx.Graphics.Vulkan
             return null;
         }
 
+        public Auto<DisposableBuffer> GetBuffer(CommandBuffer commandBuffer, BufferHandle handle, int offset, int size, bool isWrite)
+        {
+            if (TryGetBuffer(handle, out var holder))
+            {
+                return holder.GetBuffer(commandBuffer, offset, size, isWrite);
+            }
+
+            return null;
+        }
+
         public Auto<DisposableBuffer> GetBufferI8ToI16(CommandBufferScoped cbs, BufferHandle handle, int offset, int size)
         {
             if (TryGetBuffer(handle, out var holder))

Ryujinx.Graphics.Vulkan/CacheByRange.cs

@@ -25,6 +25,11 @@ namespace Ryujinx.Graphics.Vulkan
             return other is I8ToI16CacheKey;
         }
 
+        public void SetBuffer(Auto<DisposableBuffer> buffer)
+        {
+            _buffer = buffer;
+        }
+
         public void Dispose()
         {
             _gd.PipelineInternal.DirtyIndexBuffer(_buffer);
@@ -160,6 +165,44 @@ namespace Ryujinx.Graphics.Vulkan
             }
         }
 
+        public void ClearRange(int offset, int size)
+        {
+            if (_ranges != null && _ranges.Count > 0)
+            {
+                int end = offset + size;
+
+                List<ulong> toRemove = null;
+
+                foreach (KeyValuePair<ulong, List<Entry>> range in _ranges)
+                {
+                    (int rOffset, int rSize) = UnpackRange(range.Key);
+
+                    int rEnd = rOffset + rSize;
+
+                    if (rEnd > offset && rOffset < end)
+                    {
+                        List<Entry> entries = range.Value;
+
+                        foreach (Entry entry in entries)
+                        {
+                            entry.Key.Dispose();
+                            entry.Value.Dispose();
+                        }
+
+                        (toRemove ??= new List<ulong>()).Add(range.Key);
+                    }
+                }
+
+                if (toRemove != null)
+                {
+                    foreach (ulong range in toRemove)
+                    {
+                        _ranges.Remove(range);
+                    }
+                }
+            }
+        }
+
         private List<Entry> GetEntries(int offset, int size)
         {
             if (_ranges == null)
@@ -184,6 +227,11 @@ namespace Ryujinx.Graphics.Vulkan
             return (uint)offset | ((ulong)size << 32);
         }
 
+        private static (int offset, int size) UnpackRange(ulong range)
+        {
+            return ((int)range, (int)(range >> 32));
+        }
+
         public void Dispose()
         {
             Clear();

Ryujinx.Graphics.Vulkan/PipelineBase.cs

@@ -21,7 +21,7 @@ namespace Ryujinx.Graphics.Vulkan
 
         protected readonly AutoFlushCounter AutoFlush;
 
-        private PipelineDynamicState _dynamicState;
+        protected PipelineDynamicState DynamicState;
         private PipelineState _newState;
         private bool _stateDirty;
         private GAL.PrimitiveTopology _topology;
@@ -150,7 +150,7 @@ namespace Ryujinx.Graphics.Vulkan
         {
             EndRenderPass();
 
-            var dst = Gd.BufferManager.GetBuffer(CommandBuffer, destination, true).Get(Cbs, offset, size).Value;
+            var dst = Gd.BufferManager.GetBuffer(CommandBuffer, destination, offset, size, true).Get(Cbs, offset, size).Value;
 
             BufferHolder.InsertBufferBarrier(
                 Gd,
@@ -238,8 +238,8 @@ namespace Ryujinx.Graphics.Vulkan
         {
             EndRenderPass();
 
-            var src = Gd.BufferManager.GetBuffer(CommandBuffer, source, false);
-            var dst = Gd.BufferManager.GetBuffer(CommandBuffer, destination, true);
+            var src = Gd.BufferManager.GetBuffer(CommandBuffer, source, srcOffset, size, false);
+            var dst = Gd.BufferManager.GetBuffer(CommandBuffer, destination, dstOffset, size, true);
 
             BufferHolder.Copy(Gd, Cbs, src, dst, srcOffset, dstOffset, size);
         }
@@ -388,7 +388,7 @@ namespace Ryujinx.Graphics.Vulkan
                 var oldDepthTestEnable = _newState.DepthTestEnable;
                 var oldDepthWriteEnable = _newState.DepthWriteEnable;
                 var oldTopology = _newState.Topology;
-                var oldViewports = _dynamicState.Viewports;
+                var oldViewports = DynamicState.Viewports;
                 var oldViewportsCount = _newState.ViewportsCount;
 
                 _newState.CullMode = CullModeFlags.CullModeNone;
@@ -411,9 +411,9 @@ namespace Ryujinx.Graphics.Vulkan
                 _newState.DepthWriteEnable = oldDepthWriteEnable;
                 _newState.Topology = oldTopology;
 
-                _dynamicState.Viewports = oldViewports;
-                _dynamicState.ViewportsCount = (int)oldViewportsCount;
-                _dynamicState.SetViewportsDirty();
+                DynamicState.Viewports = oldViewports;
+                DynamicState.ViewportsCount = (int)oldViewportsCount;
+                DynamicState.SetViewportsDirty();
 
                 _newState.ViewportsCount = oldViewportsCount;
                 SignalStateChange();
@@ -448,8 +448,13 @@ namespace Ryujinx.Graphics.Vulkan
             ResumeTransformFeedbackInternal();
             DrawCount++;
 
-            var buffer = Gd.BufferManager.GetBuffer(CommandBuffer, indirectBuffer.Handle, true).Get(Cbs, indirectBuffer.Offset, indirectBuffer.Size).Value;
-            var countBuffer = Gd.BufferManager.GetBuffer(CommandBuffer, parameterBuffer.Handle, true).Get(Cbs, parameterBuffer.Offset, parameterBuffer.Size).Value;
+            var buffer = Gd.BufferManager
+                .GetBuffer(CommandBuffer, indirectBuffer.Handle, indirectBuffer.Offset, indirectBuffer.Size, true)
+                .Get(Cbs, indirectBuffer.Offset, indirectBuffer.Size).Value;
+
+            var countBuffer = Gd.BufferManager
+                .GetBuffer(CommandBuffer, parameterBuffer.Handle, parameterBuffer.Offset, parameterBuffer.Size, true)
+                .Get(Cbs, parameterBuffer.Offset, parameterBuffer.Size).Value;
 
             Gd.DrawIndirectCountApi.CmdDrawIndirectCount(
                 CommandBuffer,
@@ -478,8 +483,13 @@ namespace Ryujinx.Graphics.Vulkan
             ResumeTransformFeedbackInternal();
             DrawCount++;
 
-            var buffer = Gd.BufferManager.GetBuffer(CommandBuffer, indirectBuffer.Handle, true).Get(Cbs, indirectBuffer.Offset, indirectBuffer.Size).Value;
-            var countBuffer = Gd.BufferManager.GetBuffer(CommandBuffer, parameterBuffer.Handle, true).Get(Cbs, parameterBuffer.Offset, parameterBuffer.Size).Value;
+            var buffer = Gd.BufferManager
+                .GetBuffer(CommandBuffer, indirectBuffer.Handle, parameterBuffer.Offset, parameterBuffer.Size, true)
+                .Get(Cbs, indirectBuffer.Offset, indirectBuffer.Size).Value;
+
+            var countBuffer = Gd.BufferManager
+                .GetBuffer(CommandBuffer, parameterBuffer.Handle, parameterBuffer.Offset, parameterBuffer.Size, true)
+                .Get(Cbs, parameterBuffer.Offset, parameterBuffer.Size).Value;
 
             Gd.DrawIndirectCountApi.CmdDrawIndexedIndirectCount(
                 CommandBuffer,
@@ -535,7 +545,7 @@ namespace Ryujinx.Graphics.Vulkan
 
         public void SetDepthBias(PolygonModeMask enables, float factor, float units, float clamp)
         {
-            _dynamicState.SetDepthBias(factor, units, clamp);
+            DynamicState.SetDepthBias(factor, units, clamp);
 
             _newState.DepthBiasEnable = enables != 0;
             SignalStateChange();
@@ -753,10 +763,10 @@ namespace Ryujinx.Graphics.Vulkan
                 var offset = new Offset2D(region.X, region.Y);
                 var extent = new Extent2D((uint)region.Width, (uint)region.Height);
 
-                _dynamicState.SetScissor(i, new Rect2D(offset, extent));
+                DynamicState.SetScissor(i, new Rect2D(offset, extent));
             }
 
-            _dynamicState.ScissorsCount = count;
+            DynamicState.ScissorsCount = count;
 
             _newState.ScissorsCount = (uint)count;
             SignalStateChange();
@@ -764,7 +774,7 @@ namespace Ryujinx.Graphics.Vulkan
 
         public void SetStencilTest(StencilTestDescriptor stencilTest)
         {
-            _dynamicState.SetStencilMasks(
+            DynamicState.SetStencilMasks(
                 (uint)stencilTest.BackFuncMask,
                 (uint)stencilTest.BackMask,
                 (uint)stencilTest.BackFuncRef,
@@ -813,7 +823,8 @@ namespace Ryujinx.Graphics.Vulkan
 
                 if (range.Handle != BufferHandle.Null)
                 {
-                    _transformFeedbackBuffers[i] = new BufferState(Gd.BufferManager.GetBuffer(CommandBuffer, range.Handle, true), range.Offset, range.Size);
+                    _transformFeedbackBuffers[i] = 
+                        new BufferState(Gd.BufferManager.GetBuffer(CommandBuffer, range.Handle, range.Offset, range.Size, true), range.Offset, range.Size);
                     _transformFeedbackBuffers[i].BindTransformFeedbackBuffer(Gd, Cbs, (uint)i);
                 }
                 else
@@ -975,7 +986,7 @@ namespace Ryujinx.Graphics.Vulkan
             {
                 var viewport = viewports[i];
 
-                _dynamicState.SetViewport(i, new Silk.NET.Vulkan.Viewport(
+                DynamicState.SetViewport(i, new Silk.NET.Vulkan.Viewport(
                     viewport.Region.X,
                     viewport.Region.Y,
                     viewport.Region.Width == 0f ? 1f : viewport.Region.Width,
@@ -984,7 +995,7 @@ namespace Ryujinx.Graphics.Vulkan
                     Clamp(viewport.DepthFar)));
             }
 
-            _dynamicState.ViewportsCount = count;
+            DynamicState.ViewportsCount = count;
 
             float disableTransformF = disableTransform ? 1.0f : 0.0f;
             if (SupportBufferUpdater.Data.ViewportInverse.W != disableTransformF || disableTransform)
@@ -1063,7 +1074,7 @@ namespace Ryujinx.Graphics.Vulkan
             _vertexBuffersDirty = ulong.MaxValue >> (64 - _vertexBuffers.Length);
 
             _descriptorSetUpdater.SignalCommandBufferChange();
-            _dynamicState.ForceAllDirty();
+            DynamicState.ForceAllDirty();
             _currentPipelineHandle = 0;
         }
 
@@ -1201,7 +1212,7 @@ namespace Ryujinx.Graphics.Vulkan
 
         private void RecreatePipelineIfNeeded(PipelineBindPoint pbp)
         {
-            _dynamicState.ReplayIfDirty(Gd.Api, CommandBuffer);
+            DynamicState.ReplayIfDirty(Gd.Api, CommandBuffer);
 
             // Commit changes to the support buffer before drawing.
             SupportBufferUpdater.Commit();

Ryujinx.Graphics.Vulkan/PipelineFull.cs

@@ -204,6 +204,8 @@ namespace Ryujinx.Graphics.Vulkan
             }
 
             SignalCommandBufferChange();
+
+            DynamicState.ReplayIfDirty(Gd.Api, CommandBuffer);
         }
 
         public void FlushCommandsImpl()

Ryujinx.Graphics.Vulkan/StagingBuffer.cs

@@ -87,7 +87,7 @@ namespace Ryujinx.Graphics.Vulkan
         private void PushDataImpl(CommandBufferScoped cbs, BufferHolder dst, int dstOffset, ReadOnlySpan<byte> data)
         {
             var srcBuffer = _buffer.GetBuffer();
-            var dstBuffer = dst.GetBuffer();
+            var dstBuffer = dst.GetBuffer(cbs.CommandBuffer, dstOffset, data.Length, true);
 
             int offset = _freeOffset;
             int capacity = BufferSize - offset;

Ryujinx.HLE/HOS/Services/Nv/NvDrvServices/NvMap/NvMapDeviceFile.cs

@@ -3,7 +3,6 @@ using Ryujinx.Common.Logging;
 using Ryujinx.Graphics.Gpu.Memory;
 using Ryujinx.Memory;
 using System;
-using System.Collections.Concurrent;
 
 namespace Ryujinx.HLE.HOS.Services.Nv.NvDrvServices.NvMap
 {
@@ -11,13 +10,10 @@ namespace Ryujinx.HLE.HOS.Services.Nv.NvDrvServices.NvMap
     {
         private const int FlagNotFreedYet = 1;
 
-        private static ConcurrentDictionary<ulong, IdDictionary> _maps = new ConcurrentDictionary<ulong, IdDictionary>();
+        private static NvMapIdDictionary _maps = new NvMapIdDictionary();
 
         public NvMapDeviceFile(ServiceCtx context, IVirtualMemoryManager memory, ulong owner) : base(context, owner)
         {
-            IdDictionary dict = _maps.GetOrAdd(Owner, (key) => new IdDictionary());
-
-            dict.Add(0, new NvMapHandle());
         }
 
         public override NvInternalResult Ioctl(NvIoctl command, Span<byte> arguments)
@@ -232,19 +228,12 @@ namespace Ryujinx.HLE.HOS.Services.Nv.NvDrvServices.NvMap
 
         private int CreateHandleFromMap(NvMapHandle map)
         {
-            IdDictionary dict = _maps.GetOrAdd(Owner, (key) => new IdDictionary());
-
-            return dict.Add(map);
+            return _maps.Add(map);
         }
 
         private static bool DeleteMapWithHandle(ulong pid, int handle)
         {
-            if (_maps.TryGetValue(pid, out IdDictionary dict))
-            {
-                return dict.Delete(handle) != null;
-            }
-
-            return false;
+            return _maps.Delete(handle) != null;
         }
 
         public static void IncrementMapRefCount(ulong pid, int handle)
@@ -277,12 +266,7 @@ namespace Ryujinx.HLE.HOS.Services.Nv.NvDrvServices.NvMap
 
         public static NvMapHandle GetMapFromHandle(ulong pid, int handle)
         {
-            if (_maps.TryGetValue(pid, out IdDictionary dict))
-            {
-                return dict.GetData<NvMapHandle>(handle);
-            }
-
-            return null;
+            return _maps.Get(handle);
         }
     }
 }

Ryujinx.HLE/HOS/Services/Nv/NvDrvServices/NvMap/Types/NvMapIdDictionary.cs

@@ -0,0 +1,61 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Threading;
+
+namespace Ryujinx.HLE.HOS.Services.Nv.NvDrvServices.NvMap
+{
+    class NvMapIdDictionary
+    {
+        private readonly ConcurrentDictionary<int, NvMapHandle> _nvmapHandles;
+        private int _id;
+
+        public ICollection<NvMapHandle> Values => _nvmapHandles.Values;
+
+        public NvMapIdDictionary()
+        {
+            _nvmapHandles = new ConcurrentDictionary<int, NvMapHandle>();
+        }
+
+        public int Add(NvMapHandle handle)
+        {
+            int id = Interlocked.Add(ref _id, 4);
+
+            if (id != 0 && _nvmapHandles.TryAdd(id, handle))
+            {
+                return id;
+            }
+
+            throw new InvalidOperationException("NvMap ID overflow.");
+        }
+
+        public NvMapHandle Get(int id)
+        {
+            if (_nvmapHandles.TryGetValue(id, out NvMapHandle handle))
+            {
+                return handle;
+            }
+
+            return null;
+        }
+
+        public NvMapHandle Delete(int id)
+        {
+            if (_nvmapHandles.TryRemove(id, out NvMapHandle handle))
+            {
+                return handle;
+            }
+
+            return null;
+        }
+
+        public ICollection<NvMapHandle> Clear()
+        {
+            ICollection<NvMapHandle> values = _nvmapHandles.Values;
+
+            _nvmapHandles.Clear();
+
+            return values;
+        }
+    }
+}