Support compute uniform buffers emulated with global memory (#924)

3 files changed, 92 insertions, 13 deletions

diff --git a/Ryujinx.Graphics.Shader/Translation/GlobalMemory.cs b/Ryujinx.Graphics.Shader/Translation/GlobalMemory.cs
index a442357d..75bd9ddf 100644
--- a/Ryujinx.Graphics.Shader/Translation/GlobalMemory.cs
+++ b/Ryujinx.Graphics.Shader/Translation/GlobalMemory.cs
@@ -11,6 +11,11 @@ namespace Ryujinx.Graphics.Shader.Translation
 
         public const int StorageDescsSize = StorageDescSize * StorageMaxCount;
 
+        public const int UbeBaseOffset = 0x98; // In words.
+        public const int UbeMaxCount   = 9;
+        public const int UbeDescsSize  = StorageDescSize * UbeMaxCount;
+        public const int UbeFirstCbuf  = 8;
+
         public static bool UsesGlobalMemory(Instruction inst)
         {
             return (inst.IsAtomic() && IsGlobalMr(inst)) ||
@@ -30,17 +35,16 @@ namespace Ryujinx.Graphics.Shader.Translation
 
         public static int GetStorageBaseCbOffset(ShaderStage stage)
         {
-            switch (stage)
+            return stage switch
             {
-                case ShaderStage.Compute:                return StorageDescsBaseOffset + 2 * StorageDescsSize;
-                case ShaderStage.Vertex:                 return StorageDescsBaseOffset;
-                case ShaderStage.TessellationControl:    return StorageDescsBaseOffset + 1 * StorageDescsSize;
-                case ShaderStage.TessellationEvaluation: return StorageDescsBaseOffset + 2 * StorageDescsSize;
-                case ShaderStage.Geometry:               return StorageDescsBaseOffset + 3 * StorageDescsSize;
-                case ShaderStage.Fragment:               return StorageDescsBaseOffset + 4 * StorageDescsSize;
-            }
-
-            return 0;
+                ShaderStage.Compute                => StorageDescsBaseOffset + 2 * StorageDescsSize,
+                ShaderStage.Vertex                 => StorageDescsBaseOffset,
+                ShaderStage.TessellationControl    => StorageDescsBaseOffset + 1 * StorageDescsSize,
+                ShaderStage.TessellationEvaluation => StorageDescsBaseOffset + 2 * StorageDescsSize,
+                ShaderStage.Geometry               => StorageDescsBaseOffset + 3 * StorageDescsSize,
+                ShaderStage.Fragment               => StorageDescsBaseOffset + 4 * StorageDescsSize,
+                _ => 0
+            };
         }
     }
 }
\ No newline at end of file
diff --git a/Ryujinx.Graphics.Shader/Translation/Lowering.cs b/Ryujinx.Graphics.Shader/Translation/Lowering.cs
index 1ee21e0a..99aea26e 100644
--- a/Ryujinx.Graphics.Shader/Translation/Lowering.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Lowering.cs
@@ -81,7 +81,7 @@ namespace Ryujinx.Graphics.Shader.Translation
 
             Operand alignMask = Const(-config.QueryInfo(QueryInfoName.StorageBufferOffsetAlignment));
 
-            Operand baseAddrTrunc = PrependOperation(Instruction.BitwiseAnd,    sbBaseAddrLow, Const(-64));
+            Operand baseAddrTrunc = PrependOperation(Instruction.BitwiseAnd,    sbBaseAddrLow, alignMask);
             Operand byteOffset    = PrependOperation(Instruction.Subtract,      addrLow, baseAddrTrunc);
             Operand wordOffset    = PrependOperation(Instruction.ShiftRightU32, byteOffset, Const(2));
 
diff --git a/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs b/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs
index 8efd2c52..7988ef6c 100644
--- a/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs
@@ -31,8 +31,27 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
 
                         if (storageIndex >= 0)
                         {
+                            // Storage buffers are implemented using global memory access.
+                            // If we know from where the base address of the access is loaded,
+                            // we can guess which storage buffer it is accessing.
+                            // We can then replace the global memory access with a storage
+                            // buffer access.
                             node = ReplaceGlobalWithStorage(node, config, storageIndex);
                         }
+                        else if (config.Stage == ShaderStage.Compute && operation.Inst == Instruction.LoadGlobal)
+                        {
+                            // Here we effectively try to replace a LDG instruction with LDC.
+                            // The hardware only supports a limited amount of constant buffers
+                            // so NVN "emulates" more constant buffers using global memory access.
+                            // Here we try to replace the global access back to a constant buffer
+                            // load.
+                            storageIndex = SearchForStorageBase(asgOperation, UbeBaseOffset, UbeBaseOffset + UbeDescsSize);
+
+                            if (storageIndex >= 0)
+                            {
+                                node = ReplaceLdgWithLdc(node, config, storageIndex);
+                            }
+                        }
                     }
                 }
             }
@@ -42,8 +61,6 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
         {
             Operation operation = (Operation)node.Value;
 
-            Operation storageOp;
-
             Operand GetStorageOffset()
             {
                 Operand addrLow = operation.GetSource(0);
@@ -80,6 +97,8 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
                 sources[index] = operation.GetSource(index);
             }
 
+            Operation storageOp;
+
             if (operation.Inst.IsAtomic())
             {
                 Instruction inst = (operation.Inst & ~Instruction.MrMask) | Instruction.MrStorage;
@@ -109,6 +128,62 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
             return node;
         }
 
+        private static LinkedListNode<INode> ReplaceLdgWithLdc(LinkedListNode<INode> node, ShaderConfig config, int storageIndex)
+        {
+            Operation operation = (Operation)node.Value;
+
+            Operand GetCbufOffset()
+            {
+                Operand addrLow = operation.GetSource(0);
+
+                Operand baseAddrLow = Cbuf(0, UbeBaseOffset + storageIndex * StorageDescSize);
+
+                Operand baseAddrTrunc = Local();
+
+                Operand alignMask = Const(-config.QueryInfo(QueryInfoName.StorageBufferOffsetAlignment));
+
+                Operation andOp = new Operation(Instruction.BitwiseAnd, baseAddrTrunc, baseAddrLow, alignMask);
+
+                node.List.AddBefore(node, andOp);
+
+                Operand byteOffset = Local();
+                Operand wordOffset = Local();
+
+                Operation subOp = new Operation(Instruction.Subtract,      byteOffset, addrLow, baseAddrTrunc);
+                Operation shrOp = new Operation(Instruction.ShiftRightU32, wordOffset, byteOffset, Const(2));
+
+                node.List.AddBefore(node, subOp);
+                node.List.AddBefore(node, shrOp);
+
+                return wordOffset;
+            }
+
+            Operand[] sources = new Operand[operation.SourcesCount];
+
+            sources[0] = Const(UbeFirstCbuf + storageIndex);
+            sources[1] = GetCbufOffset();
+
+            for (int index = 2; index < operation.SourcesCount; index++)
+            {
+                sources[index] = operation.GetSource(index);
+            }
+
+            Operation ldcOp = new Operation(Instruction.LoadConstant, operation.Dest, sources);
+
+            for (int index = 0; index < operation.SourcesCount; index++)
+            {
+                operation.SetSource(index, null);
+            }
+
+            LinkedListNode<INode> oldNode = node;
+
+            node = node.List.AddBefore(node, ldcOp);
+
+            node.List.Remove(oldNode);
+
+            return node;
+        }
+
         private static int SearchForStorageBase(Operation operation, int sbStart, int sbEnd)
         {
             Queue<Operation> assignments = new Queue<Operation>();