JIT: Move ApplyRelocations to CodeCache

FEXCore/Source/Interface/Context/Context.h

@@ -80,6 +80,22 @@ public:
   void InitiateCacheGeneration() override {
     IsGeneratingCache = true;
   }
+
+  /**
+   * Applies a set of FEX relocations to the given code section.
+   *
+   * FEX relocations describe runtime-dependencies of FEX-generated code.
+   * When loading a code cache, they are used to move cached code to the
+   * dynamically chosen base address of the guest binary.
+   *
+   * Note that FEX relocations are unrelated to ELF/PE relocations.
+   *
+   * @param GuestDelta Guest address offset to apply to RIP-relative data
+   *
+   * @return Returns true on success
+   */
+  [[nodiscard]]
+  bool ApplyCodeRelocations(uint64_t GuestDelta, std::span<std::byte> Code, std::span<const CPU::Relocation> Relocations);
 };
 
 class ContextImpl final : public FEXCore::Context::Context, public CPU::CodeBufferManager {

FEXCore/Source/Interface/Core/ArchHelpers/Arm64Emitter.h

@@ -105,9 +105,12 @@ constexpr ARMEmitter::PRegister PRED_TMP_32B = ARMEmitter::PReg::p7;
 // This class contains common emitter utility functions that can
 // be used by both Arm64 JIT and ARM64 Dispatcher
 class Arm64Emitter : public ARMEmitter::Emitter {
-protected:
+public:
   Arm64Emitter(FEXCore::Context::ContextImpl* ctx, void* EmissionPtr = nullptr, size_t size = 0);
 
+  void LoadConstant(ARMEmitter::Size s, ARMEmitter::Register Reg, uint64_t Constant, bool NOPPad = false);
+
+protected:
   FEXCore::Context::ContextImpl* EmitterCTX;
 
   std::span<const ARMEmitter::Register> StaticRegisters {};
@@ -117,8 +120,6 @@ protected:
   std::span<const ARMEmitter::VRegister> GeneralFPRegisters {};
   uint32_t PairRegisters = 0;
 
-  void LoadConstant(ARMEmitter::Size s, ARMEmitter::Register Reg, uint64_t Constant, bool NOPPad = false);
-
   void FillSpecialRegs(ARMEmitter::Register TmpReg, ARMEmitter::Register TmpReg2, bool SetFIZ, bool SetPredRegs);
 
   // Correlate an ARM register back to an x86 register index.

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

@@ -10,18 +10,18 @@ $end_info$
 
 #include <FEXCore/Core/Thunks.h>
 
-namespace FEXCore::CPU {
-
-uint64_t Arm64JITCore::GetNamedSymbolLiteral(FEXCore::CPU::RelocNamedSymbolLiteral::NamedSymbol Op) {
+namespace FEXCore::Context {
+static uint64_t GetNamedSymbolLiteral(ContextImpl& CTX, FEXCore::CPU::RelocNamedSymbolLiteral::NamedSymbol Op) {
   switch (Op) {
   case FEXCore::CPU::RelocNamedSymbolLiteral::NamedSymbol::SYMBOL_LITERAL_EXITFUNCTION_LINKER:
-    return ThreadState->CurrentFrame->Pointers.Common.ExitFunctionLinker;
-    break;
-  default: ERROR_AND_DIE_FMT("Unknown named symbol literal: {}", static_cast<uint32_t>(Op)); break;
-  }
-  return ~0ULL;
-}
+    return CTX.Dispatcher->GetExitFunctionLinkerAddress();
 
+  default: ERROR_AND_DIE_FMT("Unknown named symbol literal: {}", static_cast<uint32_t>(Op));
+  }
+}
+} // namespace FEXCore::Context
+
+namespace FEXCore::CPU {
 void Arm64JITCore::InsertNamedThunkRelocation(ARMEmitter::Register Reg, const IR::SHA256Sum& Sum) {
   Relocation MoveABI {};
   MoveABI.NamedThunkMove.Header = {.Offset = GetCursorOffset(), .Type = FEXCore::CPU::RelocationTypes::RELOC_NAMED_THUNK_MOVE};
@@ -35,7 +35,7 @@ void Arm64JITCore::InsertNamedThunkRelocation(ARMEmitter::Register Reg, const IR
 }
 
 Arm64JITCore::NamedSymbolLiteralPair Arm64JITCore::InsertNamedSymbolLiteral(FEXCore::CPU::RelocNamedSymbolLiteral::NamedSymbol Op) {
-  uint64_t Pointer = GetNamedSymbolLiteral(Op);
+  uint64_t Pointer = GetNamedSymbolLiteral(*CTX, Op);
 
   NamedSymbolLiteralPair Lit {
     .Lit = Pointer,
@@ -97,58 +97,51 @@ void Arm64JITCore::InsertGuestRIPMove(ARMEmitter::Register Reg, uint64_t Constan
   LoadConstant(ARMEmitter::Size::i64Bit, Reg, Constant, false);
   Relocations.emplace_back(MoveABI);
 }
+} // namespace FEXCore::CPU
 
-bool Arm64JITCore::ApplyRelocations(uint64_t GuestEntry, std::span<std::byte> Code, std::span<const FEXCore::CPU::Relocation> Relocations) {
-  const auto OrigBase = GetBufferBase();
-  const auto OrigSize = GetBufferSize();
-  const auto OrigOffset = GetCursorOffset();
-
-  SetBuffer(reinterpret_cast<std::uint8_t*>(Code.data()), Code.size_bytes());
-  for (auto& Reloc : Relocations) {
-    SetCursorOffset(Reloc.Header.Offset);
+namespace FEXCore::Context {
+bool CodeCache::ApplyCodeRelocations(uint64_t GuestEntry, std::span<std::byte> Code,
+                                     std::span<const FEXCore::CPU::Relocation> EntryRelocations) {
+  CPU::Arm64Emitter Emitter(&CTX, Code.data(), Code.size_bytes());
+  for (size_t j = 0; j < EntryRelocations.size(); ++j) {
+    const FEXCore::CPU::Relocation& Reloc = EntryRelocations[j];
+    Emitter.SetCursorOffset(Reloc.Header.Offset);
 
     switch (Reloc.Header.Type) {
     case FEXCore::CPU::RelocationTypes::RELOC_NAMED_SYMBOL_LITERAL: {
-      uint64_t Pointer = GetNamedSymbolLiteral(Reloc.NamedSymbolLiteral.Symbol);
-
       // Generate a literal so we can place it
-      dc64(Pointer);
+      uint64_t Pointer = GetNamedSymbolLiteral(CTX, Reloc.NamedSymbolLiteral.Symbol);
+      Emitter.dc64(Pointer);
       break;
     }
     case FEXCore::CPU::RelocationTypes::RELOC_NAMED_THUNK_MOVE: {
-      uint64_t Pointer = reinterpret_cast<uint64_t>(EmitterCTX->ThunkHandler->LookupThunk(Reloc.NamedThunkMove.Symbol));
+      uint64_t Pointer = reinterpret_cast<uint64_t>(CTX.ThunkHandler->LookupThunk(Reloc.NamedThunkMove.Symbol));
       if (Pointer == ~0ULL) {
         return false;
       }
 
-      LoadConstant(ARMEmitter::Size::i64Bit, ARMEmitter::Register(Reloc.NamedThunkMove.RegisterIndex), Pointer, true);
+      Emitter.LoadConstant(ARMEmitter::Size::i64Bit, ARMEmitter::Register(Reloc.NamedThunkMove.RegisterIndex), Pointer, true);
       break;
     }
     case FEXCore::CPU::RelocationTypes::RELOC_GUEST_RIP_LITERAL: {
-      dc64(GuestEntry + Reloc.GuestRIP.GuestRIP);
+      Emitter.dc64(GuestEntry + Reloc.GuestRIP.GuestRIP);
       break;
     }
     case FEXCore::CPU::RelocationTypes::RELOC_GUEST_RIP_MOVE: {
-      // XXX: Reenable once the JIT Object Cache is upstream
-      // XXX: Should spin the relocation list, create a list of guest RIP moves, and ask for them all once, reduces lock contention.
-      uint64_t Pointer = ~0ULL; // EmitterCTX->JITObjectCache->FindRelocatedRIP(Reloc->GuestRIP.GuestRIP);
-      if (Pointer == ~0ULL) {
-        SetBuffer(OrigBase, OrigSize);
-        SetCursorOffset(OrigOffset);
-        return false;
-      }
-
-      LoadConstant(ARMEmitter::Size::i64Bit, ARMEmitter::Register(Reloc.GuestRIP.RegisterIndex), Pointer, true);
+      uint64_t Pointer = Reloc.GuestRIP.GuestRIP + GuestEntry;
+      Emitter.LoadConstant(ARMEmitter::Size::i64Bit, ARMEmitter::Register(Reloc.GuestRIP.RegisterIndex), Pointer, true);
       break;
     }
+
+    default: ERROR_AND_DIE_FMT("Unknown relocation type {}", ToUnderlying(Reloc.Header.Type));
     }
   }
 
-  SetBuffer(OrigBase, OrigSize);
-  SetCursorOffset(OrigOffset);
-
   return true;
 }
+} // namespace FEXCore::Context
+
+namespace FEXCore::CPU {
 
 fextl::vector<FEXCore::CPU::Relocation> Arm64JITCore::TakeRelocations(uint64_t GuestBaseAddress) {
   // Rebase relocations to library base address

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

@@ -536,8 +536,6 @@ private:
    * @name Relocations
    * @{ */
 
-  uint64_t GetNamedSymbolLiteral(FEXCore::CPU::RelocNamedSymbolLiteral::NamedSymbol Op);
-
   /**
    * @brief A literal pair relocation object for named symbol literals
    */
@@ -591,9 +589,6 @@ private:
 
   fextl::vector<FEXCore::CPU::Relocation> Relocations;
 
-  ///< Relocation code loading
-  bool ApplyRelocations(uint64_t GuestEntry, std::span<std::byte> Code, std::span<const FEXCore::CPU::Relocation>);
-
   /**
    * Returns any relocations generated since the last call to TakeRelocations.
    *