SharedCodeBufferManager: Allocate JIT space atomically.

This removes the fairly long lived lock that the buffer allocator held
while doing significantly more work than intended while holding that
lock.

As the first step towards moving over to the atomic bitmap allocator,
change this to be atomic to closer match what the new allocator is
doing. Since we are just doing linear allocations, this is an easy
convert and should give a good stutter improvement.

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>&) {};