Merge pull request #5823 from Sonicadvance1/202

SharedCodeBufferManager: Allocate JIT space atomically.

FEXCore/Source/Interface/Core/CodeCache.cpp

@@ -314,7 +314,7 @@ bool CodeCache::SaveData(Core::InternalThreadState& Thread, int fd, const Execut
   std::ranges::copy(GIT_HASH, header.FEXVersion);
   header.NumBlocks = LookupCache.BlockList.size();
   header.NumCodePages = LookupCache.CodePages.size();
-  header.CodeBufferSize = FEXCore::AlignUp(CTX.LatestOffset, Utils::FEX_PAGE_SIZE);
+  header.CodeBufferSize = FEXCore::AlignUp(CTX.GetAllocatedSize(), Utils::FEX_PAGE_SIZE);
   header.NumRelocations = Relocations.size();
   header.SerializedBaseAddress = SerializedBaseAddress;
   ::write(fd, &header, sizeof(header));
@@ -361,7 +361,7 @@ bool CodeCache::SaveData(Core::InternalThreadState& Thread, int fd, const Execut
   }
 
   // Dump the host code (relocated for position-independent serialization)
-  std::span CodeBufferData(reinterpret_cast<std::byte*>(CodeBuffer->Ptr), reinterpret_cast<std::byte*>(CodeBuffer->Ptr) + CTX.LatestOffset);
+  std::span CodeBufferData(reinterpret_cast<std::byte*>(CodeBuffer->Ptr), reinterpret_cast<std::byte*>(CodeBuffer->Ptr) + CTX.GetAllocatedSize());
   if (!ApplyCodeRelocations(SerializedBaseAddress, CodeBufferData, Relocations, 0, true)) {
     LOGMAN_THROW_A_FMT(false, "Failed to apply code relocations");
     return false;
@@ -446,10 +446,10 @@ void CodeCache::Validate(const ExecutableFileSectionInfo& Section, fextl::set<ui
   }));
   (void)ApplyCodeRelocations(Section.FileStartVA, CodeBufferRangeRef, NewRelocations, 0, false);
 
-  if (ValidationCTX->LatestOffset <= CodeBufferRangeRef.size()) {
+  if (ValidationCTX->GetAllocatedSize() <= CodeBufferRangeRef.size()) {
     // Reference compilation produced fewer bytes than our cache, so validation is going to fail.
     // Make sure we don't output any garbage bytes though.
-    CodeBufferRangeRef = CodeBufferRangeRef.subspan(0, ValidationCTX->LatestOffset);
+    CodeBufferRangeRef = CodeBufferRangeRef.subspan(0, ValidationCTX->GetAllocatedSize());
   }
 
   auto [Mismatch, _] = std::mismatch(CodeBufferRangeRef.begin(), CodeBufferRangeRef.end(), CachedCode.begin());
@@ -500,7 +500,7 @@ void CodeCache::Validate(const ExecutableFileSectionInfo& Section, fextl::set<ui
 
   // Reset Context state for next validation
   ValidationThread->LookupCache->ClearCache(ValidationThread->LookupCache->AcquireWriteLock());
-  ValidationCTX->LatestOffset = 0;
+  ValidationCTX->CodeBufferOffset = ValidationCTX->CodeBufferBase;
 
   LogMan::Msg::IFmt("            successfully validated cache");
 }

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

@@ -815,6 +815,33 @@ void Arm64JITCore::EmitEntryPoint(ARMEmitter::BackwardLabel& HeaderLabel, bool C
   EmitSuspendInterruptCheck();
 }
 
+
+SharedCodeBufferManager::CodeBufferAllocation Arm64JITCore::AllocateCodeBufferInSharedCache(size_t Size) {
+  SharedCodeBufferManager::CodeBufferAllocation AllocatedInfo {};
+  LOGMAN_THROW_A_FMT(CurrentCodeBuffer->LookupCache.get() == ThreadState->LookupCache->Shared, "INVARIANT VIOLATED: SharedLookupCache "
+                                                                                               "doesn't match up!\n");
+  // Bring CodeBuffer up to date
+  if (auto Prev = CheckCodeBufferUpdate()) {
+    Allocator::VirtualDontNeed(ThreadState->CallRetStackBase, FEXCore::Core::InternalThreadState::CALLRET_STACK_SIZE);
+    auto lk = ThreadState->LookupCache->AcquireWriteLock();
+    ThreadState->LookupCache->ChangeGuestToHostMapping(*Prev, *CurrentCodeBuffer->LookupCache, lk);
+  }
+
+  // Attempt to allocate a buffer from the SharedCodeBuffers.
+  while (AllocatedInfo.BufferAllocationOffset == nullptr) {
+    AllocatedInfo = SharedCodeBuffers.AtomicAllocateBuffer(Size);
+
+    if (AllocatedInfo.BufferAllocationOffset == nullptr) {
+      // If it didn't fit then clear the buffer and try again.
+      // This has the possibility of migrating the SharedCodeBuffer. See above in `Arm64JITCore::ClearCache()`
+      CTX->ClearCodeCache(ThreadState);
+      continue;
+    }
+  }
+
+  return AllocatedInfo;
+}
+
 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");
@@ -1069,34 +1096,19 @@ CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size
     SetCursorOffset(PrevCur);
   }
 
-  // Migrate the compile output from temporary storage to the actual CodeBuffer.
-  // This can block progress in other compiling threads, so the duration of the lock should be as small as possible.
-  {
-    auto CodeBufferLock = std::unique_lock {SharedCodeBuffers.CodeBufferWriteMutex};
+  // Make sure code is 16B aligned on the tail.
+  Align16B();
 
+  // Migrate the compile output from temporary storage to the actual CodeBuffer.
+  {
     // Query size of generated code
     const auto TempSize = GetCursorOffset();
+    LOGMAN_THROW_A_FMT(TempSize % 16 == 0, "Needs to be 16B aligned!");
 
-    // Bring CodeBuffer up to date
-    {
-      LOGMAN_THROW_A_FMT(CurrentCodeBuffer->LookupCache.get() == ThreadState->LookupCache->Shared, "INVARIANT VIOLATED: SharedLookupCache "
-                                                                                                   "doesn't match up!\n");
-      if (auto Prev = CheckCodeBufferUpdate()) {
-        Allocator::VirtualDontNeed(ThreadState->CallRetStackBase, FEXCore::Core::InternalThreadState::CALLRET_STACK_SIZE);
-        auto lk = ThreadState->LookupCache->AcquireWriteLock();
-        ThreadState->LookupCache->ChangeGuestToHostMapping(*Prev, *CurrentCodeBuffer->LookupCache, lk);
-      }
+    auto AllocatedInfo = AllocateCodeBufferInSharedCache(TempSize);
 
-      // NOTE: 16-byte alignment of the new cursor offset must be preserved for block linking records
-      SetBuffer(CurrentCodeBuffer->Ptr, CurrentCodeBuffer->AllocatedSize);
-      SetCursorOffset(SharedCodeBuffers.LatestOffset);
-      Align16B();
-      if ((GetCursorOffset() + TempSize) > CurrentCodeBuffer->UsableSize()) {
-        CTX->ClearCodeCache(ThreadState);
-      }
-
-      SharedCodeBuffers.LatestOffset = GetCursorOffset();
-    }
+    // NOTE: 16-byte alignment of the new cursor offset must be preserved for block linking records
+    SetBuffer(AllocatedInfo.BufferAllocationOffset, Size);
 
     // Adjust host addresses
     const auto Delta = GetCursorAddress<uint8_t*>() - CodeData.BlockBegin;
@@ -1106,15 +1118,16 @@ CPUBackend::CompiledCode Arm64JITCore::CompileCode(uint64_t Entry, uint64_t Size
     }
     CodeBegin += Delta;
 
+    // Offset the relocations based on how far forward they moved from the temp buffer to the new buffer.
+    // TODO: Relocations should instead be relocated based on the block entrypoint instead of the codebuffer base.
+    // This would make this relocation movement here get deleted and then `CodeCache::HandleRelocations` can just handle it.
+    const size_t AllocationOffset = AllocatedInfo.BufferAllocationOffset - AllocatedInfo.BufferBase;
     for (std::size_t Idx = PrevNumAllocations; Idx != Relocations.size(); ++Idx) {
-      Relocations[Idx].Header.Offset += SharedCodeBuffers.LatestOffset;
+      Relocations[Idx].Header.Offset += AllocationOffset;
     }
 
     // Copy over CodeBuffer contents
     memcpy(GetCursorAddress<uint8_t*>(), TempCodeBuffer, TempSize);
-    SetCursorOffset(SharedCodeBuffers.LatestOffset + TempSize);
-
-    SharedCodeBuffers.LatestOffset = GetCursorOffset();
   }
 
   TempCodeBufferAllocator.DelayedDisownBuffer();

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

@@ -624,6 +624,8 @@ private:
 
   void EmitEntryPoint(ARMEmitter::BackwardLabel& HeaderLabel, bool CheckTF);
 
+  [[nodiscard]] SharedCodeBufferManager::CodeBufferAllocation AllocateCodeBufferInSharedCache(size_t Size);
+
 #define DEF_OP(x) void Op_##x(IR::IROp_Header const* IROp, IR::Ref Node)
 
   ///< Unhandled handler

FEXCore/Source/Interface/Core/SharedCodeBufferManager.cpp

@@ -66,7 +66,9 @@ fextl::shared_ptr<CodeBuffer> SharedCodeBufferManager::AllocateNew(size_t Size)
   auto Buffer = fextl::make_shared<CodeBuffer>(Size);
 
   Latest = Buffer;
-  LatestOffset = 0;
+  CodeBufferBase = Buffer->Ptr;
+  CodeBufferEnd = Buffer->Ptr + Buffer->UsableSize();
+  CodeBufferOffset = CodeBufferBase;
 
   OnCodeBufferAllocated(Buffer);
 

FEXCore/Source/Interface/Core/SharedCodeBufferManager.h

@@ -8,6 +8,7 @@ $end_info$
 #pragma once
 
 #include <FEXCore/fextl/memory.h>
+#include <FEXCore/Utils/MathUtils.h>
 #include <FEXCore/Utils/SignalScopeGuards.h>
 #include <FEXCore/Utils/TypeDefines.h>
 
@@ -63,10 +64,48 @@ public:
   fextl::shared_ptr<CodeBuffer> StartMaximalCodeBuffer();
 
   // Write offset into the latest CodeBuffer
-  std::size_t LatestOffset {};
+  uint8_t* CodeBufferBase {};
+  uint8_t* CodeBufferEnd {};
+  std::atomic<uint8_t*> CodeBufferOffset {};
 
-  // Protects writes to the latest CodeBuffer and changes to LatestOffset
-  FEXCore::ForkableUniqueMutex CodeBufferWriteMutex;
+  // Atomically allocate a fixed size buffer out of the current allocated codebuffer.
+  // Lockless because it's just a linear allocator.
+  struct CodeBufferAllocation {
+    const uint8_t* BufferBase;
+    uint8_t* BufferAllocationOffset;
+  };
+
+  CodeBufferAllocation AtomicAllocateBuffer(size_t Size) {
+    Size = FEXCore::AlignUp(Size, 16);
+    LOGMAN_THROW_A_FMT(reinterpret_cast<uintptr_t>(CodeBufferOffset.load()) % 16 == 0, "Buffer needs to always be 16B aligned!");
+
+    auto ExpectedOffset = CodeBufferOffset.load(std::memory_order_relaxed);
+    auto DesiredOffset = ExpectedOffset + Size;
+
+    if (DesiredOffset > CodeBufferEnd) {
+      // Couldn't fit.
+      return {};
+    }
+
+    while (!CodeBufferOffset.compare_exchange_strong(ExpectedOffset, DesiredOffset)) {
+      DesiredOffset = ExpectedOffset + Size;
+
+      if (DesiredOffset > CodeBufferEnd) {
+        // Couldn't fit.
+        return {};
+      }
+    }
+
+    // Managed to fit.
+    return {
+      .BufferBase = CodeBufferBase,
+      .BufferAllocationOffset = ExpectedOffset,
+    };
+  }
+
+  size_t GetAllocatedSize() const {
+    return CodeBufferOffset - CodeBufferBase;
+  }
 
   virtual void OnCodeBufferAllocated(const std::shared_ptr<CodeBuffer>&) {};