FEXCore/SharedCodeBufferManager: Pivot what tracks memory allocations

It's soon going to change how these buffers are managed, where the
CodeBuffer is going to manage its own allocations soon once it changes
over to the bitmap allocator. Additionally the Manager class is actually
going to do proper management, pooling, and invalidation handling.

Split the task preemptively before we switch to the bitmap allocator to
reduce churn. A little change in the CodeCache where it needs to query
the codebuffer directly rather than the context, but fairly safe.

Shouldn't be any real behaviour change.

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.GetAllocatedSize(), Utils::FEX_PAGE_SIZE);
+  header.CodeBufferSize = FEXCore::AlignUp(CodeBuffer->GetAllocatedSize(), Utils::FEX_PAGE_SIZE);
   header.NumRelocations = Relocations.size();
   header.SerializedBaseAddress = SerializedBaseAddress;
   ::write(fd, &header, sizeof(header));
@@ -361,7 +361,8 @@ 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.GetAllocatedSize());
+  std::span CodeBufferData(reinterpret_cast<std::byte*>(CodeBuffer->Ptr),
+                           reinterpret_cast<std::byte*>(CodeBuffer->Ptr) + CodeBuffer->GetAllocatedSize());
   if (!ApplyCodeRelocations(SerializedBaseAddress, CodeBufferData, Relocations, 0, true)) {
     LOGMAN_THROW_A_FMT(false, "Failed to apply code relocations");
     return false;
@@ -446,10 +447,10 @@ void CodeCache::Validate(const ExecutableFileSectionInfo& Section, fextl::set<ui
   }));
   (void)ApplyCodeRelocations(Section.FileStartVA, CodeBufferRangeRef, NewRelocations, 0, false);
 
-  if (ValidationCTX->GetAllocatedSize() <= CodeBufferRangeRef.size()) {
+  if (NewCodeBuffer->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->GetAllocatedSize());
+    CodeBufferRangeRef = CodeBufferRangeRef.subspan(0, NewCodeBuffer->GetAllocatedSize());
   }
 
   auto [Mismatch, _] = std::mismatch(CodeBufferRangeRef.begin(), CodeBufferRangeRef.end(), CachedCode.begin());
@@ -500,7 +501,7 @@ void CodeCache::Validate(const ExecutableFileSectionInfo& Section, fextl::set<ui
 
   // Reset Context state for next validation
   ValidationThread->LookupCache->ClearCache(ValidationThread->LookupCache->AcquireWriteLock());
-  ValidationCTX->CodeBufferOffset = ValidationCTX->CodeBufferBase;
+  NewCodeBuffer->CodeBufferOffset = NewCodeBuffer->Ptr;
 
   LogMan::Msg::IFmt("            successfully validated cache");
 }

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

@@ -816,8 +816,8 @@ void Arm64JITCore::EmitEntryPoint(ARMEmitter::BackwardLabel& HeaderLabel, bool C
 }
 
 
-SharedCodeBufferManager::CodeBufferAllocation Arm64JITCore::AllocateCodeBufferInSharedCache(size_t Size) {
-  SharedCodeBufferManager::CodeBufferAllocation AllocatedInfo {};
+CodeBuffer::CodeBufferAllocation Arm64JITCore::AllocateCodeBufferInSharedCache(size_t Size) {
+  CodeBuffer::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
@@ -829,7 +829,7 @@ SharedCodeBufferManager::CodeBufferAllocation Arm64JITCore::AllocateCodeBufferIn
 
   // Attempt to allocate a buffer from the SharedCodeBuffers.
   while (AllocatedInfo.BufferAllocationOffset == nullptr) {
-    AllocatedInfo = SharedCodeBuffers.AtomicAllocateBuffer(Size);
+    AllocatedInfo = CurrentCodeBuffer->AtomicAllocateBuffer(Size);
 
     if (AllocatedInfo.BufferAllocationOffset == nullptr) {
       // If it didn't fit then clear the buffer and try again.

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

@@ -627,7 +627,7 @@ private:
 
   void EmitEntryPoint(ARMEmitter::BackwardLabel& HeaderLabel, bool CheckTF);
 
-  [[nodiscard]] SharedCodeBufferManager::CodeBufferAllocation AllocateCodeBufferInSharedCache(size_t Size);
+  [[nodiscard]] CodeBuffer::CodeBufferAllocation AllocateCodeBufferInSharedCache(size_t Size);
 
 #define DEF_OP(x) void Op_##x(IR::IROp_Header const* IROp, IR::Ref Node)
 

FEXCore/Source/Interface/Core/SharedCodeBufferManager.cpp

@@ -34,6 +34,9 @@ CodeBuffer::CodeBuffer(size_t Size)
   FEXCore::Allocator::VirtualTHPControl(Ptr, Size, FEXCore::Allocator::THPControl::Enable);
 
   LookupCache = fextl::make_unique<GuestToHostMap>();
+
+  CodeBufferEnd = Ptr + UsableSize();
+  CodeBufferOffset = Ptr;
 }
 
 CodeBuffer::~CodeBuffer() {
@@ -66,9 +69,6 @@ fextl::shared_ptr<CodeBuffer> SharedCodeBufferManager::AllocateNew(size_t Size)
   auto Buffer = fextl::make_shared<CodeBuffer>(Size);
 
   Latest = Buffer;
-  CodeBufferBase = Buffer->Ptr;
-  CodeBufferEnd = Buffer->Ptr + Buffer->UsableSize();
-  CodeBufferOffset = CodeBufferBase;
 
   OnCodeBufferAllocated(Buffer);
 

FEXCore/Source/Interface/Core/SharedCodeBufferManager.h

@@ -22,8 +22,12 @@ struct GuestToHostMap;
 namespace FEXCore::CPU {
 struct CodeBuffer {
   uint8_t* Ptr;
+  uint8_t* CodeBufferEnd;
   size_t AllocatedSize; // including guard page; see UsableSize()
 
+  // Code buffer allocation information.
+  std::atomic<uint8_t*> CodeBufferOffset {};
+
   fextl::unique_ptr<GuestToHostMap> LookupCache;
 
   CodeBuffer(size_t Size);
@@ -38,35 +42,6 @@ struct CodeBuffer {
   size_t UsableSize() const {
     return AllocatedSize - FEXCore::Utils::FEX_PAGE_SIZE;
   }
-};
-
-/**
- * A manager that coordinates access to the CodeBuffer used for compiling new code across threads.
- *
- * The CodeBuffer is managed as a partially persistent data structure:
- * - Exactly one CodeBuffer is now designated as "active", which means data can be appended to it
- * - Lossy modifications to the active CodeBuffer will not invalidate any data in use by other threads (which is what enables save CodeBuffer sharing across threads)
- * - Instead, such lossy modifications trigger a new "version" of the data in the modifying thread. Old versions of the CodeBuffer persist as read-only data for use by the other threads.
- * - The other threads can update their version of the CodeBuffer. This will decrease the reference count and eventually trigger deallocation of the old version
- */
-class SharedCodeBufferManager {
-public:
-  // Get the CodeBuffer that was most recently allocated.
-  // This is the only CodeBuffer that data may be written to.
-  fextl::shared_ptr<CodeBuffer> GetLatest();
-
-  // Allocate a new CodeBuffer with geometric growth up to an internal maximum.
-  // Subsequent calls to GetLatest will point to the returned buffer.
-  fextl::shared_ptr<CodeBuffer> StartLargerCodeBuffer();
-
-  // Allocate a new CodeBuffer with maximum internal size.
-  // Subsequent calls to GetLatest will point to the returned buffer.
-  fextl::shared_ptr<CodeBuffer> StartMaximalCodeBuffer();
-
-  // Write offset into the latest CodeBuffer
-  uint8_t* CodeBufferBase {};
-  uint8_t* CodeBufferEnd {};
-  std::atomic<uint8_t*> CodeBufferOffset {};
 
   // Atomically allocate a fixed size buffer out of the current allocated codebuffer.
   // Lockless because it's just a linear allocator.
@@ -98,14 +73,40 @@ public:
 
     // Managed to fit.
     return {
-      .BufferBase = CodeBufferBase,
+      .BufferBase = Ptr,
       .BufferAllocationOffset = ExpectedOffset,
     };
   }
 
   size_t GetAllocatedSize() const {
-    return CodeBufferOffset - CodeBufferBase;
+    return CodeBufferOffset - Ptr;
   }
+};
+
+/**
+ * A manager that coordinates access to the CodeBuffer used for compiling new code across threads.
+ *
+ * The CodeBuffer is managed as a partially persistent data structure:
+ * - Exactly one CodeBuffer is now designated as "active", which means data can be appended to it
+ * - Lossy modifications to the active CodeBuffer will not invalidate any data in use by other threads (which is what enables save CodeBuffer sharing across threads)
+ * - Instead, such lossy modifications trigger a new "version" of the data in the modifying thread. Old versions of the CodeBuffer persist as read-only data for use by the other threads.
+ * - The other threads can update their version of the CodeBuffer. This will decrease the reference count and eventually trigger deallocation of the old version
+ */
+class SharedCodeBufferManager {
+public:
+  virtual ~SharedCodeBufferManager() = default;
+
+  // Get the CodeBuffer that was most recently allocated.
+  // This is the only CodeBuffer that data may be written to.
+  fextl::shared_ptr<CodeBuffer> GetLatest();
+
+  // Allocate a new CodeBuffer with geometric growth up to an internal maximum.
+  // Subsequent calls to GetLatest will point to the returned buffer.
+  fextl::shared_ptr<CodeBuffer> StartLargerCodeBuffer();
+
+  // Allocate a new CodeBuffer with maximum internal size.
+  // Subsequent calls to GetLatest will point to the returned buffer.
+  fextl::shared_ptr<CodeBuffer> StartMaximalCodeBuffer();
 
   virtual void OnCodeBufferAllocated(const std::shared_ptr<CodeBuffer>&) {};