From 19d5ac1318832fef830ae771be4fb9b6dea79863 Mon Sep 17 00:00:00 2001
From: IndecisiveTurtle <47210458+raphaelthegreat@users.noreply.github.com>
Date: Sun, 25 Aug 2024 14:21:37 +0300
Subject: [PATCH] spirv: Only specific storage image format for atomics

---
 .../backend/spirv/emit_spirv.cpp              |  1 +
 .../backend/spirv/spirv_emit_context.cpp      |  2 +-
 .../ir/passes/resource_tracking_pass.cpp      | 46 +++++++++----------
 src/shader_recompiler/runtime_info.h          |  1 +
 4 files changed, 24 insertions(+), 26 deletions(-)

diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.cpp b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
index 460537b7..ca49047c 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
@@ -187,6 +187,7 @@ void DefineEntryPoint(const IR::Program& program, EmitContext& ctx, Id main) {
     ctx.AddCapability(spv::Capability::Int64);
     if (info.has_storage_images) {
         ctx.AddCapability(spv::Capability::StorageImageExtendedFormats);
+        ctx.AddCapability(spv::Capability::StorageImageWriteWithoutFormat);
     }
     switch (program.info.stage) {
     case Stage::Compute: {
diff --git a/src/shader_recompiler/backend/spirv/spirv_emit_context.cpp b/src/shader_recompiler/backend/spirv/spirv_emit_context.cpp
index 7c28cdef..ee472667 100644
--- a/src/shader_recompiler/backend/spirv/spirv_emit_context.cpp
+++ b/src/shader_recompiler/backend/spirv/spirv_emit_context.cpp
@@ -447,7 +447,7 @@ spv::ImageFormat GetFormat(const AmdGpu::Image& image) {
 
 Id ImageType(EmitContext& ctx, const ImageResource& desc, Id sampled_type) {
     const auto image = ctx.info.ReadUd<AmdGpu::Image>(desc.sgpr_base, desc.dword_offset);
-    const auto format = desc.is_storage ? GetFormat(image) : spv::ImageFormat::Unknown;
+    const auto format = desc.is_atomic ? GetFormat(image) : spv::ImageFormat::Unknown;
     const u32 sampled = desc.is_storage ? 2 : 1;
     switch (desc.type) {
     case AmdGpu::ImageType::Color1D:
diff --git a/src/shader_recompiler/ir/passes/resource_tracking_pass.cpp b/src/shader_recompiler/ir/passes/resource_tracking_pass.cpp
index ace6a37d..abdb97b9 100644
--- a/src/shader_recompiler/ir/passes/resource_tracking_pass.cpp
+++ b/src/shader_recompiler/ir/passes/resource_tracking_pass.cpp
@@ -143,20 +143,8 @@ IR::Type BufferDataType(const IR::Inst& inst, AmdGpu::NumberFormat num_format) {
     }
 }
 
-bool IsImageInstruction(const IR::Inst& inst) {
+bool IsImageAtomicInstruction(const IR::Inst& inst) {
     switch (inst.GetOpcode()) {
-    case IR::Opcode::ImageSampleExplicitLod:
-    case IR::Opcode::ImageSampleImplicitLod:
-    case IR::Opcode::ImageSampleDrefExplicitLod:
-    case IR::Opcode::ImageSampleDrefImplicitLod:
-    case IR::Opcode::ImageFetch:
-    case IR::Opcode::ImageGather:
-    case IR::Opcode::ImageGatherDref:
-    case IR::Opcode::ImageQueryDimensions:
-    case IR::Opcode::ImageQueryLod:
-    case IR::Opcode::ImageGradient:
-    case IR::Opcode::ImageRead:
-    case IR::Opcode::ImageWrite:
     case IR::Opcode::ImageAtomicIAdd32:
     case IR::Opcode::ImageAtomicSMin32:
     case IR::Opcode::ImageAtomicUMin32:
@@ -178,20 +166,27 @@ bool IsImageStorageInstruction(const IR::Inst& inst) {
     switch (inst.GetOpcode()) {
     case IR::Opcode::ImageWrite:
     case IR::Opcode::ImageRead:
-    case IR::Opcode::ImageAtomicIAdd32:
-    case IR::Opcode::ImageAtomicSMin32:
-    case IR::Opcode::ImageAtomicUMin32:
-    case IR::Opcode::ImageAtomicSMax32:
-    case IR::Opcode::ImageAtomicUMax32:
-    case IR::Opcode::ImageAtomicInc32:
-    case IR::Opcode::ImageAtomicDec32:
-    case IR::Opcode::ImageAtomicAnd32:
-    case IR::Opcode::ImageAtomicOr32:
-    case IR::Opcode::ImageAtomicXor32:
-    case IR::Opcode::ImageAtomicExchange32:
         return true;
     default:
-        return false;
+        return IsImageAtomicInstruction(inst);
+    }
+}
+
+bool IsImageInstruction(const IR::Inst& inst) {
+    switch (inst.GetOpcode()) {
+    case IR::Opcode::ImageSampleExplicitLod:
+    case IR::Opcode::ImageSampleImplicitLod:
+    case IR::Opcode::ImageSampleDrefExplicitLod:
+    case IR::Opcode::ImageSampleDrefImplicitLod:
+    case IR::Opcode::ImageFetch:
+    case IR::Opcode::ImageGather:
+    case IR::Opcode::ImageGatherDref:
+    case IR::Opcode::ImageQueryDimensions:
+    case IR::Opcode::ImageQueryLod:
+    case IR::Opcode::ImageGradient:
+        return true;
+    default:
+        return IsImageStorageInstruction(inst);
     }
 }
 
@@ -537,6 +532,7 @@ void PatchImageInstruction(IR::Block& block, IR::Inst& inst, Info& info, Descrip
         .nfmt = static_cast<AmdGpu::NumberFormat>(image.GetNumberFmt()),
         .is_storage = IsImageStorageInstruction(inst),
         .is_depth = bool(inst_info.is_depth),
+        .is_atomic = IsImageAtomicInstruction(inst),
     });
 
     // Read sampler sharp. This doesn't exist for IMAGE_LOAD/IMAGE_STORE instructions
diff --git a/src/shader_recompiler/runtime_info.h b/src/shader_recompiler/runtime_info.h
index 9dde9f26..41ab7157 100644
--- a/src/shader_recompiler/runtime_info.h
+++ b/src/shader_recompiler/runtime_info.h
@@ -110,6 +110,7 @@ struct ImageResource {
     AmdGpu::NumberFormat nfmt;
     bool is_storage;
     bool is_depth;
+    bool is_atomic{};
 
     u64 GetKey(const Info& info) const {
         const auto sharp = GetTsharp(info);