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.

FEXCore/Source/Interface/Core/JIT/Arm64Relocations.cpp

@@ -63,7 +63,7 @@ void Arm64JITCore::PlaceNamedSymbolLiteral(NamedSymbolLiteralPair& Lit) {
   auto CurrentCursor = GetCursorAddress<uint8_t*>();
   Lit.MoveABI.NamedSymbolLiteral.Offset = CurrentCursor - CodeData.BlockBegin;
 
-  Bind(&Lit.Loc);
+  BindOrRestart(&Lit.Loc);
   dc64(Lit.Lit);
   Relocations.emplace_back(Lit.MoveABI);
 }

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) {

FEXCore/Source/Interface/Core/JIT/JIT.cpp

@@ -28,6 +28,7 @@ $end_info$
 #include <FEXCore/Utils/CompilerDefs.h>
 #include <FEXCore/Utils/EnumUtils.h>
 #include <FEXCore/Utils/LogManager.h>
+#include <FEXCore/Utils/LongJump.h>
 #include <FEXCore/Utils/Profiler.h>
 #include <FEXCore/Utils/Telemetry.h>
 #include <FEXCore/Utils/TypeDefines.h>
@@ -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<RestartOptions::RestartOptionControl>(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<uint8_t*>());
         DebugData->GuestOpcodes.push_back({BlockIROp->GuestEntryOffset, GetCursorAddress<uint8_t*>() - 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<uint64_t>();
-    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.

FEXCore/Source/Interface/Core/JIT/JITClass.h

@@ -23,6 +23,7 @@ $end_info$
 #include <FEXCore/fextl/memory.h>
 #include <FEXCore/fextl/string.h>
 #include <FEXCore/fextl/vector.h>
+#include <FEXCore/Utils/LongJump.h>
 
 #include <CodeEmitter/Emitter.h>
 
@@ -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<uint64_t>(), 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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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<typename T>
+  requires (std::is_same_v<T, ARMEmitter::ForwardLabel> || std::is_same_v<T, ARMEmitter::BackwardLabel> ||
+            std::is_same_v<T, ARMEmitter::BiDirectionalLabel> || std::is_same_v<T, ARMEmitter::ForwardLabel::Reference>)
+  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 {};