1 files changed, 199 insertions, 0 deletions
diff --git a/src/shader_recompiler/ir_opt/texture_pass.cpp b/src/shader_recompiler/ir_opt/texture_pass.cpp
new file mode 100644
index 000000000..80e4ad6a9
--- /dev/null
+++ b/src/shader_recompiler/ir_opt/texture_pass.cpp
@@ -0,0 +1,199 @@
+// Copyright 2021 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#include <optional>
+
+#include <boost/container/flat_set.hpp>
+#include <boost/container/small_vector.hpp>
+
+#include "shader_recompiler/environment.h"
+#include "shader_recompiler/frontend/ir/basic_block.h"
+#include "shader_recompiler/frontend/ir/ir_emitter.h"
+#include "shader_recompiler/ir_opt/passes.h"
+#include "shader_recompiler/shader_info.h"
+
+namespace Shader::Optimization {
+namespace {
+struct ConstBufferAddr {
+    u32 index;
+    u32 offset;
+};
+
+struct TextureInst {
+    ConstBufferAddr cbuf;
+    IR::Inst* inst;
+    IR::Block* block;
+};
+
+using TextureInstVector = boost::container::small_vector<TextureInst, 24>;
+
+using VisitedBlocks = boost::container::flat_set<IR::Block*, std::less<IR::Block*>,
+                                                 boost::container::small_vector<IR::Block*, 2>>;
+
+IR::Opcode IndexedInstruction(const IR::Inst& inst) {
+    switch (inst.Opcode()) {
+    case IR::Opcode::BindlessImageSampleImplicitLod:
+    case IR::Opcode::BoundImageSampleImplicitLod:
+        return IR::Opcode::ImageSampleImplicitLod;
+    case IR::Opcode::BoundImageSampleExplicitLod:
+    case IR::Opcode::BindlessImageSampleExplicitLod:
+        return IR::Opcode::ImageSampleExplicitLod;
+    case IR::Opcode::BoundImageSampleDrefImplicitLod:
+    case IR::Opcode::BindlessImageSampleDrefImplicitLod:
+        return IR::Opcode::ImageSampleDrefImplicitLod;
+    case IR::Opcode::BoundImageSampleDrefExplicitLod:
+    case IR::Opcode::BindlessImageSampleDrefExplicitLod:
+        return IR::Opcode::ImageSampleDrefExplicitLod;
+    default:
+        return IR::Opcode::Void;
+    }
+}
+
+bool IsBindless(const IR::Inst& inst) {
+    switch (inst.Opcode()) {
+    case IR::Opcode::BindlessImageSampleImplicitLod:
+    case IR::Opcode::BindlessImageSampleExplicitLod:
+    case IR::Opcode::BindlessImageSampleDrefImplicitLod:
+    case IR::Opcode::BindlessImageSampleDrefExplicitLod:
+        return true;
+    case IR::Opcode::BoundImageSampleImplicitLod:
+    case IR::Opcode::BoundImageSampleExplicitLod:
+    case IR::Opcode::BoundImageSampleDrefImplicitLod:
+    case IR::Opcode::BoundImageSampleDrefExplicitLod:
+        return false;
+    default:
+        throw InvalidArgument("Invalid opcode {}", inst.Opcode());
+    }
+}
+
+bool IsTextureInstruction(const IR::Inst& inst) {
+    return IndexedInstruction(inst) != IR::Opcode::Void;
+}
+
+std::optional<ConstBufferAddr> Track(IR::Block* block, const IR::Value& value,
+                                     VisitedBlocks& visited) {
+    if (value.IsImmediate()) {
+        // Immediates can't be a storage buffer
+        return std::nullopt;
+    }
+    const IR::Inst* const inst{value.InstRecursive()};
+    if (inst->Opcode() == IR::Opcode::GetCbuf) {
+        const IR::Value index{inst->Arg(0)};
+        const IR::Value offset{inst->Arg(1)};
+        if (!index.IsImmediate()) {
+            // Reading a bindless texture from variable indices is valid
+            // but not supported here at the moment
+            return std::nullopt;
+        }
+        if (!offset.IsImmediate()) {
+            // TODO: Support arrays of textures
+            return std::nullopt;
+        }
+        return ConstBufferAddr{
+            .index{index.U32()},
+            .offset{offset.U32()},
+        };
+    }
+    // Reversed loops are more likely to find the right result
+    for (size_t arg = inst->NumArgs(); arg--;) {
+        IR::Block* inst_block{block};
+        if (inst->Opcode() == IR::Opcode::Phi) {
+            // If we are going through a phi node, mark the current block as visited
+            visited.insert(block);
+            // and skip already visited blocks to avoid looping forever
+            IR::Block* const phi_block{inst->PhiBlock(arg)};
+            if (visited.contains(phi_block)) {
+                // Already visited, skip
+                continue;
+            }
+            inst_block = phi_block;
+        }
+        const std::optional storage_buffer{Track(inst_block, inst->Arg(arg), visited)};
+        if (storage_buffer) {
+            return *storage_buffer;
+        }
+    }
+    return std::nullopt;
+}
+
+TextureInst MakeInst(Environment& env, IR::Block* block, IR::Inst& inst) {
+    ConstBufferAddr addr;
+    if (IsBindless(inst)) {
+        VisitedBlocks visited;
+        const std::optional<ConstBufferAddr> track_addr{Track(block, IR::Value{&inst}, visited)};
+        if (!track_addr) {
+            throw NotImplementedException("Failed to track bindless texture constant buffer");
+        }
+        addr = *track_addr;
+    } else {
+        addr = ConstBufferAddr{
+            .index{env.TextureBoundBuffer()},
+            .offset{inst.Arg(0).U32()},
+        };
+    }
+    return TextureInst{
+        .cbuf{addr},
+        .inst{&inst},
+        .block{block},
+    };
+}
+
+class Descriptors {
+public:
+    explicit Descriptors(TextureDescriptors& descriptors_) : descriptors{descriptors_} {}
+
+    u32 Add(const TextureDescriptor& descriptor) {
+        // TODO: Handle arrays
+        auto it{std::ranges::find_if(descriptors, [&descriptor](const TextureDescriptor& existing) {
+            return descriptor.cbuf_index == existing.cbuf_index &&
+                   descriptor.cbuf_offset == existing.cbuf_offset &&
+                   descriptor.type == existing.type;
+        })};
+        if (it != descriptors.end()) {
+            return static_cast<u32>(std::distance(descriptors.begin(), it));
+        }
+        descriptors.push_back(descriptor);
+        return static_cast<u32>(descriptors.size()) - 1;
+    }
+
+private:
+    TextureDescriptors& descriptors;
+};
+} // Anonymous namespace
+
+void TexturePass(Environment& env, IR::Program& program) {
+    TextureInstVector to_replace;
+    for (IR::Function& function : program.functions) {
+        for (IR::Block* const block : function.post_order_blocks) {
+            for (IR::Inst& inst : block->Instructions()) {
+                if (!IsTextureInstruction(inst)) {
+                    continue;
+                }
+                to_replace.push_back(MakeInst(env, block, inst));
+            }
+        }
+    }
+    // Sort instructions to visit textures by constant buffer index, then by offset
+    std::ranges::sort(to_replace, [](const auto& lhs, const auto& rhs) {
+        return lhs.cbuf.offset < rhs.cbuf.offset;
+    });
+    std::stable_sort(to_replace.begin(), to_replace.end(), [](const auto& lhs, const auto& rhs) {
+        return lhs.cbuf.index < rhs.cbuf.index;
+    });
+    Descriptors descriptors{program.info.texture_descriptors};
+    for (TextureInst& texture_inst : to_replace) {
+        // TODO: Handle arrays
+        IR::Inst* const inst{texture_inst.inst};
+        const u32 index{descriptors.Add(TextureDescriptor{
+            .type{inst->Flags<IR::TextureInstInfo>().type},
+            .cbuf_index{texture_inst.cbuf.index},
+            .cbuf_offset{texture_inst.cbuf.offset},
+            .count{1},
+        })};
+        inst->ReplaceOpcode(IndexedInstruction(*inst));
+        inst->SetArg(0, IR::Value{index});
+    }
+}
+
+} // namespace Shader::Optimization