From 324d0528d3326be3412699095294de937db4b5e5 Mon Sep 17 00:00:00 2001
From: gdk <gab.dark.100@gmail.com>
Date: Wed, 6 Apr 2022 21:04:45 -0300
Subject: [PATCH] SPIR-V: Fix TextureSize for MS and Buffer sampler types

---
 .../CodeGen/Spirv/CodeGenContext.cs           |  1 +
 .../CodeGen/Spirv/Declarations.cs             |  1 +
 .../CodeGen/Spirv/Instructions.cs             | 23 ++++++++++++++-----
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/Ryujinx.Graphics.Shader/CodeGen/Spirv/CodeGenContext.cs b/Ryujinx.Graphics.Shader/CodeGen/Spirv/CodeGenContext.cs
index cdfef6ba0..e6ef4933e 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Spirv/CodeGenContext.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Spirv/CodeGenContext.cs
@@ -25,6 +25,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
         public Instruction SharedMemory { get; set; }
         public Instruction InputsArray { get; set; }
         public Instruction OutputsArray { get; set; }
+        public Dictionary<TextureMeta, SamplerType> SamplersTypes { get; } = new Dictionary<TextureMeta, SamplerType>();
         public Dictionary<TextureMeta, (Instruction, Instruction, Instruction)> Samplers { get; } = new Dictionary<TextureMeta, (Instruction, Instruction, Instruction)>();
         public Dictionary<TextureMeta, (Instruction, Instruction)> Images { get; } = new Dictionary<TextureMeta, (Instruction, Instruction)>();
         public Dictionary<int, Instruction> Inputs { get; } = new Dictionary<int, Instruction>();
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Spirv/Declarations.cs b/Ryujinx.Graphics.Shader/CodeGen/Spirv/Declarations.cs
index 573c9feeb..f6c81a3a2 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Spirv/Declarations.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Spirv/Declarations.cs
@@ -213,6 +213,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
                 var sampledImageVariable = context.Variable(sampledImagePointerType, StorageClass.UniformConstant);
 
                 context.Samplers.Add(meta, (imageType, sampledImageType, sampledImageVariable));
+                context.SamplersTypes.Add(meta, descriptor.Type);
 
                 context.Name(sampledImageVariable, $"{GetStagePrefix(context.Config.Stage)}_tex{nameSuffix}");
                 context.Decorate(sampledImageVariable, Decoration.DescriptorSet, (LiteralInteger)setIndex);
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Spirv/Instructions.cs b/Ryujinx.Graphics.Shader/CodeGen/Spirv/Instructions.cs
index e10df30d4..db1bdd49f 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Spirv/Instructions.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Spirv/Instructions.cs
@@ -1662,10 +1662,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
                 index = context.GetS32(texOp.GetSource(0));
             }
 
-            int lodSrcIndex = isBindless || isIndexed ? 1 : 0;
-
-            var lod = context.GetS32(operation.GetSource(lodSrcIndex));
-
             var meta = new TextureMeta(texOp.CbufSlot, texOp.Handle, texOp.Format);
 
             (var imageType, var sampledImageType, var sampledImageVariable) = context.Samplers[meta];
@@ -1679,9 +1675,24 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
             }
             else
             {
-                int components = (texOp.Type & SamplerType.Mask) == SamplerType.TextureCube ? 2 : texOp.Type.GetDimensions();
+                var type = context.SamplersTypes[meta];
+                bool hasLod = !type.HasFlag(SamplerType.Multisample) && type != SamplerType.TextureBuffer;
+
+                int components = (type & SamplerType.Mask) == SamplerType.TextureCube ? 2 : type.GetDimensions();
                 var resultType = components == 1 ? context.TypeS32() : context.TypeVector(context.TypeS32(), components);
-                var result = context.ImageQuerySizeLod(resultType, image, lod);
+
+                SpvInstruction result;
+
+                if (hasLod)
+                {
+                    int lodSrcIndex = isBindless || isIndexed ? 1 : 0;
+                    var lod = context.GetS32(operation.GetSource(lodSrcIndex));
+                    result = context.ImageQuerySizeLod(resultType, image, lod);
+                }
+                else
+                {
+                    result = context.ImageQuerySize(resultType, image);
+                }
 
                 if (components != 1)
                 {