JIT: Rewrite block linking to support direct ExitFunction calls

The constraints introduced by shared code buffers make supporting
calls with the previous layout impossible. The main additional constraint
imposed by call-ret that if a host location is ever pushed onto the
call-ret stack, then it must forever be a valid jump target. While
this is reasonable in the: unlinked, direct linked, unlinked,
direct linked case; it's almost impossible to achieve in the: unlinked,
indirect linked, unlinked, direct linked case while ensuring
all backpatching cases are valid with the current approach.

To solve this introduce an additional layer of indirection, jump thunks,
these are emitted at the end of a multiblock and are used to handle the
two cases of calling the initial linker, and calling an indirect linked
block. Initially at the ExitFunction location a branch/call to a unique
jump thunk will be emitted, which will have the code layout:
00: b 0x8
04: br TMP1
08: ldr TMP1, <Shared exit linker>
0c: blr TMP1
10: HostCode
18: GuestRIP
20: CallerOffset

If a direct link can be performed, then the initial branch/call to the
jump thunk can be linked/unlinked to point to the jump thunk in a
single 32-bit atomic operation. For an indirect link, the HostCode
member is updated with a 64 bit atomic operation, and then a 32 bit
atomic operation is used to replace the branch at 00 with a load of
HostCode. Indirect unlinks are done by placing back the b 0x8 at 00.

Safety:
(1)
Sequential link (e.g. one waiting to lock, one locked and linking):
Linking is idempotent, would just rewrite the same data atomically.

(2)
Simultaneous link or simultaneous delink:
Impossible due to LookupCache locking.

(3)
Simultaneous link and execute:
(3.1)
Direct link: Either the direct link is observed at the thunk
callsite, or it is not observed and the linker is entered - this is
then just (1).

(3.2)
Indirect link: Either the branch at 00 in the thunk is observed
to be replaced with an ldr, in which case the modified HostCode
must be observed due to the cache flush. Alternatively the branch
replacement isn't observed and it's just (1).

(4)
Simultaneous unlink and execute:
(4.1)
Direct link: Either the jump to the jump thunk is seen, which must
be in its base unlinked state with the branch at 00 as that would
be inserted by any previous indirect unlink. In such a case the
linker would just be entered, giving (5). Alternatively the modified
jump isn't seen and it calls the original host code (which is fine).

(4.2)
Indirect link: If an ldr is seen at 00, then the rest of that sequence
will function fine as HostCode is left untouched. If a branch is seen
at 00, then it will just call the linker giving (5).

(5)
Sequential unlink then link:
Unlinking restores the callsite and jump thunk to their original
contents (aside from a modified HostCode). Linking then works as
usual.

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<auto Fn>
-  static uint64_t ThreadExitFunctionLink(FEXCore::Core::CpuStateFrame* Frame, ExitFunctionLinkData* Record) {
-    auto Thread = Frame->Thread;
-    auto lk = GuardSignalDeferringSection<std::shared_lock>(static_cast<ContextImpl*>(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) {

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 <GuestReturnRIP, HostReturnPC> 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, <Shared exit linker>
+      //    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, <Shared exit linker>
+      //    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, <Shared exit linker>
+      //    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, <Shared exit linker>
+      //    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<ARMEmitter::IndexType::PRE>(CallReturnAddressReg, TMP1, REG_CALLRET_SP, -0x10);
+        } else {
+          stp<ARMEmitter::IndexType::PRE>(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

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<uintptr_t>(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<uint8_t*>(&BranchInst), 4);
+  if (Call) {
+    BranchEmit.bl(BranchOffset);
+  } else {
+    BranchEmit.b(BranchOffset);
+  }
+
+  std::atomic_ref<uint32_t>(*reinterpret_cast<uint32_t*>(CallerAddress)).store(BranchInst, std::memory_order::relaxed);
+  ARMEmitter::Emitter::ClearICache(reinterpret_cast<void*>(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<uintptr_t>(Record) - 0x10;
+  uint32_t BranchInst = 0;
+  ARMEmitter::Emitter BranchEmit(reinterpret_cast<uint8_t*>(&BranchInst), 4);
+  BranchEmit.b(0x8);
 
+  std::atomic_ref<uint32_t>(*reinterpret_cast<uint32_t*>(JumpThunkStartAddress)).store(BranchInst, std::memory_order::relaxed);
+  ARMEmitter::Emitter::ClearICache(reinterpret_cast<void*>(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<Arm64JITCore*>(Thread->CPUBackend.get())->CurrentCodeBuffer;
+      HostCode = static_cast<Context::ContextImpl*>(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<uintptr_t>(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<uint32_t>(*reinterpret_cast<uint32_t*>(branch)).store(Patch, std::memory_order::relaxed);
-    ARMEmitter::Emitter::ClearICache((void*)branch, 4);
+  uint32_t ExpectedKnownCallMarkerInst = 0;
+  ARMEmitter::Emitter ExpectedKnownCallMarkerEmit(reinterpret_cast<uint8_t*>(&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<uint32_t*>(KnownCallMarkerAddr);
+  if (ARMEmitter::Emitter::IsInt26(BranchOffset)) {
+    // Directly patch the callsite with the appropriate branch instruction.
+    uint32_t BranchInst = 0;
+    ARMEmitter::Emitter BranchEmit(reinterpret_cast<uint8_t*>(&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<uint32_t>(*reinterpret_cast<uint32_t*>(CallerAddress)).store(BranchInst, std::memory_order::relaxed);
+    ARMEmitter::Emitter::ClearICache(reinterpret_cast<void*>(CallerAddress), 4);
   } else {
-    // fallback case - do a soft-er link by patching the pointer
-    std::atomic_ref<uint64_t>(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<uint64_t>(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<uint8_t*>(&LdrInst), 4);
+    LdrEmit.ldr(TMP1, reinterpret_cast<uint64_t>(&Record->HostCode) - JumpThunkStartAddress);
+    std::atomic_ref<uint32_t>(*reinterpret_cast<uint32_t*>(JumpThunkStartAddress)).store(LdrInst, std::memory_order::relaxed);
+    ARMEmitter::Emitter::ClearICache(reinterpret_cast<void*>(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<uint64_t>(CTX->SyscallHandler);
       Common.SyscallHandlerFunc = PMF.GetVTableEntry(CTX->SyscallHandler);
     }
-    Common.ExitFunctionLink = reinterpret_cast<uintptr_t>(&Context::ContextImpl::ThreadExitFunctionLink<Arm64JITCore_ExitFunctionLink>);
+    Common.ExitFunctionLink = reinterpret_cast<uintptr_t>(&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<uint64_t>();
+    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<uint8_t*>() - CodeBegin;
 

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<IR::NodeID, ARMEmitter::BiDirectionalLabel> JumpTargets;
   fextl::map<IR::NodeID, ARMEmitter::BiDirectionalLabel> CallReturnTargets;
 
+  struct PendingJumpThunk {
+    uint64_t CallerAddress;
+    uint64_t GuestRIP;
+    ARMEmitter::ForwardLabel Label;
+  };
+  fextl::vector<PendingJumpThunk> PendingJumpThunks;
+
   Utils::PoolBufferWithTimedRetirement<uint8_t*, 5000, 500> 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<uint64_t>(), 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 {};