[REFACTOR][OrcJIT] Add SlabPoolMemoryManager — growable per-session pool

Stage B of the slab-pool refactor. `ArenaJITLinkMemoryManager` is
replaced by `SlabPoolMemoryManager`, which holds a vector of `Slab`s
per session and grows on demand instead of pre-reserving one giant
arena.

Behavior changes:

- **Default capacity**: per-slab `slab_size` drops from 1 GB to 64 MB.
  Typical ML JIT graphs are well under 10 MB; the first slab now
  reserves 64 MB instead of 1 GB, dramatically reducing VA footprint
  for small workloads on memory-constrained hosts.
- **Growth on demand**: when no existing slab can fit a graph, a fresh
  slab is mmap'd at `slab_size` bytes and appended to the pool.
  Sessions that previously failed with "pool exhausted" for cumulative
  allocations beyond 1 GB now succeed transparently.
- **Oversize path**: a single graph whose footprint exceeds
  `slab_size - 2 * kCommitGranularity` gets its own dedicated slab
  sized to fit it, rounded to the commit granularity. One graph per
  oversize slab; the slab becomes available to other allocations
  after the graph is freed but usually isn't reused (sized tightly).

New infrastructure:

- `SlabPoolExhaustedError` — retriable error class. Emitted by
  `Slab::bumpAllocate` when the requested region exceeds the pool
  limit; caught by `SlabPoolMemoryManager::allocate` to fall through
  to the next slab or mmap a new one. Other errors (mmap, mprotect,
  JITLink) keep their existing types and are propagated to the caller
  without retry.
- `Slab::computeGraphFootprint(G, page_size)` — static helper that
  pre-computes per-pool byte totals for a graph. The pool manager
  uses it to make the normal-vs-oversize decision without first
  attempting a normal allocation.
- `classifyOverflowSections(G)` — file-scope helper extracted from
  `Slab::allocate`; also used by `computeGraphFootprint`. Keeps the
  two entry points consistent on which sections go to the overflow
  (separate-mmap) path.

Concurrency:

`pool_mu_` guards only the `slabs_` vector. It is dropped before
calling into `Slab::allocate` or the caller's `OnAllocated` callback,
because LLJIT materialization frequently invokes nested lookups from
inside those callbacks — a coarse lock here deadlocks. Existing slab
pointers are stable across concurrent grows (Stage B never removes
slabs), so a snapshot taken under the lock is safe to iterate
afterwards.

Initial-slab retry: the session constructor still halves capacity on
mmap failure, now down to `kMinSlabSize = 8 MB`. Subsequent grows use
exactly `slab_size_` with no retry — mmap failures propagate.

Tests:

- `test_arena.py`: `_ARENA_SIZE` bumped from 16 MB → 256 MB so the
  co-location tests continue to exercise single-slab invariants. The
  existing overflow-section test's contiguous-region assertion still
  passes because 256 MB is enough for one slab.
- `test_basic.py`: 3 new parametrized tests (×2 C/C++ variants = 6
  cases) under the "Slab-pool growth" section —
  `test_pool_grows_under_small_slab` (16 libs, 8 MB slab, pool must
  grow), `test_small_slab_recycles_after_drop` (32-iter load/drop
  exercises free-list within a slab), `test_pool_survives_mixed_load_drop_create`
  (interleaved paths).

Verification: `pytest addons/tvm_ffi_orcjit/tests` — 70 passed,
3 skipped (was 64 + 3).

Scope: Stage C (warm-slab eviction + real munmap of drained slabs)
remains a planned follow-up. Drained slabs in Stage B stay mapped
until the session is destroyed, same as today's single arena.

Author: cyx-6

addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py

@@ -66,18 +66,19 @@ def __init__(self, orc_rt_path: str | None = None, slab_size: int = 0) -> None:
         Args:
             orc_rt_path: Optional path to the liborc_rt library. If not provided,
                         it will be automatically discovered using clang.
-            slab_size: Slab capacity in bytes for the JIT memory manager.
+            slab_size: Per-slab capacity in bytes for the JIT memory manager.
                        Linux only — ignored on macOS and Windows, where the
                        slab allocator is compiled out.
-                       0 = arch default (1 GB; falls back to smaller sizes
-                       down to 256 MB under RLIMIT_AS / container limits),
-                       >0 = custom size, <0 = disable slab allocator (LLJIT
-                       uses its default scattered-mmap allocator).
-
-                       Stage A of the slab-pool refactor: one Slab per session,
-                       so this is also the total arena capacity. Later stages
-                       will grow session memory by adding slabs of this size
-                       on demand.
+                       0 = arch default (64 MB; initial slab halves on mmap
+                       failure down to 8 MB under RLIMIT_AS / container
+                       limits), >0 = custom size, <0 = disable slab allocator
+                       (LLJIT uses its default scattered-mmap allocator).
+
+                       The session holds a growable pool of slabs: a fresh
+                       slab is mmap'd on demand when no existing one can fit
+                       a graph, and graphs larger than slab_size go to a
+                       dedicated oversize slab sized to fit. Drained slabs
+                       stay mapped until the session is destroyed.
 
         """
         if orc_rt_path is None:

addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.cc

@@ -19,52 +19,162 @@
 
 /*!
  * \file orcjit_memory_manager.cc
- * \brief Thin wrapper delegating JITLinkMemoryManager ops to a Slab.
+ * \brief Growable per-session pool of `Slab`s.
  */
 #include "orcjit_memory_manager.h"
 
 #ifdef __linux__
 
+#include <llvm/Support/Alignment.h>
 #include <llvm/Support/Error.h>
 #include <llvm/Support/FormatVariadic.h>
 
 #include <algorithm>
+#include <optional>
+#include <utility>
 
 namespace tvm {
 namespace ffi {
 namespace orcjit {
 
 using llvm::Error;
+using llvm::Expected;
 
-ArenaJITLinkMemoryManager::ArenaJITLinkMemoryManager(std::size_t page_size,
-                                                     std::size_t slab_capacity) {
-  // Try requested capacity, halve on failure down to a minimum floor.
-  // The floor is the smaller of kMinSlabCapacity and the requested size,
-  // so explicit small slabs (e.g. 16 MB for tests) are honoured.
-  // mmap(PROT_NONE | MAP_NORESERVE) can still fail under RLIMIT_AS or
-  // extreme VA fragmentation.
-  std::size_t floor = std::min(slab_capacity, kMinSlabCapacity);
-  std::size_t cap = slab_capacity;
+SlabPoolMemoryManager::SlabPoolMemoryManager(std::size_t page_size, std::size_t slab_size)
+    : page_size_(page_size), slab_size_(slab_size) {
+  // Reserve the initial slab.  Halving retry only applies here: if the
+  // very first mmap fails (RLIMIT_AS, container limits), we halve the
+  // requested size down to kMinSlabSize before giving up.  Subsequent
+  // slabs added during allocate() use exactly slab_size_ and propagate
+  // errors on mmap failure.
+  std::size_t floor = std::min(slab_size_, kMinSlabSize);
+  std::size_t cap = slab_size_;
   while (cap >= floor) {
-    auto slab = std::make_unique<Slab>(page_size, cap);
+    auto slab = std::make_unique<Slab>(page_size_, cap);
     if (slab->isValid()) {
-      slab_ = std::move(slab);
+      // Pin the actual initial-slab size to whatever we succeeded with.
+      // If RLIMIT_AS forced us to 8 MB, we keep 8 MB as the working slab
+      // size; growing later at 64 MB would just fail again.
+      slab_size_ = cap;
+      slabs_.push_back(std::move(slab));
       return;
     }
     cap /= 2;
   }
-  llvm::report_fatal_error("ArenaJITLinkMemoryManager: failed to reserve at least " +
+  llvm::report_fatal_error("SlabPoolMemoryManager: failed to reserve at least " +
                            llvm::Twine(floor / (1024 * 1024)) + " MB of virtual address space");
 }
 
-void ArenaJITLinkMemoryManager::allocate(const llvm::jitlink::JITLinkDylib* /*JD*/,
-                                         llvm::jitlink::LinkGraph& G,
-                                         OnAllocatedFunction OnAllocated) {
-  slab_->allocate(G, std::move(OnAllocated));
+std::unique_ptr<Slab> SlabPoolMemoryManager::createSlab(std::size_t capacity) {
+  auto slab = std::make_unique<Slab>(page_size_, capacity);
+  if (!slab->isValid()) return nullptr;
+  return slab;
 }
 
-void ArenaJITLinkMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
-                                           OnDeallocatedFunction OnDeallocated) {
+void SlabPoolMemoryManager::allocate(const llvm::jitlink::JITLinkDylib* /*JD*/,
+                                     llvm::jitlink::LinkGraph& G,
+                                     OnAllocatedFunction OnAllocated) {
+  // Step 1: pre-compute footprint to decide normal vs oversize.
+  auto footprint = Slab::computeGraphFootprint(G, page_size_);
+  std::size_t total = footprint.total();
+
+  // Step 2: conservative usable-per-slab estimate.  The dual-pool
+  // midpoint split means a graph cannot use the entire slab — one
+  // pool's cursor is bounded at midpoint, the other at
+  // exec_bump_limit.  2 MB of slack covers midpoint alignment.  A
+  // false-positive oversize just costs one extra mmap sized to fit,
+  // never a crash.
+  std::size_t usable = slab_size_ > 2 * Slab::kCommitGranularity
+                           ? slab_size_ - 2 * Slab::kCommitGranularity
+                           : slab_size_ / 2;
+
+  // `pool_mu_` only protects the slabs_ vector itself.  We never hold
+  // it across a call into Slab::allocate or a user callback: the
+  // LLJIT linker will frequently invoke nested lookups (which trigger
+  // recursive allocate() calls via materialization) from inside
+  // OnAllocated, and a coarse lock here would deadlock.  Slabs we've
+  // seen in a snapshot are guaranteed to outlive this call because
+  // Stage B never removes slabs from the pool.
+  using AllocResult = Expected<std::unique_ptr<InFlightAlloc>>;
+
+  // Step 3: oversize path — one graph per dedicated slab.
+  if (total > usable) {
+    std::size_t needed = total + 2 * Slab::kCommitGranularity;
+    std::size_t cap = llvm::alignTo(needed, Slab::kCommitGranularity);
+    if (cap < slab_size_) cap = slab_size_;
+    auto slab = createSlab(cap);
+    if (!slab) {
+      OnAllocated(llvm::make_error<llvm::StringError>(
+          "SlabPoolMemoryManager: mmap failed for oversize slab of " +
+              llvm::formatv("{0:x}", cap).str() + " bytes",
+          llvm::inconvertibleErrorCode()));
+      return;
+    }
+    Slab* raw = slab.get();
+    {
+      std::lock_guard<std::mutex> lock(pool_mu_);
+      slabs_.push_back(std::move(slab));
+    }
+    raw->allocate(G, std::move(OnAllocated));
+    return;
+  }
+
+  // Step 4: first-fit over existing slabs.  Take a snapshot of raw
+  // pointers under the lock, then iterate without holding it.
+  // Slab::allocate is synchronous (invokes its callback inline on
+  // every code path), so we observe the result via a captured
+  // std::optional that the callback fills before the call returns.
+  std::vector<Slab*> snapshot;
+  {
+    std::lock_guard<std::mutex> lock(pool_mu_);
+    snapshot.reserve(slabs_.size());
+    for (auto& s : slabs_) snapshot.push_back(s.get());
+  }
+  for (Slab* slab : snapshot) {
+    std::optional<AllocResult> observed;
+    slab->allocate(G, [&](AllocResult R) { observed.emplace(std::move(R)); });
+    AllocResult result = std::move(*observed);
+    if (result) {
+      OnAllocated(std::move(result));
+      return;
+    }
+    Error E = result.takeError();
+    if (E.isA<SlabPoolExhaustedError>()) {
+      // Retriable: this graph didn't fit in this slab.  Try next.
+      llvm::consumeError(std::move(E));
+      continue;
+    }
+    // Terminal error (mmap, mprotect, JITLink, BasicLayout).
+    OnAllocated(std::move(E));
+    return;
+  }
+
+  // Step 5: no existing slab fits.  Mmap a new normal-size slab.
+  // Another thread may have added slabs meanwhile; we don't re-scan
+  // (would duplicate work under contention).  Concurrent creates
+  // would at worst make the pool grow faster than strictly necessary
+  // — never incorrect.
+  auto slab = createSlab(slab_size_);
+  if (!slab) {
+    OnAllocated(llvm::make_error<llvm::StringError>(
+        "SlabPoolMemoryManager: mmap failed for new slab of " +
+            llvm::formatv("{0:x}", slab_size_).str() + " bytes",
+        llvm::inconvertibleErrorCode()));
+    return;
+  }
+  Slab* raw = slab.get();
+  {
+    std::lock_guard<std::mutex> lock(pool_mu_);
+    slabs_.push_back(std::move(slab));
+  }
+  // A fresh slab must fit any graph we've already decided is in-range
+  // (step 2 + step 3 gate).  If somehow it doesn't, the error is
+  // propagated through — not retried.
+  raw->allocate(G, std::move(OnAllocated));
+}
+
+void SlabPoolMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
+                                       OnDeallocatedFunction OnDeallocated) {
   Error DeallocErr = Error::success();
   for (auto& Alloc : Allocs) {
     auto* FA = Alloc.release().toPtr<FinalizedAllocInfo*>();

addons/tvm_ffi_orcjit/src/ffi/orcjit_memory_manager.h

@@ -19,44 +19,28 @@
 
 /*!
  * \file orcjit_memory_manager.h
- * \brief JITLinkMemoryManager backed by one fixed-size Slab (Stage A).
+ * \brief Per-session growable slab pool (Stage B).
  *
- * Holds a single `Slab` (see orcjit_slab.h) and delegates all
- * `JITLinkMemoryManager` operations to it.  The capacity-negotiation
- * retry loop (halve-on-mmap-failure) lives here because it is a
- * caller/kernel negotiation, not something a Slab itself should do.
+ * `SlabPoolMemoryManager` implements `JITLinkMemoryManager` on top of a
+ * per-session `std::vector<std::unique_ptr<Slab>>`.  On each `allocate`
+ * it picks the first `Slab` that can fit the graph; if none do, it
+ * `mmap`s a new fixed-size (`slab_size`) slab and appends it.  Graphs
+ * larger than a single normal slab go to the oversize path — a
+ * dedicated `Slab` sized to fit that one graph.
  *
- * In later stages this class is superseded by a `SlabPoolMemoryManager`
- * that owns multiple Slabs per session and grows on demand.  The public
- * API surface (one memory manager per LLJIT, constructed with a
- * page-size and a slab-capacity) stays the same.
+ * ## Lifecycle (Stage B)
  *
- * ## GOTPCRELX relaxation workaround
- *
- * The arena triggers a latent bug in LLVM JITLink's
- * `optimizeGOTAndStubAccesses()` (x86_64.cpp).  That pass relaxes
- * `call *foo@GOTPCREL(%rip)` (ff 15) → `addr32 call foo` (67 e8) and
- * sets the edge kind to `Pointer32` (absolute 32-bit address).  However
- * the `call rel32` instruction is always **PC-relative** — the `67`
- * prefix is just padding — so the fixup should be PC-relative too
- * (matching the static linker's `R_X86_64_PC32`).
+ * Once a slab is added to the pool it stays mapped until the pool
+ * (and its enclosing session) is destroyed.  Individual graphs are
+ * deallocated via `FA->owner->deallocateOne(...)`, returning bytes to
+ * the slab's free list, but the slab's VA reservation is not reclaimed.
+ * Stage C will add warm-slab eviction that munmaps drained slabs.
  *
- * The bug is latent because the relaxation only fires when the target
- * address fits in 32 bits (`isUInt<32>`).  On PIE executables every
- * resolved symbol is at a high address, so the guard is never true and
- * the relaxation never runs.  On **non-PIE** executables the PLT
- * entries for libc functions (malloc, free, …) live near 0x400000, the
- * guard passes, and the wrong fixup produces a garbage displacement →
- * SIGSEGV during ORC-runtime teardown.
+ * ## GOTPCRELX relaxation workaround
  *
- * `GOTPCRELXFixPlugin` in llvm_patches/gotpcrelx_fix.cc works around
- * this: a PreFixupPass that runs *after* `optimizeGOTAndStubAccesses`
- * detects `Pointer32` edges on `67 e8` / `e9` instructions and either
- *   (a) converts to `BranchPCRel32` when the PC-relative displacement
- *       fits in int32, or
- *   (b) reverts the relaxation entirely — restores the `ff 15` /
- *       `ff 25` opcode bytes and retargets the edge to the GOT entry
- *       with `PCRel32` + addend 0.
+ * Unchanged from Stage A — see `llvm_patches/gotpcrelx_fix.cc`.  The
+ * plugin is added per-session to the `ObjectLinkingLayer` alongside
+ * this memory manager and is orthogonal to pool growth.
  */
 #ifndef TVM_FFI_ORCJIT_ORCJIT_MEMORY_MANAGER_H_
 #define TVM_FFI_ORCJIT_ORCJIT_MEMORY_MANAGER_H_
@@ -65,48 +49,67 @@
 
 #include <cstddef>
 #include <memory>
+#include <mutex>
+#include <vector>
 
 #include "orcjit_slab.h"
 
 namespace tvm {
 namespace ffi {
 namespace orcjit {
 
-/*! \brief JITLink memory manager backed by a single `Slab`.
- *
- *  Reserves `slab_capacity` bytes of contiguous VA at construction time,
- *  halving the request down to `kMinSlabCapacity` if the initial `mmap`
- *  fails (RLIMIT_AS, container limits).
+/*!
+ * \brief `JITLinkMemoryManager` backed by a growable pool of `Slab`s.
  *
- *  `allocate` and `deallocate` forward to the underlying `Slab`.
+ * The constructor reserves one initial slab (halving its capacity down
+ * to `kMinSlabSize` if `mmap` fails under RLIMIT_AS).  Subsequent
+ * slabs are added on demand by `allocate()` and reserved at exactly
+ * `slab_size_` bytes each — no retry, no halving, errors propagate.
  */
-class ArenaJITLinkMemoryManager : public llvm::jitlink::JITLinkMemoryManager {
+class SlabPoolMemoryManager : public llvm::jitlink::JITLinkMemoryManager {
  public:
-  // Default slab capacity: 1 GB on both architectures.  Well within the
-  // PC-relative relocation limit (x86_64 ±2 GB, AArch64 ±4 GB) so
-  // cross-section fixups always fit; large enough to cover typical ML
-  // JIT workloads without oversubscribing virtual address space on
-  // memory-constrained hosts (containers, CI runners).
-  static constexpr std::size_t kDefaultSlabCapacity_x86_64 = std::size_t{1} << 30;   // 1 GB
-  static constexpr std::size_t kDefaultSlabCapacity_AArch64 = std::size_t{1} << 30;  // 1 GB
-  static constexpr std::size_t kMinSlabCapacity = std::size_t{256} << 20;            // 256 MB floor
-
-  explicit ArenaJITLinkMemoryManager(std::size_t page_size, std::size_t slab_capacity);
-  ~ArenaJITLinkMemoryManager() override = default;
-
-  ArenaJITLinkMemoryManager(const ArenaJITLinkMemoryManager&) = delete;
-  ArenaJITLinkMemoryManager& operator=(const ArenaJITLinkMemoryManager&) = delete;
-  ArenaJITLinkMemoryManager(ArenaJITLinkMemoryManager&&) = delete;
-  ArenaJITLinkMemoryManager& operator=(ArenaJITLinkMemoryManager&&) = delete;
+  // Default per-slab capacity.  64 MB is above the p99 size of typical
+  // ML JIT graphs (single-kernel bindings, fused kernels), below the
+  // PC-relative relocation limit, and a multiple of the 2 MB THP
+  // granule.  Small enough that a pinned slab only wastes 64 MB of RSS.
+  static constexpr std::size_t kDefaultSlabSize = std::size_t{64} << 20;  // 64 MB
+
+  // Lower bound on initial-slab reservation.  If the first `mmap`
+  // fails and halving drops below this, the constructor aborts.
+  // 8 MB is enough for a minimal JITDylib setup under very tight
+  // RLIMIT_AS.
+  static constexpr std::size_t kMinSlabSize = std::size_t{8} << 20;  // 8 MB
+
+  explicit SlabPoolMemoryManager(std::size_t page_size, std::size_t slab_size);
+  ~SlabPoolMemoryManager() override = default;
+
+  SlabPoolMemoryManager(const SlabPoolMemoryManager&) = delete;
+  SlabPoolMemoryManager& operator=(const SlabPoolMemoryManager&) = delete;
+  SlabPoolMemoryManager(SlabPoolMemoryManager&&) = delete;
+  SlabPoolMemoryManager& operator=(SlabPoolMemoryManager&&) = delete;
 
   void allocate(const llvm::jitlink::JITLinkDylib* JD, llvm::jitlink::LinkGraph& G,
                 OnAllocatedFunction OnAllocated) override;
 
   void deallocate(std::vector<FinalizedAlloc> Allocs,
                   OnDeallocatedFunction OnDeallocated) override;
 
+  /*! \brief Number of slabs currently held (test introspection). */
+  std::size_t numSlabs() const {
+    std::lock_guard<std::mutex> lock(pool_mu_);
+    return slabs_.size();
+  }
+
  private:
-  std::unique_ptr<Slab> slab_;
+  /*! \brief Reserve a fresh slab at exactly \p capacity bytes.  Returns
+   *         nullptr on mmap failure (caller reports the error). */
+  std::unique_ptr<Slab> createSlab(std::size_t capacity);
+
+  std::size_t page_size_;
+  std::size_t slab_size_;
+
+  mutable std::mutex pool_mu_;
+  std::vector<std::unique_ptr<Slab>> slabs_;
 };
 
 }  // namespace orcjit