14 files changed, 172 insertions, 21 deletions

diff --git a/src/shader_recompiler/CMakeLists.txt b/src/shader_recompiler/CMakeLists.txt
index 5f3868bfe..7f8dc8eed 100644
--- a/src/shader_recompiler/CMakeLists.txt
+++ b/src/shader_recompiler/CMakeLists.txt
@@ -126,6 +126,7 @@ add_library(shader_recompiler STATIC
     frontend/maxwell/translate/impl/texture_fetch_swizzled.cpp
     frontend/maxwell/translate/impl/texture_gather_swizzled.cpp
     frontend/maxwell/translate/impl/texture_gather.cpp
+    frontend/maxwell/translate/impl/texture_query.cpp
     frontend/maxwell/translate/impl/vote.cpp
     frontend/maxwell/translate/impl/warp_shuffle.cpp
     frontend/maxwell/translate/translate.cpp
diff --git a/src/shader_recompiler/backend/spirv/emit_context.cpp b/src/shader_recompiler/backend/spirv/emit_context.cpp
index 50793b5bf..c2d13f97c 100644
--- a/src/shader_recompiler/backend/spirv/emit_context.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_context.cpp
@@ -244,8 +244,9 @@ void EmitContext::DefineTextures(const Info& info, u32& binding) {
         if (desc.count != 1) {
             throw NotImplementedException("Array of textures");
         }
-        const Id type{TypeSampledImage(ImageType(*this, desc))};
-        const Id pointer_type{TypePointer(spv::StorageClass::UniformConstant, type)};
+        const Id image_type{ImageType(*this, desc)};
+        const Id sampled_type{TypeSampledImage(image_type)};
+        const Id pointer_type{TypePointer(spv::StorageClass::UniformConstant, sampled_type)};
         const Id id{AddGlobalVariable(pointer_type, spv::StorageClass::UniformConstant)};
         Decorate(id, spv::Decoration::Binding, binding);
         Decorate(id, spv::Decoration::DescriptorSet, 0U);
@@ -254,7 +255,8 @@ void EmitContext::DefineTextures(const Info& info, u32& binding) {
             // TODO: Pass count info
             textures.push_back(TextureDefinition{
                 .id{id},
-                .type{type},
+                .sampled_type{sampled_type},
+                .image_type{image_type},
             });
         }
         binding += desc.count;
diff --git a/src/shader_recompiler/backend/spirv/emit_context.h b/src/shader_recompiler/backend/spirv/emit_context.h
index 5ed815c06..0cb411a0e 100644
--- a/src/shader_recompiler/backend/spirv/emit_context.h
+++ b/src/shader_recompiler/backend/spirv/emit_context.h
@@ -31,7 +31,8 @@ private:
 
 struct TextureDefinition {
     Id id;
-    Id type;
+    Id sampled_type;
+    Id image_type;
 };
 
 struct UniformDefinitions {
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.cpp b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
index cee72f50d..4bed16e7b 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
@@ -126,10 +126,10 @@ Id DefineMain(EmitContext& ctx, IR::Program& program) {
     return main;
 }
 
-void DefineEntryPoint(Environment& env, EmitContext& ctx, Id main) {
+void DefineEntryPoint(Environment& env, const IR::Program& program, EmitContext& ctx, Id main) {
     const std::span interfaces(ctx.interfaces.data(), ctx.interfaces.size());
     spv::ExecutionModel execution_model{};
-    switch (env.ShaderStage()) {
+    switch (program.stage) {
     case Shader::Stage::Compute: {
         const std::array<u32, 3> workgroup_size{env.WorkgroupSize()};
         execution_model = spv::ExecutionModel::GLCompute;
@@ -143,6 +143,9 @@ void DefineEntryPoint(Environment& env, EmitContext& ctx, Id main) {
     case Shader::Stage::Fragment:
         execution_model = spv::ExecutionModel::Fragment;
         ctx.AddExecutionMode(main, spv::ExecutionMode::OriginUpperLeft);
+        if (program.info.stores_frag_depth) {
+            ctx.AddExecutionMode(main, spv::ExecutionMode::DepthReplacing);
+        }
         break;
     default:
         throw NotImplementedException("Stage {}", env.ShaderStage());
@@ -235,6 +238,7 @@ void SetupCapabilities(const Profile& profile, const Info& info, EmitContext& ct
     }
     // TODO: Track this usage
     ctx.AddCapability(spv::Capability::ImageGatherExtended);
+    ctx.AddCapability(spv::Capability::ImageQuery);
 }
 
 Id PhiArgDef(EmitContext& ctx, IR::Inst* inst, size_t index) {
@@ -267,7 +271,7 @@ std::vector<u32> EmitSPIRV(const Profile& profile, Environment& env, IR::Program
                            u32& binding) {
     EmitContext ctx{profile, program, binding};
     const Id main{DefineMain(ctx, program)};
-    DefineEntryPoint(env, ctx, main);
+    DefineEntryPoint(env, program, ctx, main);
     if (profile.support_float_controls) {
         ctx.AddExtension("SPV_KHR_float_controls");
         SetupDenormControl(profile, program, ctx, main);
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.h b/src/shader_recompiler/backend/spirv/emit_spirv.h
index 4da1f3707..b82b16e9d 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.h
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.h
@@ -343,6 +343,7 @@ Id EmitBindlessImageSampleDrefExplicitLod(EmitContext&);
 Id EmitBindlessImageGather(EmitContext&);
 Id EmitBindlessImageGatherDref(EmitContext&);
 Id EmitBindlessImageFetch(EmitContext&);
+Id EmitBindlessImageQueryDimensions(EmitContext&);
 Id EmitBoundImageSampleImplicitLod(EmitContext&);
 Id EmitBoundImageSampleExplicitLod(EmitContext&);
 Id EmitBoundImageSampleDrefImplicitLod(EmitContext&);
@@ -350,6 +351,7 @@ Id EmitBoundImageSampleDrefExplicitLod(EmitContext&);
 Id EmitBoundImageGather(EmitContext&);
 Id EmitBoundImageGatherDref(EmitContext&);
 Id EmitBoundImageFetch(EmitContext&);
+Id EmitBoundImageQueryDimensions(EmitContext&);
 Id EmitImageSampleImplicitLod(EmitContext& ctx, IR::Inst* inst, const IR::Value& index, Id coords,
                               Id bias_lc, Id offset);
 Id EmitImageSampleExplicitLod(EmitContext& ctx, IR::Inst* inst, const IR::Value& index, Id coords,
@@ -364,6 +366,7 @@ Id EmitImageGatherDref(EmitContext& ctx, IR::Inst* inst, const IR::Value& index,
                        const IR::Value& offset, const IR::Value& offset2, Id dref);
 Id EmitImageFetch(EmitContext& ctx, IR::Inst* inst, const IR::Value& index, Id coords, Id offset,
                   Id lod, Id ms);
+Id EmitImageQueryDimensions(EmitContext& ctx, IR::Inst* inst, const IR::Value& index, Id lod);
 Id EmitVoteAll(EmitContext& ctx, Id pred);
 Id EmitVoteAny(EmitContext& ctx, Id pred);
 Id EmitVoteEqual(EmitContext& ctx, Id pred);
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_image.cpp b/src/shader_recompiler/backend/spirv/emit_spirv_image.cpp
index b6e9d3c0c..3ea0011aa 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv_image.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_image.cpp
@@ -91,7 +91,15 @@ private:
 Id Texture(EmitContext& ctx, const IR::Value& index) {
     if (index.IsImmediate()) {
         const TextureDefinition def{ctx.textures.at(index.U32())};
-        return ctx.OpLoad(def.type, def.id);
+        return ctx.OpLoad(def.sampled_type, def.id);
+    }
+    throw NotImplementedException("Indirect texture sample");
+}
+
+Id TextureImage(EmitContext& ctx, const IR::Value& index) {
+    if (index.IsImmediate()) {
+        const TextureDefinition def{ctx.textures.at(index.U32())};
+        return ctx.OpImage(def.image_type, ctx.OpLoad(def.sampled_type, def.id));
     }
     throw NotImplementedException("Indirect texture sample");
 }
@@ -149,6 +157,10 @@ Id EmitBindlessImageFetch(EmitContext&) {
     throw LogicError("Unreachable instruction");
 }
 
+Id EmitBindlessImageQueryDimensions(EmitContext&) {
+    throw LogicError("Unreachable instruction");
+}
+
 Id EmitBoundImageSampleImplicitLod(EmitContext&) {
     throw LogicError("Unreachable instruction");
 }
@@ -177,6 +189,10 @@ Id EmitBoundImageFetch(EmitContext&) {
     throw LogicError("Unreachable instruction");
 }
 
+Id EmitBoundImageQueryDimensions(EmitContext&) {
+    throw LogicError("Unreachable instruction");
+}
+
 Id EmitImageSampleImplicitLod(EmitContext& ctx, IR::Inst* inst, const IR::Value& index, Id coords,
                               Id bias_lc, Id offset) {
     const auto info{inst->Flags<IR::TextureInstInfo>()};
@@ -241,4 +257,34 @@ Id EmitImageFetch(EmitContext& ctx, IR::Inst* inst, const IR::Value& index, Id c
                 Texture(ctx, index), coords, operands.Mask(), operands.Span());
 }
 
+Id EmitImageQueryDimensions(EmitContext& ctx, IR::Inst* inst, const IR::Value& index, Id lod) {
+    const auto info{inst->Flags<IR::TextureInstInfo>()};
+    const Id image{TextureImage(ctx, index)};
+    const Id zero{ctx.u32_zero_value};
+    const auto mips{[&] { return ctx.OpImageQueryLevels(ctx.U32[1], image); }};
+    switch (info.type) {
+    case TextureType::Color1D:
+    case TextureType::Shadow1D:
+        return ctx.OpCompositeConstruct(ctx.U32[4], ctx.OpImageQuerySizeLod(ctx.U32[1], image, lod),
+                                        zero, zero, mips());
+    case TextureType::ColorArray1D:
+    case TextureType::Color2D:
+    case TextureType::ColorCube:
+    case TextureType::ShadowArray1D:
+    case TextureType::Shadow2D:
+    case TextureType::ShadowCube:
+        return ctx.OpCompositeConstruct(ctx.U32[4], ctx.OpImageQuerySizeLod(ctx.U32[2], image, lod),
+                                        zero, mips());
+    case TextureType::ColorArray2D:
+    case TextureType::Color3D:
+    case TextureType::ColorArrayCube:
+    case TextureType::ShadowArray2D:
+    case TextureType::Shadow3D:
+    case TextureType::ShadowArrayCube:
+        return ctx.OpCompositeConstruct(ctx.U32[4], ctx.OpImageQuerySizeLod(ctx.U32[3], image, lod),
+                                        mips());
+    }
+    throw LogicError("Unspecified image type {}", info.type.Value());
+}
+
 } // namespace Shader::Backend::SPIRV
diff --git a/src/shader_recompiler/environment.h b/src/shader_recompiler/environment.h
index 6dec4b255..0c62c1c54 100644
--- a/src/shader_recompiler/environment.h
+++ b/src/shader_recompiler/environment.h
@@ -4,6 +4,7 @@
 
 #include "common/common_types.h"
 #include "shader_recompiler/program_header.h"
+#include "shader_recompiler/shader_info.h"
 #include "shader_recompiler/stage.h"
 
 namespace Shader {
@@ -14,6 +15,8 @@ public:
 
     [[nodiscard]] virtual u64 ReadInstruction(u32 address) = 0;
 
+    [[nodiscard]] virtual TextureType ReadTextureType(u32 cbuf_index, u32 cbuf_offset) = 0;
+
     [[nodiscard]] virtual u32 TextureBoundBuffer() const = 0;
 
     [[nodiscard]] virtual std::array<u32, 3> WorkgroupSize() const = 0;
diff --git a/src/shader_recompiler/frontend/ir/ir_emitter.cpp b/src/shader_recompiler/frontend/ir/ir_emitter.cpp
index 0296f8773..f281c023f 100644
--- a/src/shader_recompiler/frontend/ir/ir_emitter.cpp
+++ b/src/shader_recompiler/frontend/ir/ir_emitter.cpp
@@ -1493,6 +1493,12 @@ Value IREmitter::ImageFetch(const Value& handle, const Value& coords, const Valu
     return Inst(op, Flags{info}, handle, coords, offset, lod, multisampling);
 }
 
+Value IREmitter::ImageQueryDimension(const Value& handle, const IR::U32& lod) {
+    const Opcode op{handle.IsImmediate() ? Opcode::BoundImageQueryDimensions
+                                         : Opcode::BindlessImageQueryDimensions};
+    return Inst(op, handle, lod);
+}
+
 U1 IREmitter::VoteAll(const U1& value) {
     return Inst<U1>(Opcode::VoteAll, value);
 }
diff --git a/src/shader_recompiler/frontend/ir/ir_emitter.h b/src/shader_recompiler/frontend/ir/ir_emitter.h
index 446fd7785..771c186d4 100644
--- a/src/shader_recompiler/frontend/ir/ir_emitter.h
+++ b/src/shader_recompiler/frontend/ir/ir_emitter.h
@@ -239,6 +239,7 @@ public:
                                                  const F32& dref, const F32& lod,
                                                  const Value& offset, const F32& lod_clamp,
                                                  TextureInstInfo info);
+    [[nodiscard]] Value ImageQueryDimension(const Value& handle, const IR::U32& lod);
 
     [[nodiscard]] Value ImageGather(const Value& handle, const Value& coords, const Value& offset,
                                     const Value& offset2, TextureInstInfo info);
diff --git a/src/shader_recompiler/frontend/ir/opcodes.inc b/src/shader_recompiler/frontend/ir/opcodes.inc
index e12b92c47..5d7462d76 100644
--- a/src/shader_recompiler/frontend/ir/opcodes.inc
+++ b/src/shader_recompiler/frontend/ir/opcodes.inc
@@ -356,6 +356,7 @@ OPCODE(BindlessImageSampleDrefExplicitLod,                  F32,            U32,
 OPCODE(BindlessImageGather,                                 F32x4,          U32,            Opaque,         Opaque,         Opaque,                         )
 OPCODE(BindlessImageGatherDref,                             F32x4,          U32,            Opaque,         Opaque,         Opaque,         F32,            )
 OPCODE(BindlessImageFetch,                                  F32x4,          U32,            Opaque,         U32,            U32,                            )
+OPCODE(BindlessImageQueryDimensions,                        U32x4,          U32,            U32,                                                            )
 
 OPCODE(BoundImageSampleImplicitLod,                         F32x4,          U32,            Opaque,         Opaque,         Opaque,                         )
 OPCODE(BoundImageSampleExplicitLod,                         F32x4,          U32,            Opaque,         Opaque,         Opaque,                         )
@@ -364,6 +365,7 @@ OPCODE(BoundImageSampleDrefExplicitLod,                     F32,            U32,
 OPCODE(BoundImageGather,                                    F32x4,          U32,            Opaque,         Opaque,         Opaque,                         )
 OPCODE(BoundImageGatherDref,                                F32x4,          U32,            Opaque,         Opaque,         Opaque,         F32,            )
 OPCODE(BoundImageFetch,                                     F32x4,          U32,            Opaque,         U32,            U32,                            )
+OPCODE(BoundImageQueryDimensions,                           U32x4,          U32,            U32,                                                            )
 
 OPCODE(ImageSampleImplicitLod,                              F32x4,          U32,            Opaque,         Opaque,         Opaque,                         )
 OPCODE(ImageSampleExplicitLod,                              F32x4,          U32,            Opaque,         Opaque,         Opaque,                         )
@@ -372,6 +374,7 @@ OPCODE(ImageSampleDrefExplicitLod,                          F32,            U32,
 OPCODE(ImageGather,                                         F32x4,          U32,            Opaque,         Opaque,         Opaque,                         )
 OPCODE(ImageGatherDref,                                     F32x4,          U32,            Opaque,         Opaque,         Opaque,         F32,            )
 OPCODE(ImageFetch,                                          F32x4,          U32,            Opaque,         U32,            U32,                            )
+OPCODE(ImageQueryDimensions,                                U32x4,          U32,            U32,                                                            )
 
 // Warp operations
 OPCODE(VoteAll,                                             U1,             U1,                                                                             )
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp
index 788765c21..96ee2e741 100644
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp
@@ -373,14 +373,6 @@ void TranslatorVisitor::TXD_b(u64) {
     ThrowNotImplemented(Opcode::TXD_b);
 }
 
-void TranslatorVisitor::TXQ(u64) {
-    ThrowNotImplemented(Opcode::TXQ);
-}
-
-void TranslatorVisitor::TXQ_b(u64) {
-    ThrowNotImplemented(Opcode::TXQ_b);
-}
-
 void TranslatorVisitor::VABSDIFF(u64) {
     ThrowNotImplemented(Opcode::VABSDIFF);
 }
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/texture_query.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/texture_query.cpp
new file mode 100644
index 000000000..e8ea8faeb
--- /dev/null
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/texture_query.cpp
@@ -0,0 +1,76 @@
+// Copyright 2021 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#include <optional>
+
+#include "common/bit_field.h"
+#include "common/common_types.h"
+#include "shader_recompiler/frontend/ir/modifiers.h"
+#include "shader_recompiler/frontend/maxwell/translate/impl/impl.h"
+
+namespace Shader::Maxwell {
+namespace {
+enum class Mode : u64 {
+    Dimension = 1,
+    TextureType = 2,
+    SamplePos = 5,
+};
+
+IR::Value Query(TranslatorVisitor& v, const IR::U32& handle, Mode mode, IR::Reg src_reg) {
+    switch (mode) {
+    case Mode::Dimension: {
+        const IR::U32 lod{v.X(src_reg)};
+        return v.ir.ImageQueryDimension(handle, lod);
+    }
+    case Mode::TextureType:
+    case Mode::SamplePos:
+    default:
+        throw NotImplementedException("Mode {}", mode);
+    }
+}
+
+void Impl(TranslatorVisitor& v, u64 insn, std::optional<u32> cbuf_offset) {
+    union {
+        u64 raw;
+        BitField<49, 1, u64> nodep;
+        BitField<0, 8, IR::Reg> dest_reg;
+        BitField<8, 8, IR::Reg> src_reg;
+        BitField<22, 3, Mode> mode;
+        BitField<31, 4, u64> mask;
+    } const txq{insn};
+
+    IR::Reg src_reg{txq.src_reg};
+    IR::U32 handle;
+    if (cbuf_offset) {
+        handle = v.ir.Imm32(*cbuf_offset);
+    } else {
+        handle = v.X(src_reg);
+        ++src_reg;
+    }
+    const IR::Value query{Query(v, handle, txq.mode, src_reg)};
+    IR::Reg dest_reg{txq.dest_reg};
+    for (int element = 0; element < 4; ++element) {
+        if (((txq.mask >> element) & 1) == 0) {
+            continue;
+        }
+        v.X(dest_reg, IR::U32{v.ir.CompositeExtract(query, element)});
+        ++dest_reg;
+    }
+}
+} // Anonymous namespace
+
+void TranslatorVisitor::TXQ(u64 insn) {
+    union {
+        u64 raw;
+        BitField<36, 13, u64> cbuf_offset;
+    } const txq{insn};
+
+    Impl(*this, insn, static_cast<u32>(txq.cbuf_offset));
+}
+
+void TranslatorVisitor::TXQ_b(u64 insn) {
+    Impl(*this, insn, std::nullopt);
+}
+
+} // namespace Shader::Maxwell
diff --git a/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp b/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
index 6fe06fda8..80ca8db26 100644
--- a/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
+++ b/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
@@ -365,7 +365,8 @@ void VisitUsages(Info& info, IR::Inst& inst) {
     case IR::Opcode::ImageSampleDrefImplicitLod:
     case IR::Opcode::ImageSampleDrefExplicitLod:
     case IR::Opcode::ImageGather:
-    case IR::Opcode::ImageGatherDref: {
+    case IR::Opcode::ImageGatherDref:
+    case IR::Opcode::ImageQueryDimensions: {
         const TextureType type{inst.Flags<IR::TextureInstInfo>().type};
         info.uses_sampled_1d |= type == TextureType::Color1D || type == TextureType::ColorArray1D ||
                                 type == TextureType::Shadow1D || type == TextureType::ShadowArray1D;
diff --git a/src/shader_recompiler/ir_opt/texture_pass.cpp b/src/shader_recompiler/ir_opt/texture_pass.cpp
index 0167dd06e..dfacf848f 100644
--- a/src/shader_recompiler/ir_opt/texture_pass.cpp
+++ b/src/shader_recompiler/ir_opt/texture_pass.cpp
@@ -54,6 +54,9 @@ IR::Opcode IndexedInstruction(const IR::Inst& inst) {
     case IR::Opcode::BindlessImageFetch:
     case IR::Opcode::BoundImageFetch:
         return IR::Opcode::ImageFetch;
+    case IR::Opcode::BoundImageQueryDimensions:
+    case IR::Opcode::BindlessImageQueryDimensions:
+        return IR::Opcode::ImageQueryDimensions;
     default:
         return IR::Opcode::Void;
     }
@@ -68,6 +71,7 @@ bool IsBindless(const IR::Inst& inst) {
     case IR::Opcode::BindlessImageGather:
     case IR::Opcode::BindlessImageGatherDref:
     case IR::Opcode::BindlessImageFetch:
+    case IR::Opcode::BindlessImageQueryDimensions:
         return true;
     case IR::Opcode::BoundImageSampleImplicitLod:
     case IR::Opcode::BoundImageSampleExplicitLod:
@@ -76,6 +80,7 @@ bool IsBindless(const IR::Inst& inst) {
     case IR::Opcode::BoundImageGather:
     case IR::Opcode::BoundImageGatherDref:
     case IR::Opcode::BoundImageFetch:
+    case IR::Opcode::BoundImageQueryDimensions:
         return false;
     default:
         throw InvalidArgument("Invalid opcode {}", inst.Opcode());
@@ -198,13 +203,20 @@ void TexturePass(Environment& env, IR::Program& program) {
     for (TextureInst& texture_inst : to_replace) {
         // TODO: Handle arrays
         IR::Inst* const inst{texture_inst.inst};
+        inst->ReplaceOpcode(IndexedInstruction(*inst));
+
+        const auto& cbuf{texture_inst.cbuf};
+        auto flags{inst->Flags<IR::TextureInstInfo>()};
+        if (inst->Opcode() == IR::Opcode::ImageQueryDimensions) {
+            flags.type.Assign(env.ReadTextureType(cbuf.index, cbuf.offset));
+            inst->SetFlags(flags);
+        }
         const u32 index{descriptors.Add(TextureDescriptor{
-            .type{inst->Flags<IR::TextureInstInfo>().type},
-            .cbuf_index{texture_inst.cbuf.index},
-            .cbuf_offset{texture_inst.cbuf.offset},
+            .type{flags.type},
+            .cbuf_index{cbuf.index},
+            .cbuf_offset{cbuf.offset},
             .count{1},
         })};
-        inst->ReplaceOpcode(IndexedInstruction(*inst));
         inst->SetArg(0, IR::Value{index});
     }
 }