[llvm-jitlink] Add optional slab allocator for testing locality optimizations.

The llvm-jitlink utility now accepts a '-slab-allocate <size>' option. If given,
llvm-jitlink will use a slab-based memory manager rather than the default
InProcessMemoryManager. Using a slab allocator will allow reliable testing of
future locality based optimizations (e.g. PLT and GOT elimination) in JITLink.

The <size> argument is a number, optionally followed by a units specifier (Kb,
Mb, or Gb). If the units are not given then the number is assumed to be in Kb.

llvm-svn: 371244

tools/llvm-jitlink/llvm-jitlink.cpp

@@ -31,6 +31,7 @@
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Process.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/Timer.h"
@@ -99,6 +100,13 @@ static cl::opt<bool> ShowTimes("show-times",
                                cl::desc("Show times for llvm-jitlink phases"),
                                cl::init(false));
 
+static cl::opt<std::string> SlabAllocateSizeString(
+    "slab-allocate",
+    cl::desc("Allocate from a slab of the given size "
+             "(allowable suffixes: Kb, Mb, Gb. default = "
+             "Kb)"),
+    cl::init(""));
+
 static cl::opt<bool> ShowRelocatedSectionContents(
     "show-relocated-section-contents",
     cl::desc("show section contents after fixups have been applied"),
@@ -221,7 +229,175 @@ static void dumpSectionContents(raw_ostream &OS, AtomGraph &G) {
   }
 }
 
-Session::Session(Triple TT) : ObjLayer(ES, MemMgr), TT(std::move(TT)) {
+class JITLinkSlabAllocator final : public JITLinkMemoryManager {
+public:
+  static Expected<std::unique_ptr<JITLinkSlabAllocator>>
+  Create(uint64_t SlabSize) {
+    Error Err = Error::success();
+    std::unique_ptr<JITLinkSlabAllocator> Allocator(
+        new JITLinkSlabAllocator(SlabSize, Err));
+    if (Err)
+      return std::move(Err);
+    return std::move(Allocator);
+  }
+
+  Expected<std::unique_ptr<JITLinkMemoryManager::Allocation>>
+  allocate(const SegmentsRequestMap &Request) override {
+
+    using AllocationMap = DenseMap<unsigned, sys::MemoryBlock>;
+
+    // Local class for allocation.
+    class IPMMAlloc : public Allocation {
+    public:
+      IPMMAlloc(AllocationMap SegBlocks) : SegBlocks(std::move(SegBlocks)) {}
+      MutableArrayRef<char> getWorkingMemory(ProtectionFlags Seg) override {
+        assert(SegBlocks.count(Seg) && "No allocation for segment");
+        return {static_cast<char *>(SegBlocks[Seg].base()),
+                SegBlocks[Seg].allocatedSize()};
+      }
+      JITTargetAddress getTargetMemory(ProtectionFlags Seg) override {
+        assert(SegBlocks.count(Seg) && "No allocation for segment");
+        return reinterpret_cast<JITTargetAddress>(SegBlocks[Seg].base());
+      }
+      void finalizeAsync(FinalizeContinuation OnFinalize) override {
+        OnFinalize(applyProtections());
+      }
+      Error deallocate() override {
+        for (auto &KV : SegBlocks)
+          if (auto EC = sys::Memory::releaseMappedMemory(KV.second))
+            return errorCodeToError(EC);
+        return Error::success();
+      }
+
+    private:
+      Error applyProtections() {
+        for (auto &KV : SegBlocks) {
+          auto &Prot = KV.first;
+          auto &Block = KV.second;
+          if (auto EC = sys::Memory::protectMappedMemory(Block, Prot))
+            return errorCodeToError(EC);
+          if (Prot & sys::Memory::MF_EXEC)
+            sys::Memory::InvalidateInstructionCache(Block.base(),
+                                                    Block.allocatedSize());
+        }
+        return Error::success();
+      }
+
+      AllocationMap SegBlocks;
+    };
+
+    AllocationMap Blocks;
+
+    for (auto &KV : Request) {
+      auto &Seg = KV.second;
+
+      if (Seg.getContentAlignment() > PageSize)
+        return make_error<StringError>("Cannot request higher than page "
+                                       "alignment",
+                                       inconvertibleErrorCode());
+
+      if (PageSize % Seg.getContentAlignment() != 0)
+        return make_error<StringError>("Page size is not a multiple of "
+                                       "alignment",
+                                       inconvertibleErrorCode());
+
+      uint64_t ZeroFillStart =
+          alignTo(Seg.getContentSize(), Seg.getZeroFillAlignment());
+      uint64_t SegmentSize = ZeroFillStart + Seg.getZeroFillSize();
+
+      // Round segment size up to page boundary.
+      SegmentSize = (SegmentSize + PageSize - 1) & ~(PageSize - 1);
+
+      // Take segment bytes from the front of the slab.
+      void *SlabBase = SlabRemaining.base();
+      uint64_t SlabRemainingSize = SlabRemaining.allocatedSize();
+
+      if (SegmentSize > SlabRemainingSize)
+        return make_error<StringError>("Slab allocator out of memory",
+                                       inconvertibleErrorCode());
+
+      sys::MemoryBlock SegMem(SlabBase, SegmentSize);
+      SlabRemaining =
+          sys::MemoryBlock(reinterpret_cast<char *>(SlabBase) + SegmentSize,
+                           SlabRemainingSize - SegmentSize);
+
+      // Zero out the zero-fill memory.
+      memset(static_cast<char *>(SegMem.base()) + ZeroFillStart, 0,
+             Seg.getZeroFillSize());
+
+      // Record the block for this segment.
+      Blocks[KV.first] = std::move(SegMem);
+    }
+    return std::unique_ptr<InProcessMemoryManager::Allocation>(
+        new IPMMAlloc(std::move(Blocks)));
+  }
+
+private:
+  JITLinkSlabAllocator(uint64_t SlabSize, Error &Err) {
+    ErrorAsOutParameter _(&Err);
+
+    PageSize = sys::Process::getPageSizeEstimate();
+
+    if (!isPowerOf2_64(PageSize)) {
+      Err = make_error<StringError>("Page size is not a power of 2",
+                                    inconvertibleErrorCode());
+      return;
+    }
+
+    // Round slab request up to page size.
+    SlabSize = (SlabSize + PageSize - 1) & ~(PageSize - 1);
+
+    const sys::Memory::ProtectionFlags ReadWrite =
+        static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                  sys::Memory::MF_WRITE);
+
+    std::error_code EC;
+    SlabRemaining =
+        sys::Memory::allocateMappedMemory(SlabSize, nullptr, ReadWrite, EC);
+
+    if (EC) {
+      Err = errorCodeToError(EC);
+      return;
+    }
+  }
+
+  sys::MemoryBlock SlabRemaining;
+  uint64_t PageSize = 0;
+};
+
+Expected<uint64_t> getSlabAllocSize(StringRef SizeString) {
+  SizeString = SizeString.trim();
+
+  uint64_t Units = 1024;
+
+  if (SizeString.endswith_lower("kb"))
+    SizeString = SizeString.drop_back(2).rtrim();
+  else if (SizeString.endswith_lower("mb")) {
+    Units = 1024 * 1024;
+    SizeString = SizeString.drop_back(2).rtrim();
+  } else if (SizeString.endswith_lower("gb")) {
+    Units = 1024 * 1024 * 1024;
+    SizeString = SizeString.drop_back(2).rtrim();
+  }
+
+  uint64_t SlabSize = 0;
+  if (SizeString.getAsInteger(10, SlabSize))
+    return make_error<StringError>("Invalid numeric format for slab size",
+                                   inconvertibleErrorCode());
+
+  return SlabSize * Units;
+}
+
+static std::unique_ptr<jitlink::JITLinkMemoryManager> createMemoryManager() {
+  if (!SlabAllocateSizeString.empty()) {
+    auto SlabSize = ExitOnErr(getSlabAllocSize(SlabAllocateSizeString));
+    return ExitOnErr(JITLinkSlabAllocator::Create(SlabSize));
+  }
+  return std::make_unique<jitlink::InProcessMemoryManager>();
+}
+
+Session::Session(Triple TT)
+    : MemMgr(createMemoryManager()), ObjLayer(ES, *MemMgr), TT(std::move(TT)) {
 
   /// Local ObjectLinkingLayer::Plugin class to forward modifyPassConfig to the
   /// Session.

tools/llvm-jitlink/llvm-jitlink.h

@@ -26,7 +26,7 @@ namespace llvm {
 
 struct Session {
   orc::ExecutionSession ES;
-  jitlink::InProcessMemoryManager MemMgr;
+  std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr;
   orc::ObjectLinkingLayer ObjLayer;
   std::vector<orc::JITDylib *> JDSearchOrder;
   Triple TT;