Make sure attributes used on subsequent shader stages are initialized (#2538)

4 files changed, 105 insertions, 17 deletions

diff --git a/Ryujinx.Graphics.Shader/Translation/EmitterContext.cs b/Ryujinx.Graphics.Shader/Translation/EmitterContext.cs
index 49a89374..5cdd5c0a 100644
--- a/Ryujinx.Graphics.Shader/Translation/EmitterContext.cs
+++ b/Ryujinx.Graphics.Shader/Translation/EmitterContext.cs
@@ -15,6 +15,8 @@ namespace Ryujinx.Graphics.Shader.Translation
 
         public bool IsNonMain { get; }
 
+        public int OperationsCount => _operations.Count;
+
         private readonly IReadOnlyDictionary<ulong, int> _funcs;
         private readonly List<Operation> _operations;
         private readonly Dictionary<ulong, Operand> _labels;
@@ -200,6 +202,7 @@ namespace Ryujinx.Graphics.Shader.Translation
 
                     if (target.Enabled)
                     {
+                        Config.SetOutputUserAttribute(rtIndex);
                         regIndexBase += 4;
                     }
                 }
diff --git a/Ryujinx.Graphics.Shader/Translation/ShaderConfig.cs b/Ryujinx.Graphics.Shader/Translation/ShaderConfig.cs
index 3e7be582..c7704c2b 100644
--- a/Ryujinx.Graphics.Shader/Translation/ShaderConfig.cs
+++ b/Ryujinx.Graphics.Shader/Translation/ShaderConfig.cs
@@ -41,6 +41,10 @@ namespace Ryujinx.Graphics.Shader.Translation
 
         private readonly TranslationCounts _counts;
 
+        public int UsedInputAttributes { get; private set; }
+        public int UsedOutputAttributes { get; private set; }
+        public int PassthroughAttributes { get; private set; }
+
         private int _usedConstantBuffers;
         private int _usedStorageBuffers;
         private int _usedStorageBuffersWrite;
@@ -170,6 +174,8 @@ namespace Ryujinx.Graphics.Shader.Translation
 
             TextureHandlesForCache.UnionWith(other.TextureHandlesForCache);
 
+            UsedInputAttributes |= other.UsedInputAttributes;
+            UsedOutputAttributes |= other.UsedOutputAttributes;
             _usedConstantBuffers |= other._usedConstantBuffers;
             _usedStorageBuffers |= other._usedStorageBuffers;
             _usedStorageBuffersWrite |= other._usedStorageBuffersWrite;
@@ -191,6 +197,28 @@ namespace Ryujinx.Graphics.Shader.Translation
             }
         }
 
+        public void SetInputUserAttribute(int index)
+        {
+            UsedInputAttributes |= 1 << index;
+        }
+
+        public void SetOutputUserAttribute(int index)
+        {
+            UsedOutputAttributes |= 1 << index;
+        }
+
+        public void MergeOutputUserAttributes(int mask)
+        {
+            if (GpPassthrough)
+            {
+                PassthroughAttributes = mask & ~UsedOutputAttributes;
+            }
+            else
+            {
+                UsedOutputAttributes |= mask;
+            }
+        }
+
         public void SetClipDistanceWritten(int index)
         {
             ClipDistancesWritten |= (byte)(1 << index);
diff --git a/Ryujinx.Graphics.Shader/Translation/Translator.cs b/Ryujinx.Graphics.Shader/Translation/Translator.cs
index 685b6a20..f1e92d7c 100644
--- a/Ryujinx.Graphics.Shader/Translation/Translator.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Translator.cs
@@ -5,6 +5,7 @@ using Ryujinx.Graphics.Shader.StructuredIr;
 using Ryujinx.Graphics.Shader.Translation.Optimizations;
 using System;
 using System.Collections.Generic;
+using System.Numerics;
 
 using static Ryujinx.Graphics.Shader.IntermediateRepresentation.OperandHelper;
 
@@ -120,24 +121,17 @@ namespace Ryujinx.Graphics.Shader.Translation
             Block[][] cfg;
             ulong maxEndAddress = 0;
 
-            bool hasBindless;
-
             if ((options.Flags & TranslationFlags.Compute) != 0)
             {
                 config = new ShaderConfig(gpuAccessor, options, counts);
 
-                cfg = Decoder.Decode(gpuAccessor, address, out hasBindless);
+                cfg = Decoder.Decode(config, address);
             }
             else
             {
                 config = new ShaderConfig(new ShaderHeader(gpuAccessor, address), gpuAccessor, options, counts);
 
-                cfg = Decoder.Decode(gpuAccessor, address + HeaderSize, out hasBindless);
-            }
-
-            if (hasBindless)
-            {
-                config.SetUsedFeature(FeatureFlags.Bindless);
+                cfg = Decoder.Decode(config, address + HeaderSize);
             }
 
             for (int funcIndex = 0; funcIndex < cfg.Length; funcIndex++)
@@ -151,7 +145,7 @@ namespace Ryujinx.Graphics.Shader.Translation
                         maxEndAddress = block.EndAddress;
                     }
 
-                    if (!hasBindless)
+                    if (!config.UsedFeatures.HasFlag(FeatureFlags.Bindless))
                     {
                         for (int index = 0; index < block.OpCodes.Count; index++)
                         {
@@ -169,8 +163,10 @@ namespace Ryujinx.Graphics.Shader.Translation
             return cfg;
         }
 
-        internal static FunctionCode[] EmitShader(Block[][] cfg, ShaderConfig config)
+        internal static FunctionCode[] EmitShader(Block[][] cfg, ShaderConfig config, bool initializeOutputs, out int initializationOperations)
         {
+            initializationOperations = 0;
+
             Dictionary<ulong, int> funcIds = new Dictionary<ulong, int>();
 
             for (int funcIndex = 0; funcIndex < cfg.Length; funcIndex++)
@@ -184,6 +180,12 @@ namespace Ryujinx.Graphics.Shader.Translation
             {
                 EmitterContext context = new EmitterContext(config, funcIndex != 0, funcIds);
 
+                if (initializeOutputs && funcIndex == 0)
+                {
+                    EmitOutputsInitialization(context, config);
+                    initializationOperations = context.OperationsCount;
+                }
+
                 for (int blkIndex = 0; blkIndex < cfg[funcIndex].Length; blkIndex++)
                 {
                     Block block = cfg[funcIndex][blkIndex];
@@ -201,6 +203,39 @@ namespace Ryujinx.Graphics.Shader.Translation
             return funcs.ToArray();
         }
 
+        private static void EmitOutputsInitialization(EmitterContext context, ShaderConfig config)
+        {
+            // Compute has no output attributes, and fragment is the last stage, so we
+            // don't need to initialize outputs on those stages.
+            if (config.Stage == ShaderStage.Compute || config.Stage == ShaderStage.Fragment)
+            {
+                return;
+            }
+
+            void InitializeOutput(int baseAttr)
+            {
+                for (int c = 0; c < 4; c++)
+                {
+                    context.Copy(Attribute(baseAttr + c * 4), ConstF(c == 3 ? 1f : 0f));
+                }
+            }
+
+            if (config.Stage == ShaderStage.Vertex)
+            {
+                InitializeOutput(AttributeConsts.PositionX);
+            }
+
+            int usedAttribtes = context.Config.UsedOutputAttributes;
+            while (usedAttribtes != 0)
+            {
+                int index = BitOperations.TrailingZeroCount(usedAttribtes);
+
+                InitializeOutput(AttributeConsts.UserAttributeBase + index * 16);
+
+                usedAttribtes &= ~(1 << index);
+            }
+        }
+
         private static void EmitOps(EmitterContext context, Block block)
         {
             for (int opIndex = 0; opIndex < block.OpCodes.Count; opIndex++)
diff --git a/Ryujinx.Graphics.Shader/Translation/TranslatorContext.cs b/Ryujinx.Graphics.Shader/Translation/TranslatorContext.cs
index ab74d039..47cf0ac8 100644
--- a/Ryujinx.Graphics.Shader/Translation/TranslatorContext.cs
+++ b/Ryujinx.Graphics.Shader/Translation/TranslatorContext.cs
@@ -38,7 +38,7 @@ namespace Ryujinx.Graphics.Shader.Translation
                    operand.Value < AttributeConsts.UserAttributeEnd;
         }
 
-        private static FunctionCode[] Combine(FunctionCode[] a, FunctionCode[] b)
+        private static FunctionCode[] Combine(FunctionCode[] a, FunctionCode[] b, int aStart)
         {
             // Here we combine two shaders.
             // For shader A:
@@ -57,7 +57,7 @@ namespace Ryujinx.Graphics.Shader.Translation
 
             Operand lblB = Label();
 
-            for (int index = 0; index < a[0].Code.Length; index++)
+            for (int index = aStart; index < a[0].Code.Length; index++)
             {
                 Operation operation = a[0].Code[index];
 
@@ -103,7 +103,17 @@ namespace Ryujinx.Graphics.Shader.Translation
 
                         if (temp != null)
                         {
-                            operation.SetSource(srcIndex, temp);
+                            // TODO: LoadAttribute should accept any integer value as first argument,
+                            // then we don't need special case here. Right now it expects the first
+                            // operand to be of type "attribute".
+                            if ((operation.Inst & Instruction.Mask) == Instruction.LoadAttribute)
+                            {
+                                operation.TurnIntoCopy(temp);
+                            }
+                            else
+                            {
+                                operation.SetSource(srcIndex, temp);
+                            }
                         }
                     }
                 }
@@ -126,13 +136,25 @@ namespace Ryujinx.Graphics.Shader.Translation
             return output;
         }
 
-        public ShaderProgram Translate(out ShaderProgramInfo shaderProgramInfo, TranslatorContext other = null)
+        public ShaderProgram Translate(
+            out ShaderProgramInfo shaderProgramInfo,
+            TranslatorContext nextStage = null,
+            TranslatorContext other = null)
         {
-            FunctionCode[] code = EmitShader(_cfg, _config);
+            if (nextStage != null)
+            {
+                _config.MergeOutputUserAttributes(nextStage._config.UsedInputAttributes);
+            }
+
+            FunctionCode[] code = EmitShader(_cfg, _config, initializeOutputs: other == null, out _);
 
             if (other != null)
             {
-                code = Combine(EmitShader(other._cfg, other._config), code);
+                other._config.MergeOutputUserAttributes(_config.UsedOutputAttributes);
+
+                FunctionCode[] otherCode = EmitShader(other._cfg, other._config, initializeOutputs: true, out int aStart);
+
+                code = Combine(otherCode, code, aStart);
 
                 _config.InheritFrom(other._config);
             }