rpcsx/rpcs3/Emu/CPU/Backends/AArch64/AArch64JIT.cpp

#include "stdafx.h"

#include "AArch64JIT.h"
#include "AArch64ASM.h"

LOG_CHANNEL(jit_log, "JIT");

#define STDOUT_DEBUG 0

#define DPRINT1(...)\
    do {\
        printf(__VA_ARGS__);\
        printf("\n");\
        fflush(stdout);\
    } while (0)

#if STDOUT_DEBUG
#define DPRINT DPRINT1
#else
#define DPRINT jit_log.trace
#endif

namespace aarch64
{
    using instruction_info_t = GHC_frame_preservation_pass::instruction_info_t;
    using function_info_t = GHC_frame_preservation_pass::function_info_t;

    GHC_frame_preservation_pass::GHC_frame_preservation_pass(const config_t& configuration)
        : m_config(configuration)
    {}

    void GHC_frame_preservation_pass::reset()
    {
        m_visited_functions.clear();
    }

    void GHC_frame_preservation_pass::force_tail_call_terminators(llvm::Function& f)
    {
        // GHC functions are not call-stack preserving and can therefore never return if they make any external calls at all.
        // Replace every terminator clause with a tail call explicitly. This is already done for X64 to work, but better safe than sorry.
        for (auto& bb : f)
        {
            auto bit = bb.begin(), prev = bb.end();
            for (; bit != bb.end(); prev = bit, ++bit)
            {
                if (prev == bb.end())
                {
                    continue;
                }

                if (llvm::isa<llvm::ReturnInst>(&*bit))
                {
                    if (auto ci = llvm::dyn_cast<llvm::CallInst>(&*prev))
                    {
                        // This is a "ret" that is coming after a "call" to another funciton.
                        // Enforce that it must be a tail call.
                        if (!ci->isTailCall())
                        {
                            ci->setTailCall();
                        }
                    }
                }
            }
        }
    }

    function_info_t GHC_frame_preservation_pass::preprocess_function(const llvm::Function& f)
    {
        function_info_t result{};
        result.instruction_count = f.getInstructionCount();

        // Blanket exclusions. Stubs or dispatchers that do not compute anything themselves.
        if (f.getName() == "__spu-null")
        {
            // Don't waste the effort processing this stub. It has no points of concern
            result.num_external_calls = 1;
            return result;
        }

        if (m_config.use_stack_frames)
        {
            // Stack frame estimation. SPU code can be very long and consumes several KB of stack.
            u32 stack_frame_size = 128u;
            // Actual ratio is usually around 1:4
            const u32 expected_compiled_instr_count = f.getInstructionCount() * 4;
            // Because GHC doesn't preserve stack (all stack is scratch), we know we'll start to spill once we go over the number of actual regs.
            // We use a naive allocator that just assumes each instruction consumes a register slot. We "spill" every 32 instructions.
            // FIXME: Aggressive spill is only really a thing with vector operations. We can detect those instead.
            // A proper fix is to port this to a MF pass, but I have PTSD from working at MF level.
            const u32 spill_pages = (expected_compiled_instr_count + 127u) / 128u;
            stack_frame_size *= std::min(spill_pages, 32u); // 128 to 4k dynamic. It is unlikely that any frame consumes more than 4096 bytes

            result.stack_frame_size = stack_frame_size;
        }

        result.instruction_count = f.getInstructionCount();
        result.num_external_calls = 0;

        // The LR is not spared by LLVM in cases where there is a lot of spilling.
        // This is much easier to manage with a custom LLVM branch as we can just mark X30 as off-limits as a GPR.
        // This is another thing to be moved to a MachineFunction pass. Ideally we should check the instruction stream for writes to LR and reload it on exit.
        // For now, assume it is dirtied if the function is of any reasonable length.
        result.clobbers_x30 = result.instruction_count > 32;
        result.is_leaf = true;

        for (auto& bb : f)
        {
            for (auto& inst : bb)
            {
                if (auto ci = llvm::dyn_cast<llvm::CallInst>(&inst))
                {
                    if (llvm::isa<llvm::InlineAsm>(ci->getCalledOperand()))
                    {
                        // Inline ASM blocks are ignored
                        continue;
                    }

                    result.num_external_calls++;
                    if (ci->isTailCall())
                    {
                        // This is not a leaf if it has at least one exit point / terminator that is not a return instruction.
                        result.is_leaf = false;
                    }
                    else
                    {
                        // Returning calls always clobber x30
                        result.clobbers_x30 = true;
                    }
                }
            }
        }

        return result;
    }

    instruction_info_t GHC_frame_preservation_pass::decode_instruction(const llvm::Function& f, const llvm::Instruction* i)
    {
        instruction_info_t result{};
        if (auto ci = llvm::dyn_cast<llvm::CallInst>(i))
        {
            // Watch out for injected ASM blocks...
            if (llvm::isa<llvm::InlineAsm>(ci->getCalledOperand()))
            {
                // Not a real call. This is just an insert of inline asm
                return result;
            }

            result.is_call_inst = true;
            result.is_returning = true;
            result.preserve_stack = !ci->isTailCall();
            result.callee = ci->getCalledFunction();
            result.is_tail_call = ci->isTailCall();

            if (!result.callee)
            {
                // Indirect call (call from raw value).
                result.is_indirect = true;
                result.callee_is_GHC = ci->getCallingConv() == llvm::CallingConv::GHC;
                result.callee_name = "__indirect_call";
            }
            else
            {
                result.callee_is_GHC = result.callee->getCallingConv() == llvm::CallingConv::GHC;
                result.callee_name = result.callee->getName().str();
            }
            return result;
        }

        if (auto bi = llvm::dyn_cast<llvm::BranchInst>(i))
        {
            // More likely to jump out via an unconditional...
            if (!bi->isConditional())
            {
                ensure(bi->getNumSuccessors() == 1);
                auto targetbb = bi->getSuccessor(0);

                result.callee = targetbb->getParent();
                result.callee_name = result.callee->getName().str();
                result.is_call_inst = result.callee_name != f.getName();
            }

            return result;
        }

        if (auto bi = llvm::dyn_cast<llvm::IndirectBrInst>(i))
        {
            // Very unlikely to be the same function. Can be considered a function exit.
            ensure(bi->getNumDestinations() == 1);
            auto targetbb = ensure(bi->getSuccessor(0)); // This is guaranteed to fail but I've yet to encounter this

            result.callee = targetbb->getParent();
            result.callee_name = result.callee->getName().str();
            result.is_call_inst = result.callee_name != f.getName();
            return result;
        }

        if (auto bi = llvm::dyn_cast<llvm::CallBrInst>(i))
        {
            ensure(bi->getNumSuccessors() == 1);
            auto targetbb = bi->getSuccessor(0);

            result.callee = targetbb->getParent();
            result.callee_name = result.callee->getName().str();
            result.is_call_inst = result.callee_name != f.getName();
            return result;
        }

        if (auto bi = llvm::dyn_cast<llvm::InvokeInst>(i))
        {
            ensure(bi->getNumSuccessors() == 2);
            auto targetbb = bi->getSuccessor(0);

            result.callee = targetbb->getParent();
            result.callee_name = result.callee->getName().str();
            result.is_call_inst = result.callee_name != f.getName();
            return result;
        }

        return result;
    }

    gpr GHC_frame_preservation_pass::get_base_register_for_call(const std::string& callee_name, gpr default_reg)
    {
        // We go over the base_register_lookup table and find the first matching pattern
        for (const auto& pattern : m_config.base_register_lookup)
        {
            if (callee_name.starts_with(pattern.first))
            {
                return pattern.second;
            }
        }

        return default_reg;
    }

    void GHC_frame_preservation_pass::run(llvm::IRBuilder<>* irb, llvm::Function& f)
    {
        if (f.getCallingConv() != llvm::CallingConv::GHC)
        {
            // If we're not doing GHC, the calling conv will have stack fixup on its own via prologue/epilogue
            return;
        }

        if (f.getInstructionCount() == 0)
        {
            // Nothing to do. Happens with placeholder functions such as branch patchpoints
            return;
        }

        const auto this_name = f.getName().str();
        if (m_visited_functions.find(this_name) != m_visited_functions.end())
        {
            // Already processed. Only useful when recursing which is currently not used.
            DPRINT("Function %s was already processed. Skipping.\n", this_name.c_str());
            return;
        }

        if (this_name != "__spu-null") // This name is meaningless and doesn't uniquely identify a function
        {
            m_visited_functions.insert(this_name);
        }

        if (m_config.exclusion_callback && m_config.exclusion_callback(this_name))
        {
            // Function is explicitly excluded
            return;
        }

        // Preprocessing.
        auto function_info = preprocess_function(f);
        if (function_info.num_external_calls == 0 && function_info.stack_frame_size == 0)
        {
            // No stack frame injection and no external calls to patch up. This is a leaf function, nothing to do.
            DPRINT("Ignoring function %s", this_name.c_str());
            return;
        }

        // Force tail calls on all terminators
        force_tail_call_terminators(f);

        // Check for leaves
        if (function_info.is_leaf && !m_config.use_stack_frames)
        {
            // Sanity check. If this function had no returning calls, it should have been omitted from processing.
            ensure(function_info.clobbers_x30, "Function has no terminator and no non-tail calls but was allowed for frame processing!");
            DPRINT("Function %s is a leaf.", this_name.c_str());
            process_leaf_function(irb, f);
            return;
        }

        // Asm snippets for patching stack frame
        ASMBlock frame_prologue, frame_epilogue;

        if (function_info.stack_frame_size > 0)
        {
            // NOTE: The stack frame here is purely optional, we can pre-allocate scratch on the gateway.
            // However, that is an optimization for another time, this helps make debugging easier.
            frame_prologue.sub(sp, sp, UASM::Imm(function_info.stack_frame_size));
            frame_epilogue.add(sp, sp, UASM::Imm(function_info.stack_frame_size));

            // Emit the frame prologue. We use a BB here for extra safety as it solves the problem of backwards jumps re-executing the prologue.
            auto functionStart = &f.front();
            auto prologueBB = llvm::BasicBlock::Create(f.getContext(), "", &f, functionStart);
            irb->SetInsertPoint(prologueBB, prologueBB->begin());
            frame_prologue.insert(irb, f.getContext());
            irb->CreateBr(functionStart);
        }

        // Now we start processing
        bool terminator_found = false;
        for (auto& bb : f)
        {
            for (auto bit = bb.begin(); bit != bb.end();)
            {
                const auto instruction_info = decode_instruction(f, &(*bit));
                if (!instruction_info.is_call_inst)
                {
                    ++bit;
                    continue;
                }

                std::string callee_name = "__unknown";
                if (const auto cf = instruction_info.callee)
                {
                    callee_name = cf->getName().str();
                    if (cf->hasFnAttribute(llvm::Attribute::AlwaysInline) || callee_name.starts_with("llvm."))
                    {
                        // Always inlined call. Likely inline Asm. Skip
                        ++bit;
                        continue;
                    }

                    // Technically We should also ignore any host functions linked in, usually starting with ppu_ or spu_ prefix.
                    // However, there is not much guarantee that those are safe with only rare exceptions, and it doesn't hurt to patch the frame around them that much anyway.
                }

                if (instruction_info.preserve_stack)
                {
                    // Non-tail call. If we have a stack allocated, we preserve it across the call
                    ++bit;
                    continue;
                }

                ensure(instruction_info.is_tail_call);
                terminator_found = true;

                // Now we patch the call if required. For normal calls that 'return' (i.e calls to C/C++ ABI), we do not patch them as they will manage the stack themselves (callee-managed)
                bit = patch_tail_call(irb, f, bit, instruction_info, function_info, frame_epilogue);

                // Next
                if (bit != bb.end())
                {
                    ++bit;
                }
            }
        }

        if (!terminator_found)
        {
            // If we got here, we must be using stack frames.
            ensure(function_info.is_leaf && function_info.stack_frame_size > 0, "Leaf function was processed without using stack frames!");

            // We want to insert a frame cleanup at the tail at every return instruction we find.
            for (auto& bb : f)
            {
                for (auto& i : bb)
                {
                    if (is_ret_instruction(&i))
                    {
                        irb->SetInsertPoint(&i);
                        frame_epilogue.insert(irb, f.getContext());
                    }
                }
            }
        }
    }

    llvm::BasicBlock::iterator
    GHC_frame_preservation_pass::patch_tail_call(
        llvm::IRBuilder<>* irb,
        llvm::Function& f,
        llvm::BasicBlock::iterator where,
        const instruction_info_t& instruction_info,
        const function_info_t& function_info,
        const UASM& frame_epilogue)
    {
        auto ci = llvm::dyn_cast<llvm::CallInst>(where);
        irb->SetInsertPoint(ensure(ci));

        const auto this_name = f.getName().str();

        // Insert breadcrumb info before the call
        // WARNING: This can corrupt the call because LLVM somehow ignores the clobbered register during a call instruction for some reason
        // In case of a blr on x27..x29 you can end up corrupting the binary, but it is invaluable for debugging.
        // Debug frames are disabled in shipping code so this is not a big deal.
        if (m_config.debug_info)
        {
            // Call-chain tracing
            ASMBlock c;
            c.mov(x29, x28);
            c.mov(x28, x27);
            c.adr(x27, UASM::Reg(pc));
            c.insert(irb, f.getContext());
        }

        // Clean up any injected frames before the call
        if (function_info.stack_frame_size > 0)
        {
            frame_epilogue.insert(irb, f.getContext());
        }

        // Insert the next piece after the call, before the ret
        ++where;
        ensure(llvm::isa<llvm::ReturnInst>(where));
        irb->SetInsertPoint(llvm::dyn_cast<llvm::Instruction>(where));

        if (instruction_info.callee_is_GHC &&                    // Calls to C++ ABI will always return
            !instruction_info.is_indirect &&                     // We don't know enough when calling indirectly to know if we'll return or not
            !is_faux_function(instruction_info.callee_name))     // Ignore branch patch-points and imposter functions. Their behavior is unreliable.
        {
            // We're making a one-way call. This branch shouldn't even bother linking as it will never return here.
            ASMBlock c;
            c.brk(0x99);
            c.insert(irb, f.getContext());
            return where;
        }

        // Patch the return path. No GHC call shall ever return to another. If we reach the function endpoint, immediately abort to GW
        auto thread_base_reg = get_base_register_for_call(f.getName().str());
        auto arg_index = static_cast<int>(thread_base_reg) - static_cast<int>(x19);
        ASMBlock c;

        auto thread_arg = ensure(f.getArg(arg_index)); // Guaranteed to hold our original 'thread'
        c.mov(x30, UASM::Var(thread_arg));
        c.ldr(x30, x30, UASM::Imm(m_config.hypervisor_context_offset));
        c.insert(irb, f.getContext());

        // Next
        return where;
    }

    bool GHC_frame_preservation_pass::is_ret_instruction(const llvm::Instruction* i)
    {
        if (llvm::isa<llvm::ReturnInst>(i))
        {
            return true;
        }

        // Check for inline asm invoking "ret". This really shouldn't be a thing, but it is present in SPULLVMRecompiler for some reason.
        if (auto ci = llvm::dyn_cast<llvm::CallInst>(i))
        {
            if (auto asm_ = llvm::dyn_cast<llvm::InlineAsm>(ci->getCalledOperand()))
            {
                if (asm_->getAsmString() == "ret")
                {
                    return true;
                }
            }
        }

        return false;
    }

    bool GHC_frame_preservation_pass::is_inlined_call(const llvm::CallInst* ci)
    {
        const auto callee = ci->getCalledFunction();
        if (!callee)
        {
            // Indirect BLR
            return false;
        }

        const std::string callee_name = callee->getName().str();
        if (callee_name.starts_with("llvm."))
        {
            // Intrinsic
            return true;
        }

        if (callee->hasFnAttribute(llvm::Attribute::AlwaysInline))
        {
            // Assume LLVM always obeys this
            return true;
        }

        return false;
    }

    bool GHC_frame_preservation_pass::is_faux_function(const std::string& function_name)
    {
        // Is it a branch patch-point?
        if (function_name.find("-pp-") != umax)
        {
            return true;
        }

        // Now we search the known imposters list
        if (m_config.faux_function_list.empty())
        {
            return false;
        }

        const auto& x = m_config.faux_function_list;
        return std::find(x.begin(), x.end(), function_name) != x.end();
    }

    void GHC_frame_preservation_pass::process_leaf_function(llvm::IRBuilder<>* irb, llvm::Function& f)
    {
        for (auto &bb : f)
        {
            for (auto bit = bb.begin(); bit != bb.end();)
            {
                auto i = llvm::dyn_cast<llvm::Instruction>(bit);
                if (!is_ret_instruction(i))
                {
                    ++bit;
                    continue;
                }

                // Insert sequence before the return
                irb->SetInsertPoint(llvm::dyn_cast<llvm::Instruction>(bit));

                if (m_config.debug_info)
                {
                    // We need to save the chain return point.
                    ASMBlock c;
                    c.mov(x29, x28);
                    c.mov(x28, x27);
                    c.adr(x27, UASM::Reg(pc));
                    c.insert(irb, f.getContext());
                }

                // Now we need to reload LR. We abuse the function's caller arg set for this to avoid messing with regs too much
                auto thread_base_reg = get_base_register_for_call(f.getName().str());
                auto arg_index = static_cast<int>(thread_base_reg) - static_cast<int>(x19);
                ASMBlock c;

                auto thread_arg = ensure(f.getArg(arg_index)); // Guaranteed to hold our original 'thread'
                c.mov(x30, UASM::Var(thread_arg));
                c.ldr(x30, x30, UASM::Imm(m_config.hypervisor_context_offset));
                c.insert(irb, f.getContext());

                if (bit != bb.end())
                {
                    ++bit;
                }
            }
        }
    }
}