shader: compute texture handles read from storage buffers for each draw

Switch 2 shaders can read a texture handle from a storage buffer (Switch
Sports' scenery: a per-object entry of a table whose pointer is in the driver
constant buffer). The bindless handle search looked through the load into its
address and returned the table pointer's constant buffer word as the handle,
which named an arbitrary descriptor; the game repoints that one at transient
render targets every few seconds, so the island and menu lighting flipped
between correct, dark and red.

The search no longer looks into memory loads. A handle it cannot trace is now
built into a small program over constant buffer words, storage buffer words and
integer operations (Shader::Info::handle_programs), which the pipelines evaluate
when they bind the texture. Handles a single evaluation cannot give (one that
changes in a loop) keep the previous search.

1 files changed, 426 insertions, 6 deletions

diff --git a/src/shader_recompiler/ir_opt/texture_pass.cpp b/src/shader_recompiler/ir_opt/texture_pass.cpp
index fae69ab91b..f413cf108a 100644
--- a/src/shader_recompiler/ir_opt/texture_pass.cpp
+++ b/src/shader_recompiler/ir_opt/texture_pass.cpp
@@ -5,6 +5,10 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
 #include <algorithm>
+#include <queue>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
 #include <bit>
 #include <optional>
 #include "common/container/unordered_map.h"
@@ -340,15 +344,20 @@ static inline u32 GetTextureHandleCached(Environment& env, const ConstBufferAddr
 }
 
 // Cached variants of existing helpers
+// A handle computed per draw has no descriptor to read at translation time: assume an ordinary
+// floating-point 2D texture, the instruction's own type deciding the rest.
 static inline TextureType ReadTextureTypeCached(Environment& env, const ConstBufferAddr& cbuf) {
+    if (cbuf.handle_program != ~0U) return TextureType::Color2D;
     return env.ReadTextureType(GetTextureHandleCached(env, cbuf));
 }
 static inline TexturePixelFormat ReadTexturePixelFormatCached(Environment& env,
                                                                 const ConstBufferAddr& cbuf) {
+    if (cbuf.handle_program != ~0U) return TexturePixelFormat::A8B8G8R8_UNORM;
     return env.ReadTexturePixelFormat(GetTextureHandleCached(env, cbuf));
 }
 static inline bool IsTexturePixelFormatIntegerCached(Environment& env,
                                                         const ConstBufferAddr& cbuf) {
+    if (cbuf.handle_program != ~0U) return false;
     return env.IsTexturePixelFormatInteger(GetTextureHandleCached(env, cbuf));
 }
 
@@ -366,12 +375,387 @@ static inline std::optional<ConstBufferAddr> TrackCached(const IR::Value& v, Env
 
 std::optional<ConstBufferAddr> TryGetConstBuffer(const IR::Inst* inst, Environment& env, const HostTranslateInfo& host_info, const IR::Block* block);
 
+bool IsMemoryLoad(IR::Opcode opcode) {
+    switch (opcode) {
+    case IR::Opcode::LoadGlobalU8: case IR::Opcode::LoadGlobalS8: case IR::Opcode::LoadGlobalU16:
+    case IR::Opcode::LoadGlobalS16: case IR::Opcode::LoadGlobal32: case IR::Opcode::LoadGlobal64:
+    case IR::Opcode::LoadGlobal128: case IR::Opcode::LoadStorageU8: case IR::Opcode::LoadStorageS8:
+    case IR::Opcode::LoadStorageU16: case IR::Opcode::LoadStorageS16: case IR::Opcode::LoadStorage32:
+    case IR::Opcode::LoadStorage64: case IR::Opcode::LoadStorage128: case IR::Opcode::LoadLocal:
+    case IR::Opcode::LoadSharedU8: case IR::Opcode::LoadSharedS8: case IR::Opcode::LoadSharedU16:
+    case IR::Opcode::LoadSharedS16: case IR::Opcode::LoadSharedU32: case IR::Opcode::LoadSharedU64:
+    case IR::Opcode::LoadSharedU128:
+        return true;
+    default:
+        return false;
+    }
+}
+
 std::optional<ConstBufferAddr> Track(const IR::Value& value, Environment& env, const HostTranslateInfo& host_info, const IR::Block* block) {
-    return IR::BreadthFirstSearch(value, [&env, &host_info, block](const IR::Inst* inst) {
-        return TryGetConstBuffer(inst, env, host_info, block);
-    }, [&](IR::Value operand) { return ResolveConditionalPhi(operand, block); });
+    // A value loaded from memory is not a constant buffer word, whatever its address is built
+    // from: the search does not look into a load's operands (a handle read from a storage
+    // buffer would otherwise be taken from the buffer's address in the constant buffer).
+    const IR::Value start{ResolveConditionalPhi(value, block)};
+    if (start.IsImmediate()) {
+        return std::nullopt;
+    }
+    std::unordered_set<const IR::Inst*> visited;
+    std::queue<const IR::Inst*> queue;
+    queue.push(start.InstRecursive());
+    visited.insert(start.InstRecursive());
+    while (!queue.empty()) {
+        const IR::Inst* const inst{queue.front()};
+        queue.pop();
+        if (const std::optional result = TryGetConstBuffer(inst, env, host_info, block)) {
+            return result;
+        }
+        if (IsMemoryLoad(inst->GetOpcode())) {
+            continue;
+        }
+        for (size_t arg = inst->NumArgs(); arg--;) {
+            const IR::Value arg_value{ResolveConditionalPhi(inst->Arg(arg), block)};
+            if (arg_value.IsImmediate()) {
+                continue;
+            }
+            const IR::Inst* const arg_inst{arg_value.InstRecursive()};
+            if (visited.insert(arg_inst).second) {
+                queue.push(arg_inst);
+            }
+        }
+    }
+    return std::nullopt;
+}
+
+// Builds a handle program (shader_info.h) for a value computed from constant buffer and storage
+// buffer words; false when the value depends on anything else.
+thread_local IR::Opcode g_handle_failure{IR::Opcode::Void}; // the first operation not supported
+
+bool IsUndefined(IR::Opcode opcode) {
+    return opcode == IR::Opcode::UndefU1 || opcode == IR::Opcode::UndefU8 ||
+           opcode == IR::Opcode::UndefU16 || opcode == IR::Opcode::UndefU32 ||
+           opcode == IR::Opcode::UndefU64;
+}
+
+// Values already emitted by the program being built, by the slot that keeps them.
+thread_local std::unordered_map<const IR::Inst*, u32>* g_handle_slots{};
+
+bool BuildHandleNode(const IR::Value& value, std::vector<u32>& code, int& budget,
+                     const IR::Block* block);
+
+bool BuildHandleProgram(const IR::Value& value, std::vector<u32>& code, int& budget,
+                        const IR::Block* block) {
+    using Op = HandleOp;
+    if (value.IsImmediate()) {
+        return BuildHandleNode(value, code, budget, block);
+    }
+    std::unordered_map<const IR::Inst*, u32> local_slots;
+    const bool outermost = g_handle_slots == nullptr;
+    if (outermost) {
+        g_handle_slots = &local_slots;
+    }
+    const IR::Inst* const inst{value.InstRecursive()};
+    bool result;
+    if (const auto it = g_handle_slots->find(inst); it != g_handle_slots->end()) {
+        if (it->second == ~0U) {
+            // Reached again while building it: a loop carries the value, which no single
+            // evaluation can give.
+            g_handle_failure = IR::Opcode::Phi;
+            result = false;
+        } else {
+            code.push_back(static_cast<u32>(Op::Load));
+            code.push_back(it->second);
+            result = true;
+        }
+    } else {
+        g_handle_slots->emplace(inst, ~0U);
+        const size_t start = code.size();
+        result = BuildHandleNode(value, code, budget, block);
+        g_handle_slots->erase(inst);
+        if (!result) {
+            code.resize(start);
+        }
+        if (result && inst->UseCount() > 1 && g_handle_slots->size() < 256) {
+            const u32 slot = static_cast<u32>(
+                std::ranges::count_if(*g_handle_slots, [](const auto& entry) { return entry.second != ~0U; }));
+            g_handle_slots->emplace(inst, slot);
+            code.push_back(static_cast<u32>(Op::Tee));
+            code.push_back(slot);
+        }
+    }
+    if (outermost) {
+        g_handle_slots = nullptr;
+    }
+    return result;
 }
 
+bool BuildHandleNode(const IR::Value& value, std::vector<u32>& code, int& budget,
+                     const IR::Block* block) {
+    if (--budget < 0) {
+        return false;
+    }
+    using Op = HandleOp;
+    const auto emit = [&](Op op) { code.push_back(static_cast<u32>(op)); };
+    if (value.IsImmediate()) {
+        emit(Op::Immediate);
+        code.push_back(value.Type() == IR::Type::U1 ? (value.U1() ? 1U : 0U) : value.U32());
+        return true;
+    }
+    const IR::Inst* const inst{value.InstRecursive()};
+    const auto build = [&](size_t arg) { return BuildHandleProgram(inst->Arg(arg), code, budget, block); };
+    const auto binary = [&](Op op) {
+        if (!build(0) || !build(1)) return false;
+        emit(op);
+        return true;
+    };
+    const auto unary = [&](Op op) {
+        if (!build(0)) return false;
+        emit(op);
+        return true;
+    };
+    // A word of a constant buffer read or of a storage buffer load, at byte offset +extra.
+    const auto load_word = [&](const IR::Inst* source, u32 extra) {
+        const auto add_extra = [&] {
+            if (extra != 0) {
+                emit(Op::Immediate);
+                code.push_back(extra);
+                emit(Op::Add);
+            }
+        };
+        switch (source->GetOpcode()) {
+        case IR::Opcode::GetCbufU32:
+        case IR::Opcode::GetCbufU32x2: {
+            const IR::Value index{source->Arg(0)};
+            const IR::Value offset{source->Arg(1)};
+            if (!index.IsImmediate()) return false;
+            if (offset.IsImmediate()) {
+                emit(Op::Cbuf);
+                code.push_back(index.U32());
+                code.push_back(offset.U32() + extra);
+                return true;
+            }
+            if (!BuildHandleProgram(offset, code, budget, block)) return false;
+            add_extra();
+            emit(Op::CbufDynamic);
+            code.push_back(index.U32());
+            return true;
+        }
+        case IR::Opcode::LoadStorage32:
+        case IR::Opcode::LoadStorage64:
+        case IR::Opcode::LoadStorage128: {
+            const IR::Value binding{source->Arg(0)};
+            if (!binding.IsImmediate()) return false;
+            if (!BuildHandleProgram(source->Arg(1), code, budget, block)) return false;
+            add_extra();
+            emit(Op::Storage);
+            code.push_back(binding.U32());
+            return true;
+        }
+        default:
+            return false;
+        }
+    };
+    switch (inst->GetOpcode()) {
+    case IR::Opcode::Identity:
+    case IR::Opcode::ConditionRef:
+    case IR::Opcode::ShuffleIndex: // a value made uniform (R2UR): the same in every lane
+    case IR::Opcode::VoteAny:      // of a condition on per-draw words only, the same in every lane
+    case IR::Opcode::VoteAll:
+        return build(0);
+    case IR::Opcode::GetCbufU32:
+    case IR::Opcode::LoadStorage32:
+        return load_word(inst, 0);
+    case IR::Opcode::CompositeExtractU32x2: {
+        const IR::Value element{inst->Arg(1)};
+        if (!element.IsImmediate() || inst->Arg(0).IsImmediate()) return false;
+        return load_word(inst->Arg(0).InstRecursive(), element.U32() * 4);
+    }
+    case IR::Opcode::Phi: {
+        // Keep the incoming value that is defined; of two, the one made uniform (the
+        // register-to-uniform move a conditional block makes for the handle).
+        const IR::Value resolved{ResolveConditionalPhi(value, block)};
+        if (resolved.IsImmediate() || resolved.InstRecursive() != inst) {
+            return BuildHandleProgram(resolved, code, budget, block);
+        }
+        std::optional<size_t> chosen;
+        for (size_t arg = 0; arg < inst->NumArgs(); ++arg) {
+            const IR::Value incoming{inst->Arg(arg)};
+            if (!incoming.IsImmediate() && IsUndefined(incoming.InstRecursive()->GetOpcode())) {
+                continue;
+            }
+            IR::Value probe{incoming};
+            bool shuffled = false;
+            for (int depth = 0; depth < 8 && !probe.IsImmediate(); ++depth) {
+                const IR::Inst* const probe_inst{probe.InstRecursive()};
+                if (probe_inst->GetOpcode() == IR::Opcode::ShuffleIndex) {
+                    shuffled = true;
+                    break;
+                }
+                if (probe_inst->GetOpcode() != IR::Opcode::Identity) break;
+                probe = probe_inst->Arg(0);
+            }
+            if (!chosen || shuffled) chosen = arg;
+        }
+        if (chosen && build(*chosen)) {
+            return true;
+        }
+        for (size_t arg = 0; arg < inst->NumArgs(); ++arg) {
+            const IR::Value incoming{inst->Arg(arg)};
+            if (arg == chosen || (!incoming.IsImmediate() && IsUndefined(incoming.InstRecursive()->GetOpcode()))) {
+                continue;
+            }
+            if (build(arg)) {
+                return true;
+            }
+        }
+        return false;
+    }
+    case IR::Opcode::IAdd32: return binary(Op::Add);
+    case IR::Opcode::ISub32: return binary(Op::Sub);
+    case IR::Opcode::IMul32: return binary(Op::Mul);
+    case IR::Opcode::BitwiseAnd32: return binary(Op::And);
+    case IR::Opcode::BitwiseOr32: return binary(Op::Or);
+    case IR::Opcode::BitwiseXor32: return binary(Op::Xor);
+    case IR::Opcode::ShiftLeftLogical32: return binary(Op::Shl);
+    case IR::Opcode::ShiftRightLogical32: return binary(Op::Shr);
+    case IR::Opcode::ShiftRightArithmetic32: return binary(Op::Sar);
+    case IR::Opcode::INeg32: return unary(Op::Neg);
+    case IR::Opcode::BitwiseNot32: return unary(Op::Not);
+    case IR::Opcode::UMin32: return binary(Op::UMin);
+    case IR::Opcode::UMax32: return binary(Op::UMax);
+    case IR::Opcode::SMin32: return binary(Op::SMin);
+    case IR::Opcode::SMax32: return binary(Op::SMax);
+    case IR::Opcode::SLessThan: return binary(Op::SLt);
+    case IR::Opcode::SLessThanEqual: return binary(Op::SLe);
+    case IR::Opcode::SGreaterThan: return binary(Op::SGt);
+    case IR::Opcode::SGreaterThanEqual: return binary(Op::SGe);
+    case IR::Opcode::ULessThan: return binary(Op::ULt);
+    case IR::Opcode::ULessThanEqual: return binary(Op::ULe);
+    case IR::Opcode::UGreaterThan: return binary(Op::UGt);
+    case IR::Opcode::UGreaterThanEqual: return binary(Op::UGe);
+    case IR::Opcode::IEqual: return binary(Op::Eq);
+    case IR::Opcode::INotEqual: return binary(Op::Ne);
+    case IR::Opcode::LogicalAnd: return binary(Op::LAnd);
+    case IR::Opcode::LogicalOr: return binary(Op::LOr);
+    case IR::Opcode::LogicalXor: return binary(Op::LXor);
+    case IR::Opcode::LogicalNot: return unary(Op::LNot);
+    case IR::Opcode::GetZeroFromOp: // the flag a producing operation sets on its result
+        if (!build(0)) return false;
+        emit(Op::Immediate);
+        code.push_back(0);
+        emit(Op::Eq);
+        return true;
+    case IR::Opcode::GetSignFromOp:
+        if (!build(0)) return false;
+        emit(Op::Immediate);
+        code.push_back(0);
+        emit(Op::SLt);
+        return true;
+    case IR::Opcode::GetCarryFromOp: {
+        // The carry of a 32-bit addition: the sum below an operand.
+        const IR::Value op{inst->Arg(0)};
+        if (op.IsImmediate() || op.InstRecursive()->GetOpcode() != IR::Opcode::IAdd32) {
+            g_handle_failure = inst->GetOpcode();
+            return false;
+        }
+        const IR::Inst* const add{op.InstRecursive()};
+        if (!BuildHandleProgram(op, code, budget, block) ||
+            !BuildHandleProgram(add->Arg(0), code, budget, block)) {
+            return false;
+        }
+        emit(Op::ULt);
+        return true;
+    }
+    case IR::Opcode::BitCastU32F32:
+    case IR::Opcode::BitCastF32U32:
+        return build(0);
+    case IR::Opcode::UndefU1:
+    case IR::Opcode::UndefU32:
+        // Read before any write on some path: no value to compute, take zero.
+        emit(Op::Immediate);
+        code.push_back(0);
+        return true;
+    case IR::Opcode::BitFieldUExtract: {
+        // (base >> offset) & ((1 << count) - 1), the shift by 32 giving the all-ones mask.
+        const auto mask = [&] {
+            emit(Op::Immediate);
+            code.push_back(1);
+            if (!build(2)) return false;
+            emit(Op::Shl);
+            emit(Op::Immediate);
+            code.push_back(1);
+            emit(Op::Sub);
+            return true;
+        };
+        if (!build(0) || !build(1)) return false;
+        emit(Op::Shr);
+        if (!mask()) return false;
+        emit(Op::And);
+        return true;
+    }
+    case IR::Opcode::BitFieldSExtract: {
+        // (base << (32 - offset - count)) >> (32 - count), arithmetic.
+        if (!build(0)) return false;
+        emit(Op::Immediate);
+        code.push_back(32);
+        if (!build(1)) return false;
+        emit(Op::Sub);
+        if (!build(2)) return false;
+        emit(Op::Sub);
+        emit(Op::Shl);
+        emit(Op::Immediate);
+        code.push_back(32);
+        if (!build(2)) return false;
+        emit(Op::Sub);
+        emit(Op::Sar);
+        return true;
+    }
+    case IR::Opcode::BitFieldInsert: {
+        // (base & ~mask) | ((insert << offset) & mask), mask = ((1 << count) - 1) << offset.
+        const auto mask = [&] {
+            emit(Op::Immediate);
+            code.push_back(1);
+            if (!build(3)) return false;
+            emit(Op::Shl);
+            emit(Op::Immediate);
+            code.push_back(1);
+            emit(Op::Sub);
+            if (!build(2)) return false;
+            emit(Op::Shl);
+            return true;
+        };
+        if (!build(0) || !mask()) return false;
+        emit(Op::Not);
+        emit(Op::And);
+        if (!build(1) || !build(2)) return false;
+        emit(Op::Shl);
+        if (!mask()) return false;
+        emit(Op::And);
+        emit(Op::Or);
+        return true;
+    }
+    case IR::Opcode::UClamp32:
+    case IR::Opcode::SClamp32: {
+        const bool is_signed = inst->GetOpcode() == IR::Opcode::SClamp32;
+        if (!build(0) || !build(1)) return false;
+        emit(is_signed ? Op::SMax : Op::UMax);
+        if (!build(2)) return false;
+        emit(is_signed ? Op::SMin : Op::UMin);
+        return true;
+    }
+    case IR::Opcode::SelectU1:
+    case IR::Opcode::SelectU32:
+        if (!build(0) || !build(1) || !build(2)) return false;
+        emit(Op::Select);
+        return true;
+    default:
+        g_handle_failure = inst->GetOpcode();
+        return false;
+    }
+}
+
+// Programs built while making the texture instructions of the pass running on this thread.
+thread_local std::vector<std::vector<u32>>* g_handle_programs{};
+
 std::optional<u32> TryGetConstant(IR::Value& value, Environment& env) {
     const IR::Inst* inst = value.InstRecursive();
     if (inst->GetOpcode() != IR::Opcode::GetCbufU32) {
@@ -525,10 +909,44 @@ TextureInst MakeInst(Environment& env, IR::Block* block, IR::Inst& inst, const H
     if (IsBindless(inst)) {
         const std::optional<ConstBufferAddr> track_addr{TrackCached(inst.Arg(0), env, host_info, block)};
 
-        if (!track_addr) {
-            throw NotImplementedException("Failed to track bindless texture constant buffer");
-        } else {
+        if (track_addr) {
             addr = *track_addr;
+        } else {
+            // Read from a storage buffer (a handle table): computed for each draw.
+            std::vector<u32> code;
+            int budget = 4096;
+            g_handle_failure = IR::Opcode::Void;
+            if (!g_handle_programs || !BuildHandleProgram(inst.Arg(0), code, budget, block)) {
+                // Not computable for a whole draw (a handle that changes in a loop): keep the
+                // earlier search through loads, which at least gives the shader a texture.
+                const std::optional<ConstBufferAddr> legacy{IR::BreadthFirstSearch(inst.Arg(0),
+                    [&env, &host_info, block](const IR::Inst* candidate) {
+                        return TryGetConstBuffer(candidate, env, host_info, block);
+                    }, [&](IR::Value operand) { return ResolveConditionalPhi(operand, block); })};
+                if (!legacy) {
+                    throw NotImplementedException("Failed to track bindless texture constant buffer (stopped at {}, budget {})",
+                                                  g_handle_failure, budget);
+                }
+                return TextureInst{.cbuf = *legacy, .inst = &inst, .block = block};
+            }
+            auto& programs{*g_handle_programs};
+            auto it{std::ranges::find(programs, code)};
+            if (it == programs.end()) {
+                programs.push_back(std::move(code));
+                it = programs.end() - 1;
+            }
+            addr = ConstBufferAddr{
+                .index = 0,
+                .offset = 0,
+                .shift_left = 0,
+                .secondary_index = 0,
+                .secondary_offset = 0,
+                .secondary_shift_left = 0,
+                .dynamic_offset = {},
+                .count = 1,
+                .has_secondary = false,
+                .handle_program = static_cast<u32>(it - programs.begin()),
+            };
         }
     } else {
         addr = ConstBufferAddr{
@@ -724,6 +1142,7 @@ void TexturePass(Environment& env, IR::Program& program, const HostTranslateInfo
     env.handle_cache.clear();
     env.track_cache.clear();
     TextureInstVector to_replace;
+    g_handle_programs = &program.info.handle_programs;
     for (IR::Block* const block : program.post_order_blocks) {
         env.track_cache.clear();
         for (IR::Inst& inst : block->Instructions()) {
@@ -881,6 +1300,7 @@ void TexturePass(Environment& env, IR::Program& program, const HostTranslateInfo
                     .secondary_shift_left = cbuf.secondary_shift_left,
                     .count = count,
                     .size_shift = size_shift,
+                    .handle_program = cbuf.handle_program,
                 });
                 flags.is_integer.Assign(is_integer ? 1 : 0);
             }