Implement BFI, BRK, FLO, FSWZADD, PBK, SHFL and TXD shader instructions, misc. fixes

3 files changed, 140 insertions, 1 deletions

diff --git a/Ryujinx.Graphics.Shader/Translation/EmitterContextInsts.cs b/Ryujinx.Graphics.Shader/Translation/EmitterContextInsts.cs
index 7d64e7ca..58a37b52 100644
--- a/Ryujinx.Graphics.Shader/Translation/EmitterContextInsts.cs
+++ b/Ryujinx.Graphics.Shader/Translation/EmitterContextInsts.cs
@@ -6,6 +6,11 @@ namespace Ryujinx.Graphics.Shader.Translation
 {
     static class EmitterContextInsts
     {
+        public static Operand BitCount(this EmitterContext context, Operand a)
+        {
+            return context.Add(Instruction.BitCount, Local(), a);
+        }
+
         public static Operand BitfieldExtractS32(this EmitterContext context, Operand a, Operand b, Operand c)
         {
             return context.Add(Instruction.BitfieldExtractS32, Local(), a, b, c);
@@ -106,6 +111,16 @@ namespace Ryujinx.Graphics.Shader.Translation
             return context.Add(Instruction.EndPrimitive);
         }
 
+        public static Operand FindFirstSetS32(this EmitterContext context, Operand a)
+        {
+            return context.Add(Instruction.FindFirstSetS32, Local(), a);
+        }
+
+        public static Operand FindFirstSetU32(this EmitterContext context, Operand a)
+        {
+            return context.Add(Instruction.FindFirstSetU32, Local(), a);
+        }
+
         public static Operand FPAbsNeg(this EmitterContext context, Operand a, bool abs, bool neg)
         {
             return context.FPNegate(context.FPAbsolute(a, abs), neg);
@@ -256,6 +271,11 @@ namespace Ryujinx.Graphics.Shader.Translation
             return context.Add(Instruction.Truncate, Local(), a);
         }
 
+        public static Operand FPSwizzleAdd(this EmitterContext context, Operand a, Operand b, int mask)
+        {
+            return context.Add(Instruction.SwizzleAdd, Local(), a, b, Const(mask));
+        }
+
         public static Operand IAbsNeg(this EmitterContext context, Operand a, bool abs, bool neg)
         {
             return context.INegate(context.IAbsolute(a, abs), neg);
@@ -418,6 +438,26 @@ namespace Ryujinx.Graphics.Shader.Translation
             return context.Add(Instruction.ShiftRightU32, Local(), a, b);
         }
 
+        public static Operand Shuffle(this EmitterContext context, Operand a, Operand b, Operand c)
+        {
+            return context.Add(Instruction.Shuffle, Local(), a, b, c);
+        }
+
+        public static Operand ShuffleDown(this EmitterContext context, Operand a, Operand b, Operand c)
+        {
+            return context.Add(Instruction.ShuffleDown, Local(), a, b, c);
+        }
+
+        public static Operand ShuffleUp(this EmitterContext context, Operand a, Operand b, Operand c)
+        {
+            return context.Add(Instruction.ShuffleUp, Local(), a, b, c);
+        }
+
+        public static Operand ShuffleXor(this EmitterContext context, Operand a, Operand b, Operand c)
+        {
+            return context.Add(Instruction.ShuffleXor, Local(), a, b, c);
+        }
+
         public static Operand StoreGlobal(this EmitterContext context, Operand a, Operand b)
         {
             return context.Add(Instruction.StoreGlobal, null, a, b);
diff --git a/Ryujinx.Graphics.Shader/Translation/Optimizations/ConstantFolding.cs b/Ryujinx.Graphics.Shader/Translation/Optimizations/ConstantFolding.cs
index d64579b7..97852ac1 100644
--- a/Ryujinx.Graphics.Shader/Translation/Optimizations/ConstantFolding.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Optimizations/ConstantFolding.cs
@@ -21,6 +21,10 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
                     EvaluateBinary(operation, (x, y) => x + y);
                     break;
 
+                case Instruction.BitCount:
+                    EvaluateUnary(operation, (x) => BitCount(x));
+                    break;
+
                 case Instruction.BitwiseAnd:
                     EvaluateBinary(operation, (x, y) => x & y);
                     break;
@@ -208,6 +212,21 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
             return true;
         }
 
+        private static int BitCount(int value)
+        {
+            int count = 0;
+
+            for (int bit = 0; bit < 32; bit++)
+            {
+                if (value.Extract(bit))
+                {
+                    count++;
+                }
+            }
+
+            return count;
+        }
+
         private static void BitfieldExtractS32(Operation operation)
         {
             int value = GetBitfieldExtractValue(operation);
diff --git a/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs b/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs
index d5e57546..22d794a4 100644
--- a/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs
@@ -1,5 +1,6 @@
 using Ryujinx.Graphics.Shader.IntermediateRepresentation;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 
 namespace Ryujinx.Graphics.Shader.Translation.Optimizations
@@ -59,7 +60,8 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
 
                                 modified = true;
                             }
-                            else if (operation.Inst == Instruction.PackHalf2x16 && PropagatePack(operation))
+                            else if ((operation.Inst == Instruction.PackHalf2x16 && PropagatePack(operation)) ||
+                                     (operation.Inst == Instruction.ShuffleXor   && MatchDdxOrDdy(operation)))
                             {
                                 if (operation.Dest.UseOps.Count == 0)
                                 {
@@ -135,6 +137,84 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
             return modified;
         }
 
+        public static bool MatchDdxOrDdy(Operation operation)
+        {
+            // It's assumed that "operation.Inst" is ShuffleXor,
+            // that should be checked before calling this method.
+            Debug.Assert(operation.Inst == Instruction.ShuffleXor);
+
+            bool modified = false;
+
+            Operand src2 = operation.GetSource(1);
+            Operand src3 = operation.GetSource(2);
+
+            if (src2.Type != OperandType.Constant || (src2.Value != 1 && src2.Value != 2))
+            {
+                return false;
+            }
+
+            if (src3.Type != OperandType.Constant || src3.Value != 0x1c03)
+            {
+                return false;
+            }
+
+            bool isDdy = src2.Value == 2;
+            bool isDdx = !isDdy;
+
+            // We can replace any use by a FSWZADD with DDX/DDY, when
+            // the following conditions are true:
+            // - The mask should be 0b10100101 for DDY, or 0b10011001 for DDX.
+            // - The first source operand must be the shuffle output.
+            // - The second source operand must be the shuffle first source operand.
+            INode[] uses = operation.Dest.UseOps.ToArray();
+
+            foreach (INode use in uses)
+            {
+                if (!(use is Operation test))
+                {
+                    continue;
+                }
+
+                if (!(use is Operation useOp) || useOp.Inst != Instruction.SwizzleAdd)
+                {
+                    continue;
+                }
+
+                Operand fswzaddSrc1 = useOp.GetSource(0);
+                Operand fswzaddSrc2 = useOp.GetSource(1);
+                Operand fswzaddSrc3 = useOp.GetSource(2);
+
+                if (fswzaddSrc1 != operation.Dest)
+                {
+                    continue;
+                }
+
+                if (fswzaddSrc2 != operation.GetSource(0))
+                {
+                    continue;
+                }
+
+                if (fswzaddSrc3.Type != OperandType.Constant)
+                {
+                    continue;
+                }
+
+                int mask = fswzaddSrc3.Value;
+
+                if ((isDdx && mask != 0b10011001) ||
+                    (isDdy && mask != 0b10100101))
+                {
+                    continue;
+                }
+
+                useOp.TurnInto(isDdx ? Instruction.Ddx : Instruction.Ddy, fswzaddSrc2);
+
+                modified = true;
+            }
+
+            return modified;
+        }
+
         private static void RemoveNode(BasicBlock block, LinkedListNode<INode> llNode)
         {
             // Remove a node from the nodes list, and also remove itself