diff --git a/FEXCore/Source/Interface/Context/Context.h b/FEXCore/Source/Interface/Context/Context.h index 9342cd394..f7ef33bfa 100644 --- a/FEXCore/Source/Interface/Context/Context.h +++ b/FEXCore/Source/Interface/Context/Context.h @@ -59,8 +59,9 @@ namespace Validation { namespace FEXCore::Context { struct FEX_PACKED ExitFunctionLinkData { - uint64_t HostBranch; + uint64_t HostCode; uint64_t GuestRIP; + int64_t CallerOffset; }; struct CustomIRResult { @@ -251,14 +252,6 @@ public: static bool ThreadRemoveCodeEntry(FEXCore::Core::InternalThreadState* Thread, uint64_t GuestRIP); - template - static uint64_t ThreadExitFunctionLink(FEXCore::Core::CpuStateFrame* Frame, ExitFunctionLinkData* Record) { - auto Thread = Frame->Thread; - auto lk = GuardSignalDeferringSection(static_cast(Thread->CTX)->CodeInvalidationMutex, Thread); - - return Fn(Frame, Record); - } - // Wrapper which takes CpuStateFrame instead of InternalThreadState and unique_locks CodeInvalidationMutex // Must be called from owning thread static void ThreadRemoveCodeEntryFromJit(FEXCore::Core::CpuStateFrame* Frame, uint64_t GuestRIP) { diff --git a/FEXCore/Source/Interface/Core/JIT/BranchOps.cpp b/FEXCore/Source/Interface/Core/JIT/BranchOps.cpp index ef074f746..1381edbcf 100644 --- a/FEXCore/Source/Interface/Core/JIT/BranchOps.cpp +++ b/FEXCore/Source/Interface/Core/JIT/BranchOps.cpp @@ -65,17 +65,91 @@ DEF_OP(ExitFunction) { br(TMP2); } else { #endif - // Align to 16 byte to allow atomic patching of the following 16 byte - // of code (excluding the RIP data) on platforms that support LSE2 - Align16B(); + // In order to support direct branches without constantly hitting the L1 cache, we emit a call to a block linker, + // this will compile the branch target block when it is hit and replace the branch to the linker at the callsite + // with a direct branch to the destination block. Upon invalidation of the target block the backpatch is undone. + // + // In addition, to avoid needing to lookup in the cache for returns and any indirect branch prediction penalty, + // a shadow stack of pairs is maintained, acting as a first level cache for any + // return operations. As the guest may not balance calls and returns exactly, an exception handler is expected to + // be installed by the frontend, to reset the shadow stack to the middle of its valid bounds on overflow/underflow. + // This shadow stack is also cleared on block invalidation operations or codebuffer switches, to ensure all pointed-to + // host code is always valid. + + // This code will be backpatched by Arm64JITCore_ExitFunctionLink, below is an enumeration of all the possible cases. + // Jump thunks are emitted in JIT.cpp after compilation of the entire multiblock. + // + // Call with known return block - unlinked + // 00: adr TMP1, 0xC + // 04: stp RetReg, TMP1, [SpReg, -0x10]! + // 08: bl JmpThunk00 + // JmpThunk00: + // 00: b 0x8 + // 04: br TMP1 + // 08: ldr TMP1, + // 0c: blr TMP1 + // 10: HostCode + // 18: GuestRIP + // 20: CallerOffset + // + // Call with known return block after backpatching - linked in branch immediate range + // 00: adr TMP1, 0xC + // 04: stp RetReg, TMP1, [SpReg, -0x10]! + // 08: bl HostCode - MODIFIED + // + // Call with known return block after backpatching - linked out of range + // 00: adr TMP1, 0xC + // 04: stp RetReg, TMP1, [SpReg, -0x10]! + // 08: bl JmpThunk00 + // JmpThunk00: + // 00: ldr TMP1, 0x10 - MODIFIED 2nd + // 04: br TMP1 + // 08: ldr TMP1, + // 0c: blr TMP1 + // 10: HostCode - MODIFIED 1st + // 18: GuestRIP + // 20: CallerOffset + // + // Jump - unlinked + // 00: b JmpThunk00 + // JmpThunk00: + // 00: b 0x8 + // 04: br TMP1 + // 08: ldr TMP1, + // 0c: blr TMP1 + // 10: HostCode + // 18: GuestRIP + // 20: CallerOffset + // + // Jump after backpatching - linked in branch immediate range + // 00: b HostCode - MODIFIED + // + // Jump after backpatching - linked out of range + // 00: b JmpThunk00 + // JmpThunk00: + // 00: ldr TMP1, 0x10 - MODIFIED 2nd + // 04: br TMP1 + // 08: ldr TMP1, + // 0c: blr TMP1 + // 10: HostCode - MODIFIED 1st + // 18: GuestRIP + // 20: CallerOffset ARMEmitter::ForwardLabel l_BranchHost; - ldr(TMP1, &l_BranchHost); - blr(TMP1); + ARMEmitter::ForwardLabel l_CallReturn; + if (Op->Hint == IR::BranchHint::Call) { + if (!Op->CallReturnBlock.IsInvalid()) { + auto CallReturnAddressReg = GetReg(Op->CallReturnAddress).X(); + PendingCallReturnTargetLabel = &CallReturnTargets.try_emplace(Op->CallReturnBlock.ID()).first->second; + adr(TMP1, &l_CallReturn); + stp(CallReturnAddressReg, TMP1, REG_CALLRET_SP, -0x10); + } else { + stp(ARMEmitter::XReg::zr, ARMEmitter::XReg::zr, REG_CALLRET_SP, -0x10); + } + } - Bind(&l_BranchHost); - dc64(ThreadState->CurrentFrame->Pointers.Common.ExitFunctionLinker); - dc64(NewRIP); + EmitLinkedBranch(NewRIP, Op->Hint == IR::BranchHint::Call); + Bind(&l_CallReturn); #ifdef _M_ARM_64EC } #endif diff --git a/FEXCore/Source/Interface/Core/JIT/JIT.cpp b/FEXCore/Source/Interface/Core/JIT/JIT.cpp index 47aadf4a1..e0e4438bd 100644 --- a/FEXCore/Source/Interface/Core/JIT/JIT.cpp +++ b/FEXCore/Source/Interface/Core/JIT/JIT.cpp @@ -495,58 +495,109 @@ void Arm64JITCore::Op_Unhandled(const IR::IROp_Header* IROp, IR::Ref Node) { } } -static void DirectBlockDelinker(FEXCore::Core::CpuStateFrame* Frame, FEXCore::Context::ExitFunctionLinkData* Record) { - // Emit new 16 bytes of code to a temporary patch, then atomically apply it - __uint128_t Patch; - ARMEmitter::Emitter emit((uint8_t*)&Patch, sizeof(Patch)); - emit.ldr(TMP1, 8); // PC-relative value pointing to constant after blr - emit.blr(TMP1); - emit.dc64(Frame->Pointers.Common.ExitFunctionLinker); +static void DirectBlockDelinker(FEXCore::Core::CpuStateFrame* Frame, FEXCore::Context::ExitFunctionLinkData* Record, bool Call) { + uintptr_t JumpThunkStartAddress = reinterpret_cast(Record) - 0x10; + uintptr_t CallerAddress = JumpThunkStartAddress + Record->CallerOffset; + auto BranchOffset = JumpThunkStartAddress / 4 - CallerAddress / 4; - auto branch = reinterpret_cast<__uint128_t*>((uintptr_t)Record - 8); - std::atomic_ref<__uint128_t>(*branch).store(Patch, std::memory_order::relaxed); - ARMEmitter::Emitter::ClearICache((void*)branch, sizeof(*branch)); + // Replace the patched callsite with a branch to the jump thunk. + uint32_t BranchInst = 0; + ARMEmitter::Emitter BranchEmit(reinterpret_cast(&BranchInst), 4); + if (Call) { + BranchEmit.bl(BranchOffset); + } else { + BranchEmit.b(BranchOffset); + } + + std::atomic_ref(*reinterpret_cast(CallerAddress)).store(BranchInst, std::memory_order::relaxed); + ARMEmitter::Emitter::ClearICache(reinterpret_cast(CallerAddress), 4); } -static uint64_t Arm64JITCore_ExitFunctionLink(FEXCore::Core::CpuStateFrame* Frame, FEXCore::Context::ExitFunctionLinkData* Record) { - auto Thread = Frame->Thread; - auto Lock = Thread->LookupCache->AcquireLock(); +static void IndirectBlockDelinker(FEXCore::Core::CpuStateFrame* Frame, FEXCore::Context::ExitFunctionLinkData* Record) { + uintptr_t JumpThunkStartAddress = reinterpret_cast(Record) - 0x10; + uint32_t BranchInst = 0; + ARMEmitter::Emitter BranchEmit(reinterpret_cast(&BranchInst), 4); + BranchEmit.b(0x8); + std::atomic_ref(*reinterpret_cast(JumpThunkStartAddress)).store(BranchInst, std::memory_order::relaxed); + ARMEmitter::Emitter::ClearICache(reinterpret_cast(JumpThunkStartAddress), 4); + + // No need to reset HostCode here as the exit linker pointer is stored separately, and if the block is relinked it will be updated. +} + +uint64_t Arm64JITCore::ExitFunctionLink(FEXCore::Core::CpuStateFrame* Frame, FEXCore::Context::ExitFunctionLinkData* Record) { + auto Thread = Frame->Thread; bool TFSet = Thread->CurrentFrame->State.flags[X86State::RFLAG_TF_RAW_LOC]; uintptr_t HostCode {}; auto GuestRip = Record->GuestRIP; - if (!TFSet) { - HostCode = Thread->LookupCache->FindBlock(GuestRip); - } - - if (TFSet || !HostCode) { + if (TFSet) { // If TF is set, the cache must be skipped as different code needs to be generated. Frame->State.rip = GuestRip; return Frame->Pointers.Common.DispatcherLoopTop; + } else { + HostCode = Thread->LookupCache->FindBlock(GuestRip); + if (!HostCode) { + // Hold a reference to the code buffer, to avoid linking unmapped code if compilation triggers a recreation. + auto CodeBuffer = static_cast(Thread->CPUBackend.get())->CurrentCodeBuffer; + HostCode = static_cast(Thread->CTX)->CompileBlock(Frame, GuestRip, 0); + if (Thread->LookupCache->Shared != CodeBuffer->LookupCache.get()) { + return HostCode; + } + } } - uintptr_t branch = (uintptr_t)(Record)-8; - LOGMAN_THROW_A_FMT((branch % 16) == 0, "Incorrect alignment for block linking record"); + // See ExitFunction in BranchOps.cpp for an assembly level view of the handled cases. + uintptr_t JumpThunkStartAddress = reinterpret_cast(Record) - 0x10; + uintptr_t CallerAddress = JumpThunkStartAddress + Record->CallerOffset; + auto BranchOffset = HostCode / 4 - CallerAddress / 4; - auto offset = HostCode / 4 - branch / 4; - if (ARMEmitter::Emitter::IsInt26(offset)) { - // This is the optimal case, where the target can be encoded in a single instruction. - // Atomically patch the code with a relative branch. - const uint32_t Patch = (0b0001'01 << 26) | (offset & ((1u << 26) - 1)); - std::atomic_ref(*reinterpret_cast(branch)).store(Patch, std::memory_order::relaxed); - ARMEmitter::Emitter::ClearICache((void*)branch, 4); + uint32_t ExpectedKnownCallMarkerInst = 0; + ARMEmitter::Emitter ExpectedKnownCallMarkerEmit(reinterpret_cast(&ExpectedKnownCallMarkerInst), 4); + ExpectedKnownCallMarkerEmit.adr(TMP1, 0xC); + + // Lock here is necessary to prevent simultaneous linking and delinking + auto lk = Thread->LookupCache->AcquireLock(); + + // For non-calls, this would extend into the block's code, however that's fine as an out-of-range adr would never + // be generated avoiding any false positives. + uintptr_t KnownCallMarkerAddr = CallerAddress - 0x8; + uint32_t KnownCallMarkerInst = *reinterpret_cast(KnownCallMarkerAddr); + if (ARMEmitter::Emitter::IsInt26(BranchOffset)) { + // Directly patch the callsite with the appropriate branch instruction. + uint32_t BranchInst = 0; + ARMEmitter::Emitter BranchEmit(reinterpret_cast(&BranchInst), 4); + + if (KnownCallMarkerInst == ExpectedKnownCallMarkerInst) { + BranchEmit.bl(BranchOffset); + Thread->LookupCache->AddBlockLink(GuestRip, Record, [](FEXCore::Core::CpuStateFrame* Frame, FEXCore::Context::ExitFunctionLinkData* Record) { + DirectBlockDelinker(Frame, Record, true); + }); + } else { + BranchEmit.b(BranchOffset); + Thread->LookupCache->AddBlockLink(GuestRip, Record, [](FEXCore::Core::CpuStateFrame* Frame, FEXCore::Context::ExitFunctionLinkData* Record) { + DirectBlockDelinker(Frame, Record, false); + }); + } + + std::atomic_ref(*reinterpret_cast(CallerAddress)).store(BranchInst, std::memory_order::relaxed); + ARMEmitter::Emitter::ClearICache(reinterpret_cast(CallerAddress), 4); } else { - // fallback case - do a soft-er link by patching the pointer - std::atomic_ref(Record->HostBranch).store(HostCode, std::memory_order::seq_cst); + // This case is common between calls and jumps as the thunk callsite can be left untouched. + std::atomic_ref(Record->HostCode).store(HostCode, std::memory_order::seq_cst); #ifdef _M_ARM_64 // Make memory write visible to other threads reading the same location - asm volatile("dc cvau, %0; dsb ish" : : "r"(Record->HostBranch) :); + asm volatile("dc cvau, %0; dsb ish" : : "r"(Record->HostCode) :); #endif - } - // Add de-linking handler - Thread->LookupCache->AddBlockLink(GuestRip, Record, DirectBlockDelinker); + uint32_t LdrInst = 0; + ARMEmitter::Emitter LdrEmit(reinterpret_cast(&LdrInst), 4); + LdrEmit.ldr(TMP1, reinterpret_cast(&Record->HostCode) - JumpThunkStartAddress); + std::atomic_ref(*reinterpret_cast(JumpThunkStartAddress)).store(LdrInst, std::memory_order::relaxed); + ARMEmitter::Emitter::ClearICache(reinterpret_cast(JumpThunkStartAddress), 4); + + Thread->LookupCache->AddBlockLink(GuestRip, Record, IndirectBlockDelinker); + } return HostCode; } @@ -598,7 +649,7 @@ Arm64JITCore::Arm64JITCore(FEXCore::Context::ContextImpl* ctx, FEXCore::Core::In Common.SyscallHandlerObj = reinterpret_cast(CTX->SyscallHandler); Common.SyscallHandlerFunc = PMF.GetVTableEntry(CTX->SyscallHandler); } - Common.ExitFunctionLink = reinterpret_cast(&Context::ContextImpl::ThreadExitFunctionLink); + Common.ExitFunctionLink = reinterpret_cast(&Arm64JITCore::ExitFunctionLink); // Platform Specific auto& AArch64 = ThreadState->CurrentFrame->Pointers.AArch64; @@ -767,6 +818,7 @@ CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size JumpTargets.clear(); CallReturnTargets.clear(); + PendingJumpThunks.clear(); uint32_t SSACount = IR->GetSSACount(); this->Entry = Entry; @@ -890,6 +942,30 @@ CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size } PendingTargetLabel = nullptr; + ARMEmitter::ForwardLabel l_ExitLink; + for (auto& PendingJumpThunk : PendingJumpThunks) { + // Align as 64-bit atomics are used on the HostCode field. + Align(8); + + ARMEmitter::ForwardLabel l_DoLink; + uint64_t ThunkAddress = GetCursorAddress(); + Bind(&PendingJumpThunk.Label); + b(&l_DoLink); + br(TMP1); + Bind(&l_DoLink); + ldr(TMP1, &l_ExitLink); + blr(TMP1); + + // This is a ExitFunctionLinkData struct + Bind(&l_ExitLink); + dc64(0); // HostCode + dc64(PendingJumpThunk.GuestRIP); // GuestRIP + dc64(PendingJumpThunk.CallerAddress - ThunkAddress); // CallerOffset + } + + Bind(&l_ExitLink); + dc64(ThreadState->CurrentFrame->Pointers.Common.ExitFunctionLinker); + // CodeSize not including the header or tail data. const uint64_t CodeOnlySize = GetCursorAddress() - CodeBegin; diff --git a/FEXCore/Source/Interface/Core/JIT/JITClass.h b/FEXCore/Source/Interface/Core/JIT/JITClass.h index 64d1ab733..4b3edaf63 100644 --- a/FEXCore/Source/Interface/Core/JIT/JITClass.h +++ b/FEXCore/Source/Interface/Core/JIT/JITClass.h @@ -31,6 +31,10 @@ namespace FEXCore::Core { struct InternalThreadState; } +namespace FEXCore::Context { +struct ExitFunctionLinkData; +} + namespace FEXCore::CPU { class Arm64JITCore final : public CPUBackend, public Arm64Emitter { public: @@ -66,8 +70,17 @@ private: fextl::map JumpTargets; fextl::map CallReturnTargets; + struct PendingJumpThunk { + uint64_t CallerAddress; + uint64_t GuestRIP; + ARMEmitter::ForwardLabel Label; + }; + fextl::vector PendingJumpThunks; + Utils::PoolBufferWithTimedRetirement TempAllocator; + static uint64_t ExitFunctionLink(FEXCore::Core::CpuStateFrame* Frame, FEXCore::Context::ExitFunctionLinkData* Record); + [[nodiscard]] ARMEmitter::Register GetReg(IR::PhysicalRegister Reg) const { LOGMAN_THROW_A_FMT(Reg.Class == IR::GPRFixedClass.Val || Reg.Class == IR::GPRClass.Val, "Unexpected Class: {}", Reg.Class); @@ -292,6 +305,17 @@ private: uint32_t End; }; + void EmitLinkedBranch(uint64_t GuestRIP, bool Call) { + PendingJumpThunks.push_back({GetCursorAddress(), GuestRIP, {}}); + auto& Thunk = PendingJumpThunks.back(); + Bind(&Thunk.Label); + if (Call) { + bl(&Thunk.Label); + } else { + b(&Thunk.Label); + } + } + // This is purely a debugging aid for developers to see if they are in JIT code space when inspecting raw memory void EmitDetectionString(); IR::RegisterAllocationPass* RAPass {};