Use vector transform feedback outputs if possible (#3832)

6 files changed, 52 insertions, 34 deletions

diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
index 418af6cb..9eb20f6f 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
@@ -10,12 +10,12 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         public StructuredFunction CurrentFunction { get; set; }
 
+        public StructuredProgramInfo Info { get; }
+
         public ShaderConfig Config { get; }
 
         public OperandManager OperandManager { get; }
 
-        private readonly StructuredProgramInfo _info;
-
         private readonly StringBuilder _sb;
 
         private int _level;
@@ -24,7 +24,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         public CodeGenContext(StructuredProgramInfo info, ShaderConfig config)
         {
-            _info = info;
+            Info = info;
             Config = config;
 
             OperandManager = new OperandManager();
@@ -72,19 +72,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         public StructuredFunction GetFunction(int id)
         {
-            return _info.Functions[id];
-        }
-
-        public TransformFeedbackOutput GetTransformFeedbackOutput(int location, int component)
-        {
-            int index = (AttributeConsts.UserAttributeBase / 4) + location * 4 + component;
-            return _info.TransformFeedbackOutputs[index];
-        }
-
-        public TransformFeedbackOutput GetTransformFeedbackOutput(int location)
-        {
-            int index = location / 4;
-            return _info.TransformFeedbackOutputs[index];
+            return Info.Functions[id];
         }
 
         private void UpdateIndentation()
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
index 91fd286d..4f2751b1 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
@@ -210,7 +210,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
                 if (context.Config.TransformFeedbackEnabled && context.Config.LastInVertexPipeline)
                 {
-                    var tfOutput = context.GetTransformFeedbackOutput(AttributeConsts.PositionX);
+                    var tfOutput = context.Info.GetTransformFeedbackOutput(AttributeConsts.PositionX);
                     if (tfOutput.Valid)
                     {
                         context.AppendLine($"layout (xfb_buffer = {tfOutput.Buffer}, xfb_offset = {tfOutput.Offset}, xfb_stride = {tfOutput.Stride}) out gl_PerVertex");
@@ -604,19 +604,45 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
             if (context.Config.TransformFeedbackEnabled && context.Config.LastInVertexPipeline)
             {
-                for (int c = 0; c < 4; c++)
+                int attrOffset = AttributeConsts.UserAttributeBase + attr * 16;
+                int components = context.Config.LastInPipeline ? context.Info.GetTransformFeedbackOutputComponents(attrOffset) : 1;
+
+                if (components > 1)
                 {
-                    char swzMask = "xyzw"[c];
+                    string type = components switch
+                    {
+                        2 => "vec2",
+                        3 => "vec3",
+                        4 => "vec4",
+                        _ => "float"
+                    };
 
                     string xfb = string.Empty;
 
-                    var tfOutput = context.GetTransformFeedbackOutput(attr, c);
+                    var tfOutput = context.Info.GetTransformFeedbackOutput(attrOffset);
                     if (tfOutput.Valid)
                     {
                         xfb = $", xfb_buffer = {tfOutput.Buffer}, xfb_offset = {tfOutput.Offset}, xfb_stride = {tfOutput.Stride}";
                     }
 
-                    context.AppendLine($"layout (location = {attr}, component = {c}{xfb}) out float {name}_{swzMask};");
+                    context.AppendLine($"layout (location = {attr}{xfb}) out {type} {name};");
+                }
+                else
+                {
+                    for (int c = 0; c < 4; c++)
+                    {
+                        char swzMask = "xyzw"[c];
+
+                        string xfb = string.Empty;
+
+                        var tfOutput = context.Info.GetTransformFeedbackOutput(attrOffset + c * 4);
+                        if (tfOutput.Valid)
+                        {
+                            xfb = $", xfb_buffer = {tfOutput.Buffer}, xfb_offset = {tfOutput.Offset}, xfb_stride = {tfOutput.Stride}";
+                        }
+
+                        context.AppendLine($"layout (location = {attr}, component = {c}{xfb}) out float {name}_{swzMask};");
+                    }
                 }
             }
             else
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
index e9dbdd2d..e1b8eb6e 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
@@ -134,7 +134,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                     if (assignment.Destination is AstOperand operand && operand.Type.IsAttribute())
                     {
                         bool perPatch = operand.Type == OperandType.AttributePerPatch;
-                        dest = OperandManager.GetOutAttributeName(operand.Value, context.Config, perPatch);
+                        dest = OperandManager.GetOutAttributeName(context, operand.Value, perPatch);
                     }
                     else
                     {
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
index 388285a8..b890b015 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
@@ -22,7 +22,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             }
             else if (node is AstOperand operand)
             {
-                return context.OperandManager.GetExpression(operand, context.Config);
+                return context.OperandManager.GetExpression(context, operand);
             }
 
             throw new ArgumentException($"Invalid node type \"{node?.GetType().Name ?? "null"}\".");
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
index 09404001..022e3a44 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
@@ -205,7 +205,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             if (src2 is AstOperand operand && operand.Type == OperandType.Constant)
             {
                 int attrOffset = baseAttr.Value + (operand.Value << 2);
-                return OperandManager.GetAttributeName(attrOffset, context.Config, perPatch: false, isOutAttr: false, indexExpr);
+                return OperandManager.GetAttributeName(context, attrOffset, perPatch: false, isOutAttr: false, indexExpr);
             }
             else
             {
@@ -332,7 +332,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             if (src2 is AstOperand operand && operand.Type == OperandType.Constant)
             {
                 int attrOffset = baseAttr.Value + (operand.Value << 2);
-                attrName = OperandManager.GetAttributeName(attrOffset, context.Config, perPatch: false, isOutAttr: true);
+                attrName = OperandManager.GetAttributeName(context, attrOffset, perPatch: false, isOutAttr: true);
             }
             else
             {
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
index 67442e5a..031b1c44 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
@@ -103,15 +103,15 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return name;
         }
 
-        public string GetExpression(AstOperand operand, ShaderConfig config)
+        public string GetExpression(CodeGenContext context, AstOperand operand)
         {
             return operand.Type switch
             {
                 OperandType.Argument => GetArgumentName(operand.Value),
-                OperandType.Attribute => GetAttributeName(operand.Value, config, perPatch: false),
-                OperandType.AttributePerPatch => GetAttributeName(operand.Value, config, perPatch: true),
+                OperandType.Attribute => GetAttributeName(context, operand.Value, perPatch: false),
+                OperandType.AttributePerPatch => GetAttributeName(context, operand.Value, perPatch: true),
                 OperandType.Constant => NumberFormatter.FormatInt(operand.Value),
-                OperandType.ConstantBuffer => GetConstantBufferName(operand, config),
+                OperandType.ConstantBuffer => GetConstantBufferName(operand, context.Config),
                 OperandType.LocalVariable => _locals[operand],
                 OperandType.Undefined => DefaultNames.UndefinedName,
                 _ => throw new ArgumentException($"Invalid operand type \"{operand.Type}\".")
@@ -153,13 +153,15 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return GetVec4Indexed(GetUbName(stage, slotExpr) + $"[{offsetExpr} >> 2]", offsetExpr + " & 3", indexElement);
         }
 
-        public static string GetOutAttributeName(int value, ShaderConfig config, bool perPatch)
+        public static string GetOutAttributeName(CodeGenContext context, int value, bool perPatch)
         {
-            return GetAttributeName(value, config, perPatch, isOutAttr: true);
+            return GetAttributeName(context, value, perPatch, isOutAttr: true);
         }
 
-        public static string GetAttributeName(int value, ShaderConfig config, bool perPatch, bool isOutAttr = false, string indexExpr = "0")
+        public static string GetAttributeName(CodeGenContext context, int value, bool perPatch, bool isOutAttr = false, string indexExpr = "0")
         {
+            ShaderConfig config = context.Config;
+
             if ((value & AttributeConsts.LoadOutputMask) != 0)
             {
                 isOutAttr = true;
@@ -192,6 +194,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             }
             else if (value >= AttributeConsts.UserAttributeBase && value < AttributeConsts.UserAttributeEnd)
             {
+                int attrOffset = value;
                 value -= AttributeConsts.UserAttributeBase;
 
                 string prefix = isOutAttr
@@ -215,14 +218,15 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                     ((config.LastInVertexPipeline && isOutAttr) ||
                     (config.Stage == ShaderStage.Fragment && !isOutAttr)))
                 {
-                    string name = $"{prefix}{(value >> 4)}_{swzMask}";
+                    int components = config.LastInPipeline ? context.Info.GetTransformFeedbackOutputComponents(attrOffset) : 1;
+                    string name = components > 1 ? $"{prefix}{(value >> 4)}" : $"{prefix}{(value >> 4)}_{swzMask}";
 
                     if (AttributeInfo.IsArrayAttributeGlsl(config.Stage, isOutAttr))
                     {
                         name += isOutAttr ? "[gl_InvocationID]" : $"[{indexExpr}]";
                     }
 
-                    return name;
+                    return components > 1 ? name + '.' + swzMask : name;
                 }
                 else
                 {