Fix shaders with mixed PBK and SSY addresses on the stack (#2329)

* Fix shaders with mixed PBK and SSY addresses on the stack

* Address PR feedback and nits

Ryujinx.Graphics.Gpu/Shader/ShaderCache.cs

@@ -36,7 +36,7 @@ namespace Ryujinx.Graphics.Gpu.Shader
         /// <summary>
         /// Version of the codegen (to be changed when codegen or guest format change).
         /// </summary>
-        private const ulong ShaderCodeGenVersion = 2317;
+        private const ulong ShaderCodeGenVersion = 2329;
 
         // Progress reporting helpers
         private volatile int _shaderCount;

Ryujinx.Graphics.Shader/Decoders/Decoder.cs

@@ -318,6 +318,12 @@ namespace Ryujinx.Graphics.Shader.Decoders
                     opCode is OpCodeExit;
         }
 
+        private enum MergeType
+        {
+            Brk = 0,
+            Sync = 1
+        }
+
         private struct PathBlockState
         {
             public Block Block { get; }
@@ -332,6 +338,7 @@ namespace Ryujinx.Graphics.Shader.Decoders
             private RestoreType _restoreType;
 
             private ulong _restoreValue;
+            private MergeType _restoreMergeType;
 
             public bool ReturningFromVisit => _restoreType != RestoreType.None;
 
@@ -340,6 +347,7 @@ namespace Ryujinx.Graphics.Shader.Decoders
                 Block             = block;
                 _restoreType      = RestoreType.None;
                 _restoreValue     = 0;
+                _restoreMergeType = default;
             }
 
             public PathBlockState(int oldStackSize)
@@ -347,20 +355,22 @@ namespace Ryujinx.Graphics.Shader.Decoders
                 Block             = null;
                 _restoreType      = RestoreType.PopPushOp;
                 _restoreValue     = (ulong)oldStackSize;
+                _restoreMergeType = default;
             }
 
-            public PathBlockState(ulong syncAddress)
+            public PathBlockState(ulong syncAddress, MergeType mergeType)
             {
                 Block             = null;
                 _restoreType      = RestoreType.PushBranchOp;
                 _restoreValue     = syncAddress;
+                _restoreMergeType = mergeType;
             }
 
-            public void RestoreStackState(Stack<ulong> branchStack)
+            public void RestoreStackState(Stack<(ulong, MergeType)> branchStack)
             {
                 if (_restoreType == RestoreType.PushBranchOp)
                 {
-                    branchStack.Push(_restoreValue);
+                    branchStack.Push((_restoreValue, _restoreMergeType));
                 }
                 else if (_restoreType == RestoreType.PopPushOp)
                 {
@@ -380,7 +390,7 @@ namespace Ryujinx.Graphics.Shader.Decoders
 
             HashSet<Block> visited = new HashSet<Block>();
 
-            Stack<ulong> branchStack = new Stack<ulong>();
+            Stack<(ulong, MergeType)> branchStack = new Stack<(ulong, MergeType)>();
 
             void Push(PathBlockState pbs)
             {
@@ -426,7 +436,9 @@ namespace Ryujinx.Graphics.Shader.Decoders
 
                     for (int index = pushOpIndex; index < pushOpsCount; index++)
                     {
-                        branchStack.Push(current.PushOpCodes[index].GetAbsoluteAddress());
+                        OpCodePush currentPushOp = current.PushOpCodes[index];
+                        MergeType pushMergeType = currentPushOp.Emitter == InstEmit.Ssy ? MergeType.Sync : MergeType.Brk;
+                        branchStack.Push((currentPushOp.GetAbsoluteAddress(), pushMergeType));
                     }
                 }
 
@@ -452,27 +464,51 @@ namespace Ryujinx.Graphics.Shader.Decoders
                 }
                 else if (current.GetLastOp() is OpCodeBranchPop op)
                 {
-                    ulong targetAddress = branchStack.Pop();
+                    MergeType popMergeType = op.Emitter == InstEmit.Sync ? MergeType.Sync : MergeType.Brk;
 
+                    bool found = true;
+                    ulong targetAddress = 0UL;
+                    MergeType mergeType;
+
+                    do
+                    {
                         if (branchStack.Count == 0)
                         {
-                        branchStack.Push(targetAddress);
+                            found = false;
+                            break;
+                        }
 
+                        (targetAddress, mergeType) = branchStack.Pop();
+
+                        // Push the target address (this will be used to push the address
+                        // back into the SSY/PBK stack when we return from that block),
+                        Push(new PathBlockState(targetAddress, mergeType));
+                    }
+                    while (mergeType != popMergeType);
+
+                    // Make sure we found the correct address,
+                    // the push and pop instruction types must match, so:
+                    // - BRK can only consume addresses pushed by PBK.
+                    // - SYNC can only consume addresses pushed by SSY.
+                    if (found)
+                    {
+                        if (branchStack.Count == 0)
+                        {
+                            // If the entire stack was consumed, then the current pop instruction
+                            // just consumed the address from out push instruction.
                             op.Targets.Add(pushOp, op.Targets.Count);
 
                             pushOp.PopOps.TryAdd(op, Local());
                         }
                         else
                         {
-                        // First we push the target address (this will be used to push the
+                            // Push the block itself into the work "queue" (well, it's a stack)
-                        // address back into the SSY/PBK stack when we return from that block),
-                        // then we push the block itself into the work "queue" (well, it's a stack)
                             // for processing.
-                        Push(new PathBlockState(targetAddress));
                             Push(new PathBlockState(blocks[targetAddress]));
                         }
                     }
                 }
             }
         }
+    }
 }

Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs

@@ -1,5 +1,4 @@
 using Ryujinx.Graphics.Shader.IntermediateRepresentation;
-using System;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;