Improve shader BRX instruction code generation (#3759)

* Improve shader BRX instruction code generation

* Shader cache version bump, add some comments and asserts

Ryujinx.Graphics.Gpu/Shader/DiskCache/DiskCacheHostStorage.cs

@@ -22,7 +22,7 @@ namespace Ryujinx.Graphics.Gpu.Shader.DiskCache
         private const ushort FileFormatVersionMajor = 1;
         private const ushort FileFormatVersionMinor = 2;
         private const uint FileFormatVersionPacked = ((uint)FileFormatVersionMajor << 16) | FileFormatVersionMinor;
-        private const uint CodeGenVersion = 3732;
+        private const uint CodeGenVersion = 3759;
 
         private const string SharedTocFileName = "shared.toc";
         private const string SharedDataFileName = "shared.data";

Ryujinx.Graphics.Shader/Decoders/Decoder.cs

@@ -377,6 +377,8 @@ namespace Ryujinx.Graphics.Shader.Decoders
 
                 if (lastOp.Name == InstName.Brx && block.Successors.Count == (hasNext ? 1 : 0))
                 {
+                    HashSet<ulong> visited = new HashSet<ulong>();
+
                     InstBrx opBrx = new InstBrx(lastOp.RawOpCode);
                     ulong baseOffset = lastOp.GetAbsoluteAddress();
 
@@ -392,12 +394,17 @@ namespace Ryujinx.Graphics.Shader.Decoders
                     for (int i = 0; i < cbOffsetsCount; i++)
                     {
                         uint targetOffset = config.ConstantBuffer1Read(cbBaseOffset + i * 4);
-                        Block target = getBlock(baseOffset + targetOffset);
+                        ulong targetAddress = baseOffset + targetOffset;
+
+                        if (visited.Add(targetAddress))
+                        {
+                            Block target = getBlock(targetAddress);
                             target.Predecessors.Add(block);
                             block.Successors.Add(target);
                         }
                     }
                 }
+            }
 
             return hasNewTarget;
         }

Ryujinx.Graphics.Shader/Instructions/InstEmitFlowControl.cs

@@ -41,20 +41,77 @@ namespace Ryujinx.Graphics.Shader.Instructions
 
             Operand address = context.IAdd(Register(op.SrcA, RegisterType.Gpr), Const(offset));
 
-            // Sorting the target addresses in descending order improves the code,
+            var targets = context.CurrBlock.Successors.Skip(startIndex);
-            // since it will always check the most distant targets first, then the
-            // near ones. This can be easily transformed into if/else statements.
-            var sortedTargets = context.CurrBlock.Successors.Skip(startIndex).OrderByDescending(x => x.Address);
 
-            Block lastTarget = sortedTargets.LastOrDefault();
+            bool allTargetsSinglePred = true;
+            int total = context.CurrBlock.Successors.Count - startIndex;
+            int count = 0;
 
-            foreach (Block possibleTarget in sortedTargets)
+            foreach (var target in targets.OrderBy(x => x.Address))
             {
-                Operand label = context.GetLabel(possibleTarget.Address);
+                if (++count < total && (target.Predecessors.Count > 1 || target.Address <= context.CurrBlock.Address))
-
-                if (possibleTarget != lastTarget)
                 {
-                    context.BranchIfTrue(label, context.ICompareEqual(address, Const((int)possibleTarget.Address)));
+                    allTargetsSinglePred = false;
+                    break;
+                }
+            }
+
+            if (allTargetsSinglePred)
+            {
+                // Chain blocks, each target block will check if the BRX target address
+                // matches its own address, if not, it jumps to the next target which will do the same check,
+                // until it reaches the last possible target, which executed unconditionally.
+                // We can only do this if the BRX block is the only predecessor of all target blocks.
+                // Additionally, this is not supported for blocks located before the current block,
+                // since it will be too late to insert a label, but this is something that can be improved
+                // in the future if necessary.
+
+                var sortedTargets = targets.OrderBy(x => x.Address);
+
+                Block currentTarget = null;
+                ulong firstTargetAddress = 0;
+
+                foreach (Block nextTarget in sortedTargets)
+                {
+                    if (currentTarget != null)
+                    {
+                        if (currentTarget.Address != nextTarget.Address)
+                        {
+                            context.SetBrxTarget(currentTarget.Address, address, (int)currentTarget.Address, nextTarget.Address);
+                        }
+                    }
+                    else
+                    {
+                        firstTargetAddress = nextTarget.Address;
+                    }
+
+                    currentTarget = nextTarget;
+                }
+
+                context.Branch(context.GetLabel(firstTargetAddress));
+            }
+            else
+            {
+                // Emit the branches sequentially.
+                // This generates slightly worse code, but should work for all cases.
+
+                var sortedTargets = targets.OrderByDescending(x => x.Address);
+                ulong lastTargetAddress = ulong.MaxValue;
+
+                count = 0;
+
+                foreach (Block target in sortedTargets)
+                {
+                    Operand label = context.GetLabel(target.Address);
+
+                    if (++count < total)
+                    {
+                        if (target.Address != lastTargetAddress)
+                        {
+                            context.BranchIfTrue(label, context.ICompareEqual(address, Const((int)target.Address)));
+                        }
+
+                        lastTargetAddress = target.Address;
                     }
                     else
                     {
@@ -62,6 +119,7 @@ namespace Ryujinx.Graphics.Shader.Instructions
                     }
                 }
             }
+        }
 
         public static void Cal(EmitterContext context)
         {

Ryujinx.Graphics.Shader/Translation/EmitterContext.cs

@@ -21,8 +21,33 @@ namespace Ryujinx.Graphics.Shader.Translation
 
         public int OperationsCount => _operations.Count;
 
+        private struct BrxTarget
+        {
+            public readonly Operand Selector;
+            public readonly int ExpectedValue;
+            public readonly ulong NextTargetAddress;
+
+            public BrxTarget(Operand selector, int expectedValue, ulong nextTargetAddress)
+            {
+                Selector = selector;
+                ExpectedValue = expectedValue;
+                NextTargetAddress = nextTargetAddress;
+            }
+        }
+
+        private class BlockLabel
+        {
+            public readonly Operand Label;
+            public BrxTarget BrxTarget;
+
+            public BlockLabel(Operand label)
+            {
+                Label = label;
+            }
+        }
+
         private readonly List<Operation> _operations;
-        private readonly Dictionary<ulong, Operand> _labels;
+        private readonly Dictionary<ulong, BlockLabel> _labels;
 
         public EmitterContext(DecodedProgram program, ShaderConfig config, bool isNonMain)
         {
@@ -30,7 +55,7 @@ namespace Ryujinx.Graphics.Shader.Translation
             Config = config;
             IsNonMain = isNonMain;
             _operations = new List<Operation>();
-            _labels = new Dictionary<ulong, Operand>();
+            _labels = new Dictionary<ulong, BlockLabel>();
 
             EmitStart();
         }
@@ -158,14 +183,40 @@ namespace Ryujinx.Graphics.Shader.Translation
 
         public Operand GetLabel(ulong address)
         {
-            if (!_labels.TryGetValue(address, out Operand label))
+            return EnsureBlockLabel(address).Label;
-            {
-                label = Label();
-
-                _labels.Add(address, label);
         }
 
-            return label;
+        public void SetBrxTarget(ulong address, Operand selector, int targetValue, ulong nextTargetAddress)
+        {
+            BlockLabel blockLabel = EnsureBlockLabel(address);
+            Debug.Assert(blockLabel.BrxTarget.Selector == null);
+            blockLabel.BrxTarget = new BrxTarget(selector, targetValue, nextTargetAddress);
+        }
+
+        public void EnterBlock(ulong address)
+        {
+            BlockLabel blockLabel = EnsureBlockLabel(address);
+
+            MarkLabel(blockLabel.Label);
+
+            BrxTarget brxTarget = blockLabel.BrxTarget;
+
+            if (brxTarget.Selector != null)
+            {
+                this.BranchIfFalse(GetLabel(brxTarget.NextTargetAddress), this.ICompareEqual(brxTarget.Selector, Const(brxTarget.ExpectedValue)));
+            }
+        }
+
+        private BlockLabel EnsureBlockLabel(ulong address)
+        {
+            if (!_labels.TryGetValue(address, out BlockLabel blockLabel))
+            {
+                blockLabel = new BlockLabel(Label());
+
+                _labels.Add(address, blockLabel);
+            }
+
+            return blockLabel;
         }
 
         public void PrepareForVertexReturn()

Ryujinx.Graphics.Shader/Translation/Translator.cs

@@ -162,7 +162,7 @@ namespace Ryujinx.Graphics.Shader.Translation
                 {
                     context.CurrBlock = block;
 
-                    context.MarkLabel(context.GetLabel(block.Address));
+                    context.EnterBlock(block.Address);
 
                     EmitOps(context, block);
                 }