2 files changed, 42 insertions, 1 deletions
diff --git a/src/Ryujinx.Graphics.Gpu/Engine/Threed/StateUpdater.cs b/src/Ryujinx.Graphics.Gpu/Engine/Threed/StateUpdater.cs
index 1c9bf1d2..87e58ead 100644
--- a/src/Ryujinx.Graphics.Gpu/Engine/Threed/StateUpdater.cs
+++ b/src/Ryujinx.Graphics.Gpu/Engine/Threed/StateUpdater.cs
@@ -23,6 +23,8 @@ namespace Ryujinx.Graphics.Gpu.Engine.Threed
         public const int PrimitiveRestartStateIndex = 12;
         public const int RenderTargetStateIndex = 27;
 
+        private const ulong MaxUnknownStorageSize = 0x100000;
+
         private readonly GpuContext _context;
         private readonly GpuChannel _channel;
         private readonly DeviceStateWithShadow<ThreedClassState> _state;
@@ -356,7 +358,19 @@ namespace Ryujinx.Graphics.Gpu.Engine.Threed
 
                     SbDescriptor sbDescriptor = _channel.MemoryManager.Physical.Read<SbDescriptor>(sbDescAddress);
 
-                    _channel.BufferManager.SetGraphicsStorageBuffer(stage, sb.Slot, sbDescriptor.PackAddress(), (uint)sbDescriptor.Size, sb.Flags);
+                    uint size;
+                    if (sb.SbCbSlot == 0)
+                    {
+                        // Only trust the SbDescriptor size if it comes from slot 0.
+                        size = (uint)sbDescriptor.Size;
+                    }
+                    else
+                    {
+                        // TODO: Use full mapped size and somehow speed up buffer sync.
+                        size = (uint)_channel.MemoryManager.GetMappedSize(sbDescriptor.PackAddress(), MaxUnknownStorageSize);
+                    }
+
+                    _channel.BufferManager.SetGraphicsStorageBuffer(stage, sb.Slot, sbDescriptor.PackAddress(), size, sb.Flags);
                 }
             }
         }
diff --git a/src/Ryujinx.Graphics.Gpu/Memory/MemoryManager.cs b/src/Ryujinx.Graphics.Gpu/Memory/MemoryManager.cs
index 0d4a41f0..c7a138c9 100644
--- a/src/Ryujinx.Graphics.Gpu/Memory/MemoryManager.cs
+++ b/src/Ryujinx.Graphics.Gpu/Memory/MemoryManager.cs
@@ -638,6 +638,33 @@ namespace Ryujinx.Graphics.Gpu.Memory
         }
 
         /// <summary>
+        /// Translates a GPU virtual address and returns the number of bytes that are mapped after it.
+        /// </summary>
+        /// <param name="va">GPU virtual address to be translated</param>
+        /// <param name="maxSize">Maximum size in bytes to scan</param>
+        /// <returns>Number of bytes, 0 if unmapped</returns>
+        public ulong GetMappedSize(ulong va, ulong maxSize)
+        {
+            if (!ValidateAddress(va))
+            {
+                return 0;
+            }
+
+            ulong startVa = va;
+            ulong endVa = va + maxSize;
+
+            ulong pte = GetPte(va);
+
+            while (pte != PteUnmapped && va < endVa)
+            {
+                va += PageSize - (va & PageMask);
+                pte = GetPte(va);
+            }
+
+            return Math.Min(maxSize, va - startVa);
+        }
+
+        /// <summary>
         /// Gets the kind of a given memory page.
         /// This might indicate the type of resource that can be allocated on the page, and also texture tiling.
         /// </summary>