JIT/arm64: Implement block linking

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

@@ -72,11 +72,14 @@ DEF_OP(ExitFunction) {
   uint64_t NewRIP;
 
   if (IsInlineConstant(Op->NewRIP, &NewRIP)) {
-    uintptr_t CacheEntry = State->BlockCache->GetL1Pointer() + (NewRIP & BlockCache::L1_ENTRIES_MASK) * 16;
-    LoadConstant(x0, CacheEntry);
+    Literal l_BranchHost{ExitFunctionLinkerAddress};
+    Literal l_BranchGuest{NewRIP};
 
-    RipReg = x2;
-    LoadConstant(RipReg, NewRIP);
+    ldr(x0, &l_BranchHost);
+    blr(x0);
+
+    place(&l_BranchHost);
+    place(&l_BranchGuest);
   } else {
     RipReg = GetReg<RA_64>(Op->Header.Args[0].ID());
     
@@ -85,17 +88,17 @@ DEF_OP(ExitFunction) {
 
     and_(x3, RipReg, BlockCache::L1_ENTRIES_MASK);
     add(x0, x0, Operand(x3, Shift::LSL, 4));
+
+    ldp(x1, x0, MemOperand(x0));
+    cmp(x0, RipReg);
+    b(&FullLookup, Condition::ne);
+    br(x1);
+
+    bind(&FullLookup);
+    LoadConstant(TMP1, AbsoluteLoopTopAddress);
+    str(RipReg, MemOperand(STATE, offsetof(FEXCore::Core::ThreadState, State.rip)));
+    br(TMP1);
   }
-
-  ldp(x1, x0, MemOperand(x0));
-  cmp(x0, RipReg);
-  b(&FullLookup, Condition::ne);
-  br(x1);
-
-  bind(&FullLookup);
-  LoadConstant(TMP1, AbsoluteLoopTopAddress);
-  str(RipReg, MemOperand(STATE, offsetof(FEXCore::Core::ThreadState, State.rip)));
-  br(TMP1);
 }
 
 DEF_OP(Jump) {

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

@@ -832,7 +832,7 @@ void *JITCore::CompileCode([[maybe_unused]] FEXCore::IR::IRListView<true> const
     // Make sure RIP is syncronized to the context
     LoadConstant(x0, HeaderOp->Entry);
     str(x0, MemOperand(STATE, offsetof(FEXCore::Core::ThreadState, State.rip)));
-    
+
     LoadConstant(x0, ThreadSharedData.InterpreterFallbackHelperAddress);
     br(x0);
   } else {
@@ -988,6 +988,35 @@ void JITCore::PopCalleeSavedRegisters() {
   }
 }
 
+uint64_t JITCore::ExitFunctionLink(JITCore *core, FEXCore::Core::InternalThreadState *Thread, uint64_t *record) {
+  auto GuestRip = record[1];
+
+  auto HostCode = Thread->BlockCache->FindBlock(GuestRip);
+
+  if (!HostCode) {
+    //printf("ExitFunctionLink: Aborting, %lX not in cache\n", GuestRip);
+    Thread->State.State.rip = GuestRip;
+    return core->AbsoluteLoopTopAddress;
+  }
+
+  uintptr_t branch = (uintptr_t)(record) - 8;
+
+  auto offset = HostCode/4 - branch/4;
+  if (IsInt26(offset)) {
+    // optimal case - can branch directly
+    // patch the code
+    vixl::aarch64::Assembler emit((uint8_t*)(branch), 24);
+    emit.b(offset);
+    emit.FinalizeCode();
+    vixl::aarch64::CPU::EnsureIAndDCacheCoherency(&record[-2], 16);
+  } else {
+    // fallback case - do a soft-er link by patching the pointer
+    record[0] = HostCode;
+  }
+  
+  return HostCode;
+}
+
 void JITCore::CreateCustomDispatch(FEXCore::Core::InternalThreadState *Thread) {
   auto OriginalBuffer = *GetBuffer();
 
@@ -1052,6 +1081,7 @@ void JITCore::CreateCustomDispatch(FEXCore::Core::InternalThreadState *Thread) {
 
   Literal l_CompileBlock {CompileBlockPtr};
   Literal l_CompileFallback {CompileFallbackPtr};
+  Literal l_ExitFunctionLink {(uintptr_t)&ExitFunctionLink};
 
   // Push all the register we need to save
   PushCalleeSavedRegisters();
@@ -1086,6 +1116,7 @@ void JITCore::CreateCustomDispatch(FEXCore::Core::InternalThreadState *Thread) {
   aarch64::Label LoopTop{};
   aarch64::Label ExitSpillSRA{};
   aarch64::Label ThreadPauseHandlerSpillSRA{};
+  aarch64::Label ExitFunctionLinker{};
 
   bind(&LoopTop);
   AbsoluteLoopTopAddress = GetLabelAddress<uint64_t>(&LoopTop);
@@ -1173,6 +1204,23 @@ void JITCore::CreateCustomDispatch(FEXCore::Core::InternalThreadState *Thread) {
     ret();
   }
 
+  {
+    bind(&ExitFunctionLinker);
+    ExitFunctionLinkerAddress = GetLabelAddress<uint64_t>(&ExitFunctionLinker);
+
+    SpillStaticRegs();
+    
+    LoadConstant(x0, (uintptr_t)this);
+    mov(x1, STATE);
+    mov(x2, lr);
+    
+    ldr(x3, &l_ExitFunctionLink);
+    blr(x3);
+
+    FillStaticRegs();
+    br(x0);
+  }
+
   aarch64::Label FallbackCore;
   // Need to create the block
   {
@@ -1314,6 +1362,7 @@ void JITCore::CreateCustomDispatch(FEXCore::Core::InternalThreadState *Thread) {
   place(&l_Sleep);
   place(&l_CompileBlock);
   place(&l_CompileFallback);
+  place(&l_ExitFunctionLink);
 
   FinalizeCode();
   uint64_t CodeEnd = Buffer->GetOffsetAddress<uint64_t>(GetCursorOffset());

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

@@ -201,12 +201,15 @@ private:
   void PushCalleeSavedRegisters();
   void PopCalleeSavedRegisters();
 
+  static uint64_t ExitFunctionLink(JITCore *core, FEXCore::Core::InternalThreadState *Thread, uint64_t *record);
+
   /**
    * @name Dispatch Helper functions
    * @{ */
   uint64_t AbsoluteLoopTopAddressFillSRA{};
   uint64_t AbsoluteLoopTopAddress{};
   uint64_t ThreadPauseHandlerAddressSpillSRA{};
+  uint64_t ExitFunctionLinkerAddress{};
   uint64_t ThreadPauseHandlerAddress{};
 
   uint64_t ThreadStopHandlerAddressSpillSRA{};