From 3ab5c23f492183ae6f5cf8f62c4239bf181d2630 Mon Sep 17 00:00:00 2001
From: gdk <gab.dark.100@gmail.com>
Date: Sat, 2 Nov 2019 23:07:21 -0300
Subject: [PATCH] Add partial support for array of samplers, and add pass to
 identify them from bindless texture accesses

---
 Ryujinx.Graphics.Gpu/Engine/Methods.cs        |  2 +-
 .../CodeGen/Glsl/Declarations.cs              | 45 ++++++++--
 .../Glsl/Instructions/InstGenMemory.cs        | 65 +++++++++-----
 .../CodeGen/Glsl/OperandManager.cs            | 16 +++-
 .../TextureOperation.cs                       | 15 +++-
 Ryujinx.Graphics.Shader/SamplerType.cs        |  5 +-
 .../StructuredIr/AstTextureOperation.cs       | 11 ++-
 .../StructuredIr/StructuredProgram.cs         |  1 +
 .../Optimizations/BindlessToIndexed.cs        | 84 +++++++++++++++++++
 .../Translation/Optimizations/Optimizer.cs    |  5 ++
 10 files changed, 210 insertions(+), 39 deletions(-)
 create mode 100644 Ryujinx.Graphics.Shader/Translation/Optimizations/BindlessToIndexed.cs

diff --git a/Ryujinx.Graphics.Gpu/Engine/Methods.cs b/Ryujinx.Graphics.Gpu/Engine/Methods.cs
index 5d15631293..b1326ec572 100644
--- a/Ryujinx.Graphics.Gpu/Engine/Methods.cs
+++ b/Ryujinx.Graphics.Gpu/Engine/Methods.cs
@@ -670,7 +670,7 @@ namespace Ryujinx.Graphics.Gpu.Engine
 
         private static Target GetTarget(SamplerType type)
         {
-            type &= ~SamplerType.Shadow;
+            type &= ~(SamplerType.Indexed | SamplerType.Shadow);
 
             switch (type)
             {
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
index 7c67bc13c2..6c4ba949d4 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
@@ -231,7 +231,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
             foreach (AstTextureOperation texOp in info.Samplers.OrderBy(x => x.Handle))
             {
-                string samplerName = OperandManager.GetSamplerName(context.Config.Stage, texOp);
+                string indexExpr = NumberFormatter.FormatInt(texOp.ArraySize);
+
+                string samplerName = OperandManager.GetSamplerName(context.Config.Stage, texOp, indexExpr);
 
                 if (!samplers.TryAdd(samplerName, texOp))
                 {
@@ -257,12 +259,25 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
                     desc = new TextureDescriptor(samplerName, texOp.Type, operand.CbufSlot, operand.CbufOffset);
                 }
+                else if ((texOp.Type & SamplerType.Indexed) != 0)
+                {
+                    for (int index = 0; index < texOp.ArraySize; index++)
+                    {
+                        string indexExpr = NumberFormatter.FormatInt(index);
+
+                        string indexedSamplerName = OperandManager.GetSamplerName(context.Config.Stage, texOp, indexExpr);
+
+                        desc = new TextureDescriptor(indexedSamplerName, texOp.Type, texOp.Handle + index * 2);
+
+                        context.TextureDescriptors.Add(desc);
+                    }
+                }
                 else
                 {
                     desc = new TextureDescriptor(samplerName, texOp.Type, texOp.Handle);
-                }
 
-                context.TextureDescriptors.Add(desc);
+                    context.TextureDescriptors.Add(desc);
+                }
             }
         }
 
@@ -272,7 +287,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
             foreach (AstTextureOperation texOp in info.Images.OrderBy(x => x.Handle))
             {
-                string imageName = OperandManager.GetImageName(context.Config.Stage, texOp);
+                string indexExpr = NumberFormatter.FormatInt(texOp.ArraySize);
+
+                string imageName = OperandManager.GetImageName(context.Config.Stage, texOp, indexExpr);
 
                 if (!images.TryAdd(imageName, texOp))
                 {
@@ -290,9 +307,25 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
                 AstTextureOperation texOp = kv.Value;
 
-                TextureDescriptor desc = new TextureDescriptor(imageName, texOp.Type, texOp.Handle);
+                if ((texOp.Type & SamplerType.Indexed) != 0)
+                {
+                    for (int index = 0; index < texOp.ArraySize; index++)
+                    {
+                        string indexExpr = NumberFormatter.FormatInt(index);
 
-                context.ImageDescriptors.Add(desc);
+                        string indexedSamplerName = OperandManager.GetSamplerName(context.Config.Stage, texOp, indexExpr);
+
+                        var desc = new TextureDescriptor(indexedSamplerName, texOp.Type, texOp.Handle + index * 2);
+
+                        context.TextureDescriptors.Add(desc);
+                    }
+                }
+                else
+                {
+                    var desc = new TextureDescriptor(imageName, texOp.Type, texOp.Handle);
+
+                    context.ImageDescriptors.Add(desc);
+                }
             }
         }
 
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
index 913cace16d..21e39fcf4f 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
@@ -15,11 +15,26 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
 
             bool isBindless = (texOp.Flags & TextureFlags.Bindless) != 0;
 
-            bool isArray = (texOp.Type & SamplerType.Array) != 0;
+            bool isArray   = (texOp.Type & SamplerType.Array)   != 0;
+            bool isIndexed = (texOp.Type & SamplerType.Indexed) != 0;
 
             string texCall = "imageStore";
 
-            string imageName = OperandManager.GetImageName(context.Config.Stage, texOp);
+            int srcIndex = isBindless ? 1 : 0;
+
+            string Src(VariableType type)
+            {
+                return GetSoureExpr(context, texOp.GetSource(srcIndex++), type);
+            }
+
+            string indexExpr = null;
+
+            if (isIndexed)
+            {
+                indexExpr = Src(VariableType.S32);
+            }
+
+            string imageName = OperandManager.GetImageName(context.Config.Stage, texOp, indexExpr);
 
             texCall += "(" + imageName;
 
@@ -34,13 +49,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                 arrayIndexElem = pCount++;
             }
 
-            int srcIndex = isBindless ? 1 : 0;
-
-            string Src(VariableType type)
-            {
-                return GetSoureExpr(context, texOp.GetSource(srcIndex++), type);
-            }
-
             void Append(string str)
             {
                 texCall += ", " + str;
@@ -174,6 +182,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             bool hasOffsets     = (texOp.Flags & TextureFlags.Offsets)     != 0;
 
             bool isArray       = (texOp.Type & SamplerType.Array)       != 0;
+            bool isIndexed     = (texOp.Type & SamplerType.Indexed)     != 0;
             bool isMultisample = (texOp.Type & SamplerType.Multisample) != 0;
             bool isShadow      = (texOp.Type & SamplerType.Shadow)      != 0;
 
@@ -209,7 +218,21 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                 texCall += "Offsets";
             }
 
-            string samplerName = OperandManager.GetSamplerName(context.Config.Stage, texOp);
+            int srcIndex = isBindless ? 1 : 0;
+
+            string Src(VariableType type)
+            {
+                return GetSoureExpr(context, texOp.GetSource(srcIndex++), type);
+            }
+
+            string indexExpr = null;
+
+            if (isIndexed)
+            {
+                indexExpr = Src(VariableType.S32);
+            }
+
+            string samplerName = OperandManager.GetSamplerName(context.Config.Stage, texOp, indexExpr);
 
             texCall += "(" + samplerName;
 
@@ -249,13 +272,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                 hasExtraCompareArg = true;
             }
 
-            int srcIndex = isBindless ? 1 : 0;
-
-            string Src(VariableType type)
-            {
-                return GetSoureExpr(context, texOp.GetSource(srcIndex++), type);
-            }
-
             void Append(string str)
             {
                 texCall += ", " + str;
@@ -395,11 +411,20 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
         {
             AstTextureOperation texOp = (AstTextureOperation)operation;
 
-            bool isBindless  = (texOp.Flags & TextureFlags.Bindless) != 0;
+            bool isBindless = (texOp.Flags & TextureFlags.Bindless) != 0;
 
-            string samplerName = OperandManager.GetSamplerName(context.Config.Stage, texOp);
+            bool isIndexed = (texOp.Type & SamplerType.Indexed) != 0;
 
-            IAstNode src0 = operation.GetSource(isBindless ? 1 : 0);
+            string indexExpr = null;
+
+            if (isIndexed)
+            {
+                indexExpr = GetSoureExpr(context, texOp.GetSource(0), VariableType.S32);
+            }
+
+            string samplerName = OperandManager.GetSamplerName(context.Config.Stage, texOp, indexExpr);
+
+            IAstNode src0 = operation.GetSource(isBindless || isIndexed ? 1 : 0);
 
             string src0Expr = GetSoureExpr(context, src0, GetSrcVarType(operation.Inst, 0));
 
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
index fe3073967c..36f76ec5fc 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
@@ -223,7 +223,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return ubName + "_" + DefaultNames.UniformNameSuffix;
         }
 
-        public static string GetSamplerName(ShaderStage stage, AstTextureOperation texOp)
+        public static string GetSamplerName(ShaderStage stage, AstTextureOperation texOp, string indexExpr)
         {
             string suffix;
 
@@ -235,16 +235,26 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             }
             else
             {
-                suffix = (texOp.Handle - 8).ToString();
+                suffix = texOp.Handle.ToString();
+
+                if ((texOp.Type & SamplerType.Indexed) != 0)
+                {
+                    suffix += $"a[{indexExpr}]";
+                }
             }
 
             return GetShaderStagePrefix(stage) + "_" + DefaultNames.SamplerNamePrefix + suffix;
         }
 
-        public static string GetImageName(ShaderStage stage, AstTextureOperation texOp)
+        public static string GetImageName(ShaderStage stage, AstTextureOperation texOp, string indexExpr)
         {
             string suffix = texOp.Handle.ToString();
 
+            if ((texOp.Type & SamplerType.Indexed) != 0)
+            {
+                suffix += $"a[{indexExpr}]";
+            }
+
             return GetShaderStagePrefix(stage) + "_" + DefaultNames.ImageNamePrefix + suffix;
         }
 
diff --git a/Ryujinx.Graphics.Shader/IntermediateRepresentation/TextureOperation.cs b/Ryujinx.Graphics.Shader/IntermediateRepresentation/TextureOperation.cs
index b96c55e880..718d2c2ec1 100644
--- a/Ryujinx.Graphics.Shader/IntermediateRepresentation/TextureOperation.cs
+++ b/Ryujinx.Graphics.Shader/IntermediateRepresentation/TextureOperation.cs
@@ -2,10 +2,10 @@ namespace Ryujinx.Graphics.Shader.IntermediateRepresentation
 {
     class TextureOperation : Operation
     {
-        public SamplerType  Type  { get; }
-        public TextureFlags Flags { get; }
+        public SamplerType  Type  { get; private set; }
+        public TextureFlags Flags { get; private set; }
 
-        public int Handle { get; }
+        public int Handle { get; private set; }
 
         public TextureOperation(
             Instruction      inst,
@@ -20,5 +20,14 @@ namespace Ryujinx.Graphics.Shader.IntermediateRepresentation
             Flags  = flags;
             Handle = handle;
         }
+
+        public void TurnIntoIndexed(int handle)
+        {
+            Type |= SamplerType.Indexed;
+
+            Flags &= ~TextureFlags.Bindless;
+
+            Handle = handle;
+        }
     }
 }
\ No newline at end of file
diff --git a/Ryujinx.Graphics.Shader/SamplerType.cs b/Ryujinx.Graphics.Shader/SamplerType.cs
index 42854c9786..5e0b776c3d 100644
--- a/Ryujinx.Graphics.Shader/SamplerType.cs
+++ b/Ryujinx.Graphics.Shader/SamplerType.cs
@@ -14,8 +14,9 @@ namespace Ryujinx.Graphics.Shader
         Mask = 0xff,
 
         Array       = 1 << 8,
-        Multisample = 1 << 9,
-        Shadow      = 1 << 10
+        Indexed     = 1 << 9,
+        Multisample = 1 << 10,
+        Shadow      = 1 << 11
     }
 
     static class SamplerTypeExtensions
diff --git a/Ryujinx.Graphics.Shader/StructuredIr/AstTextureOperation.cs b/Ryujinx.Graphics.Shader/StructuredIr/AstTextureOperation.cs
index 7261b9ff1f..d6d40732fa 100644
--- a/Ryujinx.Graphics.Shader/StructuredIr/AstTextureOperation.cs
+++ b/Ryujinx.Graphics.Shader/StructuredIr/AstTextureOperation.cs
@@ -7,19 +7,22 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
         public SamplerType  Type  { get; }
         public TextureFlags Flags { get; }
 
-        public int Handle { get; }
+        public int Handle    { get; }
+        public int ArraySize { get; }
 
         public AstTextureOperation(
             Instruction       inst,
             SamplerType       type,
             TextureFlags      flags,
             int               handle,
+            int               arraySize,
             int               compMask,
             params IAstNode[] sources) : base(inst, compMask, sources)
         {
-            Type = type;
-            Flags  = flags;
-            Handle = handle;
+            Type      = type;
+            Flags     = flags;
+            Handle    = handle;
+            ArraySize = arraySize;
         }
     }
 }
\ No newline at end of file
diff --git a/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgram.cs b/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgram.cs
index c4ffbe1ad6..ef8443ab72 100644
--- a/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgram.cs
+++ b/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgram.cs
@@ -60,6 +60,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
                     texOp.Type,
                     texOp.Flags,
                     texOp.Handle,
+                    4, // TODO: Non-hardcoded array size.
                     componentMask,
                     sources);
             }
diff --git a/Ryujinx.Graphics.Shader/Translation/Optimizations/BindlessToIndexed.cs b/Ryujinx.Graphics.Shader/Translation/Optimizations/BindlessToIndexed.cs
new file mode 100644
index 0000000000..8cb62db937
--- /dev/null
+++ b/Ryujinx.Graphics.Shader/Translation/Optimizations/BindlessToIndexed.cs
@@ -0,0 +1,84 @@
+using Ryujinx.Graphics.Shader.IntermediateRepresentation;
+using System.Collections.Generic;
+
+using static Ryujinx.Graphics.Shader.IntermediateRepresentation.OperandHelper;
+
+namespace Ryujinx.Graphics.Shader.Translation.Optimizations
+{
+    static class BindlessToIndexed
+    {
+        private const int StorageDescsBaseOffset = 0x44; // In words.
+
+        private const int UbeStorageDescsBaseOffset = 0x84; // In words.
+        private const int UbeStorageMaxCount        = 14;
+
+        private const int StorageDescSize = 4; // In words.
+        private const int StorageMaxCount = 16;
+
+        private const int StorageDescsSize  = StorageDescSize * StorageMaxCount;
+
+        public static void RunPass(BasicBlock block)
+        {
+            // We can turn a bindless texture access into a indexed access,
+            // as long the following conditions are true:
+            // - The handle is loaded using a LDC instruction.
+            // - The handle is loaded from the constant buffer with the handles (CB2 for NVN).
+            // - The load has a constant offset.
+            // The base offset of the array of handles on the constant buffer is the constant offset.
+            for (LinkedListNode<INode> node = block.Operations.First; node != null; node = node.Next)
+            {
+                if (!(node.Value is TextureOperation texOp))
+                {
+                    continue;
+                }
+
+                if ((texOp.Flags & TextureFlags.Bindless) == 0)
+                {
+                    continue;
+                }
+
+                if (!(texOp.GetSource(0).AsgOp is Operation handleAsgOp))
+                {
+                    continue;
+                }
+
+                if (handleAsgOp.Inst != Instruction.LoadConstant)
+                {
+                    continue;
+                }
+
+                Operand ldcSrc0 = handleAsgOp.GetSource(0);
+                Operand ldcSrc1 = handleAsgOp.GetSource(1);
+
+                if (ldcSrc0.Type != OperandType.Constant || ldcSrc0.Value != 2)
+                {
+                    continue;
+                }
+
+                if (!(ldcSrc1.AsgOp is Operation addOp))
+                {
+                    continue;
+                }
+
+                Operand addSrc1 = addOp.GetSource(1);
+
+                if (addSrc1.Type != OperandType.Constant)
+                {
+                    continue;
+                }
+
+                texOp.TurnIntoIndexed(addSrc1.Value);
+
+                Operand index = Local();
+
+                Operand source = addOp.GetSource(0);
+
+                Operation shrBy1 = new Operation(Instruction.ShiftRightU32, index, source, Const(1));
+
+                block.Operations.AddBefore(node, shrBy1);
+
+                texOp.SetSource(0, index);
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs b/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs
index 22d794a40d..6ee27884c2 100644
--- a/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Optimizations/Optimizer.cs
@@ -84,6 +84,11 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
                 }
             }
             while (modified);
+
+            for (int blkIndex = 0; blkIndex < blocks.Length; blkIndex++)
+            {
+                BindlessToIndexed.RunPass(blocks[blkIndex]);
+            }
         }
 
         private static void PropagateCopy(Operation copyOp)