Support res scale on images, correctly blacklist for SUST, move logic out of backend. (#1657)

* Support res scale on images, correctly blacklist for SUST, move logic
out of backend.

* Fix Typo

3 files changed, 69 insertions, 14 deletions

diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
index 85347dfd..f86e82da 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
@@ -84,7 +84,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             AppendLine("}" + suffix);
         }
 
-        public int FindTextureDescriptorIndex(AstTextureOperation texOp)
+        private int FindDescriptorIndex(List<TextureDescriptor> list, AstTextureOperation texOp)
         {
             AstOperand operand = texOp.GetSource(0) as AstOperand;
             bool bindless = (texOp.Flags & TextureFlags.Bindless) > 0;
@@ -92,13 +92,24 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             int cBufSlot = bindless ? operand.CbufSlot : 0;
             int cBufOffset = bindless ? operand.CbufOffset : 0;
 
-            return TextureDescriptors.FindIndex(descriptor =>
+            return list.FindIndex(descriptor =>
                 descriptor.Type == texOp.Type &&
                 descriptor.HandleIndex == texOp.Handle &&
+                descriptor.Format == texOp.Format &&
                 descriptor.CbufSlot == cBufSlot &&
                 descriptor.CbufOffset == cBufOffset);
         }
 
+        public int FindTextureDescriptorIndex(AstTextureOperation texOp)
+        {
+            return FindDescriptorIndex(TextureDescriptors, texOp);
+        }
+
+        public int FindImageDescriptorIndex(AstTextureOperation texOp)
+        {
+            return FindDescriptorIndex(ImageDescriptors, texOp);
+        }
+
         public StructuredFunction GetFunction(int id)
         {
             return _info.Functions[id];
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
index 0250d852..734546a0 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
@@ -509,7 +509,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             {
                 string stage = OperandManager.GetShaderStagePrefix(context.Config.Stage);
 
-                int scaleElements = context.TextureDescriptors.Count;
+                int scaleElements = context.TextureDescriptors.Count + context.ImageDescriptors.Count;
 
                 if (context.Config.Stage == ShaderStage.Fragment)
                 {
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
index 456bfc4e..f10eb101 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
@@ -47,6 +47,43 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                 texCall += ", " + str;
             }
 
+            string ApplyScaling(string vector)
+            {
+                int index = context.FindImageDescriptorIndex(texOp);
+                TextureUsageFlags flags = TextureUsageFlags.NeedsScaleValue;
+
+                if ((context.Config.Stage == ShaderStage.Fragment || context.Config.Stage == ShaderStage.Compute) &&
+                    texOp.Inst == Instruction.ImageLoad &&
+                    !isBindless &&
+                    !isIndexed)
+                {
+                    // Image scales start after texture ones.
+                    int scaleIndex = context.TextureDescriptors.Count + index;
+
+                    if (pCount == 3 && isArray)
+                    {
+                        // The array index is not scaled, just x and y.
+                        vector = "ivec3(Helper_TexelFetchScale((" + vector + ").xy, " + scaleIndex + "), (" + vector + ").z)";
+                    }
+                    else if (pCount == 2 && !isArray)
+                    {
+                        vector = "Helper_TexelFetchScale(" + vector + ", " + scaleIndex + ")";
+                    }
+                    else
+                    {
+                        flags |= TextureUsageFlags.ResScaleUnsupported;
+                    }
+                }
+                else
+                {
+                    flags |= TextureUsageFlags.ResScaleUnsupported;
+                }
+
+                context.ImageDescriptors[index] = context.ImageDescriptors[index].SetFlag(flags);
+
+                return vector;
+            }
+
             if (pCount > 1)
             {
                 string[] elems = new string[pCount];
@@ -56,7 +93,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                     elems[index] = Src(VariableType.S32);
                 }
 
-                Append("ivec" + pCount + "(" + string.Join(", ", elems) + ")");
+                Append(ApplyScaling("ivec" + pCount + "(" + string.Join(", ", elems) + ")"));
             }
             else
             {
@@ -404,28 +441,35 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                 if (intCoords)
                 {
                     int index = context.FindTextureDescriptorIndex(texOp);
+                    TextureUsageFlags flags = TextureUsageFlags.NeedsScaleValue;
 
                     if ((context.Config.Stage == ShaderStage.Fragment || context.Config.Stage == ShaderStage.Compute) &&
-                        (texOp.Flags & TextureFlags.Bindless) == 0 &&
-                        texOp.Type != SamplerType.Indexed)
+                        !isBindless &&
+                        !isIndexed)
                     {
                         if (pCount == 3 && isArray)
                         {
                             // The array index is not scaled, just x and y.
-                            return "ivec3(Helper_TexelFetchScale((" + vector + ").xy, " + index + "), (" + vector + ").z)";
+                            vector = "ivec3(Helper_TexelFetchScale((" + vector + ").xy, " + index + "), (" + vector + ").z)";
                         }
                         else if (pCount == 2 && !isArray)
                         {
-                            return "Helper_TexelFetchScale(" + vector + ", " + index + ")";
+                            vector = "Helper_TexelFetchScale(" + vector + ", " + index + ")";
                         }
-                    }
+                        else
+                        {
+                            flags |= TextureUsageFlags.ResScaleUnsupported;
+                        }
+                    } 
+                    else
+                    {
+                        // Resolution scaling cannot be applied to this texture right now.
+                        // Flag so that we know to blacklist scaling on related textures when binding them.
 
-                    // Resolution scaling cannot be applied to this texture right now.
-                    // Flag so that we know to blacklist scaling on related textures when binding them.
+                        flags |= TextureUsageFlags.ResScaleUnsupported;
+                    }
 
-                    TextureDescriptor descriptor = context.TextureDescriptors[index];
-                    descriptor.Flags |= TextureUsageFlags.ResScaleUnsupported;
-                    context.TextureDescriptors[index] = descriptor;
+                    context.TextureDescriptors[index] = context.TextureDescriptors[index].SetFlag(flags);
                 }
 
                 return vector;