Implement shader storage buffer operations using new Load/Store instructions (#4993)

* Implement storage buffer operations using new Load/Store instruction

* Extend GenerateMultiTargetStorageOp to also match access with constant offset, and log and comments

* Remove now unused code

* Catch more complex cases of global memory usage

* Shader cache version bump

* Extend global access elimination to work with more shared memory cases

* Change alignment requirement from 16 bytes to 8 bytes, handle cases where we need more than 16 storage buffers

* Tweak preferencing to catch more cases

* Enable CB0 elimination even when host storage buffer alignment is > 16 (for Intel)

* Fix storage buffer bindings

* Simplify some code

* Shader cache version bump

* Fix typo

* Extend global memory elimination to handle shared memory with multiple possible offsets and local memory

3 files changed, 40 insertions, 109 deletions

diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
index 24ea66d0..01d8a6e7 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
@@ -68,33 +68,45 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
 
                 string args = string.Empty;
 
-                for (int argIndex = 0; argIndex < arity; argIndex++)
+                if (atomic && operation.StorageKind == StorageKind.StorageBuffer)
                 {
+                    args = GenerateLoadOrStore(context, operation, isStore: false);
+
+                    AggregateType dstType = operation.Inst == Instruction.AtomicMaxS32 || operation.Inst == Instruction.AtomicMinS32
+                        ? AggregateType.S32
+                        : AggregateType.U32;
+
+                    for (int argIndex = operation.SourcesCount - arity + 2; argIndex < operation.SourcesCount; argIndex++)
+                    {
+                        args += ", " + GetSoureExpr(context, operation.GetSource(argIndex), dstType);
+                    }
+                }
+                else if (atomic && operation.StorageKind == StorageKind.SharedMemory)
+                {
+                    args = LoadShared(context, operation);
+
                     // For shared memory access, the second argument is unused and should be ignored.
                     // It is there to make both storage and shared access have the same number of arguments.
                     // For storage, both inputs are consumed when the argument index is 0, so we should skip it here.
-                    if (argIndex == 1 && (atomic || operation.StorageKind == StorageKind.SharedMemory))
-                    {
-                        continue;
-                    }
 
-                    if (argIndex != 0)
+                    for (int argIndex = 2; argIndex < arity; argIndex++)
                     {
                         args += ", ";
-                    }
 
-                    if (argIndex == 0 && atomic)
-                    {
-                        switch (operation.StorageKind)
-                        {
-                            case StorageKind.SharedMemory: args += LoadShared(context, operation); break;
-                            case StorageKind.StorageBuffer: args += LoadStorage(context, operation); break;
+                        AggregateType dstType = GetSrcVarType(inst, argIndex);
 
-                            default: throw new InvalidOperationException($"Invalid storage kind \"{operation.StorageKind}\".");
-                        }
+                        args += GetSoureExpr(context, operation.GetSource(argIndex), dstType);
                     }
-                    else
+                }
+                else
+                {
+                    for (int argIndex = 0; argIndex < arity; argIndex++)
                     {
+                        if (argIndex != 0)
+                        {
+                            args += ", ";
+                        }
+
                         AggregateType dstType = GetSrcVarType(inst, argIndex);
 
                         args += GetSoureExpr(context, operation.GetSource(argIndex), dstType);
@@ -173,9 +185,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                     case Instruction.LoadShared:
                         return LoadShared(context, operation);
 
-                    case Instruction.LoadStorage:
-                        return LoadStorage(context, operation);
-
                     case Instruction.Lod:
                         return Lod(context, operation);
 
@@ -203,15 +212,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                     case Instruction.StoreShared8:
                         return StoreShared8(context, operation);
 
-                    case Instruction.StoreStorage:
-                        return StoreStorage(context, operation);
-
-                    case Instruction.StoreStorage16:
-                        return StoreStorage16(context, operation);
-
-                    case Instruction.StoreStorage8:
-                        return StoreStorage8(context, operation);
-
                     case Instruction.TextureSample:
                         return TextureSample(context, operation);
 
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
index 6cf36a2a..f42d9898 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
@@ -85,7 +85,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             Add(Instruction.Load,                     InstType.Special);
             Add(Instruction.LoadLocal,                InstType.Special);
             Add(Instruction.LoadShared,               InstType.Special);
-            Add(Instruction.LoadStorage,              InstType.Special);
             Add(Instruction.Lod,                      InstType.Special);
             Add(Instruction.LogarithmB2,              InstType.CallUnary,      "log2");
             Add(Instruction.LogicalAnd,               InstType.OpBinaryCom,    "&&",              9);
@@ -123,9 +122,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             Add(Instruction.StoreShared,              InstType.Special);
             Add(Instruction.StoreShared16,            InstType.Special);
             Add(Instruction.StoreShared8,             InstType.Special);
-            Add(Instruction.StoreStorage,             InstType.Special);
-            Add(Instruction.StoreStorage16,           InstType.Special);
-            Add(Instruction.StoreStorage8,            InstType.Special);
             Add(Instruction.Subtract,                 InstType.OpBinary,       "-",               2);
             Add(Instruction.SwizzleAdd,               InstType.CallTernary,    HelperFunctionNames.SwizzleAdd);
             Add(Instruction.TextureSample,            InstType.Special);
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
index dfc8197b..c8084d9d 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
@@ -210,17 +210,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             return $"{arrayName}[{offsetExpr}]";
         }
 
-        public static string LoadStorage(CodeGenContext context, AstOperation operation)
-        {
-            IAstNode src1 = operation.GetSource(0);
-            IAstNode src2 = operation.GetSource(1);
-
-            string indexExpr  = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
-            string offsetExpr = GetSoureExpr(context, src2, GetSrcVarType(operation.Inst, 1));
-
-            return GetStorageBufferAccessor(indexExpr, offsetExpr, context.Config.Stage);
-        }
-
         public static string Lod(CodeGenContext context, AstOperation operation)
         {
             AstTextureOperation texOp = (AstTextureOperation)operation;
@@ -326,60 +315,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             return $"{HelperFunctionNames.StoreShared8}({offsetExpr}, {src})";
         }
 
-        public static string StoreStorage(CodeGenContext context, AstOperation operation)
-        {
-            IAstNode src1 = operation.GetSource(0);
-            IAstNode src2 = operation.GetSource(1);
-            IAstNode src3 = operation.GetSource(2);
-
-            string indexExpr  = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
-            string offsetExpr = GetSoureExpr(context, src2, GetSrcVarType(operation.Inst, 1));
-
-            AggregateType srcType = OperandManager.GetNodeDestType(context, src3);
-
-            string src = TypeConversion.ReinterpretCast(context, src3, srcType, AggregateType.U32);
-
-            string sb = GetStorageBufferAccessor(indexExpr, offsetExpr, context.Config.Stage);
-
-            return $"{sb} = {src}";
-        }
-
-        public static string StoreStorage16(CodeGenContext context, AstOperation operation)
-        {
-            IAstNode src1 = operation.GetSource(0);
-            IAstNode src2 = operation.GetSource(1);
-            IAstNode src3 = operation.GetSource(2);
-
-            string indexExpr  = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
-            string offsetExpr = GetSoureExpr(context, src2, GetSrcVarType(operation.Inst, 1));
-
-            AggregateType srcType = OperandManager.GetNodeDestType(context, src3);
-
-            string src = TypeConversion.ReinterpretCast(context, src3, srcType, AggregateType.U32);
-
-            string sb = GetStorageBufferAccessor(indexExpr, offsetExpr, context.Config.Stage);
-
-            return $"{HelperFunctionNames.StoreStorage16}({indexExpr}, {offsetExpr}, {src})";
-        }
-
-        public static string StoreStorage8(CodeGenContext context, AstOperation operation)
-        {
-            IAstNode src1 = operation.GetSource(0);
-            IAstNode src2 = operation.GetSource(1);
-            IAstNode src3 = operation.GetSource(2);
-
-            string indexExpr  = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
-            string offsetExpr = GetSoureExpr(context, src2, GetSrcVarType(operation.Inst, 1));
-
-            AggregateType srcType = OperandManager.GetNodeDestType(context, src3);
-
-            string src = TypeConversion.ReinterpretCast(context, src3, srcType, AggregateType.U32);
-
-            string sb = GetStorageBufferAccessor(indexExpr, offsetExpr, context.Config.Stage);
-
-            return $"{HelperFunctionNames.StoreStorage8}({indexExpr}, {offsetExpr}, {src})";
-        }
-
         public static string TextureSample(CodeGenContext context, AstOperation operation)
         {
             AstTextureOperation texOp = (AstTextureOperation)operation;
@@ -701,25 +636,34 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             }
         }
 
-        private static string GenerateLoadOrStore(CodeGenContext context, AstOperation operation, bool isStore)
+        public static string GenerateLoadOrStore(CodeGenContext context, AstOperation operation, bool isStore)
         {
             StorageKind storageKind = operation.StorageKind;
 
             string varName;
             AggregateType varType;
             int srcIndex = 0;
-            int inputsCount = isStore ? operation.SourcesCount - 1 : operation.SourcesCount;
+            bool isStoreOrAtomic = operation.Inst == Instruction.Store || operation.Inst.IsAtomic();
+            int inputsCount = isStoreOrAtomic ? operation.SourcesCount - 1 : operation.SourcesCount;
+
+            if (operation.Inst == Instruction.AtomicCompareAndSwap)
+            {
+                inputsCount--;
+            }
 
             switch (storageKind)
             {
                 case StorageKind.ConstantBuffer:
+                case StorageKind.StorageBuffer:
                     if (!(operation.GetSource(srcIndex++) is AstOperand bindingIndex) || bindingIndex.Type != OperandType.Constant)
                     {
                         throw new InvalidOperationException($"First input of {operation.Inst} with {storageKind} storage must be a constant operand.");
                     }
 
                     int binding = bindingIndex.Value;
-                    BufferDefinition buffer = context.Config.Properties.ConstantBuffers[binding];
+                    BufferDefinition buffer = storageKind == StorageKind.ConstantBuffer
+                        ? context.Config.Properties.ConstantBuffers[binding]
+                        : context.Config.Properties.StorageBuffers[binding];
 
                     if (!(operation.GetSource(srcIndex++) is AstOperand fieldIndex) || fieldIndex.Type != OperandType.Constant)
                     {
@@ -825,15 +769,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             return varName;
         }
 
-        private static string GetStorageBufferAccessor(string slotExpr, string offsetExpr, ShaderStage stage)
-        {
-            string sbName = OperandManager.GetShaderStagePrefix(stage);
-
-            sbName += "_" + DefaultNames.StorageNamePrefix;
-
-            return $"{sbName}[{slotExpr}].{DefaultNames.DataName}[{offsetExpr}]";
-        }
-
         private static string GetMask(int index)
         {
             return $".{"rgba".AsSpan(index, 1)}";