Add "free-list" for JIT generator objects

Summary:
Redo of D77572569. The original had a bug where upon shutdown the module state was free'd but we still go down the `jit_dealloc_gen()` path which tries to make use of it.

We now swap back the original dealloc methods for the system generators on shutdown. We also move the free list data to the module state and bump the module state's refcount so we don't free memory from underneath outstanding JIT generators.

Reviewed By: alexmalyshev

Differential Revision: D78516909

fbshipit-source-id: 5ff870968d90873bedd1b4a34edae132133f4ee9

Author: jbower-fb

cinderx/Common/log.h

@@ -79,6 +79,16 @@ std::string repr(BorrowedRef<> obj);
     }                                             \
   }
 
+#define JIT_CHECK_ONCE(COND, ...)   \
+  {                                 \
+    static bool checked = false;    \
+    if (!checked) {                 \
+      JIT_CHECK(COND, __VA_ARGS__); \
+    } else {                        \
+      checked = true;               \
+    }                               \
+  }
+
 #define JIT_ABORT(...)                                               \
   {                                                                  \
     fmt::print(stderr, "JIT: {}:{} -- Abort\n", __FILE__, __LINE__); \
@@ -97,6 +107,7 @@ std::string repr(BorrowedRef<> obj);
 #ifdef Py_DEBUG
 #define JIT_DABORT(...) JIT_ABORT(__VA_ARGS__)
 #define JIT_DCHECK(COND, ...) JIT_CHECK((COND), __VA_ARGS__)
+#define JIT_DCHECK_ONCE(COND, ...) JIT_CHECK_ONCE((COND), __VA_ARGS__)
 #else
 #define JIT_DABORT(...)     \
   if (0) {                  \
@@ -106,6 +117,10 @@ std::string repr(BorrowedRef<> obj);
   if (0) {                          \
     JIT_CHECK((COND), __VA_ARGS__); \
   }
+#define JIT_DCHECK_ONCE(COND, ...)       \
+  if (0) {                               \
+    JIT_CHECK_ONCE((COND), __VA_ARGS__); \
+  }
 #endif
 
 } // namespace jit

cinderx/Common/util.h

@@ -128,6 +128,12 @@ constexpr T roundUp(T x, size_t n) {
   return roundDown(x + n - 1, n);
 }
 
+template <typename T1, typename T2>
+  requires std::is_integral_v<T1> && std::is_integral_v<T2>
+constexpr std::common_type_t<T1, T1> ceilDiv(T1 a, T2 b) {
+  return (a + b - 1) / b;
+}
+
 constexpr int kCoFlagsAnyGenerator =
     CO_ASYNC_GENERATOR | CO_COROUTINE | CO_GENERATOR | CO_ITERABLE_COROUTINE;
 

cinderx/Jit/codegen/gen_asm.cpp

@@ -1376,13 +1376,10 @@ void NativeGenerator::generateResumeEntry() {
 #if PY_VERSION_HEX < 0x030C0000
   auto gi_jit_data_offset = offsetof(PyGenObject, gi_jit_data);
 #else
-  int python_frame_data_bytes =
-      _PyFrame_NumSlotsForCodeObject(GetFunction()->code) *
-      cinderx::getModuleState()->genType()->tp_itemsize;
-  Py_ssize_t gi_jit_data_offset =
-      cinderx::getModuleState()->genType()->tp_basicsize +
-      python_frame_data_bytes;
-
+  Py_ssize_t python_frame_slots =
+      _PyFrame_NumSlotsForCodeObject(GetFunction()->code);
+  Py_ssize_t gi_jit_data_offset = _PyObject_VAR_SIZE(
+      cinderx::getModuleState()->genType(), python_frame_slots);
 #endif
   as_->mov(jit_data_r, x86::ptr(x86::rdi, gi_jit_data_offset));
 

cinderx/Jit/generators_mm.cpp

@@ -0,0 +1,121 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#include "cinderx/python.h"
+
+#if PY_VERSION_HEX >= 0x030C0000
+
+#include "internal/pycore_object.h"
+
+#include "cinderx/Common/log.h"
+#include "cinderx/Common/util.h"
+#include "cinderx/Jit/generators_mm.h"
+#include "cinderx/module_state.h"
+#if PY_VERSION_HEX >= 0x030E0000
+#include "internal/pycore_interpframe.h"
+#endif
+
+namespace jit {
+
+JitGenFreeList:: // NOLINT(cppcoreguidelines-pro-type-member-init)
+    JitGenFreeList() {
+  Entry* next = nullptr;
+  for (size_t i = 0; i < kGenFreeListEntries; ++i) {
+    entries_[i].next = next;
+    next = &entries_[i];
+  }
+  head_ = next;
+}
+
+void* JitGenFreeList::rawAllocate() {
+  JIT_DCHECK(head_, "No free generator entries");
+  Entry* entry = head_;
+  head_ = entry->next;
+  // The memory for the free-list is backed by the module state, so bump the
+  // reference count to prevent it being free'd before all free-listed
+  // generators are.
+  Py_INCREF(cinderx::getModuleState()->module());
+  return entry->data;
+}
+
+bool JitGenFreeList::fromThisArena(void* ptr) {
+  return ptr >= &entries_ && ptr < &entries_[kGenFreeListEntries - 1] + 1;
+}
+
+void JitGenFreeList::free(PyObject* ptr) {
+  if (!fromThisArena(ptr)) {
+    PyObject_GC_Del(ptr);
+    return;
+  }
+  // Note we assert in allocate() that the "presize" of data is
+  // sizeof(PyGC_HEAD)
+  Entry* entry = reinterpret_cast<Entry*>( // NOLINT(performance-no-int-to-ptr)
+      reinterpret_cast<uintptr_t>(ptr) - sizeof(PyGC_Head));
+  JIT_DCHECK(
+      (reinterpret_cast<uintptr_t>(entry) -
+       reinterpret_cast<uintptr_t>(&entries_[0])) %
+              kGenFreeListEntrySize ==
+          0,
+      "Incorrect pointer calculation");
+  entry->next = head_;
+  head_ = entry;
+  // See comment in rawAllocate()
+  Py_DECREF(cinderx::getModuleState()->module());
+}
+
+std::pair<JitGenObject*, size_t> JitGenFreeList::allocate(
+    BorrowedRef<PyCodeObject> code,
+    uint64_t jit_data_size) {
+  BorrowedRef<PyTypeObject> gen_tp = cinderx::getModuleState()->genType();
+  // We *assume* these assertions hold in free().
+  JIT_DCHECK_ONCE(
+      _PyType_PreHeaderSize(gen_tp) == sizeof(PyGC_Head) &&
+          !_PyType_HasFeature(gen_tp, Py_TPFLAGS_PREHEADER),
+      "Unexpected pre-header setup");
+
+  // A "slot" is the size of PyObject* and we assume this just means 64 bits for
+  // purposes of sizing allocation to cover JIT data.
+  static_assert(sizeof(uint64_t) == sizeof(PyObject*));
+
+  // +1 for the pointer to JIT data (GenDataFooter*)
+  size_t slots =
+      _PyFrame_NumSlotsForCodeObject(code) + 1 + ceilDiv(jit_data_size, 8);
+  // All the generator types should be the same size.
+  size_t size = _PyObject_VAR_SIZE(gen_tp, slots);
+  size_t total_size = sizeof(PyGC_Head) + size;
+
+  bool is_coro = !!(code->co_flags & CO_COROUTINE);
+
+  if (!head_ || total_size > kGenFreeListEntrySize) {
+    JitGenObject* gen = is_coro
+        ? reinterpret_cast<JitGenObject*>(PyObject_GC_NewVar(
+              PyCoroObject, cinderx::getModuleState()->coroType(), slots))
+        : reinterpret_cast<JitGenObject*>(
+              PyObject_GC_NewVar(PyGenObject, gen_tp, slots));
+    // See comment in allocate_and_link_interpreter_frame about failure.
+    JIT_CHECK(gen != nullptr, "Failed to allocate JitGenObject");
+    return {gen, size};
+  }
+
+  void* raw = rawAllocate();
+  // Zero the pre-header, which in this case is the GC header. The
+  // The reference for this is gc_alloc() + _PyObject_GC_Link(). It
+  // would be nice if the latter were public so we could custom
+  // allocate GC'able objects.
+  // Note we are NOT bumping the GC's young generation counter here as
+  // _PyObject_GC_Link would. I argue we're not actually increasing memory
+  // pressure so this is not needed.
+  (reinterpret_cast<PyObject**>(raw))[0] = nullptr;
+  (reinterpret_cast<PyObject**>(raw))[1] = nullptr;
+  PyVarObject* op =
+      reinterpret_cast<PyVarObject*>( // NOLINT(performance-no-int-to-ptr)
+          reinterpret_cast<uintptr_t>(raw) + sizeof(PyGC_Head));
+
+  PyTypeObject* tp = is_coro ? cinderx::getModuleState()->coroType() : gen_tp;
+  _PyObject_InitVar(op, tp, slots);
+
+  return {reinterpret_cast<JitGenObject*>(op), size};
+}
+
+} // namespace jit
+
+#endif // PY_VERSION_HEX >= 0x030C0000

cinderx/Jit/generators_mm.h

@@ -0,0 +1,57 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+#pragma once
+
+#include <Python.h>
+
+#if PY_VERSION_HEX >= 0x030C0000
+
+#include "cinderx/Common/ref.h"
+#include "cinderx/Jit/generators_mm_iface.h"
+
+#include <array>
+
+namespace jit {
+
+struct JitGenObject;
+
+// These values were determined experimentally on IG's webservers by utilizing
+// the stats above. The number of outstanding requests seems to burst up to ~60k
+// on startup but then quickly settles down to around 1-2k, so 2048 entries
+// should be enough. The average size seems to be ~400 bytes with the max being
+// about 10x that. Performance experiments showed a size of 512 was a greater
+// improvement compared to 1024. Presumably the trade off in extra fixed memory
+// allocation cost on workers isn't worth it for greater sizes.
+constexpr size_t kGenFreeListEntries = 2048;
+constexpr size_t kGenFreeListEntrySize = 512;
+
+// Basically a free-list but the backing memory is pre-allocated in a single
+// block. This makes it possible to determine if the storage is from this pool
+// even after deopt by just examining a generator's pointer value.
+class JitGenFreeList : public IJitGenFreeList {
+ public:
+  JitGenFreeList();
+  ~JitGenFreeList() override = default;
+
+  std::pair<JitGenObject*, size_t> allocate(
+      BorrowedRef<PyCodeObject> code,
+      uint64_t jit_spill_words) override;
+  void free(PyObject* ptr) override;
+
+ private:
+  void* rawAllocate();
+  bool fromThisArena(void* ptr);
+
+  struct Entry {
+    union {
+      uint8_t data[kGenFreeListEntrySize];
+      Entry* next;
+    };
+  };
+
+  std::array<Entry, kGenFreeListEntries> entries_;
+  Entry* head_;
+};
+
+} // namespace jit
+
+#endif // PY_VERSION_HEX >= 0x030C0000

cinderx/Jit/generators_mm_iface.h

@@ -0,0 +1,24 @@
+// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+#pragma once
+
+#include <Python.h>
+
+#include "cinderx/Common/ref.h"
+
+namespace jit {
+
+struct JitGenObject;
+
+class IJitGenFreeList {
+ public:
+  IJitGenFreeList() = default;
+  virtual ~IJitGenFreeList() = default;
+
+  virtual std::pair<JitGenObject*, size_t> allocate(
+      BorrowedRef<PyCodeObject> code,
+      uint64_t jit_spill_words) = 0;
+  virtual void free(PyObject* ptr) = 0;
+};
+
+} // namespace jit

cinderx/Jit/generators_rt.cpp

@@ -12,7 +12,10 @@
 #include "cinderx/Jit/deopt.h"
 #include "cinderx/Jit/frame.h"
 #include "cinderx/Jit/generators_borrowed.h"
+#include "cinderx/Jit/generators_mm.h"
+#include "cinderx/Jit/generators_rt.h"
 #include "cinderx/Jit/runtime.h"
+#include "cinderx/UpstreamBorrow/borrowed.h"
 #include "cinderx/module_state.h"
 
 #include <string_view>
@@ -32,11 +35,57 @@ PyObject* JitGenObject::yieldFrom() {
 
 namespace {
 
-void jitgen_dealloc(PyObject* obj) {
-  if (!deopt_jit_gen(obj)) {
+const destructor original_gen_dealloc = PyGen_Type.tp_dealloc;
+const destructor original_coro_dealloc = PyCoro_Type.tp_dealloc;
+
+// This is mostly a copy of gen_dealloc from genobject.c but with a deopt at the
+// start, using our own memory manager for free at the end, some minor
+// tweaks for C++, and to use public APIs.
+void jitgen_dealloc(PyObject* self) {
+  if (!deopt_jit_gen(self)) {
     JIT_ABORT("Tried to dealloc a running JIT generator");
   }
-  return Py_TYPE(obj)->tp_dealloc(obj);
+
+  PyGenObject* gen = reinterpret_cast<PyGenObject*>(self);
+
+  PyObject_GC_UnTrack(gen);
+
+  if (gen->gi_weakreflist != nullptr) {
+    PyObject_ClearWeakRefs(self);
+  }
+
+  PyObject_GC_Track(self);
+
+  if (PyObject_CallFinalizerFromDealloc(self)) {
+    return; /* resurrected.  :( */
+  }
+
+  PyObject_GC_UnTrack(self);
+  if (PyAsyncGen_CheckExact(gen)) {
+    /* We have to handle this case for asynchronous generators
+       right here, because this code has to be between UNTRACK
+       and GC_Del. */
+    Py_CLEAR(reinterpret_cast<PyAsyncGenObject*>(gen)->ag_origin_or_finalizer);
+  }
+  if (gen->gi_frame_state < FRAME_CLEARED) {
+    _PyInterpreterFrame* frame = generatorFrame(gen);
+    gen->gi_frame_state = FRAME_CLEARED;
+    frame->previous = nullptr;
+    _PyFrame_ClearExceptCode(frame);
+  }
+  PyCodeObject* code = frameCode(generatorFrame(gen));
+  if (code->co_flags & CO_COROUTINE) {
+    Py_CLEAR(reinterpret_cast<PyCoroObject*>(gen)->cr_origin_or_finalizer);
+  }
+  Py_DECREF(code);
+  Py_CLEAR(gen->gi_name);
+  Py_CLEAR(gen->gi_qualname);
+  _PyErr_ClearExcState(&gen->gi_exc_state);
+#if PY_VERSION_HEX < 0x030E0000
+  Py_CLEAR(gen->gi_ci_awaiter);
+#endif
+
+  cinderx::getModuleState()->jitGenFreeList()->free(self);
 }
 
 int jitgen_traverse(PyObject* obj, visitproc visit, void* arg) {
@@ -715,6 +764,10 @@ void init_jit_genobject_type() {
   auto copy_getset = [](PyGetSetDef* src, PyGetSetDef* target) {
     int i;
     for (i = 0; src[i].name != nullptr; ++i) {
+      JIT_CHECK(
+          target[i].name != nullptr,
+          "Missing getter/setter on JIT generator: {}",
+          src[i].name);
       JIT_CHECK(
           std::string_view(target[i].name) == src[i].name,
           "Name mismatch: {} != {}",
@@ -744,6 +797,10 @@ void init_jit_genobject_type() {
   auto copy_methods = [](PyMethodDef* src, PyMethodDef* target) {
     int i;
     for (i = 0; src[i].ml_name != nullptr; ++i) {
+      JIT_CHECK(
+          target[i].ml_name != nullptr,
+          "Missing method on JIT generator: {}",
+          src[i].ml_name);
       JIT_CHECK(
           std::string_view(target[i].ml_name) == src[i].ml_name,
           "Name mismatch: {} != {}",
@@ -766,6 +823,16 @@ void init_jit_genobject_type() {
   copy_methods(PyGen_Type.tp_methods, gen_type->tp_methods);
   copy_methods(PyCoro_Type.tp_methods, coro_type->tp_methods);
 
+  cinderx::getModuleState()->setJitGenFreeList(new JitGenFreeList());
+
+  // Override dealloc so we can use a "free-list" for our objects.
+  JIT_CHECK(
+      PyGen_Type.tp_dealloc == original_gen_dealloc &&
+          PyCoro_Type.tp_dealloc == original_coro_dealloc,
+      "PyGen/Coro_Type already overridden");
+  PyGen_Type.tp_dealloc = reinterpret_cast<destructor>(jitgen_dealloc);
+  PyCoro_Type.tp_dealloc = reinterpret_cast<destructor>(jitgen_dealloc);
+
 #ifdef ENABLE_GENERATOR_AWAITER
   JIT_CHECK(
       PyCoro_Type.tp_flags & Ci_TPFLAGS_HAVE_AM_EXTRA,
@@ -776,6 +843,11 @@ void init_jit_genobject_type() {
 #endif
 }
 
+void shutdown_jit_genobject_type() {
+  PyGen_Type.tp_dealloc = original_gen_dealloc;
+  PyCoro_Type.tp_dealloc = original_coro_dealloc;
+}
+
 static PyMethodDef anextawaitable_methods[] = {
     {"send", (PyCFunction)Ci_anextawaitable_send, METH_O, ""},
     {"throw", (PyCFunction)Ci_anextawaitable_throw, METH_VARARGS, ""},