[llvm] [ORC] Add opt-in per-JITDylib colocating slab allocator (PR #207970)

[via llvm-commits]

================
@@ -48,22 +68,68 @@ class LLVM_ABI MapperJITLinkMemoryManager
   // synchronous overload
   using JITLinkMemoryManager::deallocate;
 
+  /// Restrict each JITDylib to a single reservation ("slab"). Once a
+  /// JITDylib's slab is full, further allocations for it fail rather than
+  /// silently spilling into a second, possibly out-of-range, slab.
+  void allowSingleSlab();
+
+  /// Allow each JITDylib to grow into additional reservations on demand (the
+  /// default). Objects in different slabs of the same JITDylib are not
+  /// guaranteed to be within range of each other.
+  void allowMultipleSlabs();
+
+  /// Set a custom policy, invoked before reserving an additional slab for a
+  /// JITDylib (i.e. when its existing reservations can't satisfy a request).
+  /// Returning an Error fails the triggering allocation.
+  void setOnSecondSlab(unique_function<Error()> Policy);
+
 private:
   class InFlightAlloc;
 
+  using AvailableMemoryMap = IntervalMap<ExecutorAddr, bool>;
+
+  // Returns the pool of reserved-but-not-yet-allocated ranges for the given
+  // key, creating it on first use. The key is the JITDylib when colocating
+  // per-JITDylib, otherwise nullptr (a single shared pool). Must be called with
+  // Mutex held.
+  AvailableMemoryMap &getAvailableMemory(const jitlink::JITLinkDylib *Key);
+
   std::mutex Mutex;
 
   // We reserve multiples of this from the executor address space
   size_t ReservationUnits;
 
-  // Ranges that have been reserved in executor but not yet allocated
-  using AvailableMemoryMap = IntervalMap<ExecutorAddr, bool>;
+  // When true, each JITDylib gets its own pool of reservations (so a
+  // JITDylib's objects are colocated); when false a single nullptr-keyed pool
+  // is shared by all JITDylibs.
+  bool ColocatePerJITDylib;
+
+  // Policy consulted before reserving an additional slab for a pool (i.e. when
+  // a pool already owns a reservation but none of its free ranges fit the
+  // request). Returning an Error fails the allocation. Defaults to allowing
+  // additional slabs (preserving the historical behavior).
+  unique_function<Error()> OnSecondSlab = []() -> Error {
+    return Error::success();
+  };
+
+  // Pool keys that already own at least one reservation, used to detect when a
+  // further reservation would be a "second slab" for that pool.
+  DenseSet<const jitlink::JITLinkDylib *> ReservedKeys;
----------------
mkovacevic99 wrote:

For dangling pointers: The only hook ORC gives us for "something happened to a JITDylib" is ResourceManager, but it only fires when a resource tracker gets removed — not when the JITDylib itself actually goes away. A JITDylib can create and drop trackers while staying perfectly alive, so there's no reliable way to tell "just a tracker" apart from "the whole JITDylib is gone" from inside that callback. If we used it to clear a JITDylib's pool, we'd risk wiping out a pool that's still very much in use.

As far as I can tell, actually fixing this properly would mean teaching ORC something it doesn't currently know how to say: "this JITDylib is being destroyed for real.".

Do either of you see a cleaner way to handle this that I'm missing?

https://github.com/llvm/llvm-project/pull/207970