1 files changed, 105 insertions, 4 deletions
diff --git a/src/Ryujinx.Graphics.Shader/Translation/HelperFunctionManager.cs b/src/Ryujinx.Graphics.Shader/Translation/HelperFunctionManager.cs
index 6958b86f..51a39682 100644
--- a/src/Ryujinx.Graphics.Shader/Translation/HelperFunctionManager.cs
+++ b/src/Ryujinx.Graphics.Shader/Translation/HelperFunctionManager.cs
@@ -9,13 +9,13 @@ namespace Ryujinx.Graphics.Shader.Translation
     class HelperFunctionManager
     {
         private readonly List<Function> _functionList;
-        private readonly Dictionary<HelperFunctionName, int> _functionIds;
+        private readonly Dictionary<int, int> _functionIds;
         private readonly ShaderStage _stage;
 
         public HelperFunctionManager(List<Function> functionList, ShaderStage stage)
         {
             _functionList = functionList;
-            _functionIds = new Dictionary<HelperFunctionName, int>();
+            _functionIds = new Dictionary<int, int>();
             _stage = stage;
         }
 
@@ -29,14 +29,30 @@ namespace Ryujinx.Graphics.Shader.Translation
 
         public int GetOrCreateFunctionId(HelperFunctionName functionName)
         {
-            if (_functionIds.TryGetValue(functionName, out int functionId))
+            if (_functionIds.TryGetValue((int)functionName, out int functionId))
             {
                 return functionId;
             }
 
             Function function = GenerateFunction(functionName);
             functionId = AddFunction(function);
-            _functionIds.Add(functionName, functionId);
+            _functionIds.Add((int)functionName, functionId);
+
+            return functionId;
+        }
+
+        public int GetOrCreateFunctionId(HelperFunctionName functionName, int id)
+        {
+            int key = (int)functionName | (id << 16);
+
+            if (_functionIds.TryGetValue(key, out int functionId))
+            {
+                return functionId;
+            }
+
+            Function function = GenerateFunction(functionName, id);
+            functionId = AddFunction(function);
+            _functionIds.Add(key, functionId);
 
             return functionId;
         }
@@ -140,6 +156,67 @@ namespace Ryujinx.Graphics.Shader.Translation
             return new Function(ControlFlowGraph.Create(context.GetOperations()).Blocks, "ConvertFloatToDouble", false, 1, 2);
         }
 
+        private static Function GenerateFunction(HelperFunctionName functionName, int id)
+        {
+            return functionName switch
+            {
+                HelperFunctionName.SharedAtomicMaxS32 => GenerateSharedAtomicSigned(id, isMin: false),
+                HelperFunctionName.SharedAtomicMinS32 => GenerateSharedAtomicSigned(id, isMin: true),
+                HelperFunctionName.SharedStore8 => GenerateSharedStore8(id),
+                HelperFunctionName.SharedStore16 => GenerateSharedStore16(id),
+                _ => throw new ArgumentException($"Invalid function name {functionName}")
+            };
+        }
+
+        private static Function GenerateSharedAtomicSigned(int id, bool isMin)
+        {
+            EmitterContext context = new EmitterContext();
+
+            Operand wordOffset = Argument(0);
+            Operand value = Argument(1);
+
+            Operand result = GenerateSharedAtomicCasLoop(context, wordOffset, id, (memValue) =>
+            {
+                return isMin
+                    ? context.IMinimumS32(memValue, value)
+                    : context.IMaximumS32(memValue, value);
+            });
+
+            context.Return(result);
+
+            return new Function(ControlFlowGraph.Create(context.GetOperations()).Blocks, $"SharedAtomic{(isMin ? "Min" : "Max")}_{id}", true, 2, 0);
+        }
+
+        private static Function GenerateSharedStore8(int id)
+        {
+            return GenerateSharedStore(id, 8);
+        }
+
+        private static Function GenerateSharedStore16(int id)
+        {
+            return GenerateSharedStore(id, 16);
+        }
+
+        private static Function GenerateSharedStore(int id, int bitSize)
+        {
+            EmitterContext context = new EmitterContext();
+
+            Operand offset = Argument(0);
+            Operand value = Argument(1);
+
+            Operand wordOffset = context.ShiftRightU32(offset, Const(2));
+            Operand bitOffset = GetBitOffset(context, offset);
+
+            GenerateSharedAtomicCasLoop(context, wordOffset, id, (memValue) =>
+            {
+                return context.BitfieldInsert(memValue, value, bitOffset, Const(bitSize));
+            });
+
+            context.Return();
+
+            return new Function(ControlFlowGraph.Create(context.GetOperations()).Blocks, $"SharedStore{bitSize}_{id}", false, 2, 0);
+        }
+
         private Function GenerateTexelFetchScaleFunction()
         {
             EmitterContext context = new EmitterContext();
@@ -226,5 +303,29 @@ namespace Ryujinx.Graphics.Shader.Translation
                     return context.IAdd(Const(1), index);
             }
         }
+
+        public static Operand GetBitOffset(EmitterContext context, Operand offset)
+        {
+            return context.ShiftLeft(context.BitwiseAnd(offset, Const(3)), Const(3));
+        }
+
+        private static Operand GenerateSharedAtomicCasLoop(EmitterContext context, Operand wordOffset, int id, Func<Operand, Operand> opCallback)
+        {
+            Operand lblLoopHead = Label();
+
+            context.MarkLabel(lblLoopHead);
+
+            Operand oldValue = context.Load(StorageKind.SharedMemory, id, wordOffset);
+            Operand newValue = opCallback(oldValue);
+
+            Operand casResult = context.AtomicCompareAndSwap(StorageKind.SharedMemory, id, wordOffset, oldValue, newValue);
+
+            Operand casFail = context.ICompareNotEqual(casResult, oldValue);
+
+            context.BranchIfTrue(lblLoopHead, casFail);
+
+            return oldValue;
+        }
+
     }
 }
 \ No newline at end of file