Initial support for image stores, support texture sample on compute

9 files changed, 176 insertions, 17 deletions

diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
index 322bfbf5..abfe55a5 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
@@ -12,6 +12,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
         public List<BufferDescriptor>  CBufferDescriptors { get; }
         public List<BufferDescriptor>  SBufferDescriptors { get; }
         public List<TextureDescriptor> TextureDescriptors { get; }
+        public List<TextureDescriptor> ImageDescriptors   { get; }
 
         public OperandManager OperandManager { get; }
 
@@ -28,6 +29,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             CBufferDescriptors = new List<BufferDescriptor>();
             SBufferDescriptors = new List<BufferDescriptor>();
             TextureDescriptors = new List<TextureDescriptor>();
+            ImageDescriptors   = new List<TextureDescriptor>();
 
             OperandManager = new OperandManager();
 
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
index 3644b21a..ab10d91a 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
@@ -88,6 +88,13 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                 context.AppendLine();
             }
 
+            if (info.Images.Count != 0)
+            {
+                DeclareImages(context, info);
+
+                context.AppendLine();
+            }
+
             if (context.Config.Stage != ShaderStage.Compute)
             {
                 if (info.IAttributes.Count != 0)
@@ -204,7 +211,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                     continue;
                 }
 
-                string samplerTypeName = GetSamplerTypeName(texOp.Target);
+                string samplerTypeName = GetSamplerTypeName(texOp.Type);
 
                 context.AppendLine("uniform " + samplerTypeName + " " + samplerName + ";");
             }
@@ -221,17 +228,47 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                 {
                     AstOperand operand = texOp.GetSource(0) as AstOperand;
 
-                    desc = new TextureDescriptor(samplerName, texOp.Target, operand.CbufSlot, operand.CbufOffset);
+                    desc = new TextureDescriptor(samplerName, texOp.Type, operand.CbufSlot, operand.CbufOffset);
                 }
                 else
                 {
-                    desc = new TextureDescriptor(samplerName, texOp.Target, texOp.Handle);
+                    desc = new TextureDescriptor(samplerName, texOp.Type, texOp.Handle);
                 }
 
                 context.TextureDescriptors.Add(desc);
             }
         }
 
+        private static void DeclareImages(CodeGenContext context, StructuredProgramInfo info)
+        {
+            Dictionary<string, AstTextureOperation> images = new Dictionary<string, AstTextureOperation>();
+
+            foreach (AstTextureOperation texOp in info.Images.OrderBy(x => x.Handle))
+            {
+                string imageName = OperandManager.GetImageName(context.Config.Stage, texOp);
+
+                if (!images.TryAdd(imageName, texOp))
+                {
+                    continue;
+                }
+
+                string imageTypeName = GetImageTypeName(texOp.Type);
+
+                context.AppendLine("writeonly uniform " + imageTypeName + " " + imageName + ";");
+            }
+
+            foreach (KeyValuePair<string, AstTextureOperation> kv in images)
+            {
+                string imageName = kv.Key;
+
+                AstTextureOperation texOp = kv.Value;
+
+                TextureDescriptor desc = new TextureDescriptor(imageName, texOp.Type, texOp.Handle);
+
+                context.ImageDescriptors.Add(desc);
+            }
+        }
+
         private static void DeclareInputAttributes(CodeGenContext context, StructuredProgramInfo info)
         {
             string suffix = context.Config.Stage == ShaderStage.Geometry ? "[]" : string.Empty;
@@ -284,36 +321,65 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             }
         }
 
-        private static string GetSamplerTypeName(TextureTarget type)
+        private static string GetSamplerTypeName(SamplerType type)
         {
             string typeName;
 
-            switch (type & TextureTarget.Mask)
+            switch (type & SamplerType.Mask)
             {
-                case TextureTarget.Texture1D:   typeName = "sampler1D";   break;
-                case TextureTarget.Texture2D:   typeName = "sampler2D";   break;
-                case TextureTarget.Texture3D:   typeName = "sampler3D";   break;
-                case TextureTarget.TextureCube: typeName = "samplerCube"; break;
+                case SamplerType.Texture1D:     typeName = "sampler1D";     break;
+                case SamplerType.TextureBuffer: typeName = "samplerBuffer"; break;
+                case SamplerType.Texture2D:     typeName = "sampler2D";     break;
+                case SamplerType.Texture3D:     typeName = "sampler3D";     break;
+                case SamplerType.TextureCube:   typeName = "samplerCube";   break;
 
                 default: throw new ArgumentException($"Invalid sampler type \"{type}\".");
             }
 
-            if ((type & TextureTarget.Multisample) != 0)
+            if ((type & SamplerType.Multisample) != 0)
             {
                 typeName += "MS";
             }
 
-            if ((type & TextureTarget.Array) != 0)
+            if ((type & SamplerType.Array) != 0)
             {
                 typeName += "Array";
             }
 
-            if ((type & TextureTarget.Shadow) != 0)
+            if ((type & SamplerType.Shadow) != 0)
             {
                 typeName += "Shadow";
             }
 
             return typeName;
         }
+
+        private static string GetImageTypeName(SamplerType type)
+        {
+            string typeName;
+
+            switch (type & SamplerType.Mask)
+            {
+                case SamplerType.Texture1D:     typeName = "image1D";     break;
+                case SamplerType.TextureBuffer: typeName = "imageBuffer"; break;
+                case SamplerType.Texture2D:     typeName = "image2D";     break;
+                case SamplerType.Texture3D:     typeName = "image3D";     break;
+                case SamplerType.TextureCube:   typeName = "imageCube";   break;
+
+                default: throw new ArgumentException($"Invalid sampler type \"{type}\".");
+            }
+
+            if ((type & SamplerType.Multisample) != 0)
+            {
+                typeName += "MS";
+            }
+
+            if ((type & SamplerType.Array) != 0)
+            {
+                typeName += "Array";
+            }
+
+            return typeName;
+        }
     }
 }
\ No newline at end of file
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs
index 67de3b43..a06b0cc8 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs
@@ -5,6 +5,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
         public const string LocalNamePrefix = "temp";
 
         public const string SamplerNamePrefix = "tex";
+        public const string ImageNamePrefix   = "img";
 
         public const string IAttributePrefix = "in_attr";
         public const string OAttributePrefix = "out_attr";
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
index 65246d97..b5407eb8 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
@@ -21,6 +21,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                 context.CBufferDescriptors.ToArray(),
                 context.SBufferDescriptors.ToArray(),
                 context.TextureDescriptors.ToArray(),
+                context.ImageDescriptors.ToArray(),
                 context.GetCode());
         }
 
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslProgram.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslProgram.cs
index 7807cb49..31b7f312 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslProgram.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslProgram.cs
@@ -5,6 +5,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
         public BufferDescriptor[]  CBufferDescriptors { get; }
         public BufferDescriptor[]  SBufferDescriptors { get; }
         public TextureDescriptor[] TextureDescriptors { get; }
+        public TextureDescriptor[] ImageDescriptors   { get; }
 
         public string Code { get; }
 
@@ -12,11 +13,13 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             BufferDescriptor[]  cBufferDescriptors,
             BufferDescriptor[]  sBufferDescriptors,
             TextureDescriptor[] textureDescriptors,
+            TextureDescriptor[] imageDescriptors,
             string              code)
         {
             CBufferDescriptors = cBufferDescriptors;
             SBufferDescriptors = sBufferDescriptors;
             TextureDescriptors = textureDescriptors;
+            ImageDescriptors   = imageDescriptors;
             Code               = code;
         }
     }
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
index 350de348..3bf31c16 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
@@ -87,6 +87,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             {
                 switch (inst)
                 {
+                    case Instruction.ImageStore:
+                        return InstGenMemory.ImageStore(context, operation);
+
                     case Instruction.LoadAttribute:
                         return InstGenMemory.LoadAttribute(context, operation);
 
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
index 6989e997..6dfbe61a 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
@@ -48,6 +48,8 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             Add(Instruction.ExponentB2,               InstType.CallUnary,      "exp2");
             Add(Instruction.Floor,                    InstType.CallUnary,      "floor");
             Add(Instruction.FusedMultiplyAdd,         InstType.CallTernary,    "fma");
+            Add(Instruction.ImageLoad,                InstType.Special);
+            Add(Instruction.ImageStore,               InstType.Special);
             Add(Instruction.IsNan,                    InstType.CallUnary,      "isnan");
             Add(Instruction.LoadAttribute,            InstType.Special);
             Add(Instruction.LoadConstant,             InstType.Special);
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
index 5c9ec830..f2f6ae0c 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
@@ -9,6 +9,80 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
 {
     static class InstGenMemory
     {
+        public static string ImageStore(CodeGenContext context, AstOperation operation)
+        {
+            AstTextureOperation texOp = (AstTextureOperation)operation;
+
+            bool isBindless = (texOp.Flags & TextureFlags.Bindless) != 0;
+
+            bool isArray = (texOp.Type & SamplerType.Array) != 0;
+
+            string texCall = "imageStore";
+
+            string imageName = OperandManager.GetImageName(context.Config.Stage, texOp);
+
+            texCall += "(" + imageName;
+
+            int coordsCount = texOp.Type.GetDimensions();
+
+            int pCount = coordsCount;
+
+            int arrayIndexElem = -1;
+
+            if (isArray)
+            {
+                arrayIndexElem = pCount++;
+            }
+
+            int srcIndex = isBindless ? 1 : 0;
+
+            string Src(VariableType type)
+            {
+                return GetSoureExpr(context, texOp.GetSource(srcIndex++), type);
+            }
+
+            void Append(string str)
+            {
+                texCall += ", " + str;
+            }
+
+            if (pCount > 1)
+            {
+                string[] elems = new string[pCount];
+
+                for (int index = 0; index < pCount; index++)
+                {
+                    elems[index] = Src(VariableType.S32);
+                }
+
+                Append("ivec" + pCount + "(" + string.Join(", ", elems) + ")");
+            }
+            else
+            {
+                Append(Src(VariableType.S32));
+            }
+
+            string[] cElems = new string[4];
+
+            for (int index = 0; index < 4; index++)
+            {
+                if (srcIndex < texOp.SourcesCount)
+                {
+                    cElems[index] = Src(VariableType.F32);
+                }
+                else
+                {
+                    cElems[index] = NumberFormatter.FormatFloat(0);
+                }
+            }
+
+            Append("vec4(" + string.Join(", ", cElems) + ")");
+
+            texCall += ")";
+
+            return texCall;
+        }
+
         public static string LoadAttribute(CodeGenContext context, AstOperation operation)
         {
             IAstNode src1 = operation.GetSource(0);
@@ -98,9 +172,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             bool hasOffset   = (texOp.Flags & TextureFlags.Offset)    != 0;
             bool hasOffsets  = (texOp.Flags & TextureFlags.Offsets)   != 0;
 
-            bool isArray       = (texOp.Target & TextureTarget.Array)       != 0;
-            bool isMultisample = (texOp.Target & TextureTarget.Multisample) != 0;
-            bool isShadow      = (texOp.Target & TextureTarget.Shadow)      != 0;
+            bool isArray       = (texOp.Type & SamplerType.Array)       != 0;
+            bool isMultisample = (texOp.Type & SamplerType.Multisample) != 0;
+            bool isShadow      = (texOp.Type & SamplerType.Shadow)      != 0;
 
             // This combination is valid, but not available on GLSL.
             // For now, ignore the LOD level and do a normal sample.
@@ -134,7 +208,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
 
             texCall += "(" + samplerName;
 
-            int coordsCount = texOp.Target.GetDimensions();
+            int coordsCount = texOp.Type.GetDimensions();
 
             int pCount = coordsCount;
 
@@ -147,7 +221,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
 
             // The sampler 1D shadow overload expects a
             // dummy value on the middle of the vector, who knows why...
-            bool hasDummy1DShadowElem = texOp.Target == (TextureTarget.Texture1D | TextureTarget.Shadow);
+            bool hasDummy1DShadowElem = texOp.Type == (SamplerType.Texture1D | SamplerType.Shadow);
 
             if (hasDummy1DShadowElem)
             {
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
index 88095920..fe307396 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
@@ -241,6 +241,13 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return GetShaderStagePrefix(stage) + "_" + DefaultNames.SamplerNamePrefix + suffix;
         }
 
+        public static string GetImageName(ShaderStage stage, AstTextureOperation texOp)
+        {
+            string suffix = texOp.Handle.ToString();
+
+            return GetShaderStagePrefix(stage) + "_" + DefaultNames.ImageNamePrefix + suffix;
+        }
+
         public static string GetShaderStagePrefix(ShaderStage stage)
         {
             int index = (int)stage;