Initial support for shader attribute indexing (#2546)

* Initial support for shader attribute indexing

* Support output indexing too, other improvements

* Fix order

* Address feedback

7 files changed, 117 insertions, 38 deletions

diff --git a/Ryujinx.Graphics.Shader/CodeGen/Constants.cs b/Ryujinx.Graphics.Shader/CodeGen/Constants.cs
deleted file mode 100644
index 59e9f145..00000000
--- a/Ryujinx.Graphics.Shader/CodeGen/Constants.cs
+++ /dev/null
@@ -1,9 +0,0 @@
-namespace Ryujinx.Graphics.Shader.CodeGen
-{
-    static class Constants
-    {
-        public const int MaxShaderStorageBuffers = 16;
-
-        public const int ConstantBufferSize = 0x10000; // In bytes
-    }
-}
\ No newline at end of file
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
index 2a93be32..7b26801a 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
@@ -402,14 +402,23 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         private static void DeclareInputAttributes(CodeGenContext context, StructuredProgramInfo info)
         {
-            int usedAttribtes = context.Config.UsedInputAttributes;
-            while (usedAttribtes != 0)
+            if (context.Config.UsedFeatures.HasFlag(FeatureFlags.IaIndexing))
             {
-                int index = BitOperations.TrailingZeroCount(usedAttribtes);
+                string suffix = context.Config.Stage == ShaderStage.Geometry ? "[]" : string.Empty;
 
-                DeclareInputAttribute(context, info, index);
+                context.AppendLine($"layout (location = 0) in vec4 {DefaultNames.IAttributePrefix}{suffix}[{Constants.MaxAttributes}];");
+            }
+            else
+            {
+                int usedAttributes = context.Config.UsedInputAttributes;
+                while (usedAttributes != 0)
+                {
+                    int index = BitOperations.TrailingZeroCount(usedAttributes);
+
+                    DeclareInputAttribute(context, info, index);
 
-                usedAttribtes &= ~(1 << index);
+                    usedAttributes &= ~(1 << index);
+                }
             }
         }
 
@@ -448,14 +457,21 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         private static void DeclareOutputAttributes(CodeGenContext context, StructuredProgramInfo info)
         {
-            int usedAttribtes = context.Config.UsedOutputAttributes;
-            while (usedAttribtes != 0)
+            if (context.Config.UsedFeatures.HasFlag(FeatureFlags.OaIndexing))
             {
-                int index = BitOperations.TrailingZeroCount(usedAttribtes);
+                context.AppendLine($"layout (location = 0) out vec4 {DefaultNames.OAttributePrefix}[{Constants.MaxAttributes}];");
+            }
+            else
+            {
+                int usedAttributes = context.Config.UsedOutputAttributes;
+                while (usedAttributes != 0)
+                {
+                    int index = BitOperations.TrailingZeroCount(usedAttributes);
 
-                DeclareOutputAttribute(context, index);
+                    DeclareOutputAttribute(context, index);
 
-                usedAttribtes &= ~(1 << index);
+                    usedAttributes &= ~(1 << index);
+                }
             }
         }
 
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
index 6ea700ac..2d6ede0a 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
@@ -128,7 +128,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
                     if (assignment.Destination is AstOperand operand && operand.Type == OperandType.Attribute)
                     {
-                        dest = OperandManager.GetOutAttributeName(operand, context.Config);
+                        dest = OperandManager.GetOutAttributeName(operand.Value, context.Config);
                     }
                     else
                     {
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
index d5cd0f72..41a8b942 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
@@ -161,6 +161,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                     case Instruction.PackHalf2x16:
                         return PackHalf2x16(context, operation);
 
+                    case Instruction.StoreAttribute:
+                        return StoreAttribute(context, operation);
+
                     case Instruction.StoreLocal:
                         return StoreLocal(context, operation);
 
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
index f3774a60..a4957658 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
@@ -109,6 +109,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             Add(Instruction.ShuffleXor,               InstType.CallQuaternary, HelperFunctionNames.ShuffleXor);
             Add(Instruction.Sine,                     InstType.CallUnary,      "sin");
             Add(Instruction.SquareRoot,               InstType.CallUnary,      "sqrt");
+            Add(Instruction.StoreAttribute,           InstType.Special);
             Add(Instruction.StoreLocal,               InstType.Special);
             Add(Instruction.StoreShared,              InstType.Special);
             Add(Instruction.StoreStorage,             InstType.Special);
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
index ef3b0bed..a0aec28e 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
@@ -137,15 +137,25 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
         {
             IAstNode src1 = operation.GetSource(0);
             IAstNode src2 = operation.GetSource(1);
+            IAstNode src3 = operation.GetSource(2);
 
-            if (!(src1 is AstOperand attr) || attr.Type != OperandType.Attribute)
+            if (!(src1 is AstOperand baseAttr) || baseAttr.Type != OperandType.Constant)
             {
-                throw new InvalidOperationException("First source of LoadAttribute must be a attribute.");
+                throw new InvalidOperationException($"First input of {nameof(Instruction.LoadAttribute)} must be a constant operand.");
             }
 
-            string indexExpr = GetSoureExpr(context, src2, GetSrcVarType(operation.Inst, 1));
+            string indexExpr = GetSoureExpr(context, src3, GetSrcVarType(operation.Inst, 2));
 
-            return OperandManager.GetAttributeName(attr, context.Config, isOutAttr: false, indexExpr);
+            if (src2 is AstOperand operand && operand.Type == OperandType.Constant)
+            {
+                return OperandManager.GetAttributeName(baseAttr.Value + (operand.Value << 2), context.Config, isOutAttr: false, indexExpr);
+            }
+            else
+            {
+                string attrExpr = GetSoureExpr(context, src2, GetSrcVarType(operation.Inst, 1));
+                attrExpr = Enclose(attrExpr, src2, Instruction.ShiftRightS32, isLhs: true);
+                return OperandManager.GetAttributeName(attrExpr, context.Config, isOutAttr: false, indexExpr);
+            }
         }
 
         public static string LoadConstant(CodeGenContext context, AstOperation operation)
@@ -154,16 +164,15 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             IAstNode src2 = operation.GetSource(1);
 
             string offsetExpr = GetSoureExpr(context, src2, GetSrcVarType(operation.Inst, 1));
-
             offsetExpr = Enclose(offsetExpr, src2, Instruction.ShiftRightS32, isLhs: true);
 
             var config = context.Config;
             bool indexElement = !config.GpuAccessor.QueryHostHasVectorIndexingBug();
 
-            if (src1 is AstOperand oper && oper.Type == OperandType.Constant)
+            if (src1 is AstOperand operand && operand.Type == OperandType.Constant)
             {
                 bool cbIndexable = config.UsedFeatures.HasFlag(Translation.FeatureFlags.CbIndexing);
-                return OperandManager.GetConstantBufferName(oper.Value, offsetExpr, config.Stage, cbIndexable, indexElement);
+                return OperandManager.GetConstantBufferName(operand.Value, offsetExpr, config.Stage, cbIndexable, indexElement);
             }
             else
             {
@@ -250,6 +259,34 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             return $"textureQueryLod({samplerName}, {coordsExpr}){GetMask(texOp.Index)}";
         }
 
+        public static string StoreAttribute(CodeGenContext context, AstOperation operation)
+        {
+            IAstNode src1 = operation.GetSource(0);
+            IAstNode src2 = operation.GetSource(1);
+            IAstNode src3 = operation.GetSource(2);
+
+            if (!(src1 is AstOperand baseAttr) || baseAttr.Type != OperandType.Constant)
+            {
+                throw new InvalidOperationException($"First input of {nameof(Instruction.StoreAttribute)} must be a constant operand.");
+            }
+
+            string attrName;
+
+            if (src2 is AstOperand operand && operand.Type == OperandType.Constant)
+            {
+                attrName = OperandManager.GetAttributeName(baseAttr.Value + (operand.Value << 2), context.Config, isOutAttr: true);
+            }
+            else
+            {
+                string attrExpr = GetSoureExpr(context, src2, GetSrcVarType(operation.Inst, 1));
+                attrExpr = Enclose(attrExpr, src2, Instruction.ShiftRightS32, isLhs: true);
+                attrName = OperandManager.GetAttributeName(attrExpr, context.Config, isOutAttr: true);
+            }
+
+            string value = GetSoureExpr(context, src3, GetSrcVarType(operation.Inst, 2));
+            return $"{attrName} = {value}";
+        }
+
         public static string StoreLocal(CodeGenContext context, AstOperation operation)
         {
             return StoreLocalOrShared(context, operation, DefaultNames.LocalMemoryName);
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
index 60a471eb..b9f1b4ef 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
@@ -99,7 +99,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return operand.Type switch
             {
                 OperandType.Argument => GetArgumentName(operand.Value),
-                OperandType.Attribute => GetAttributeName(operand, config),
+                OperandType.Attribute => GetAttributeName(operand.Value, config),
                 OperandType.Constant => NumberFormatter.FormatInt(operand.Value),
                 OperandType.ConstantBuffer => GetConstantBufferName(
                     operand.CbufSlot,
@@ -142,15 +142,13 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return GetVec4Indexed(GetUbName(stage, slotExpr) + $"[{offsetExpr} >> 2]", offsetExpr + " & 3", indexElement);
         }
 
-        public static string GetOutAttributeName(AstOperand attr, ShaderConfig config)
+        public static string GetOutAttributeName(int value, ShaderConfig config)
         {
-            return GetAttributeName(attr, config, isOutAttr: true);
+            return GetAttributeName(value, config, isOutAttr: true);
         }
 
-        public static string GetAttributeName(AstOperand attr, ShaderConfig config, bool isOutAttr = false, string indexExpr = "0")
+        public static string GetAttributeName(int value, ShaderConfig config, bool isOutAttr = false, string indexExpr = "0")
         {
-            int value = attr.Value;
-
             char swzMask = GetSwizzleMask((value >> 2) & 3);
 
             if (value >= AttributeConsts.UserAttributeBase && value < AttributeConsts.UserAttributeEnd)
@@ -161,7 +159,18 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                     ? DefaultNames.OAttributePrefix
                     : DefaultNames.IAttributePrefix;
 
-                if ((config.Options.Flags & TranslationFlags.Feedback) != 0)
+                if (config.UsedFeatures.HasFlag(isOutAttr ? FeatureFlags.OaIndexing : FeatureFlags.IaIndexing))
+                {
+                    string name = prefix;
+
+                    if (config.Stage == ShaderStage.Geometry && !isOutAttr)
+                    {
+                        name += $"[{indexExpr}]";
+                    }
+
+                    return name + $"[{(value >> 4)}]." + swzMask;
+                }
+                else if (config.Options.Flags.HasFlag(TranslationFlags.Feedback))
                 {
                     string name = $"{prefix}{(value >> 4)}_{swzMask}";
 
@@ -231,6 +240,20 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return isOutAttr ? "// bad_attr0x" + value.ToString("X") : "0.0";
         }
 
+        public static string GetAttributeName(string attrExpr, ShaderConfig config, bool isOutAttr = false, string indexExpr = "0")
+        {
+            string name = isOutAttr
+                ? DefaultNames.OAttributePrefix
+                : DefaultNames.IAttributePrefix;
+
+            if (config.Stage == ShaderStage.Geometry && !isOutAttr)
+            {
+                name += $"[{indexExpr}]";
+            }
+
+            return $"{name}[{attrExpr} >> 2][{attrExpr} & 3]";
+        }
+
         public static string GetUbName(ShaderStage stage, int slot, bool cbIndexable)
         {
             if (cbIndexable)
@@ -314,12 +337,20 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
         {
             if (node is AstOperation operation)
             {
-                // Load attribute basically just returns the attribute value.
-                // Some built-in attributes may have different types, so we need
-                // to return the type based on the attribute that is being read.
                 if (operation.Inst == Instruction.LoadAttribute)
                 {
-                    return GetOperandVarType((AstOperand)operation.GetSource(0));
+                    // Load attribute basically just returns the attribute value.
+                    // Some built-in attributes may have different types, so we need
+                    // to return the type based on the attribute that is being read.
+                    if (operation.GetSource(0) is AstOperand operand && operand.Type == OperandType.Constant)
+                    {
+                        if (_builtInAttributes.TryGetValue(operand.Value & ~3, out BuiltInAttribute builtInAttr))
+                        {
+                            return builtInAttr.Type;
+                        }
+                    }
+
+                    return OperandInfo.GetVarType(OperandType.Attribute);
                 }
                 else if (operation.Inst == Instruction.Call)
                 {