[ExecutionEngine][MCJIT] Add a freeGeneratedCode hook.

[sanjoy]

Summary

Add ExecutionEngine::freeGeneratedCode and implement for MCJIT.
freeGeneratedCode lets an ExecutionEngine client drop all generated code
and related state. This hook can be used to:

1. Free code no longer in use. This is useful for MCJIT clients that relocate code generated by MCJIT into space allocated by the rest of the runtime, and don't need the code buffers managed by MCJIT once this is done.
2. Reset MCJIT from a state after generateCodeForModule has been called but relocations haven't been applied to a good state ready for a new compilation. This is useful for error recovery.

[lhames]

Hi Sanjoy,

Sorry for the delay on this review.

How does this differ from destroying the execution engine altogether and creating a new one?

• Lang.

[sanjoy]

How does this differ from destroying the execution engine altogether
and creating a new one?

In terms of what it does, I do not think that calling
freeGeneratedCode is very different from destroying and re-creating
the execution engine. But creating a new execution engine, compiling
a module, and immediately destroying it for every compile[1] did not
feel like the right semantic model (it does strictly more than I need
MCJIT to do, and can paper over other lifetime issues).

If you think this create-and-destroy option is the more reasonable
than adding a freeGeneratedCode hook then I can try to get that to
work with our VM and see how that goes.

• Sanjoy

[1]: We copy out generated code into our custom code arena, so keeping
a second copy of the generated code in MCJIT's memory manager is
unnecessary once the code has been copied out.

[sanjoy]

We're not pursuing this anymore.

include/llvm/ExecutionEngine/ExecutionEngine.h

public:
   ///
   /// It is the user-level function for completing the process of making the
   /// object usable for execution.  It should be called after sections within an
   /// object have been relocated using mapSectionAddress.  When this method is
   /// called the MCJIT execution engine will reapply relocations for a loaded
   /// object.  This method has no effect for the interpeter.
   virtual void finalizeObject() {}
 
+  /// freeGeneratedCode - clear out all generated code and related state.
+  ///
+  /// Calling freeGeneratedCode resets the ExecutionEngine to a state where it
+  /// owns the modules it did before the call, but does not have any generated
+  /// code.
+  virtual void freeGeneratedCode() {
+    llvm_unreachable("freeGeneratedCode unimplemented!");
+  }
+
   /// runStaticConstructorsDestructors - This method is used to execute all of
   /// the static constructors or destructors for a program.
   ///
   /// \param isDtors - Run the destructors instead of constructors.
   virtual void runStaticConstructorsDestructors(bool isDtors);
 
   /// This method is used to execute all of the static constructors or
   /// destructors for a particular module.

include/llvm/ExecutionEngine/RuntimeDyld.h

public:
     /// Allocate a memory block of (at least) the given size suitable for data.
     /// The SectionID is a unique identifier assigned by the JIT engine, and
     /// optionally recorded by the memory manager to access a loaded section.
     virtual uint8_t *allocateDataSection(uintptr_t Size, unsigned Alignment,
                                          unsigned SectionID,
                                          StringRef SectionName,
                                          bool IsReadOnly) = 0;
 
+    /// Free all of the memory allocated by allocateDataSection and
+    /// allocateCodeSection that this memory manager currently owns.
+    virtual void freeAllocatedSections() {
+      llvm_unreachable("freeAllocatedSections unimplemented!");
+    }
+
     /// Inform the memory manager about the total amount of memory required to
     /// allocate all sections to be loaded:
     /// \p CodeSize - the total size of all code sections
     /// \p DataSizeRO - the total size of all read-only data sections
     /// \p DataSizeRW - the total size of all read-write data sections
     ///
     /// Note that by default the callback is disabled. To enable it
     /// redefine the method needsToReserveAllocationSpace to return true.
public:
 
   /// Get the target address and flags for the named symbol.
   /// This address is the one used for relocation.
   SymbolInfo getSymbol(StringRef Name) const;
 
   /// Resolve the relocations for all symbols we currently know about.
   void resolveRelocations();
 
+  /// Clear the set of relocations yet to be applied to generated code.
+  void clearPendingRelocations();
+
+  /// Clear out information held about section addresses and global symbols.
+  void clearSectionReferences();
+
   /// Map a section to its target address space value.
   /// Map the address of a JIT section as returned from the memory manager
   /// to the address in the target process as the running code will see it.
   /// This is the address which will be used for relocation resolution.
   void mapSectionAddress(const void *LocalAddress, uint64_t TargetAddress);
 
   /// Register any EH frame sections that have been loaded but not previously
   /// registered with the memory manager.  Note, RuntimeDyld is responsible

include/llvm/ExecutionEngine/SectionMemoryManager.h

public:
   /// executable code.
   ///
   /// The value of \p Alignment must be a power of two.  If \p Alignment is zero
   /// a default alignment of 16 will be used.
   uint8_t *allocateDataSection(uintptr_t Size, unsigned Alignment,
                                unsigned SectionID, StringRef SectionName,
                                bool isReadOnly) override;
 
+  /// \brief Free all of the memory allocated by \p allocateDataSection and \p
+  /// allocateCodeSection that this memory manager currently owns.
+  void freeAllocatedSections() override;
+
   /// \brief Update section-specific memory permissions and other attributes.
   ///
   /// This method is called when object loading is complete and section page
   /// permissions can be applied.  It is up to the memory manager implementation
   /// to decide whether or not to act on this method.  The memory manager will
   /// typically allocate all sections as read-write and then apply specific
   /// permissions when this method is called.  Code sections cannot be executed
   /// until this function has been called.  In addition, any cache coherency

lib/ExecutionEngine/MCJIT/MCJIT.h

class MCJIT : public ExecutionEngine {
   Function *FindFunctionNamedInModulePtrSet(const char *FnName,
                                             ModulePtrSet::iterator I,
                                             ModulePtrSet::iterator E);
 
   void runStaticConstructorsDestructorsInModulePtrSet(bool isDtors,
                                                       ModulePtrSet::iterator I,
                                                       ModulePtrSet::iterator E);
 
+  void freeGeneratedCodeInternal(bool FreeAllocatedSections);
+
 public:
-  ~MCJIT() override;
+  ~MCJIT() override { freeGeneratedCodeInternal(false); }
 
   /// @name ExecutionEngine interface implementation
   /// @{
   void addModule(std::unique_ptr<Module> M) override;
   void addObjectFile(std::unique_ptr<object::ObjectFile> O) override;
   void addObjectFile(object::OwningBinary<object::ObjectFile> O) override;
   void addArchive(object::OwningBinary<object::Archive> O) override;
   bool removeModule(Module *M) override;
 
+  void freeGeneratedCode() override { freeGeneratedCodeInternal(true); }
+
   /// FindFunctionNamed - Search all of the active modules to find the one that
   /// defines FnName.  This is very slow operation and shouldn't be used for
   /// general code.
   Function *FindFunctionNamed(const char *FnName) override;
 
   /// Sets the object manager that MCJIT should use to avoid compilation.
   void setObjectCache(ObjectCache *manager) override;
 

lib/ExecutionEngine/MCJIT/MCJIT.cpp

MCJIT::MCJIT(std::unique_ptr<Module> M, std::unique_ptr<TargetMachine> tm,
   std::unique_ptr<Module> First = std::move(Modules[0]);
   Modules.clear();
 
   OwnedModules.addModule(std::move(First));
   setDataLayout(TM->getDataLayout());
   RegisterJITEventListener(JITEventListener::createGDBRegistrationListener());
 }
 
-MCJIT::~MCJIT() {
-  MutexGuard locked(lock);
-
-  Dyld.deregisterEHFrames();
-
-  for (auto &Obj : LoadedObjects)
-    if (Obj)
-      NotifyFreeingObject(*Obj);
-
-  Archives.clear();
-}
-
 void MCJIT::addModule(std::unique_ptr<Module> M) {
   MutexGuard locked(lock);
   OwnedModules.addModule(std::move(M));
 }
 
 bool MCJIT::removeModule(Module *M) {
   MutexGuard locked(lock);
   return OwnedModules.removeModule(M);
LinkingSymbolResolver::findSymbol(const std::string &Name) {
   if (!Result && Name[0] == '_')
     Result = ParentEngine.findSymbol(Name.substr(1), false);
   if (Result)
     return Result;
   if (ParentEngine.isSymbolSearchingDisabled())
     return nullptr;
   return ClientResolver->findSymbol(Name);
 }
+
+void MCJIT::freeGeneratedCodeInternal(bool FreeAllocatedSections) {
+  MutexGuard locked(lock);
+
+  Dyld.deregisterEHFrames();
+  Dyld.clearPendingRelocations();
+  Dyld.clearSectionReferences();
+
+  for (auto &Obj : LoadedObjects)
+    if (Obj)
+      NotifyFreeingObject(*Obj);
+
+  // Not all memory managers support freeAllocatedSections, it is the client's
+  // responsibility to ensure that MemMgr->freeAllocatedSections() is
+  // implemented if it wants to call freeGeneratedCode().
+  if (FreeAllocatedSections)
+    MemMgr->freeAllocatedSections();
+
+  LoadedObjects.clear();
+  Buffers.clear();
+  Archives.clear();
+}

lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp

 // Empty out-of-line virtual destructor as the key function.
 RuntimeDyldImpl::~RuntimeDyldImpl() {}
 
 // Pin LoadedObjectInfo's vtables to this file.
 void RuntimeDyld::LoadedObjectInfo::anchor() {}
 
 namespace llvm {
 
+void RuntimeDyldImpl::clearPendingEHFrameRegistrations() {}
+
 void RuntimeDyldImpl::registerEHFrames() {}
 
 void RuntimeDyldImpl::deregisterEHFrames() {}
 
 #ifndef NDEBUG
 static void dumpSectionMemory(const SectionEntry &S, StringRef State) {
   dbgs() << "----- Contents of section " << S.Name << " " << State << " -----";
 
DEBUG(dbgs() << "Resolving relocations Section #" << i << "\t"
                  << format("%p", (uintptr_t)Addr) << "\n");
     DEBUG(dumpSectionMemory(Sections[i], "before relocations"));
     resolveRelocationList(Relocations[i], Addr);
     DEBUG(dumpSectionMemory(Sections[i], "after relocations"));
     Relocations.erase(i);
   }
 }
 
+void RuntimeDyldImpl::clearSectionReferences() {
+  MutexGuard locked(lock);
+  Sections.clear();
+  GlobalSymbolTable.clear();
+}
+
+void RuntimeDyldImpl::clearPendingRelocations() {
+  MutexGuard locked(lock);
+  Relocations.clear();
+  ExternalSymbolRelocations.clear();
+}
+
 void RuntimeDyldImpl::mapSectionAddress(const void *LocalAddress,
                                         uint64_t TargetAddress) {
   MutexGuard locked(lock);
   for (unsigned i = 0, e = Sections.size(); i != e; ++i) {
     if (Sections[i].Address == LocalAddress) {
       reassignSectionAddress(i, TargetAddress);
       return;
     }

 RuntimeDyld::SymbolInfo RuntimeDyld::getSymbol(StringRef Name) const {
   if (!Dyld)
     return nullptr;
   return Dyld->getSymbol(Name);
 }
 
 void RuntimeDyld::resolveRelocations() { Dyld->resolveRelocations(); }
 
+void RuntimeDyld::clearSectionReferences() {
+  if (Dyld) {
+    Dyld->clearPendingEHFrameRegistrations();
+    Dyld->clearSectionReferences();
+  }
+}
+
+void RuntimeDyld::clearPendingRelocations() {
+  if (Dyld)
+    Dyld->clearPendingRelocations();
+}
+
 void RuntimeDyld::reassignSectionAddress(unsigned SectionID, uint64_t Addr) {
   Dyld->reassignSectionAddress(SectionID, Addr);
 }
 
 void RuntimeDyld::mapSectionAddress(const void *LocalAddress,
                                     uint64_t TargetAddress) {
   Dyld->mapSectionAddress(LocalAddress, TargetAddress);
 }

lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.h

class RuntimeDyldELF : public RuntimeDyldImpl {
   SmallVector<SID, 2> UnregisteredEHFrameSections;
   SmallVector<SID, 2> RegisteredEHFrameSections;
 
 public:
   RuntimeDyldELF(RuntimeDyld::MemoryManager &MemMgr,
                  RuntimeDyld::SymbolResolver &Resolver);
   ~RuntimeDyldELF() override;
 
+  void clearPendingEHFrameRegistrations() override {
+    UnregisteredEHFrameSections.clear();
+  }
+
   std::unique_ptr<RuntimeDyld::LoadedObjectInfo>
   loadObject(const object::ObjectFile &O) override;
 
   void resolveRelocation(const RelocationEntry &RE, uint64_t Value) override;
   relocation_iterator
   processRelocationRef(unsigned SectionID, relocation_iterator RelI,
                        const ObjectFile &Obj,
                        ObjSectionToIDMap &ObjSectionToID,

lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h

RuntimeDyld::SymbolInfo getSymbol(StringRef Name) const {
     const auto &SymEntry = pos->second;
     uint64_t TargetAddr =
       getSectionLoadAddress(SymEntry.getSectionID()) + SymEntry.getOffset();
     return RuntimeDyld::SymbolInfo(TargetAddr, SymEntry.getFlags());
   }
 
   void resolveRelocations();
 
+  void clearSectionReferences();
+  void clearPendingRelocations();
+  virtual void clearPendingEHFrameRegistrations();
+
   void reassignSectionAddress(unsigned SectionID, uint64_t Addr);
 
   void mapSectionAddress(const void *LocalAddress, uint64_t TargetAddress);
 
   // Is the linker in an error state?
   bool hasError() { return HasError; }
 
   // Mark the error condition as handled and continue.

lib/ExecutionEngine/RuntimeDyld/RuntimeDyldMachO.h

RelocationEntry getRelocationEntry(unsigned SectionID,
     uint64_t Offset;
     RI->getOffset(Offset);
     MachO::RelocationInfoType RelType =
       static_cast<MachO::RelocationInfoType>(Obj.getAnyRelocationType(RelInfo));
 
     return RelocationEntry(SectionID, Offset, RelType, 0, IsPCRel, Size);
   }
 
+  void clearPendingEHFrameRegistrations() override {
+    UnregisteredEHFrameSections.clear();
+  }
+
   /// Construct a RelocationValueRef representing the relocation target.
   /// For Symbols in known sections, this will return a RelocationValueRef
   /// representing a (SectionID, Offset) pair.
   /// For Symbols whose section is not known, this will return a
   /// (SymbolName, Offset) pair, where the Offset is taken from the instruction
   /// immediate (held in RE.Addend).
   /// In both cases the Addend field is *NOT* fixed up to be PC-relative. That
   /// should be done by the caller where appropriate by calling makePCRel on

lib/ExecutionEngine/SectionMemoryManager.cpp

 }
 
 void SectionMemoryManager::invalidateInstructionCache() {
   for (int i = 0, e = CodeMem.AllocatedMem.size(); i != e; ++i)
     sys::Memory::InvalidateInstructionCache(CodeMem.AllocatedMem[i].base(),
                                             CodeMem.AllocatedMem[i].size());
 }
 
-SectionMemoryManager::~SectionMemoryManager() {
+void SectionMemoryManager::freeAllocatedSections() {
   for (unsigned i = 0, e = CodeMem.AllocatedMem.size(); i != e; ++i)
     sys::Memory::releaseMappedMemory(CodeMem.AllocatedMem[i]);
   for (unsigned i = 0, e = RWDataMem.AllocatedMem.size(); i != e; ++i)
     sys::Memory::releaseMappedMemory(RWDataMem.AllocatedMem[i]);
   for (unsigned i = 0, e = RODataMem.AllocatedMem.size(); i != e; ++i)
     sys::Memory::releaseMappedMemory(RODataMem.AllocatedMem[i]);
+
+  for (auto *MGroup : {&CodeMem, &RWDataMem, &RODataMem}) {
+    MGroup->AllocatedMem.clear();
+    MGroup->FreeMem.clear();
+    MGroup->Near = sys::MemoryBlock();
+  }
 }
 
+SectionMemoryManager::~SectionMemoryManager() { freeAllocatedSections(); }
+
 } // namespace llvm
 

unittests/ExecutionEngine/MCJIT/CMakeLists.txt

set(LLVM_LINK_COMPONENTS
   )
 
 set(MCJITTestsSources
   MCJITTest.cpp
   MCJITCAPITest.cpp
   MCJITMemoryManagerTest.cpp
   MCJITMultipleModuleTest.cpp
   MCJITObjectCacheTest.cpp
+  MCJITObjectCleanupTest.cpp
   )
 
 if(MSVC)
   list(APPEND MCJITTestsSources MCJITTests.def)
 endif()
 
 add_llvm_unittest(MCJITTests
   ${MCJITTestsSources}
   )
 
 if(MINGW OR CYGWIN)
   set_property(TARGET MCJITTests PROPERTY LINK_FLAGS -Wl,--export-all-symbols)
 endif()

unittests/ExecutionEngine/MCJIT/MCJITObjectCleanupTest.cpp

+//===- MCJITObjectCleanupTest.cpp - Unit tests for MCJIT object cleanup ---===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "MCJITTestBase.h"
+#include "llvm/ExecutionEngine/MCJIT.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+struct TestMemoryManager : public SectionMemoryManager {
+  bool ResolutionEnabled;
+  int SectionsAllocated;
+
+  TestMemoryManager() : ResolutionEnabled(false), SectionsAllocated(0) {}
+
+  virtual uint64_t getSymbolAddress(const std::string &Name) override {
+    if (ResolutionEnabled && Name == "external_callee")
+      return 5000;
+    return SectionMemoryManager::getSymbolAddress(Name);
+  }
+
+  uint8_t *allocateDataSection(uintptr_t Size, unsigned Alignment,
+                               unsigned SectionID, StringRef SectionName,
+                               bool IsReadOnly) override {
+    SectionsAllocated++;
+    return SectionMemoryManager::allocateDataSection(Size, Alignment, SectionID,
+                                                     SectionName, IsReadOnly);
+  }
+
+  uint8_t *allocateCodeSection(uintptr_t Size, unsigned Alignment,
+                               unsigned SectionID,
+                               StringRef SectionName) override {
+    SectionsAllocated++;
+    return SectionMemoryManager::allocateCodeSection(Size, Alignment, SectionID,
+                                                     SectionName);
+  }
+
+  virtual void freeAllocatedSections() override {
+    SectionsAllocated = 0;
+    SectionMemoryManager::freeAllocatedSections();
+  }
+};
+
+class MCJITObjectCleanupTest : public testing::Test, public MCJITTestBase {
+protected:
+  Module *SourceWithExternalSyms;
+  Module *SourceWithoutExternalSyms;
+  TestMemoryManager *TMM;
+
+public:
+  MCJITObjectCleanupTest()
+      : MCJITTestBase(new TestMemoryManager), SourceWithExternalSyms(nullptr),
+        SourceWithoutExternalSyms(nullptr), TMM(nullptr) {}
+
+  void SetUp() override {
+    TMM = (TestMemoryManager *)MM.get();
+
+    M.reset(createEmptyModule("epoch"));
+    createJIT(std::move(M));
+
+    SourceWithExternalSyms = createEmptyModule("SourceWithExternalSyms");
+    Function *Callee = insertExternalReferenceToFunction<int(void)>(
+        SourceWithExternalSyms, "external_callee");
+    insertSimpleCallFunction<int(void)>(SourceWithExternalSyms, Callee);
+
+    SourceWithoutExternalSyms = createEmptyModule("SourceWithExternalSyms");
+    insertAddFunction(SourceWithoutExternalSyms);
+  }
+};
+
+TEST_F(MCJITObjectCleanupTest, Basic) {
+  SKIP_UNSUPPORTED_PLATFORM;
+
+  TMM->ResolutionEnabled = true;
+  TheJIT->addModule(std::unique_ptr<Module>(SourceWithExternalSyms));
+  TheJIT->generateCodeForModule(SourceWithExternalSyms);
+
+  EXPECT_NE(TMM->SectionsAllocated, 0)
+      << "Generating code should have allocated some sections!";
+
+  TheJIT->removeModule(SourceWithExternalSyms);
+  TheJIT->freeGeneratedCode();
+
+  EXPECT_EQ(TMM->SectionsAllocated, 0) << "All sections should have bene freed "
+                                          "after the call to "
+                                          "freeGeneratedCode!";
+
+  TMM->ResolutionEnabled = false;
+  TheJIT->addModule(std::unique_ptr<Module>(SourceWithoutExternalSyms));
+  TheJIT->generateCodeForModule(SourceWithoutExternalSyms);
+  TheJIT->finalizeObject();
+}
+}

unittests/ExecutionEngine/MCJIT/MCJITTestBase.h

void createThreeModuleCase(std::unique_ptr<Module> &A, Function *&FA,
     Function *FAExtern_in_C = insertExternalReferenceToFunction(C.get(), FA);
     FC = insertSimpleCallFunction<int32_t(int32_t, int32_t)>(C.get(), FAExtern_in_C);
   }
 };
 
 
 class MCJITTestBase : public MCJITTestAPICommon, public TrivialModuleBuilder {
 protected:
-
-  MCJITTestBase()
-    : TrivialModuleBuilder(HostTriple)
-    , OptLevel(CodeGenOpt::None)
-    , RelocModel(Reloc::Default)
-    , CodeModel(CodeModel::Default)
-    , MArch("")
-    , MM(new SectionMemoryManager)
-  {
+  explicit MCJITTestBase(RTDyldMemoryManager *MemManager = nullptr)
+      : TrivialModuleBuilder(HostTriple), OptLevel(CodeGenOpt::None),
+        RelocModel(Reloc::Default), CodeModel(CodeModel::Default), MArch(""),
+        MM(MemManager ? MemManager : new SectionMemoryManager) {
     // The architectures below are known to be compatible with MCJIT as they
     // are copied from test/ExecutionEngine/MCJIT/lit.local.cfg and should be
     // kept in sync.
     SupportedArchs.push_back(Triple::aarch64);
     SupportedArchs.push_back(Triple::arm);
     SupportedArchs.push_back(Triple::mips);
     SupportedArchs.push_back(Triple::mipsel);
     SupportedArchs.push_back(Triple::mips64);