From bbb24d8c7e6ccc61ceda1f4223fd68feee3b2f20 Mon Sep 17 00:00:00 2001
From: gdkchan <gab.dark.100@gmail.com>
Date: Mon, 5 Dec 2022 16:11:32 -0300
Subject: [PATCH] Restrict shader storage buffer search when match fails
 (#4011)

* Restrict storage buffer search when match fails

* Shader cache version bump
---
 .../Shader/DiskCache/DiskCacheHostStorage.cs    |  2 +-
 .../Optimizations/GlobalToStorage.cs            | 17 +++++++++++++++--
 .../Translation/Optimizations/Optimizer.cs      |  6 +++++-
 Ryujinx.Graphics.Shader/Translation/Rewriter.cs |  9 ++++++++-
 .../Translation/ShaderConfig.cs                 |  9 +++++++++
 5 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/Ryujinx.Graphics.Gpu/Shader/DiskCache/DiskCacheHostStorage.cs b/Ryujinx.Graphics.Gpu/Shader/DiskCache/DiskCacheHostStorage.cs
index 5c1b638b73..d2eb5ccffb 100644
--- a/Ryujinx.Graphics.Gpu/Shader/DiskCache/DiskCacheHostStorage.cs
+++ b/Ryujinx.Graphics.Gpu/Shader/DiskCache/DiskCacheHostStorage.cs
@@ -22,7 +22,7 @@ namespace Ryujinx.Graphics.Gpu.Shader.DiskCache
         private const ushort FileFormatVersionMajor = 1;
         private const ushort FileFormatVersionMinor = 2;
         private const uint FileFormatVersionPacked = ((uint)FileFormatVersionMajor << 16) | FileFormatVersionMinor;
-        private const uint CodeGenVersion = 3957;
+        private const uint CodeGenVersion = 4011;
 
         private const string SharedTocFileName = "shared.toc";
         private const string SharedDataFileName = "shared.data";
diff --git a/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs b/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs
index 7aabcc9e6d..ec8fca1da8 100644
--- a/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs
@@ -8,14 +8,25 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
 {
     static class GlobalToStorage
     {
-        public static void RunPass(BasicBlock block, ShaderConfig config)
+        public static void RunPass(BasicBlock block, ShaderConfig config, ref int sbUseMask)
         {
             int sbStart = GetStorageBaseCbOffset(config.Stage);
-
             int sbEnd = sbStart + StorageDescsSize;
 
             for (LinkedListNode<INode> node = block.Operations.First; node != null; node = node.Next)
             {
+                for (int index = 0; index < node.Value.SourcesCount; index++)
+                {
+                    Operand src = node.Value.GetSource(index);
+
+                    int storageIndex = GetStorageIndex(src, sbStart, sbEnd);
+
+                    if (storageIndex >= 0)
+                    {
+                        sbUseMask |= 1 << storageIndex;
+                    }
+                }
+
                 if (!(node.Value is Operation operation))
                 {
                     continue;
@@ -52,6 +63,8 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
                     }
                 }
             }
+
+            config.SetAccessibleStorageBuffersMask(sbUseMask);
         }
 
         private static LinkedListNode<INode> ReplaceGlobalWithStorage(BasicBlock block, LinkedListNode<INode> node, ShaderConfig config, int storageIndex)
diff --git a/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs b/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs
index 47963eacc5..a1a2054c01 100644
--- a/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs
@@ -11,14 +11,18 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
         {
             RunOptimizationPasses(blocks);
 
+            int sbUseMask = 0;
+
             // Those passes are looking for specific patterns and only needs to run once.
             for (int blkIndex = 0; blkIndex < blocks.Length; blkIndex++)
             {
-                GlobalToStorage.RunPass(blocks[blkIndex], config);
+                GlobalToStorage.RunPass(blocks[blkIndex], config, ref sbUseMask);
                 BindlessToIndexed.RunPass(blocks[blkIndex], config);
                 BindlessElimination.RunPass(blocks[blkIndex], config);
             }
 
+            config.SetAccessibleStorageBuffersMask(sbUseMask);
+
             // Run optimizations one last time to remove any code that is now optimizable after above passes.
             RunOptimizationPasses(blocks);
         }
diff --git a/Ryujinx.Graphics.Shader/Translation/Rewriter.cs b/Ryujinx.Graphics.Shader/Translation/Rewriter.cs
index bdd9b791ba..c8dab9de3f 100644
--- a/Ryujinx.Graphics.Shader/Translation/Rewriter.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Rewriter.cs
@@ -2,6 +2,7 @@ using Ryujinx.Graphics.Shader.IntermediateRepresentation;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
+using System.Numerics;
 
 using static Ryujinx.Graphics.Shader.IntermediateRepresentation.OperandHelper;
 using static Ryujinx.Graphics.Shader.Translation.GlobalMemory;
@@ -88,8 +89,14 @@ namespace Ryujinx.Graphics.Shader.Translation
             Operand sbBaseAddrLow = Const(0);
             Operand sbSlot        = Const(0);
 
-            for (int slot = 0; slot < StorageMaxCount; slot++)
+            int sbUseMask = config.AccessibleStorageBuffersMask;
+
+            while (sbUseMask != 0)
             {
+                int slot = BitOperations.TrailingZeroCount(sbUseMask);
+
+                sbUseMask &= ~(1 << slot);
+
                 config.SetUsedStorageBuffer(slot, isWrite);
 
                 int cbOffset = GetStorageCbOffset(config.Stage, slot);
diff --git a/Ryujinx.Graphics.Shader/Translation/ShaderConfig.cs b/Ryujinx.Graphics.Shader/Translation/ShaderConfig.cs
index 12cd4cd185..85b56b51f9 100644
--- a/Ryujinx.Graphics.Shader/Translation/ShaderConfig.cs
+++ b/Ryujinx.Graphics.Shader/Translation/ShaderConfig.cs
@@ -65,6 +65,8 @@ namespace Ryujinx.Graphics.Shader.Translation
         public UInt128 NextInputAttributesComponents { get; private set; }
         public UInt128 ThisInputAttributesComponents { get; private set; }
 
+        public int AccessibleStorageBuffersMask { get; private set; }
+
         private int _usedConstantBuffers;
         private int _usedStorageBuffers;
         private int _usedStorageBuffersWrite;
@@ -98,6 +100,8 @@ namespace Ryujinx.Graphics.Shader.Translation
             GpuAccessor = gpuAccessor;
             Options     = options;
 
+            AccessibleStorageBuffersMask = (1 << GlobalMemory.StorageMaxCount) - 1;
+
             UsedInputAttributesPerPatch  = new HashSet<int>();
             UsedOutputAttributesPerPatch = new HashSet<int>();
 
@@ -400,6 +404,11 @@ namespace Ryujinx.Graphics.Shader.Translation
             UsedFeatures |= flags;
         }
 
+        public void SetAccessibleStorageBuffersMask(int mask)
+        {
+            AccessibleStorageBuffersMask = mask;
+        }
+
         public void SetUsedConstantBuffer(int slot)
         {
             _usedConstantBuffers |= 1 << slot;