From b7c7789a01e4a0657fa8042e4a7d20e55fc2c225 Mon Sep 17 00:00:00 2001 From: Ryan Houdek Date: Thu, 18 Sep 2025 15:53:43 -0700 Subject: [PATCH] FEXCore/JIT: Supports restarting JIT in case of encoding failure ARM64 branches have fairly small relative distances they can encode. These can be +-1MB, or even +-32KB. The largest relative branch is +-128MB, which we already set as an upper limit of our block JIT cache size. We have for a long time just compiled these without checking with the expectation that things just happen to work. We didn't hit the asserts so it was relatively low priority. Apparently now with Steam and a MaxInst limit of 5000, we are now hitting an assert where we are encoding too large of a range. Implement support for long jumping from anywhere in the JIT for when a long jump tries to be encoded and fails, allowing us to restart the JIT at any moment. This is implemented as a long jump when this singular feature could have gotten away with some sort of invasive check and early exit path for two reasons. For one, that would be even more invasive, effectively doing try-catch logic manually. And two, the next step is supporting JIT buffer overflow for when our block size heuristic fails. This next step will mandate longjump on SIGSEGV (with cooperative interaction with the frontend) from effectively /anywhere/ in the JIT. One of the design goals of the CodeEmitter is that every code emission function doesn't do a size remaining check to allow the compiler to do some very effective optimization of emitting code blocks to memory (and it works!). But we lose the ability to sanely size check. When writing the emitter I knew we were going to need to write this cooperative guard page handler, and we're finally at a point where it needs to be done. This will be in the next PR although. --- .../Interface/Core/JIT/Arm64Relocations.cpp | 2 +- .../Source/Interface/Core/JIT/BranchOps.cpp | 18 +- FEXCore/Source/Interface/Core/JIT/JIT.cpp | 52 +++-- FEXCore/Source/Interface/Core/JIT/JITClass.h | 205 +++++++++++++++++- 4 files changed, 243 insertions(+), 34 deletions(-) diff --git a/FEXCore/Source/Interface/Core/JIT/Arm64Relocations.cpp b/FEXCore/Source/Interface/Core/JIT/Arm64Relocations.cpp index 64c590014..1450d148c 100644 --- a/FEXCore/Source/Interface/Core/JIT/Arm64Relocations.cpp +++ b/FEXCore/Source/Interface/Core/JIT/Arm64Relocations.cpp @@ -63,7 +63,7 @@ void Arm64JITCore::PlaceNamedSymbolLiteral(NamedSymbolLiteralPair& Lit) { auto CurrentCursor = GetCursorAddress(); Lit.MoveABI.NamedSymbolLiteral.Offset = CurrentCursor - CodeData.BlockBegin; - Bind(&Lit.Loc); + BindOrRestart(&Lit.Loc); dc64(Lit.Lit); Relocations.emplace_back(Lit.MoveABI); } diff --git a/FEXCore/Source/Interface/Core/JIT/BranchOps.cpp b/FEXCore/Source/Interface/Core/JIT/BranchOps.cpp index e5534734b..b7e7723b1 100644 --- a/FEXCore/Source/Interface/Core/JIT/BranchOps.cpp +++ b/FEXCore/Source/Interface/Core/JIT/BranchOps.cpp @@ -222,7 +222,7 @@ DEF_OP(CondJump) { auto TrueTargetLabel = JumpTarget(Op->TrueBlock); if (Op->FromNZCV) { - b(MapCC(Op->Cond), TrueTargetLabel); + b_OrRestart(MapCC(Op->Cond), TrueTargetLabel); } else { uint64_t Const; const bool isConst = IsInlineConstant(Op->Cmp2, &Const); @@ -235,16 +235,16 @@ DEF_OP(CondJump) { if (Op->Cond == IR::CondClass::EQ) { LOGMAN_THROW_A_FMT(Const == 0, "CondJump: Expected 0 source"); - cbz(Size, Reg, TrueTargetLabel); + cbz_OrRestart(Size, Reg, TrueTargetLabel); } else if (Op->Cond == IR::CondClass::NEQ) { LOGMAN_THROW_A_FMT(Const == 0, "CondJump: Expected 0 source"); - cbnz(Size, Reg, TrueTargetLabel); + cbnz_OrRestart(Size, Reg, TrueTargetLabel); } else if (Op->Cond == IR::CondClass::TSTZ) { LOGMAN_THROW_A_FMT(Const < 64, "CondJump: Expected valid bit source"); - tbz(Reg, Const, TrueTargetLabel); + tbz_OrRestart(Reg, Const, TrueTargetLabel); } else if (Op->Cond == IR::CondClass::TSTNZ) { LOGMAN_THROW_A_FMT(Const < 64, "CondJump: Expected valid bit source"); - tbnz(Reg, Const, TrueTargetLabel); + tbnz_OrRestart(Reg, Const, TrueTargetLabel); } else { LOGMAN_THROW_A_FMT(false, "CondJump expected simple condition"); } @@ -355,7 +355,7 @@ DEF_OP(ValidateCode) { while (len >= Size) { LoadData(); sub(ARMEmitter::Size::i64Bit, TMP1, TMP1, TMP2); - cbnz(ARMEmitter::Size::i64Bit, TMP1, &Fail); + cbnz_OrRestart(ARMEmitter::Size::i64Bit, TMP1, &Fail); len -= Size; Offset += Size; } @@ -383,10 +383,10 @@ DEF_OP(ValidateCode) { ARMEmitter::ForwardLabel End; LoadConstant(ARMEmitter::Size::i32Bit, Dst, 0); - b(&End); - Bind(&Fail); + b_OrRestart(&End); + BindOrRestart(&Fail); LoadConstant(ARMEmitter::Size::i32Bit, Dst, 1); - Bind(&End); + BindOrRestart(&End); } DEF_OP(ThreadRemoveCodeEntry) { diff --git a/FEXCore/Source/Interface/Core/JIT/JIT.cpp b/FEXCore/Source/Interface/Core/JIT/JIT.cpp index 1adf7cf8b..cc635a611 100644 --- a/FEXCore/Source/Interface/Core/JIT/JIT.cpp +++ b/FEXCore/Source/Interface/Core/JIT/JIT.cpp @@ -28,6 +28,7 @@ $end_info$ #include #include #include +#include #include #include #include @@ -788,14 +789,14 @@ void Arm64JITCore::EmitSuspendInterruptCheck() { ARMEmitter::ForwardLabel l_NoSuspend; cbz(ARMEmitter::Size::i32Bit, TMP2, &l_NoSuspend); brk(SuspendMagic); - Bind(&l_NoSuspend); + (void)Bind(&l_NoSuspend); #endif } void Arm64JITCore::EmitEntryPoint(ARMEmitter::BackwardLabel& HeaderLabel, bool CheckTF) { // Get the address of the JITCodeHeader and store in to the core state. // Two instruction cost, each 1 cycle. - adr(TMP1, &HeaderLabel); + adr_OrRestart(TMP1, &HeaderLabel); str(TMP1, STATE, offsetof(FEXCore::Core::CPUState, InlineJITBlockHeader)); if (CheckTF) { @@ -817,16 +818,25 @@ void Arm64JITCore::EmitEntryPoint(ARMEmitter::BackwardLabel& HeaderLabel, bool C CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size, bool SingleInst, const FEXCore::IR::IRListView* IR, FEXCore::Core::DebugData* DebugData, bool CheckTF) { FEXCORE_PROFILE_SCOPED("Arm64::CompileCode"); - - JumpTargets.clear(); - CallReturnTargets.clear(); - PendingJumpThunks.clear(); - uint32_t SSACount = IR->GetSSACount(); - JumpTargets.resize(IR->GetHeader()->BlockCount, {}); - this->Entry = Entry; this->DebugData = DebugData; this->IR = IR; + RestartControl.NeedsLongJumps = false; + + switch (static_cast(FEXCore::LongJump::SetJump(RestartControl.RestartJump))) { + case RestartOptions::RestartOptionControl::Incoming: + // Nothing + break; + case RestartOptions::RestartOptionControl::NeedsLongJumps: RestartControl.NeedsLongJumps = true; break; + default: LOGMAN_MSG_A_FMT("Unhandled Arm64 restart condition!"); + } + + uint32_t SSACount = IR->GetSSACount(); + JumpTargets.clear(); + CallReturnTargets.clear(); + PendingJumpThunks.clear(); + JumpTargets.resize(IR->GetHeader()->BlockCount, {}); + CodeData.EntryPoints.clear(); // Fairly excessive buffer range to make sure we don't overflow @@ -889,7 +899,7 @@ CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size if (PendingTargetLabel->Backward.Location) { EmitSuspendInterruptCheck(); } - b(PendingTargetLabel); + b_OrRestart(PendingTargetLabel); PendingTargetLabel = nullptr; } @@ -899,14 +909,14 @@ CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size const auto IsReturnTarget = CallReturnTargets.try_emplace(Node).first; if (PendingTargetLabel) { // If there is a fallthrough branch to this block, skip over the entrypoint code. - b(Target); + b_OrRestart(Target); } else if (PendingCallReturnTargetLabel && PendingCallReturnTargetLabel != &IsReturnTarget->second) { // If we just emitted a call, but the block we're now emitting is not the return block so don't fallthrough. - b(PendingCallReturnTargetLabel); + b_OrRestart(PendingCallReturnTargetLabel); } PendingCallReturnTargetLabel = nullptr; - Bind(&IsReturnTarget->second); + BindOrRestart(&IsReturnTarget->second); CodeData.EntryPoints.emplace(BlockStartRIP, GetCursorAddress()); DebugData->GuestOpcodes.push_back({BlockIROp->GuestEntryOffset, GetCursorAddress() - CodeData.BlockBegin}); @@ -915,12 +925,12 @@ CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size if (PendingCallReturnTargetLabel) { // If there is still a pending call return target, then the block we're emitting is not the return block so don't fallthrough. - b(PendingCallReturnTargetLabel); + b_OrRestart(PendingCallReturnTargetLabel); PendingCallReturnTargetLabel = nullptr; } PendingTargetLabel = nullptr; - Bind(Target); + BindOrRestart(Target); } for (auto [CodeNode, IROp] : IR->GetCode(BlockNode)) { @@ -945,7 +955,7 @@ CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size if (PendingTargetLabel->Backward.Location) { EmitSuspendInterruptCheck(); } - b(PendingTargetLabel); + b_OrRestart(PendingTargetLabel); } PendingTargetLabel = nullptr; @@ -956,21 +966,21 @@ CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size ARMEmitter::ForwardLabel l_DoLink; uint64_t ThunkAddress = GetCursorAddress(); - Bind(&PendingJumpThunk.Label); - b(&l_DoLink); + BindOrRestart(&PendingJumpThunk.Label); + b_OrRestart(&l_DoLink); br(TMP1); - Bind(&l_DoLink); + BindOrRestart(&l_DoLink); ldr(TMP1, &l_ExitLink); blr(TMP1); // This is a ExitFunctionLinkData struct - Bind(&l_ExitLink); + BindOrRestart(&l_ExitLink); dc64(0); // HostCode dc64(PendingJumpThunk.GuestRIP); // GuestRIP dc64(PendingJumpThunk.CallerAddress - ThunkAddress); // CallerOffset } - Bind(&l_ExitLink); + BindOrRestart(&l_ExitLink); dc64(ThreadState->CurrentFrame->Pointers.Common.ExitFunctionLinker); // CodeSize not including the header or tail data. diff --git a/FEXCore/Source/Interface/Core/JIT/JITClass.h b/FEXCore/Source/Interface/Core/JIT/JITClass.h index 3ae908390..f2c86a814 100644 --- a/FEXCore/Source/Interface/Core/JIT/JITClass.h +++ b/FEXCore/Source/Interface/Core/JIT/JITClass.h @@ -23,6 +23,7 @@ $end_info$ #include #include #include +#include #include @@ -66,6 +67,15 @@ private: const bool HostSupportsRPRES {}; const bool HostSupportsAFP {}; + struct RestartOptions { + FEXCore::LongJump::JumpBuf RestartJump; + bool NeedsLongJumps {}; + enum class RestartOptionControl : uint64_t { + Incoming = 0, + NeedsLongJumps = 1, + }; + }; + RestartOptions RestartControl {}; ARMEmitter::BiDirectionalLabel* PendingTargetLabel {}; ARMEmitter::BiDirectionalLabel* PendingCallReturnTargetLabel {}; FEXCore::Context::ContextImpl* CTX {}; @@ -329,14 +339,203 @@ private: void EmitLinkedBranch(uint64_t GuestRIP, bool Call) { PendingJumpThunks.push_back({GetCursorAddress(), GuestRIP, {}}); auto& Thunk = PendingJumpThunks.back(); - Bind(&Thunk.Label); + BindOrRestart(&Thunk.Label); if (Call) { - bl(&Thunk.Label); + bl_OrRestart(&Thunk.Label); } else { - b(&Thunk.Label); + b_OrRestart(&Thunk.Label); } } + // Restart helpers + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void bl_OrRestart(T* Label) { + if (bl(Label)) [[likely]] { + return; + } + + // We can support this but currently unnecessary. + LOGMAN_MSG_A_FMT("Tried to branch larger than 128MB away!"); + // Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void b_OrRestart(T* Label) { + if (b(Label)) [[likely]] { + return; + } + + // We can support this but currently unnecessary. + LOGMAN_MSG_A_FMT("Tried to branch larger than 128MB away!"); + // Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void b_OrRestart(ARMEmitter::Condition Cond, T* Label) { + if (RestartControl.NeedsLongJumps) [[unlikely]] { + ARMEmitter::ForwardLabel Skip {}; + // Invert check and skip. + (void)b(InvertCondition(Cond), &Skip); + if (!b(Label)) [[unlikely]] { + LOGMAN_MSG_A_FMT("Tried to branch larger than 128MB away!"); + } + + (void)Bind(&Skip); + return; + } + + if (b(Cond, Label)) [[likely]] { + return; + } + + Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void cbz_OrRestart(ARMEmitter::Size s, ARMEmitter::Register rt, T* Label) { + if (RestartControl.NeedsLongJumps) [[unlikely]] { + ARMEmitter::ForwardLabel Skip {}; + // Invert check and skip. + (void)cbnz(s, rt, &Skip); + if (!b(Label)) [[unlikely]] { + LOGMAN_MSG_A_FMT("Tried to branch larger than 128MB away!"); + } + + (void)Bind(&Skip); + return; + } + + if (cbz(s, rt, Label)) [[likely]] { + return; + } + + Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void cbnz_OrRestart(ARMEmitter::Size s, ARMEmitter::Register rt, T* Label) { + if (RestartControl.NeedsLongJumps) [[unlikely]] { + ARMEmitter::ForwardLabel Skip {}; + // Invert check and skip. + (void)cbz(s, rt, &Skip); + if (!b(Label)) [[unlikely]] { + LOGMAN_MSG_A_FMT("Tried to branch larger than 128MB away!"); + } + + (void)Bind(&Skip); + return; + } + + if (cbnz(s, rt, Label)) [[likely]] { + return; + } + + Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void tbz_OrRestart(ARMEmitter::Register rt, uint32_t Bit, T* Label) { + if (RestartControl.NeedsLongJumps) [[unlikely]] { + ARMEmitter::ForwardLabel Skip {}; + // Invert check and skip. + (void)tbnz(rt, Bit, &Skip); + if (!b(Label)) [[unlikely]] { + LOGMAN_MSG_A_FMT("Tried to branch larger than 128MB away!"); + } + + (void)Bind(&Skip); + return; + } + + if (tbz(rt, Bit, Label)) [[likely]] { + return; + } + + Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void tbnz_OrRestart(ARMEmitter::Register rt, uint32_t Bit, T* Label) { + if (RestartControl.NeedsLongJumps) [[unlikely]] { + ARMEmitter::ForwardLabel Skip {}; + // Invert check and skip. + (void)tbz(rt, Bit, &Skip); + if (!b(Label)) [[unlikely]] { + LOGMAN_MSG_A_FMT("Tried to branch larger than 128MB away!"); + } + + (void)Bind(&Skip); + return; + } + + if (tbnz(rt, Bit, Label)) [[likely]] { + return; + } + + Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void adr_OrRestart(ARMEmitter::Register rd, T* Label) { + if (adr(rd, Label)) [[likely]] { + return; + } + + // We can support this but currently unnecessary. + LOGMAN_MSG_A_FMT("Long ADR currently unsupported!"); + // Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void adrp_OrRestart(ARMEmitter::Register rd, T* Label) { + if (adrp(rd, Label)) [[likely]] { + return; + } + + // We can support this but currently unnecessary. + LOGMAN_MSG_A_FMT("Long ADRP currently unsupported!"); + // Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + template + requires (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) + void BindOrRestart(T* Label) { + if (Bind(Label)) [[likely]] { + return; + } + + if (RestartControl.NeedsLongJumps) [[unlikely]] { + // This should have been caught before this point. + LOGMAN_MSG_A_FMT("Oops. Unhandled long bind."); + return; + } + + Restart(RestartOptions::RestartOptionControl::NeedsLongJumps); + } + + [[noreturn]] void Restart(RestartOptions::RestartOptionControl Control) { + FEXCore::LongJump::LongJump(RestartControl.RestartJump, FEXCore::ToUnderlying(Control)); + } + // 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 {};