Merge remote-tracking branch 'origin/nx2' into codex/upstream-loading-hotkeys

1 files changed, 428 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..8a72f565dc 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,389 @@ 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, and how deep
+// in nested BuildHandleProgram calls the build is.
+thread_local std::unordered_map<const IR::Inst*, u32> g_handle_slots;
+thread_local int g_handle_depth{};
+
+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);
+    }
+    const bool outermost = g_handle_depth++ == 0;
+    if (outermost) {
+        g_handle_slots.clear();
+    }
+    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);
+        }
+    }
+    --g_handle_depth;
+    if (outermost) {
+        g_handle_slots.clear();
+    }
+    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 +911,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 +1144,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 +1302,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);
             }