FEXCore: Have ThreadRemoveCodeEntry IR op jump to dispatcher

In order to reuse code buffers, when invalidating we will no longer be
able to invalidate code while inside the code we are invalidating.

Changes the IR operation to be a block ender itself and control jumping
to the new RIP. This way it can invalidate from the dispatcher and just
restart JIT execution.

FEXCore/Source/Interface/Core/Core.cpp

@@ -661,11 +661,10 @@ ContextImpl::GenerateIR(FEXCore::Core::InternalThreadState* Thread, uint64_t Gue
           Thread->OpDispatcher->StartNewBlock();
 
           // Generate a relocatable entry for invalidation purposes.
-          auto EntryReg = Thread->OpDispatcher->_EntrypointOffset(GPRSize, 0);
-          Thread->OpDispatcher->_ThreadRemoveCodeEntry(EntryReg);
-
-          // Exit the function at this instruction after invalidation.
-          Thread->OpDispatcher->ExitFunction(Thread->OpDispatcher->_InlineEntrypointOffset(GPRSize, InstAddress - GuestRIP));
+          auto EntryToInvalidate = Thread->OpDispatcher->_EntrypointOffset(GPRSize, 0);
+          auto NewRIP = Thread->OpDispatcher->_EntrypointOffset(GPRSize, InstAddress - GuestRIP);
+          // Invalidate and exit the function
+          Thread->OpDispatcher->_ThreadRemoveCodeEntry(EntryToInvalidate, NewRIP);
 
           auto NextOpBlock = Thread->OpDispatcher->CreateNewCodeBlockAfter(CurrentBlock);
 

FEXCore/Source/Interface/Core/DiskCache.h

@@ -244,7 +244,7 @@ namespace DiskCache {
   // The current version of the diskcache.
   // This must be changed any time codegen changes occur!
   // Be aware of the impact of changing this frequently!
-  static constexpr uint16_t FormatVersion = 26;
+  static constexpr uint16_t FormatVersion = 27;
 
   static constexpr uint32_t LOOKUP_KEY_MAX_BUCKET_DEPTH = 20;
 

FEXCore/Source/Interface/Core/Dispatcher/Dispatcher.cpp

@@ -550,6 +550,28 @@ void Dispatcher::EmitDispatcher() {
     (void)b(&LoopTop);
   }
 
+  {
+    // All dynamic and static registers are spilled coming in to this handler.
+    // It's also the end of block and RIP might have changed, so we jump directly to the top of the loop.
+    ThreadDispatchRemoveCodeEntry = GetCursorAddress<uint64_t>();
+
+    // Arguments are already in x0, x1. Just jump to the handler.
+    ldr(ARMEmitter::XReg::x2, STATE, offsetof(FEXCore::Core::CpuStateFrame, Pointers.ThreadRemoveCodeEntryFromJIT));
+    if (!CTX->Config.DisableVixlIndirectCalls) [[unlikely]] {
+      GenerateIndirectRuntimeCall<void, void*, void*>(ARMEmitter::Reg::r2);
+    } else {
+      blr(ARMEmitter::Reg::r2);
+    }
+
+    FillStaticRegs({
+      .OptionalReg = ARMEmitter::Reg::r1,
+      .OptionalReg2 = ARMEmitter::Reg::r2,
+    });
+
+    // Now go back to the regular dispatcher loop
+    (void)b(&LoopTop);
+  }
+
   auto EmitLongALUOpHandler = [&](auto R, auto Offset) {
     auto Address = GetCursorAddress<uint64_t>();
 
@@ -2664,6 +2686,7 @@ void Dispatcher::InitThreadPointers(FEXCore::Core::InternalThreadState* Thread)
     Ptrs.ThreadStopHandlerSpillSRA = ThreadStopHandlerAddressSpillSRA;
     Ptrs.ThreadPauseHandlerSpillSRA = ThreadPauseHandlerAddressSpillSRA;
     Ptrs.ThreadDispatchSyscallHandler = ThreadDispatchSyscallHandler;
+    Ptrs.ThreadDispatchRemoveCodeEntry = ThreadDispatchRemoveCodeEntry;
     Ptrs.GuestSignal_SIGILL = GuestSignal_SIGILL;
     Ptrs.GuestSignal_SIGTRAP = GuestSignal_SIGTRAP;
     Ptrs.GuestSignal_SIGSEGV = GuestSignal_SIGSEGV;

FEXCore/Source/Interface/Core/Dispatcher/Dispatcher.h

@@ -82,6 +82,7 @@ private:
   uint64_t ThreadPauseHandlerAddress {};
   uint64_t ThreadPauseHandlerAddressSpillSRA {};
   uint64_t ThreadDispatchSyscallHandler {};
+  uint64_t ThreadDispatchRemoveCodeEntry {};
   uint64_t ExitFunctionLinkerAddress {};
   uint64_t SignalHandlerReturnAddress {};
   uint64_t SignalHandlerReturnAddressRT {};

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

@@ -374,27 +374,22 @@ DEF_OP(ValidateCode) {
 DEF_OP(ThreadRemoveCodeEntry) {
   auto Op = IROp->C<IR::IROp_ThreadRemoveCodeEntry>();
 
-  // Move the entry to ABI before saving state.
-  mov(ARMEmitter::Size::i64Bit, ARMEmitter::Reg::r1, GetReg(Op->Entry));
+  SpillStaticRegs(TMP1);
 
-  PushDynamicRegs(TMP4);
-  SpillStaticRegs(TMP4);
+  // Store the new RIP to go to.
+  str(GetReg(Op->NewRIP).X(), STATE, offsetof(FEXCore::Core::CpuStateFrame, State.rip));
+
+  // Move the entry to ABI before saving state.
+  mov(ARMEmitter::Size::i64Bit, ARMEmitter::Reg::r1, GetReg(Op->EntryToInvalidate));
 
   // Arguments are passed as follows:
   // X0: Thread
-  // X1: RIP
+  // X1: RIPToInvalidate
   mov(ARMEmitter::Size::i64Bit, ARMEmitter::Reg::r0, STATE.R());
 
-  ldr(ARMEmitter::XReg::x2, STATE, offsetof(FEXCore::Core::CpuStateFrame, Pointers.ThreadRemoveCodeEntryFromJIT));
-  if (!CTX->Config.DisableVixlIndirectCalls) [[unlikely]] {
-    GenerateIndirectRuntimeCall<void, void*, void*>(ARMEmitter::Reg::r2);
-  } else {
-    blr(ARMEmitter::Reg::r2);
-  }
-  FillStaticRegs();
-
-  // Fix the stack and any values that were stepped on
-  PopDynamicRegs();
+  // Jump to the invalidate dispatch handler. We won't return after this.
+  ldr(ARMEmitter::XReg::x2, STATE, offsetof(FEXCore::Core::CpuStateFrame, Pointers.ThreadDispatchRemoveCodeEntry));
+  br(ARMEmitter::XReg::x2);
 }
 
 DEF_OP(CPUID) {

FEXCore/Source/Interface/IR/IR.json

@@ -201,7 +201,7 @@
         "DestSize": "OpSize::i64Bit"
       },
 
-      "ThreadRemoveCodeEntry GPR:$Entry": {
+      "ThreadRemoveCodeEntry GPR:$EntryToInvalidate, GPR:$NewRIP": {
         "HasSideEffects": true
       },
 

FEXCore/Source/Interface/IR/IREmitter.cpp

@@ -19,6 +19,7 @@ namespace FEXCore::IR {
 
 static bool IsFragmentExit(FEXCore::IR::IROps Op) {
   switch (Op) {
+  case OP_THREADREMOVECODEENTRY:
   case OP_SYSCALL:
   case OP_EXITFUNCTION:
   case OP_BREAK: return true;

FEXCore/include/FEXCore/Core/CoreState.h

@@ -372,6 +372,7 @@ struct JITPointers {
   uint64_t ThreadStopHandlerSpillSRA {};
   uint64_t ThreadPauseHandlerSpillSRA {};
   uint64_t ThreadDispatchSyscallHandler {};
+  uint64_t ThreadDispatchRemoveCodeEntry {};
   uint64_t GuestSignal_SIGILL {};
   uint64_t GuestSignal_SIGTRAP {};
   uint64_t GuestSignal_SIGSEGV {};