[Orc] Use IntervalMap to store free memory regions in MapperJITLinkMemoryManager

[argentite]

Summary

MapperJITLinkMemoryManager uses a free list to keep track of available
memory regions. Using an IntervalMap instead of vector allow automatic
coalescing of memory regions as they are freed.

[argentite]

Added a test

[lhames]

LGTM -- Thanks @argentite!

llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h

 //===--------------- MapperJITLinkMemoryManager.h -*- C++ -*---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Implements JITLinkMemoryManager using MemoryMapper
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_EXECUTIONENGINE_ORC_MAPPERJITLINKMEMORYMANAGER_H
 #define LLVM_EXECUTIONENGINE_ORC_MAPPERJITLINKMEMORYMANAGER_H
 
+#include "llvm/ADT/IntervalMap.h"
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/MemoryMapper.h"
 
 namespace llvm {
 namespace orc {
 
 class MapperJITLinkMemoryManager : public jitlink::JITLinkMemoryManager {
 public:

 
 private:
   class InFlightAlloc;
 
   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
-  std::vector<ExecutorAddrRange> AvailableMemory;
+  using AvailableMemoryMap = IntervalMap<ExecutorAddr, bool>;
+  AvailableMemoryMap::Allocator AMAllocator;
+  IntervalMap<ExecutorAddr, bool> AvailableMemory;
+
   // Ranges that have been reserved in executor and already allocated
   DenseMap<ExecutorAddr, ExecutorAddrDiff> UsedMemory;
 
   std::unique_ptr<MemoryMapper> Mapper;
 };
 
 } // end namespace orc
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_MAPPERJITLINKMEMORYMANAGER_H

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

private:
   MapperJITLinkMemoryManager &Parent;
   LinkGraph &G;
   ExecutorAddr AllocAddr;
   std::vector<MemoryMapper::AllocInfo::SegInfo> Segs;
 };
 
 MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(
     size_t ReservationGranularity, std::unique_ptr<MemoryMapper> Mapper)
-    : ReservationUnits(ReservationGranularity), Mapper(std::move(Mapper)) {}
+    : ReservationUnits(ReservationGranularity), AvailableMemory(AMAllocator),
+      Mapper(std::move(Mapper)) {}
 
 void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
                                           OnAllocatedFunction OnAllocated) {
   BasicLayout BL(G);
 
   // find required address space
   auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(Mapper->getPageSize());
   if (!SegsSizes) {
     OnAllocated(SegsSizes.takeError());
     return;
   }
 
   auto TotalSize = SegsSizes->total();
 
-  Mutex.lock();
-
-  // find an already reserved range that is large enough
-  ExecutorAddrRange SelectedRange{};
-  std::vector<ExecutorAddrRange>::iterator SelectedRangeIt;
-  SelectedRangeIt =
-      llvm::find_if(AvailableMemory, [TotalSize](ExecutorAddrRange Range) {
-        return TotalSize < Range.size();
-      });
-  if (SelectedRangeIt != AvailableMemory.end()) {
-    SelectedRange = *SelectedRangeIt;
-    AvailableMemory.erase(SelectedRangeIt);
-  }
-
   auto CompleteAllocation = [this, &G, BL = std::move(BL),
                              OnAllocated = std::move(OnAllocated)](
                                 Expected<ExecutorAddrRange> Result) mutable {
     if (!Result) {
       Mutex.unlock();
       return OnAllocated(Result.takeError());
     }
 
for (auto &KV : BL.segments()) {
 
       SegInfos.push_back(SI);
     }
 
     UsedMemory.insert({Result->Start, NextSegAddr - Result->Start});
 
     if (NextSegAddr < Result->End) {
       // Save the remaining memory for reuse in next allocation(s)
-      auto RemainingRange = ExecutorAddrRange(NextSegAddr, Result->End);
-      AvailableMemory.push_back(RemainingRange);
+      AvailableMemory.insert(NextSegAddr, Result->End - 1, true);
     }
     Mutex.unlock();
 
     if (auto Err = BL.apply()) {
       OnAllocated(std::move(Err));
       return;
     }
 
     OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
                                                 std::move(SegInfos)));
   };
 
+  Mutex.lock();
+
+  // find an already reserved range that is large enough
+  ExecutorAddrRange SelectedRange{};
+
+  for (AvailableMemoryMap::iterator It = AvailableMemory.begin();
+       It != AvailableMemory.end(); It++) {
+    if (It.stop() - It.start() + 1 >= TotalSize) {
+      SelectedRange = ExecutorAddrRange(It.start(), It.stop() + 1);
+      It.erase();
+      break;
+    }
+  }
+
   if (SelectedRange.empty()) { // no already reserved range was found
     auto TotalAllocation = alignTo(TotalSize, ReservationUnits);
     Mapper->reserve(TotalAllocation, std::move(CompleteAllocation));
   } else {
     CompleteAllocation(SelectedRange);
   }
 }
 
Mapper->deinitialize(Bases, [this, Allocs = std::move(Allocs),
     {
       std::lock_guard<std::mutex> Lock(Mutex);
 
       for (auto &FA : Allocs) {
         ExecutorAddr Addr = FA.getAddress();
         ExecutorAddrDiff Size = UsedMemory[Addr];
 
         UsedMemory.erase(Addr);
-        AvailableMemory.push_back({Addr, Addr + Size});
+        AvailableMemory.insert(Addr, Addr + Size - 1, true);
 
         FA.release();
       }
     }
 
     OnDeallocated(Error::success());
   });
 }
 
 } // end namespace orc
 } // end namespace llvm

llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp

-//===-------------- MapperJITLinkMemoryManagerTest.cpp ----------------===//
+//===---------------- MapperJITLinkMemoryManagerTest.cpp ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "OrcTestCommon.h"
TEST(MapperJITLinkMemoryManagerTest, InProcess) {
   EXPECT_EQ(Counter->DeinitCount, 1);
 
   auto Err3 = MemMgr->deallocate(std::move(*FA2));
   EXPECT_THAT_ERROR(std::move(Err3), Succeeded());
 
   EXPECT_EQ(Counter->DeinitCount, 2);
 }
 
+TEST(MapperJITLinkMemoryManagerTest, Coalescing) {
+  auto Mapper = cantFail(InProcessMemoryMapper::Create());
+  auto MemMgr = std::make_unique<MapperJITLinkMemoryManager>(16 * 1024 * 1024,
+                                                             std::move(Mapper));
+
+  auto SSA1 = jitlink::SimpleSegmentAlloc::Create(
+      *MemMgr, nullptr, {{jitlink::MemProt::Read, {1024, Align(1)}}});
+  EXPECT_THAT_EXPECTED(SSA1, Succeeded());
+  auto SegInfo1 = SSA1->getSegInfo(jitlink::MemProt::Read);
+  ExecutorAddr TargetAddr1(SegInfo1.Addr);
+  auto FA1 = SSA1->finalize();
+  EXPECT_THAT_EXPECTED(FA1, Succeeded());
+
+  auto SSA2 = jitlink::SimpleSegmentAlloc::Create(
+      *MemMgr, nullptr, {{jitlink::MemProt::Read, {1024, Align(1)}}});
+  EXPECT_THAT_EXPECTED(SSA2, Succeeded());
+  auto FA2 = SSA2->finalize();
+  EXPECT_THAT_EXPECTED(FA2, Succeeded());
+
+  auto Err2 = MemMgr->deallocate(std::move(*FA1));
+  EXPECT_THAT_ERROR(std::move(Err2), Succeeded());
+
+  auto Err3 = MemMgr->deallocate(std::move(*FA2));
+  EXPECT_THAT_ERROR(std::move(Err3), Succeeded());
+
+  auto SSA3 = jitlink::SimpleSegmentAlloc::Create(
+      *MemMgr, nullptr, {{jitlink::MemProt::Read, {2048, Align(1)}}});
+  EXPECT_THAT_EXPECTED(SSA3, Succeeded());
+
+  auto SegInfo3 = SSA3->getSegInfo(jitlink::MemProt::Read);
+  ExecutorAddr TargetAddr3(SegInfo3.Addr);
+
+  auto FA3 = SSA3->finalize();
+  EXPECT_THAT_EXPECTED(FA3, Succeeded());
+
+  // previous two freed 1024 blocks should be fused to form a 2048 block
+  EXPECT_EQ(TargetAddr1, TargetAddr3);
+
+  auto Err4 = MemMgr->deallocate(std::move(*FA3));
+  EXPECT_THAT_ERROR(std::move(Err4), Succeeded());
+}
+
 } // namespace