Move CodeAllocator into the ModuleState

Summary:
Requires adding a new ICodeAllocator interface, and then every user can access
it via the module state.

Tweaked the interface to not use an outparam for addCode(), it now returns a
struct of the address and the error code.

Reviewed By: DinoV

Differential Revision: D78522984

fbshipit-source-id: e7972f7eb6c5ff9446a3fbb738159f5abf758701

Author: alexmalyshev

Jit/code_allocator.cpp

@@ -2,6 +2,7 @@
 
 #include "cinderx/Jit/code_allocator.h"
 
+#include "cinderx/Common/log.h"
 #include "cinderx/Jit/config.h"
 #include "cinderx/Jit/threaded_compile.h"
 
@@ -16,6 +17,11 @@ using codegen::codeSectionFromName;
 
 namespace {
 
+#define PROPAGATE_ERROR(EXPR)                                \
+  if (asmjit::Error err = (EXPR); err != asmjit::kErrorOk) { \
+    return AllocateResult{nullptr, err};                     \
+  }
+
 // 2MiB to match Linux's huge-page size.
 constexpr size_t kAllocSize = 1024 * 1024 * 2;
 
@@ -51,26 +57,24 @@ bool setHugePages([[maybe_unused]] void* ptr, [[maybe_unused]] size_t size) {
 
 } // namespace
 
-CodeAllocator* CodeAllocator::s_global_code_allocator_ = nullptr;
-
-CodeAllocator::~CodeAllocator() {}
-
-void CodeAllocator::makeGlobalCodeAllocator() {
-  JIT_CHECK(
-      s_global_code_allocator_ == nullptr, "Global allocator already set");
+ICodeAllocator* CodeAllocator::make() {
   if (getConfig().multiple_code_sections) {
-    s_global_code_allocator_ = new MultipleSectionCodeAllocator;
+    return new MultipleSectionCodeAllocator{};
   } else if (getConfig().use_huge_pages) {
-    s_global_code_allocator_ = new CodeAllocatorCinder;
-  } else {
-    s_global_code_allocator_ = new CodeAllocator;
+    return new CodeAllocatorCinder{};
   }
+  return new CodeAllocator{};
 }
 
-void CodeAllocator::freeGlobalCodeAllocator() {
-  JIT_CHECK(s_global_code_allocator_ != nullptr, "Global allocator not set");
-  delete s_global_code_allocator_;
-  s_global_code_allocator_ = nullptr;
+AllocateResult CodeAllocator::addCode(asmjit::CodeHolder* code) {
+  void* addr = nullptr;
+  asmjit::Error error = runtime_.add(&addr, code);
+
+  if (addr != nullptr && error == asmjit::kErrorOk) {
+    used_bytes_.fetch_add(code->codeSize(), std::memory_order_relaxed);
+  }
+
+  return AllocateResult{addr, error};
 }
 
 bool CodeAllocator::contains(const void* ptr) const {
@@ -81,22 +85,26 @@ bool CodeAllocator::contains(const void* ptr) const {
       asmjit::kErrorOk;
 }
 
+size_t CodeAllocator::usedBytes() const {
+  return used_bytes_.load(std::memory_order_relaxed);
+}
+
+const asmjit::Environment& CodeAllocator::asmJitEnvironment() const {
+  return runtime_.environment();
+}
+
 CodeAllocatorCinder::~CodeAllocatorCinder() {
   for (std::span<uint8_t> alloc : allocations_) {
     JIT_CHECK(
         munmap(alloc.data(), alloc.size()) == 0, "Freeing code memory failed");
   }
 }
 
-asmjit::Error CodeAllocatorCinder::addCode(
-    void** dst,
-    asmjit::CodeHolder* code) noexcept {
+AllocateResult CodeAllocatorCinder::addCode(asmjit::CodeHolder* code) {
   ThreadedCompileSerialize guard;
 
-  *dst = nullptr;
-
-  ASMJIT_PROPAGATE(code->flatten());
-  ASMJIT_PROPAGATE(code->resolveUnresolvedLinks());
+  PROPAGATE_ERROR(code->flatten());
+  PROPAGATE_ERROR(code->resolveUnresolvedLinks());
 
   size_t max_code_size = code->codeSize();
   size_t alloc_size = ((max_code_size / kAllocSize) + 1) * kAllocSize;
@@ -114,7 +122,7 @@ asmjit::Error CodeAllocatorCinder::addCode(
     current_alloc_free_ = alloc_size;
   }
 
-  ASMJIT_PROPAGATE(code->relocateToBase(uintptr_t(current_alloc_)));
+  PROPAGATE_ERROR(code->relocateToBase(uintptr_t(current_alloc_)));
 
   size_t actual_code_size = code->codeSize();
   JIT_CHECK(actual_code_size <= max_code_size, "Code grew during relocation");
@@ -136,13 +144,13 @@ asmjit::Error CodeAllocatorCinder::addCode(
     }
   }
 
-  *dst = current_alloc_;
+  void* addr = current_alloc_;
 
   current_alloc_ += actual_code_size;
   current_alloc_free_ -= actual_code_size;
   used_bytes_ += actual_code_size;
 
-  return asmjit::kErrorOk;
+  return AllocateResult{addr, asmjit::kErrorOk};
 }
 
 bool CodeAllocatorCinder::contains(const void* ptr) const {
@@ -196,15 +204,12 @@ void MultipleSectionCodeAllocator::createSlabs() noexcept {
   code_sections_[CodeSection::kCold] = region;
 }
 
-asmjit::Error MultipleSectionCodeAllocator::addCode(
-    void** dst,
-    asmjit::CodeHolder* code) noexcept {
+AllocateResult MultipleSectionCodeAllocator::addCode(asmjit::CodeHolder* code) {
   ThreadedCompileSerialize guard;
 
   if (code_sections_.empty()) {
     createSlabs();
   }
-  *dst = nullptr;
 
   size_t potential_code_size = code->codeSize();
   used_bytes_ += potential_code_size;
@@ -216,8 +221,11 @@ asmjit::Error MultipleSectionCodeAllocator::addCode(
     JIT_LOG(
         "Not enough memory to split code across sections, falling back to "
         "normal allocation.");
-    return runtime_.add(dst, code);
+    void* addr = nullptr;
+    asmjit::Error err = runtime_.add(&addr, code);
+    return AllocateResult{addr, err};
   }
+
   // Fix up the offsets for each code section before resolving links.
   // Both the `.text` and `.addrtab` sections are written to the hot section,
   // and we need to resolve offsets between them properly.
@@ -242,16 +250,16 @@ asmjit::Error MultipleSectionCodeAllocator::addCode(
 
   // Assuming that the offsets are set properly, relocating all code to be
   // relative to the start of the hot code will ensure jumps are correct.
-  ASMJIT_PROPAGATE(code->resolveUnresolvedLinks());
-  ASMJIT_PROPAGATE(
+  PROPAGATE_ERROR(code->resolveUnresolvedLinks());
+  PROPAGATE_ERROR(
       code->relocateToBase(uintptr_t(code_sections_[CodeSection::kHot])));
 
   // We assume that the hot section of the code is non-empty. This would be
   // incorrect for a completely cold function.
   JIT_CHECK(
       code->textSection()->realSize() > 0,
       "Every function must have a non-empty hot section.");
-  *dst = code_sections_[CodeSection::kHot];
+  void* addr = code_sections_[CodeSection::kHot];
 
   for (asmjit::Section* section : code->_sections) {
     size_t buffer_size = section->bufferSize();
@@ -265,7 +273,7 @@ asmjit::Error MultipleSectionCodeAllocator::addCode(
     code_sections_[code_section] += buffer_size;
   }
 
-  return asmjit::kErrorOk;
+  return AllocateResult{addr, asmjit::kErrorOk};
 }
 
 bool MultipleSectionCodeAllocator::contains(const void* ptr) const {

Jit/code_allocator.h

@@ -2,14 +2,14 @@
 
 #pragma once
 
-#include "cinderx/Common/log.h"
+#include "cinderx/Jit/code_allocator_iface.h"
 #include "cinderx/Jit/codegen/code_section.h"
 
 #include <asmjit/asmjit.h>
 
 #include <atomic>
-#include <memory>
 #include <span>
+#include <unordered_map>
 #include <vector>
 
 namespace jit {
@@ -27,59 +27,29 @@ namespace jit {
   like accommodate memory pools with different allocation characteristics, or
   have multiple threads which might compile independently.
 */
-class CodeAllocator {
+class CodeAllocator : public ICodeAllocator {
  public:
-  virtual ~CodeAllocator();
-
-  // Get the global code allocator for this process.
-  static CodeAllocator* get() {
-    JIT_CHECK(exists(), "No global code allocator");
-    return s_global_code_allocator_;
-  }
-
-  // Check if the global code allocator has been created.
-  static bool exists() {
-    return s_global_code_allocator_ != nullptr;
-  }
+  ~CodeAllocator() override = default;
 
   // To be called once by JIT initialization after enough configuration has been
   // loaded to determine which global code allocator type to use.
-  static void makeGlobalCodeAllocator();
-
-  static void freeGlobalCodeAllocator();
-
-  size_t usedBytes() const {
-    return used_bytes_;
-  }
-
-  const asmjit::Environment& asmJitEnvironment() const {
-    return runtime_.environment();
-  }
+  static ICodeAllocator* make();
 
-  virtual asmjit::Error addCode(void** dst, asmjit::CodeHolder* code) noexcept {
-    used_bytes_ += code->codeSize();
-    return runtime_.add(dst, code);
-  }
-
-  // Check if a pointer is located within this allocator's memory.
-  virtual bool contains(const void* ptr) const;
+  AllocateResult addCode(asmjit::CodeHolder* code) override;
+  bool contains(const void* ptr) const override;
+  size_t usedBytes() const override;
+  const asmjit::Environment& asmJitEnvironment() const override;
 
  protected:
   asmjit::JitRuntime runtime_;
   std::atomic<size_t> used_bytes_{0};
-
- private:
-  static CodeAllocator* s_global_code_allocator_;
 };
 
 // A code allocator which tries to allocate all code on huge pages.
 class CodeAllocatorCinder : public CodeAllocator {
  public:
   ~CodeAllocatorCinder() override;
 
-  asmjit::Error addCode(void** dst, asmjit::CodeHolder* code) noexcept override;
-  bool contains(const void* ptr) const override;
-
   size_t lostBytes() const {
     return lost_bytes_;
   }
@@ -92,6 +62,9 @@ class CodeAllocatorCinder : public CodeAllocator {
     return huge_allocs_;
   }
 
+  AllocateResult addCode(asmjit::CodeHolder* code) override;
+  bool contains(const void* ptr) const override;
+
  private:
   // List of chunks allocated for use in deallocation
   std::vector<std::span<uint8_t>> allocations_;
@@ -114,7 +87,7 @@ class MultipleSectionCodeAllocator : public CodeAllocator {
  public:
   ~MultipleSectionCodeAllocator() override;
 
-  asmjit::Error addCode(void** dst, asmjit::CodeHolder* code) noexcept override;
+  AllocateResult addCode(asmjit::CodeHolder* code) override;
   bool contains(const void* ptr) const override;
 
  private:

Jit/code_allocator_iface.h

@@ -0,0 +1,36 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#pragma once
+
+#include <asmjit/asmjit.h>
+
+#include <cstddef>
+
+namespace jit {
+
+// Resulting address and status code after calling ICodeAllocator::addCode().
+struct AllocateResult {
+  void* addr{nullptr};
+  asmjit::Error error{asmjit::kErrorNotInitialized};
+};
+
+// Interface for allocating memory for compiled code generated by the JIT.
+class ICodeAllocator {
+ public:
+  virtual ~ICodeAllocator() = default;
+
+  // Add a compiled code object to the allocator.  Get its new address on
+  // success, or nullptr and an error on failure.
+  virtual AllocateResult addCode(asmjit::CodeHolder* code) = 0;
+
+  // Check if a pointer is located within this allocator's memory.
+  virtual bool contains(const void* ptr) const = 0;
+
+  // Get the number of bytes of memory used by the allocator.
+  virtual size_t usedBytes() const = 0;
+
+  // Get the asmjit environment used by this allocator.
+  virtual const asmjit::Environment& asmJitEnvironment() const = 0;
+};
+
+} // namespace jit

Jit/codegen/gen_asm.cpp

@@ -20,7 +20,6 @@
 #include "cinderx/Common/py-portability.h"
 #include "cinderx/Common/util.h"
 #include "cinderx/Interpreter/interpreter.h"
-#include "cinderx/Jit/code_allocator.h"
 #include "cinderx/Jit/codegen/autogen.h"
 #include "cinderx/Jit/codegen/code_section.h"
 #include "cinderx/Jit/codegen/gen_asm_utils.h"
@@ -303,17 +302,16 @@ void* finalizeCode(asmjit::x86::Builder& builder, std::string_view name) {
         DebugUtils::errorAsString(err))};
   }
 
-  void* result = nullptr;
-  if (auto err = CodeAllocator::get()->addCode(&result, builder.code());
-      err != kErrorOk) {
+  ICodeAllocator* code_allocator = cinderx::getModuleState()->codeAllocator();
+  AllocateResult result = code_allocator->addCode(builder.code());
+  if (result.error != kErrorOk) {
     throw std::runtime_error{fmt::format(
         "Failed to add generated code for {} to asmjit runtime, got error code "
         "{}",
         name,
-        DebugUtils::errorAsString(err))};
+        DebugUtils::errorAsString(result.error))};
   }
-
-  return result;
+  return result.addr;
 }
 
 // Generate the final stage trampoline that is responsible for finishing
@@ -323,7 +321,8 @@ void* generateDeoptTrampoline(bool generator_mode) {
       generator_mode ? "deopt_trampoline_generators" : "deopt_trampoline";
 
   CodeHolder code;
-  ASM_CHECK(code.init(CodeAllocator::get()->asmJitEnvironment()), name);
+  ICodeAllocator* code_allocator = cinderx::getModuleState()->codeAllocator();
+  ASM_CHECK(code.init(code_allocator->asmJitEnvironment()), name);
   x86::Builder a(&code);
   Annotations annot;
 
@@ -511,7 +510,8 @@ void* generateDeoptTrampoline(bool generator_mode) {
 
 void* generateFailedDeferredCompileTrampoline() {
   CodeHolder code;
-  code.init(CodeAllocator::get()->asmJitEnvironment());
+  ICodeAllocator* code_allocator = cinderx::getModuleState()->codeAllocator();
+  code.init(code_allocator->asmJitEnvironment());
   x86::Builder a(&code);
   Annotations annot;
 
@@ -660,7 +660,8 @@ void* NativeGenerator::getVectorcallEntry() {
   JIT_CHECK(as_ == nullptr, "x86::Builder should not have been initialized.");
 
   CodeHolder code;
-  code.init(CodeAllocator::get()->asmJitEnvironment());
+  ICodeAllocator* code_allocator = cinderx::getModuleState()->codeAllocator();
+  code.init(code_allocator->asmJitEnvironment());
   ThrowableErrorHandler eh;
   code.setErrorHandler(&eh);