ir_pass: Implement shared memory lowering pass

* NVIDIA doesn't like using shared mem in fragment shader and softlocks driver

CMakeLists.txt

@@ -421,6 +421,7 @@ set(SHADER_RECOMPILER src/shader_recompiler/exception.h
                       src/shader_recompiler/ir/passes/dead_code_elimination_pass.cpp
                       src/shader_recompiler/ir/passes/identity_removal_pass.cpp
                       src/shader_recompiler/ir/passes/ir_passes.h
+                      src/shader_recompiler/ir/passes/lower_shared_mem_to_registers.cpp
                       src/shader_recompiler/ir/passes/resource_tracking_pass.cpp
                       src/shader_recompiler/ir/passes/shader_info_collection_pass.cpp
                       src/shader_recompiler/ir/passes/ssa_rewrite_pass.cpp

src/shader_recompiler/ir/passes/ir_passes.h

@@ -14,5 +14,6 @@ void DeadCodeEliminationPass(IR::Program& program);
 void ConstantPropagationPass(IR::BlockList& program);
 void ResourceTrackingPass(IR::Program& program);
 void CollectShaderInfoPass(IR::Program& program);
+void LowerSharedMemToRegisters(IR::Program& program);
 
 } // namespace Shader::Optimization

src/shader_recompiler/ir/passes/lower_shared_mem_to_registers.cpp

@@ -0,0 +1,39 @@
+// SPDX-FileCopyrightText: Copyright 2024 shadPS4 Emulator Project
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+#include <boost/container/small_vector.hpp>
+#include "shader_recompiler/ir/program.h"
+
+namespace Shader::Optimization {
+
+void LowerSharedMemToRegisters(IR::Program& program) {
+    boost::container::small_vector<IR::Inst*, 8> ds_writes;
+    Info& info{program.info};
+    for (IR::Block* const block : program.blocks) {
+        for (IR::Inst& inst : block->Instructions()) {
+            const auto opcode = inst.GetOpcode();
+            if (opcode == IR::Opcode::WriteSharedU32 || opcode == IR::Opcode::WriteSharedU64) {
+                ds_writes.emplace_back(&inst);
+                continue;
+            }
+            if (opcode == IR::Opcode::LoadSharedU32 || opcode == IR::Opcode::LoadSharedU64) {
+                // Search for write instruction with same offset
+                const IR::Inst* prod = inst.Arg(0).InstRecursive();
+                const auto it = std::ranges::find_if(ds_writes, [&](const IR::Inst* write) {
+                    const IR::Inst* write_prod = write->Arg(0).InstRecursive();
+                    return write_prod->Arg(1).U32() == prod->Arg(1).U32() &&
+                           write_prod->Arg(0) == prod->Arg(0);
+                });
+                ASSERT(it != ds_writes.end());
+                // Replace data read with value written.
+                inst.ReplaceUsesWith((*it)->Arg(1));
+            }
+        }
+    }
+    // We should have eliminated everything. Invalidate data write instructions.
+    for (const auto inst : ds_writes) {
+        inst->Invalidate();
+    }
+}
+
+} // namespace Shader::Optimization

src/shader_recompiler/recompiler.cpp

@@ -58,6 +58,9 @@ IR::Program TranslateProgram(Common::ObjectPool<IR::Inst>& inst_pool,
     Shader::Optimization::SsaRewritePass(program.post_order_blocks);
     Shader::Optimization::ResourceTrackingPass(program);
     Shader::Optimization::ConstantPropagationPass(program.post_order_blocks);
+    if (program.info.stage != Stage::Compute) {
+        Shader::Optimization::LowerSharedMemToRegisters(program);
+    }
     Shader::Optimization::IdentityRemovalPass(program.blocks);
     Shader::Optimization::DeadCodeEliminationPass(program);
     Shader::Optimization::CollectShaderInfoPass(program);