Shader: Add fallback for LDG from "ube" buffer ranges. (#4027)

We have a conversion from LDG on the compute shader to a special constant buffer binding that's used to exceed hardware limits on compute, but it was only running if the byte offset could be identified. The fallback that checks all of the bindings at runtime only checks the storage buffers.

This PR adds checking ube ranges to the LoadGlobal fallback. This extends the changes in #4011 to only check ube entries which are accessed by the shader.

Fixes particles affected by the wind in The Legend of Zelda: Breath of the Wild. May fix other weird issues with compute shaders in some games.

Try a bunch of games and drivers to make sure they don't blow up loading constants willynilly from searchable buffers.

1 files changed, 72 insertions, 16 deletions

diff --git a/Ryujinx.Graphics.Shader/Translation/Rewriter.cs b/Ryujinx.Graphics.Shader/Translation/Rewriter.cs
index c4d2c5d9..0c3c4a57 100644
--- a/Ryujinx.Graphics.Shader/Translation/Rewriter.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Rewriter.cs
@@ -1,3 +1,4 @@
+using Ryujinx.Graphics.Shader.Decoders;
 using Ryujinx.Graphics.Shader.IntermediateRepresentation;
 using System.Collections.Generic;
 using System.Diagnostics;
@@ -89,29 +90,31 @@ namespace Ryujinx.Graphics.Shader.Translation
                 return local;
             }
 
+            Operand PrependExistingOperation(Operation operation)
+            {
+                Operand local = Local();
+
+                operation.Dest = local;
+                node.List.AddBefore(node, operation);
+
+                return local;
+            }
+
             Operand addrLow  = operation.GetSource(0);
             Operand addrHigh = operation.GetSource(1);
 
             Operand sbBaseAddrLow = Const(0);
             Operand sbSlot        = Const(0);
 
-            int sbUseMask = config.AccessibleStorageBuffersMask;
+            Operand alignMask = Const(-config.GpuAccessor.QueryHostStorageBufferOffsetAlignment());
 
-            while (sbUseMask != 0)
+            Operand BindingRangeCheck(int cbOffset, out Operand baseAddrLow)
             {
-                int slot = BitOperations.TrailingZeroCount(sbUseMask);
-
-                sbUseMask &= ~(1 << slot);
-
-                config.SetUsedStorageBuffer(slot, isWrite);
-
-                int cbOffset = GetStorageCbOffset(config.Stage, slot);
-
-                Operand baseAddrLow  = Cbuf(0, cbOffset);
+                baseAddrLow = Cbuf(0, cbOffset);
                 Operand baseAddrHigh = Cbuf(0, cbOffset + 1);
-                Operand size         = Cbuf(0, cbOffset + 2);
+                Operand size = Cbuf(0, cbOffset + 2);
 
-                Operand offset = PrependOperation(Instruction.Subtract,       addrLow, baseAddrLow);
+                Operand offset = PrependOperation(Instruction.Subtract, addrLow, baseAddrLow);
                 Operand borrow = PrependOperation(Instruction.CompareLessU32, addrLow, baseAddrLow);
 
                 Operand inRangeLow = PrependOperation(Instruction.CompareLessU32, offset, size);
@@ -120,7 +123,22 @@ namespace Ryujinx.Graphics.Shader.Translation
 
                 Operand inRangeHigh = PrependOperation(Instruction.CompareEqual, addrHighBorrowed, baseAddrHigh);
 
-                Operand inRange = PrependOperation(Instruction.BitwiseAnd, inRangeLow, inRangeHigh);
+                return PrependOperation(Instruction.BitwiseAnd, inRangeLow, inRangeHigh);
+            }
+
+            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);
+
+                Operand inRange = BindingRangeCheck(cbOffset, out Operand baseAddrLow);
 
                 sbBaseAddrLow = PrependOperation(Instruction.ConditionalSelect, inRange, baseAddrLow, sbBaseAddrLow);
                 sbSlot        = PrependOperation(Instruction.ConditionalSelect, inRange, Const(slot), sbSlot);
@@ -128,8 +146,6 @@ namespace Ryujinx.Graphics.Shader.Translation
 
             if (config.AccessibleStorageBuffersMask != 0)
             {
-                Operand alignMask = Const(-config.GpuAccessor.QueryHostStorageBufferOffsetAlignment());
-
                 Operand baseAddrTrunc = PrependOperation(Instruction.BitwiseAnd, sbBaseAddrLow, alignMask);
                 Operand byteOffset    = PrependOperation(Instruction.Subtract, addrLow, baseAddrTrunc);
 
@@ -178,6 +194,46 @@ namespace Ryujinx.Graphics.Shader.Translation
                 storageOp = new Operation(Instruction.Copy, operation.Dest, Const(0));
             }
 
+            if (operation.Inst == Instruction.LoadGlobal)
+            {
+                int cbeUseMask = config.AccessibleConstantBuffersMask;
+
+                while (cbeUseMask != 0)
+                {
+                    int slot = BitOperations.TrailingZeroCount(cbeUseMask);
+                    int cbSlot = UbeFirstCbuf + slot;
+
+                    cbeUseMask &= ~(1 << slot);
+
+                    config.SetUsedConstantBuffer(cbSlot);
+
+                    Operand previousResult = PrependExistingOperation(storageOp);
+
+                    int cbOffset = GetConstantUbeOffset(slot);
+
+                    Operand inRange = BindingRangeCheck(cbOffset, out Operand baseAddrLow);
+
+                    Operand baseAddrTruncConst = PrependOperation(Instruction.BitwiseAnd, baseAddrLow, alignMask);
+                    Operand byteOffsetConst = PrependOperation(Instruction.Subtract, addrLow, baseAddrTruncConst);
+
+                    Operand cbIndex = PrependOperation(Instruction.ShiftRightU32, byteOffsetConst, Const(2));
+
+                    Operand[] sourcesCb = new Operand[operation.SourcesCount];
+
+                    sourcesCb[0] = Const(cbSlot);
+                    sourcesCb[1] = cbIndex;
+
+                    for (int index = 2; index < operation.SourcesCount; index++)
+                    {
+                        sourcesCb[index] = operation.GetSource(index);
+                    }
+
+                    Operand ldcResult = PrependOperation(Instruction.LoadConstant, sourcesCb);
+
+                    storageOp = new Operation(Instruction.ConditionalSelect, operation.Dest, inRange, ldcResult, previousResult);
+                }
+            }
+
             for (int index = 0; index < operation.SourcesCount; index++)
             {
                 operation.SetSource(index, null);