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

[Lang Hames via llvm-commits]

================
@@ -179,4 +179,184 @@ TEST(MapperJITLinkMemoryManagerTest, Coalescing) {
   EXPECT_THAT_ERROR(std::move(Err4), Succeeded());
 }
 
+// A colocating (slab) allocator places all objects inside a single reservation,
+// so any two of them are close enough for a 32-bit PC-relative cross-object
+// reference: an x86-64 REL32 is a signed 32-bit displacement, i.e. the target
+// must be within +/-2GB of the reference. The default InProcessMemoryManager
+// does NOT guarantee this -- it allocates each object independently, so two
+// objects can land more than 2GB apart and a direct cross-object call/branch
+// would fail to relocate ("out of range"). This test verifies that the
+// Mapper-based slab allocator keeps separately-allocated objects within REL32
+// range.
+TEST(MapperJITLinkMemoryManagerTest, Colocation) {
+  auto Mapper = cantFail(InProcessMemoryMapper::Create());
+  auto MemMgr = std::make_unique<MapperJITLinkMemoryManager>(16 * 1024 * 1024,
+                                                             std::move(Mapper));
+  auto SSP = std::make_shared<SymbolStringPool>();
+
+  // Allocate several "objects", the way separate input files would be added.
+  constexpr unsigned NumObjects = 8;
+  SmallVector<JITLinkMemoryManager::FinalizedAlloc> Allocs;
+  uint64_t MinAddr = ~uint64_t(0), MaxAddr = 0;
+
+  for (unsigned I = 0; I != NumObjects; ++I) {
+    auto SSA = jitlink::SimpleSegmentAlloc::Create(
+        *MemMgr, SSP, Triple("x86_64-apple-darwin"), nullptr,
+        {{MemProt::Read, {4096, Align(16)}}});
+    EXPECT_THAT_EXPECTED(SSA, Succeeded());
+
+    uint64_t Addr =
+        ExecutorAddr(SSA->getSegInfo(MemProt::Read).Addr).getValue();
+    if (Addr < MinAddr)
+      MinAddr = Addr;
+    if (Addr > MaxAddr)
+      MaxAddr = Addr;
+
+    auto FA = SSA->finalize();
+    EXPECT_THAT_EXPECTED(FA, Succeeded());
+    Allocs.push_back(std::move(*FA));
+  }
+
+  // All objects sit inside the one ~16MB reservation, so the span between the
+  // lowest and highest is far below the 2GB REL32 limit.
+  constexpr uint64_t REL32Limit = uint64_t(1) << 31; // 2GB
+  EXPECT_LT(MaxAddr - MinAddr, REL32Limit);
+
+  for (auto &FA : Allocs)
+    EXPECT_THAT_ERROR(MemMgr->deallocate(std::move(FA)), Succeeded());
+}
+
+// With per-JITDylib colocation enabled, each JITDylib draws from its own pool
+// of reservations: a second JITDylib's allocation does NOT reuse the first
+// JITDylib's leftover space, so it triggers a fresh reservation. With the
+// default (single shared pool) the second JITDylib reuses the first's leftover,
+// so only one reservation is made. Counting reservations is a deterministic way
+// to observe that objects are being grouped per-JITDylib.
+TEST(MapperJITLinkMemoryManagerTest, ColocationPerJITDylib) {
+  JITLinkDylib JDA("A"), JDB("B");
+  auto SSP = std::make_shared<SymbolStringPool>();
+
+  // Allocate a small object in the given JITDylib and finalize it.
+  auto AllocSmall = [&](MapperJITLinkMemoryManager &MemMgr,
+                        const JITLinkDylib *JD) {
+    auto SSA = jitlink::SimpleSegmentAlloc::Create(
+        MemMgr, SSP, Triple("x86_64-apple-darwin"), JD,
+        {{MemProt::Read, {4096, Align(16)}}});
+    EXPECT_THAT_EXPECTED(SSA, Succeeded());
+    auto FA = SSA->finalize();
+    EXPECT_THAT_EXPECTED(FA, Succeeded());
+    return cantFail(std::move(FA));
+  };
+
+  // Per-JITDylib colocation ON: JD B cannot reuse JD A's reservation, so a
+  // second reservation is made.
+  {
+    auto Mapper = std::make_unique<CounterMapper>(
+        cantFail(InProcessMemoryMapper::Create()));
+    auto *Counter = static_cast<CounterMapper *>(Mapper.get());
+    MapperJITLinkMemoryManager MemMgr(16 * 1024 * 1024, std::move(Mapper),
+                                      /*ColocatePerJITDylib=*/true);
+
+    auto FA_A = AllocSmall(MemMgr, &JDA);
+    EXPECT_EQ(Counter->ReserveCount, 1);
+    auto FA_B = AllocSmall(MemMgr, &JDB);
+    EXPECT_EQ(Counter->ReserveCount, 2); // separate reservation for JD B
+
+    EXPECT_THAT_ERROR(MemMgr.deallocate(std::move(FA_A)), Succeeded());
+    EXPECT_THAT_ERROR(MemMgr.deallocate(std::move(FA_B)), Succeeded());
+  }
+
+  // Default (single shared pool): JD B reuses JD A's leftover, so only one
+  // reservation is made.
+  {
+    auto Mapper = std::make_unique<CounterMapper>(
+        cantFail(InProcessMemoryMapper::Create()));
+    auto *Counter = static_cast<CounterMapper *>(Mapper.get());
+    MapperJITLinkMemoryManager MemMgr(16 * 1024 * 1024, std::move(Mapper));
+
+    auto FA_A = AllocSmall(MemMgr, &JDA);
+    EXPECT_EQ(Counter->ReserveCount, 1);
+    auto FA_B = AllocSmall(MemMgr, &JDB);
+    EXPECT_EQ(Counter->ReserveCount, 1); // reused JD A's reservation
----------------
lhames wrote:

Ditto here: assert that allocated addresses are within `reservationUnits()` of one another.

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