[JIT] Towards a working small memory model

[davide]

Summary

Hi Lang. This patch implements the idea we previously discussed on IRC, introducing --preallocate so that we can get memory upfront and use that in small memory model tests. to It's currently unfinished because we don't keep track of FunctionMemory/DataMemory (it depends on my next review). I hope this was what you had in mind. Comments welcome.

[lhames]

Hi David,

Excellent! This is exactly the kind of thing I was looking for - thank you.

What do you mean 'don't keep track of function memory / data memory' though? We may run in to issues for platforms that don't allow us to allocate RWX sections, in which case we'd have to make that distinction, but if we can get RWX memory I think we can just use one pool.

• Lang.

tools/llvm-rtdyld/llvm-rtdyld.cpp

193–200 ↗

(On Diff #37045)

allocateFromSlab needs to respect the requested alignment. There's a RoundUpToAlignment function in MathExtras.h that you can use for this.

377–384 ↗

(On Diff #37045)

Can you add the preallocate call to linkAndVerify? We'll want it there too so we can write relocation verifier tests for the small code model.

[davide]

New patch addressing your comments. To answer your question by "not keeping track of memory" I meant that my previous patch didn't push_back allocated data/code section to {Function,Data}Memory vector. I think it was wrong, because e.g., invalidateInstructionCache wouldn't have been called on those regions. I hope my reasoning is correct.
Thanks!

[davide]

Very gentle ping

[lhames]

Ping appreciated. :)

This looks good to me.

llvm/trunk/tools/llvm-rtdyld/llvm-rtdyld.cpp

MCPU("mcpu",
      cl::init(""));
 
 static cl::list<std::string>
 CheckFiles("check",
            cl::desc("File containing RuntimeDyld verifier checks."),
            cl::ZeroOrMore);
 
 static cl::opt<uint64_t>
+PreallocMemory("preallocate",
+              cl::desc("Allocate memory upfront rather than on-demand"),
+              cl::init(0));
+
+static cl::opt<uint64_t>
 TargetAddrStart("target-addr-start",
                 cl::desc("For -verify only: start of phony target address "
                          "range."),
                 cl::init(4096), // Start at "page 1" - no allocating at "null".
                 cl::Hidden);
 
 static cl::opt<uint64_t>
 TargetAddrEnd("target-addr-end",
RuntimeDyld::SymbolInfo findSymbol(const std::string &Name) override {
 
     return RTDyldMemoryManager::findSymbol(Name);
   }
 
   void registerEHFrames(uint8_t *Addr, uint64_t LoadAddr,
                         size_t Size) override {}
   void deregisterEHFrames(uint8_t *Addr, uint64_t LoadAddr,
                           size_t Size) override {}
+
+  void preallocateSlab(uint64_t Size) {
+    std::string Err;
+    sys::MemoryBlock MB = sys::Memory::AllocateRWX(Size, nullptr, &Err);
+    if (!MB.base())
+      report_fatal_error("Can't allocate enough memory: " + Err);
+
+    PreallocSlab = MB;
+    UsePreallocation = true;
+    SlabSize = Size;
+  }
+
+  uint8_t *allocateFromSlab(uintptr_t Size, unsigned Alignment, bool isCode) {
+    Size = RoundUpToAlignment(Size, Alignment);
+    if (CurrentSlabOffset + Size > SlabSize)
+      report_fatal_error("Can't allocate enough memory. Tune --preallocate");
+
+    uintptr_t OldSlabOffset = CurrentSlabOffset;
+    sys::MemoryBlock MB((void *)OldSlabOffset, Size);
+    if (isCode)
+      FunctionMemory.push_back(MB);
+    else
+      DataMemory.push_back(MB);
+    CurrentSlabOffset += Size;
+    return (uint8_t*)OldSlabOffset;
+  }
+
 private:
   std::map<std::string, uint64_t> DummyExterns;
+  sys::MemoryBlock PreallocSlab;
+  bool UsePreallocation = false;
+  uintptr_t SlabSize = 0;
+  uintptr_t CurrentSlabOffset = 0;
 };
 
 uint8_t *TrivialMemoryManager::allocateCodeSection(uintptr_t Size,
                                                    unsigned Alignment,
                                                    unsigned SectionID,
                                                    StringRef SectionName) {
+  if (UsePreallocation)
+    return allocateFromSlab(Size, Alignment, true /* isCode */);
+
   std::string Err;
   sys::MemoryBlock MB = sys::Memory::AllocateRWX(Size, nullptr, &Err);
   if (!MB.base())
     report_fatal_error("MemoryManager allocation failed: " + Err);
   FunctionMemory.push_back(MB);
   return (uint8_t*)MB.base();
 }
 
 uint8_t *TrivialMemoryManager::allocateDataSection(uintptr_t Size,
                                                    unsigned Alignment,
                                                    unsigned SectionID,
                                                    StringRef SectionName,
                                                    bool IsReadOnly) {
+  if (UsePreallocation)
+    return allocateFromSlab(Size, Alignment, false /* isCode */);
+
   std::string Err;
   sys::MemoryBlock MB = sys::Memory::AllocateRWX(Size, nullptr, &Err);
   if (!MB.base())
     report_fatal_error("MemoryManager allocation failed: " + Err);
   DataMemory.push_back(MB);
   return (uint8_t*)MB.base();
 }
 
for (const auto &P : SymAddr) {
         }
       }
     }
   }
 
   return 0;
 }
 
+static void doPreallocation(TrivialMemoryManager &MemMgr) {
+  // Allocate a slab of memory upfront, if required. This is used if
+  // we want to test small code models.
+  if (static_cast<intptr_t>(PreallocMemory) < 0)
+    report_fatal_error("Pre-allocated bytes of memory must be a positive integer.");
+
+  // FIXME: Limit the amount of memory that can be preallocated?
+  if (PreallocMemory != 0)
+    MemMgr.preallocateSlab(PreallocMemory);
+}
+
 static int executeInput() {
   // Load any dylibs requested on the command line.
   loadDylibs();
 
   // Instantiate a dynamic linker.
   TrivialMemoryManager MemMgr;
+  doPreallocation(MemMgr);
   RuntimeDyld Dyld(MemMgr, MemMgr);
 
   // FIXME: Preserve buffers until resolveRelocations time to work around a bug
   //        in RuntimeDyldELF.
   // This fixme should be fixed ASAP. This is a very brittle workaround.
   std::vector<std::unique_ptr<MemoryBuffer>> InputBuffers;
 
   // If we don't have any input files, read from stdin.
static int linkAndVerify() {
   std::unique_ptr<MCInstPrinter> InstPrinter(
       TheTarget->createMCInstPrinter(Triple(TripleName), 0, *MAI, *MII, *MRI));
 
   // Load any dylibs requested on the command line.
   loadDylibs();
 
   // Instantiate a dynamic linker.
   TrivialMemoryManager MemMgr;
+  doPreallocation(MemMgr);
   RuntimeDyld Dyld(MemMgr, MemMgr);
   Dyld.setProcessAllSections(true);
   RuntimeDyldChecker Checker(Dyld, Disassembler.get(), InstPrinter.get(),
                              llvm::dbgs());
 
   // FIXME: Preserve buffers until resolveRelocations time to work around a bug
   //        in RuntimeDyldELF.
   // This fixme should be fixed ASAP. This is a very brittle workaround.