[vk] fix LM3 and TL:LTD multiple NonUniform annotations (#3997)

should fix AMD

Signed-off-by: lizzie lizzie@eden-emu.dev

Reviewed-on: https://git.eden-emu.dev/eden-emu/eden/pulls/3997
Reviewed-by: CamilleLaVey <camillelavey99@gmail.com>
Reviewed-by: MaranBr <maranbr@eden-emu.dev>

4 files changed, 30 insertions, 44 deletions

diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.cpp b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
index fc102cadac..0447dfe90d 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
@@ -462,6 +462,13 @@ void SetupCapabilities(const Profile& profile, const Info& info, EmitContext& ct
     ctx.AddCapability(spv::Capability::ImageGatherExtended);
     ctx.AddCapability(spv::Capability::ImageQuery);
     ctx.AddCapability(spv::Capability::SampledBuffer);
+    // TODO: this usage needs to be tracked properly
+    if (ctx.profile.support_sampled_image_array_nonuniform_indexing) {
+        if (ctx.profile.supported_spirv < 0x00010400)
+            ctx.AddExtension("SPV_EXT_descriptor_indexing");
+        ctx.AddCapability(spv::Capability::ShaderNonUniform);
+        ctx.AddCapability(spv::Capability::SampledImageArrayNonUniformIndexing);
+    }
 }
 
 void PatchPhiNodes(IR::Program& program, EmitContext& ctx) {
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_image.cpp b/src/shader_recompiler/backend/spirv/emit_spirv_image.cpp
index 62a82dce9d..ad98973e29 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv_image.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_image.cpp
@@ -14,37 +14,10 @@
 
 namespace Shader::Backend::SPIRV {
 namespace {
-class DescriptorIndex {
-public:
-    explicit DescriptorIndex(EmitContext& ctx, const IR::Value& index)
-        : id{index.IsImmediate() ? ctx.Const(index.U32()) : ctx.Def(index)},
-          is_non_uniform{ctx.profile.support_sampled_image_array_nonuniform_indexing &&
-                         !index.IsImmediate()} {
-        if (!is_non_uniform) {
-            return;
-        }
-        if (ctx.profile.supported_spirv < 0x00010400) {
-            ctx.AddExtension("SPV_EXT_descriptor_indexing");
-        }
-        ctx.AddCapability(spv::Capability::ShaderNonUniform);
-        ctx.AddCapability(spv::Capability::SampledImageArrayNonUniformIndexing);
-        Decorate(ctx, id);
-    }
-
-    Id Value() const {
-        return id;
-    }
 
-    void Decorate(EmitContext& ctx, Id object) const {
-        if (is_non_uniform) {
-            ctx.Decorate(object, spv::Decoration::NonUniform);
-        }
-    }
-
-private:
-    Id id;
-    bool is_non_uniform;
-};
+[[nodiscard]] bool IsNonUniformDescriptor(EmitContext& ctx, const IR::Value& index) noexcept {
+    return ctx.profile.support_sampled_image_array_nonuniform_indexing && !index.IsImmediate();
+}
 
 class ImageOperands {
 public:
@@ -221,11 +194,13 @@ private:
 Id Texture(EmitContext& ctx, IR::TextureInstInfo info, [[maybe_unused]] const IR::Value& index) {
     const TextureDefinition& def{ctx.textures.at(info.descriptor_index)};
     if (def.count > 1) {
-        const DescriptorIndex idx{ctx, index};
-        const Id pointer{ctx.OpAccessChain(def.pointer_type, def.id, idx.Value())};
-        idx.Decorate(ctx, pointer);
+        auto const idx = index.IsImmediate() ? ctx.Const(index.U32()) : ctx.Def(index);
+        if (!ctx.non_uniform_ids.contains(idx.value) && IsNonUniformDescriptor(ctx, index)) {
+            ctx.Decorate(idx, spv::Decoration::NonUniform);
+            ctx.non_uniform_ids.insert(idx.value);
+        }
+        const Id pointer{ctx.OpAccessChain(def.pointer_type, def.id, idx)};
         const Id object{ctx.OpLoad(def.sampled_type, pointer)};
-        idx.Decorate(ctx, object);
         return object;
     } else {
         return ctx.OpLoad(def.sampled_type, def.id);
@@ -244,13 +219,14 @@ Id TextureImage(EmitContext& ctx, IR::TextureInstInfo info, const IR::Value& ind
     } else {
         const TextureDefinition& def{ctx.textures.at(info.descriptor_index)};
         if (def.count > 1) {
-            const DescriptorIndex idx{ctx, index};
-            const Id ptr{ctx.OpAccessChain(def.pointer_type, def.id, idx.Value())};
-            idx.Decorate(ctx, ptr);
-            const Id object{ctx.OpLoad(def.sampled_type, ptr)};
-            idx.Decorate(ctx, object);
-            const Id image{ctx.OpImage(def.image_type, object)};
-            idx.Decorate(ctx, image);
+            auto const idx = index.IsImmediate() ? ctx.Const(index.U32()) : ctx.Def(index);
+            if (!ctx.non_uniform_ids.contains(idx.value) && IsNonUniformDescriptor(ctx, index)) {
+                ctx.Decorate(idx, spv::Decoration::NonUniform);
+                ctx.non_uniform_ids.insert(idx.value);
+            }
+            const Id ptr = ctx.OpAccessChain(def.pointer_type, def.id, idx);
+            const Id object = ctx.OpLoad(def.sampled_type, ptr);
+            const Id image = ctx.OpImage(def.image_type, object);
             return image;
         }
         return ctx.OpImage(def.image_type, ctx.OpLoad(def.sampled_type, def.id));
diff --git a/src/shader_recompiler/backend/spirv/spirv_emit_context.h b/src/shader_recompiler/backend/spirv/spirv_emit_context.h
index de56809a98..55385de0b0 100644
--- a/src/shader_recompiler/backend/spirv/spirv_emit_context.h
+++ b/src/shader_recompiler/backend/spirv/spirv_emit_context.h
@@ -9,6 +9,7 @@
 #include <array>
 
 #include <sirit/sirit.h>
+#include <ankerl/unordered_dense.h>
 
 #include "shader_recompiler/backend/bindings.h"
 #include "shader_recompiler/frontend/ir/program.h"
@@ -366,6 +367,9 @@ public:
     Id load_const_func_u32x2{};
     Id load_const_func_u32x4{};
 
+    // Sirit::Id doesn't play nice with *::set<>
+    ankerl::unordered_dense::set<u32> non_uniform_ids;
+
 private:
     void DefineCommonTypes(const Info& info);
     void DefineCommonConstants();
diff --git a/src/shader_recompiler/ir_opt/texture_pass.cpp b/src/shader_recompiler/ir_opt/texture_pass.cpp
index 55b6f7213c..bf757da169 100644
--- a/src/shader_recompiler/ir_opt/texture_pass.cpp
+++ b/src/shader_recompiler/ir_opt/texture_pass.cpp
@@ -733,9 +733,8 @@ void TexturePass(Environment& env, IR::Program& program, const HostTranslateInfo
             break;
         }
         u32 index;
-        const u32 size_shift{cbuf.count > 1 ? DynamicDescriptorSizeShift(cbuf.dynamic_offset)
-                                            : DESCRIPTOR_SIZE_SHIFT};
-        u32 count{cbuf.count};
+        u32 size_shift = cbuf.count > 1 ? DynamicDescriptorSizeShift(cbuf.dynamic_offset) : DESCRIPTOR_SIZE_SHIFT;
+        u32 count = cbuf.count;
         switch (inst->GetOpcode()) {
         case IR::Opcode::ImageRead:
         case IR::Opcode::ImageAtomicIAdd32: