Re-enabling MCJIT object caching with memory leak fixed

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@180575 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/ExecutionEngine/ExecutionEngine.h

@@ -44,6 +44,7 @@ class JITMemoryManager;
 class MachineCodeInfo;
 class Module;
 class MutexGuard;
+class ObjectCache;
 class DataLayout;
 class Triple;
 class Type;
@@ -373,6 +374,12 @@ class ExecutionEngine {
   virtual void RegisterJITEventListener(JITEventListener *) {}
   virtual void UnregisterJITEventListener(JITEventListener *) {}
 
+  /// Sets the pre-compiled object cache.  The ownership of the ObjectCache is
+  /// not changed.  Supported by MCJIT but not JIT.
+  virtual void setObjectCache(ObjectCache *) {
+    llvm_unreachable("No support for an object cache");
+  }
+
   /// DisableLazyCompilation - When lazy compilation is off (the default), the
   /// JIT will eagerly compile every function reachable from the argument to
   /// getPointerToFunction.  If lazy compilation is turned on, the JIT will only

include/llvm/ExecutionEngine/ObjectCache.h

@@ -0,0 +1,54 @@
+//===-- ObjectCache.h - Class definition for the ObjectCache -----C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_EXECUTIONENGINE_OBJECTCACHE_H
+#define LLVM_LIB_EXECUTIONENGINE_OBJECTCACHE_H
+
+#include "llvm/Support/MemoryBuffer.h"
+
+namespace llvm {
+
+class Module;
+
+/// This is the base ObjectCache type which can be provided to an
+/// ExecutionEngine for the purpose of avoiding compilation for Modules that
+/// have already been compiled and an object file is available.
+class ObjectCache {
+public:
+  ObjectCache() { }
+
+  virtual ~ObjectCache() { }
+
+  /// notifyObjectCompiled - Provides a pointer to compiled code for Module M.
+  virtual void notifyObjectCompiled(const Module *M, const MemoryBuffer *Obj) = 0;
+
+  /// getObjectCopy - Returns a pointer to a newly allocated MemoryBuffer that
+  /// contains the object which corresponds with Module M, or 0 if an object is
+  /// not available. The caller owns the MemoryBuffer returned by this function.
+  MemoryBuffer* getObjectCopy(const Module* M) {
+    const MemoryBuffer* Obj = getObject(M);
+    if (Obj)
+      return MemoryBuffer::getMemBufferCopy(Obj->getBuffer());
+    else
+      return 0;
+  }
+
+protected:
+  /// getObject - Returns a pointer to a MemoryBuffer that contains an object
+  /// that corresponds with Module M, or 0 if an object is not available.
+  /// The pointer returned by this function is not suitable for loading because
+  /// the memory is read-only and owned by the ObjectCache. To retrieve an
+  /// owning pointer to a MemoryBuffer (which is suitable for calling
+  /// RuntimeDyld::loadObject() with) use getObjectCopy() instead.
+  virtual const MemoryBuffer* getObject(const Module* M) = 0;
+};
+
+}
+
+#endif

lib/ExecutionEngine/MCJIT/MCJIT.cpp

@@ -52,7 +52,7 @@ ExecutionEngine *MCJIT::createJIT(Module *M,
 MCJIT::MCJIT(Module *m, TargetMachine *tm, RTDyldMemoryManager *MM,
              bool AllocateGVsWithCode)
   : ExecutionEngine(m), TM(tm), Ctx(0), MemMgr(MM), Dyld(MM),
-    isCompiled(false), M(m)  {
+    IsLoaded(false), M(m), ObjCache(0)  {
 
   setDataLayout(TM->getDataLayout());
 }
@@ -64,7 +64,11 @@ MCJIT::~MCJIT() {
   delete TM;
 }
 
-void MCJIT::emitObject(Module *m) {
+void MCJIT::setObjectCache(ObjectCache* NewCache) {
+  ObjCache = NewCache;
+}
+
+ObjectBufferStream* MCJIT::emitObject(Module *m) {
   /// Currently, MCJIT only supports a single module and the module passed to
   /// this function call is expected to be the contained module.  The module
   /// is passed as a parameter here to prepare for multiple module support in
@@ -77,30 +81,66 @@ void MCJIT::emitObject(Module *m) {
   // FIXME: Track compilation state on a per-module basis when multiple modules
   //        are supported.
   // Re-compilation is not supported
-  if (isCompiled)
-    return;
+  assert(!IsLoaded);
 
   PassManager PM;
 
   PM.add(new DataLayout(*TM->getDataLayout()));
 
   // The RuntimeDyld will take ownership of this shortly
-  OwningPtr<ObjectBufferStream> Buffer(new ObjectBufferStream());
+  OwningPtr<ObjectBufferStream> CompiledObject(new ObjectBufferStream());
 
   // Turn the machine code intermediate representation into bytes in memory
   // that may be executed.
-  if (TM->addPassesToEmitMC(PM, Ctx, Buffer->getOStream(), false)) {
+  if (TM->addPassesToEmitMC(PM, Ctx, CompiledObject->getOStream(), false)) {
     report_fatal_error("Target does not support MC emission!");
   }
 
   // Initialize passes.
   PM.run(*m);
   // Flush the output buffer to get the generated code into memory
-  Buffer->flush();
+  CompiledObject->flush();
+
+  // If we have an object cache, tell it about the new object.
+  // Note that we're using the compiled image, not the loaded image (as below).
+  if (ObjCache) {
+    // MemoryBuffer is a thin wrapper around the actual memory, so it's OK
+    // to create a temporary object here and delete it after the call.
+    OwningPtr<MemoryBuffer> MB(CompiledObject->getMemBuffer());
+    ObjCache->notifyObjectCompiled(m, MB.get());
+  }
+
+  return CompiledObject.take();
+}
+
+void MCJIT::loadObject(Module *M) {
+
+  // Get a thread lock to make sure we aren't trying to load multiple times
+  MutexGuard locked(lock);
+
+  // FIXME: Track compilation state on a per-module basis when multiple modules
+  //        are supported.
+  // Re-compilation is not supported
+  if (IsLoaded)
+    return;
+
+  OwningPtr<ObjectBuffer> ObjectToLoad;
+  // Try to load the pre-compiled object from cache if possible
+  if (0 != ObjCache) {
+    OwningPtr<MemoryBuffer> PreCompiledObject(ObjCache->getObjectCopy(M));
+    if (0 != PreCompiledObject.get())
+      ObjectToLoad.reset(new ObjectBuffer(PreCompiledObject.take()));
+  }
+
+  // If the cache did not contain a suitable object, compile the object
+  if (!ObjectToLoad) {
+    ObjectToLoad.reset(emitObject(M));
+    assert(ObjectToLoad.get() && "Compilation did not produce an object.");
+  }
 
   // Load the object into the dynamic linker.
   // handing off ownership of the buffer
-  LoadedObject.reset(Dyld.loadObject(Buffer.take()));
+  LoadedObject.reset(Dyld.loadObject(ObjectToLoad.take()));
   if (!LoadedObject)
     report_fatal_error(Dyld.getErrorString());
 
@@ -113,7 +153,7 @@ void MCJIT::emitObject(Module *m) {
   NotifyObjectEmitted(*LoadedObject);
 
   // FIXME: Add support for per-module compilation state
-  isCompiled = true;
+  IsLoaded = true;
 }
 
 // FIXME: Add a parameter to identify which object is being finalized when
@@ -122,10 +162,10 @@ void MCJIT::emitObject(Module *m) {
 // protection in the interface.
 void MCJIT::finalizeObject() {
   // If the module hasn't been compiled, just do that.
-  if (!isCompiled) {
-    // If the call to Dyld.resolveRelocations() is removed from emitObject()
+  if (!IsLoaded) {
+    // If the call to Dyld.resolveRelocations() is removed from loadObject()
     // we'll need to do that here.
-    emitObject(M);
+    loadObject(M);
 
     // Set page permissions.
     MemMgr->applyPermissions();
@@ -151,8 +191,8 @@ void *MCJIT::getPointerToFunction(Function *F) {
   // dies.
 
   // FIXME: Add support for per-module compilation state
-  if (!isCompiled)
-    emitObject(M);
+  if (!IsLoaded)
+    loadObject(M);
 
   if (F->isDeclaration() || F->hasAvailableExternallyLinkage()) {
     bool AbortOnFailure = !F->hasExternalWeakLinkage();
@@ -284,8 +324,8 @@ GenericValue MCJIT::runFunction(Function *F,
 void *MCJIT::getPointerToNamedFunction(const std::string &Name,
                                        bool AbortOnFailure) {
   // FIXME: Add support for per-module compilation state
-  if (!isCompiled)
-    emitObject(M);
+  if (!IsLoaded)
+    loadObject(M);
 
   if (!isSymbolSearchingDisabled() && MemMgr) {
     void *ptr = MemMgr->getPointerToNamedFunction(Name, false);

lib/ExecutionEngine/MCJIT/MCJIT.h

@@ -12,6 +12,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
+#include "llvm/ExecutionEngine/ObjectCache.h"
 #include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include "llvm/PassManager.h"
 
@@ -34,16 +35,23 @@ class MCJIT : public ExecutionEngine {
   SmallVector<JITEventListener*, 2> EventListeners;
 
   // FIXME: Add support for multiple modules
-  bool isCompiled;
+  bool IsLoaded;
   Module *M;
   OwningPtr<ObjectImage> LoadedObject;
 
+  // An optional ObjectCache to be notified of compiled objects and used to
+  // perform lookup of pre-compiled code to avoid re-compilation.
+  ObjectCache *ObjCache;
+
 public:
   ~MCJIT();
 
   /// @name ExecutionEngine interface implementation
   /// @{
 
+  /// Sets the object manager that MCJIT should use to avoid compilation.
+  virtual void setObjectCache(ObjectCache *manager);
+
   virtual void finalizeObject();
 
   virtual void *getPointerToBasicBlock(BasicBlock *BB);
@@ -102,7 +110,9 @@ class MCJIT : public ExecutionEngine {
   /// this function call is expected to be the contained module.  The module
   /// is passed as a parameter here to prepare for multiple module support in 
   /// the future.
-  void emitObject(Module *M);
+  ObjectBufferStream* emitObject(Module *M);
+
+  void loadObject(Module *M);
 
   void NotifyObjectEmitted(const ObjectImage& Obj);
   void NotifyFreeingObject(const ObjectImage& Obj);

unittests/ExecutionEngine/MCJIT/CMakeLists.txt

@@ -10,6 +10,7 @@ set(LLVM_LINK_COMPONENTS
 set(MCJITTestsSources
   MCJITTest.cpp
   MCJITMemoryManagerTest.cpp
+  MCJITObjectCacheTest.cpp
   )
 
 if(MSVC)