Initial support for shader attribute indexing (#2546)

* Initial support for shader attribute indexing

* Support output indexing too, other improvements

* Fix order

* Address feedback

3 files changed, 48 insertions, 7 deletions

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);