Automatically do the equivalent of freeMachineCodeForFunction(F) when F is

being destroyed. This allows users to run global optimizations like globaldce
even after some functions have been jitted.

This patch also removes the Function* parameter to
JITEventListener::NotifyFreeingMachineCode() since it can cause that to be
called when the Function is partially destroyed. This change will be even more
helpful later when I think we'll want to allow machine code to actually outlive
its Function.


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

include/llvm/ExecutionEngine/ExecutionEngine.h

@@ -263,9 +263,8 @@ public:
   /// getPointerToFunction - The different EE's represent function bodies in
   /// different ways.  They should each implement this to say what a function
   /// pointer should look like.  When F is destroyed, the ExecutionEngine will
-  /// remove its global mapping but will not yet free its machine code.  Call
-  /// freeMachineCodeForFunction(F) explicitly to do that.  Note that global
-  /// optimizations can destroy Functions without notifying the ExecutionEngine.
+  /// remove its global mapping and free any machine code.  Be sure no threads
+  /// are running inside F when that happens.
   ///
   virtual void *getPointerToFunction(Function *F) = 0;
 

include/llvm/ExecutionEngine/JITEventListener.h

@@ -63,8 +63,11 @@ public:
   /// NotifyFreeingMachineCode - This is called inside of
   /// freeMachineCodeForFunction(), after the global mapping is removed, but
   /// before the machine code is returned to the allocator.  OldPtr is the
-  /// address of the machine code.
-  virtual void NotifyFreeingMachineCode(const Function &F, void *OldPtr) {}
+  /// address of the machine code and will be the same as the Code parameter to
+  /// a previous NotifyFunctionEmitted call.  The Function passed to
+  /// NotifyFunctionEmitted may have been destroyed by the time of the matching
+  /// NotifyFreeingMachineCode call.
+  virtual void NotifyFreeingMachineCode(void *OldPtr) {}
 };
 
 // This returns NULL if support isn't available.

lib/ExecutionEngine/JIT/JIT.cpp

@@ -556,10 +556,10 @@ void JIT::NotifyFunctionEmitted(
   }
 }
 
-void JIT::NotifyFreeingMachineCode(const Function &F, void *OldPtr) {
+void JIT::NotifyFreeingMachineCode(void *OldPtr) {
   MutexGuard locked(lock);
   for (unsigned I = 0, S = EventListeners.size(); I < S; ++I) {
-    EventListeners[I]->NotifyFreeingMachineCode(F, OldPtr);
+    EventListeners[I]->NotifyFreeingMachineCode(OldPtr);
   }
 }
 

lib/ExecutionEngine/JIT/JIT.h

@@ -183,7 +183,7 @@ public:
   void NotifyFunctionEmitted(
       const Function &F, void *Code, size_t Size,
       const JITEvent_EmittedFunctionDetails &Details);
-  void NotifyFreeingMachineCode(const Function &F, void *OldPtr);
+  void NotifyFreeingMachineCode(void *OldPtr);
 
 private:
   static JITCodeEmitter *createEmitter(JIT &J, JITMemoryManager *JMM,

lib/ExecutionEngine/JIT/JITEmitter.cpp

@@ -582,17 +582,24 @@ namespace {
     JITEvent_EmittedFunctionDetails EmissionDetails;
 
     struct EmittedCode {
-      void *FunctionBody;
+      void *FunctionBody;  // Beginning of the function's allocation.
+      void *Code;  // The address the function's code actually starts at.
       void *ExceptionTable;
-      EmittedCode() : FunctionBody(0), ExceptionTable(0) {}
+      EmittedCode() : FunctionBody(0), Code(0), ExceptionTable(0) {}
     };
-    DenseMap<const Function *, EmittedCode> EmittedFunctions;
+    struct EmittedFunctionConfig : public ValueMapConfig<const Function*> {
+      typedef JITEmitter *ExtraData;
+      static void onDelete(JITEmitter *, const Function*);
+      static void onRAUW(JITEmitter *, const Function*, const Function*);
+    };
+    ValueMap<const Function *, EmittedCode,
+             EmittedFunctionConfig> EmittedFunctions;
 
     // CurFnStubUses - For a given Function, a vector of stubs that it
     // references.  This facilitates the JIT detecting that a stub is no
     // longer used, so that it may be deallocated.
-    DenseMap<const Function *, SmallVector<void*, 1> > CurFnStubUses;
-    
+    DenseMap<AssertingVH<const Function>, SmallVector<void*, 1> > CurFnStubUses;
+
     // StubFnRefs - For a given pointer to a stub, a set of Functions which
     // reference the stub.  When the count of a stub's references drops to zero,
     // the stub is unused.
@@ -606,7 +613,8 @@ namespace {
 
   public:
     JITEmitter(JIT &jit, JITMemoryManager *JMM, TargetMachine &TM)
-        : SizeEstimate(0), Resolver(jit), MMI(0), CurFn(0) {
+        : SizeEstimate(0), Resolver(jit), MMI(0), CurFn(0),
+          EmittedFunctions(this) {
       MemMgr = JMM ? JMM : JITMemoryManager::CreateDefaultMemManager();
       if (jit.getJITInfo().needsGOT()) {
         MemMgr->AllocateGOT();
@@ -1062,6 +1070,7 @@ void JITEmitter::startFunction(MachineFunction &F) {
   // About to start emitting the machine code for the function.
   emitAlignment(std::max(F.getFunction()->getAlignment(), 8U));
   TheJIT->updateGlobalMapping(F.getFunction(), CurBufferPtr);
+  EmittedFunctions[F.getFunction()].Code = CurBufferPtr;
 
   MBBLocations.clear();
 
@@ -1285,12 +1294,15 @@ void JITEmitter::retryWithMoreMemory(MachineFunction &F) {
 
 /// deallocateMemForFunction - Deallocate all memory for the specified
 /// function body.  Also drop any references the function has to stubs.
+/// May be called while the Function is being destroyed inside ~Value().
 void JITEmitter::deallocateMemForFunction(const Function *F) {
-  DenseMap<const Function *, EmittedCode>::iterator Emitted =
-    EmittedFunctions.find(F);
+  ValueMap<const Function *, EmittedCode, EmittedFunctionConfig>::iterator
+    Emitted = EmittedFunctions.find(F);
   if (Emitted != EmittedFunctions.end()) {
     MemMgr->deallocateFunctionBody(Emitted->second.FunctionBody);
     MemMgr->deallocateExceptionTable(Emitted->second.ExceptionTable);
+    TheJIT->NotifyFreeingMachineCode(Emitted->second.Code);
+
     EmittedFunctions.erase(Emitted);
   }
 
@@ -1519,6 +1531,17 @@ uintptr_t JITEmitter::getJumpTableEntryAddress(unsigned Index) const {
   return (uintptr_t)((char *)JumpTableBase + Offset);
 }
 
+void JITEmitter::EmittedFunctionConfig::onDelete(
+  JITEmitter *Emitter, const Function *F) {
+  Emitter->deallocateMemForFunction(F);
+}
+void JITEmitter::EmittedFunctionConfig::onRAUW(
+  JITEmitter *, const Function*, const Function*) {
+  llvm_unreachable("The JIT doesn't know how to handle a"
+                   " RAUW on a value it has emitted.");
+}
+
+
 //===----------------------------------------------------------------------===//
 //  Public interface to this file
 //===----------------------------------------------------------------------===//
@@ -1657,13 +1680,9 @@ void JIT::updateDlsymStubTable() {
 /// freeMachineCodeForFunction - release machine code memory for given Function.
 ///
 void JIT::freeMachineCodeForFunction(Function *F) {
-
   // Delete translation for this from the ExecutionEngine, so it will get
   // retranslated next time it is used.
-  void *OldPtr = updateGlobalMapping(F, 0);
-
-  if (OldPtr)
-    TheJIT->NotifyFreeingMachineCode(*F, OldPtr);
+  updateGlobalMapping(F, 0);
 
   // Free the actual memory for the function body and related stuff.
   assert(isa<JITEmitter>(JCE) && "Unexpected MCE?");

lib/ExecutionEngine/JIT/OProfileJITEventListener.cpp

@@ -147,13 +147,13 @@ void OProfileJITEventListener::NotifyFunctionEmitted(
   }
 }
 
-// Removes the to-be-deleted function from the symbol table.
-void OProfileJITEventListener::NotifyFreeingMachineCode(
-    const Function &F, void *FnStart) {
+// Removes the being-deleted function from the symbol table.
+void OProfileJITEventListener::NotifyFreeingMachineCode(void *FnStart) {
   assert(FnStart && "Invalid function pointer");
   if (op_unload_native_code(Agent, reinterpret_cast<uint64_t>(FnStart)) == -1) {
-    DEBUG(errs() << "Failed to tell OProfile about unload of native function "
-                 << F.getName() << " at " << FnStart << "\n");
+    DEBUG(errs()
+          << "Failed to tell OProfile about unload of native function at "
+          << FnStart << "\n");
   }
 }
 

unittests/ExecutionEngine/JIT/JITEventListenerTest.cpp

@@ -37,7 +37,6 @@ struct FunctionEmittedEvent {
 };
 struct FunctionFreedEvent {
   unsigned Index;
-  const Function *F;
   void *Code;
 };
 
@@ -56,8 +55,8 @@ struct RecordingJITEventListener : public JITEventListener {
     EmittedEvents.push_back(Event);
   }
 
-  virtual void NotifyFreeingMachineCode(const Function &F, void *OldPtr) {
-    FunctionFreedEvent Event = {NextIndex++, &F, OldPtr};
+  virtual void NotifyFreeingMachineCode(void *OldPtr) {
+    FunctionFreedEvent Event = {NextIndex++, OldPtr};
     FreedEvents.push_back(Event);
   }
 };
@@ -116,11 +115,9 @@ TEST_F(JITEventListenerTest, Simple) {
       << " contain some bytes.";
 
   EXPECT_EQ(2U, Listener.FreedEvents[0].Index);
-  EXPECT_EQ(F1, Listener.FreedEvents[0].F);
   EXPECT_EQ(F1_addr, Listener.FreedEvents[0].Code);
 
   EXPECT_EQ(3U, Listener.FreedEvents[1].Index);
-  EXPECT_EQ(F2, Listener.FreedEvents[1].F);
   EXPECT_EQ(F2_addr, Listener.FreedEvents[1].Code);
 
   F1->eraseFromParent();
@@ -164,7 +161,6 @@ TEST_F(JITEventListenerTest, MultipleListenersDontInterfere) {
       << " contain some bytes.";
 
   EXPECT_EQ(1U, Listener1.FreedEvents[0].Index);
-  EXPECT_EQ(F2, Listener1.FreedEvents[0].F);
   EXPECT_EQ(F2_addr, Listener1.FreedEvents[0].Code);
 
   // Listener 2.
@@ -186,7 +182,6 @@ TEST_F(JITEventListenerTest, MultipleListenersDontInterfere) {
       << " contain some bytes.";
 
   EXPECT_EQ(2U, Listener2.FreedEvents[0].Index);
-  EXPECT_EQ(F2, Listener2.FreedEvents[0].F);
   EXPECT_EQ(F2_addr, Listener2.FreedEvents[0].Code);
 
   // Listener 3.
@@ -201,7 +196,6 @@ TEST_F(JITEventListenerTest, MultipleListenersDontInterfere) {
       << " contain some bytes.";
 
   EXPECT_EQ(1U, Listener3.FreedEvents[0].Index);
-  EXPECT_EQ(F2, Listener3.FreedEvents[0].F);
   EXPECT_EQ(F2_addr, Listener3.FreedEvents[0].Code);
 
   F1->eraseFromParent();
@@ -228,7 +222,6 @@ TEST_F(JITEventListenerTest, MatchesMachineCodeInfo) {
   EXPECT_EQ(MCI.size(), Listener.EmittedEvents[0].Size);
 
   EXPECT_EQ(1U, Listener.FreedEvents[0].Index);
-  EXPECT_EQ(F, Listener.FreedEvents[0].F);
   EXPECT_EQ(F_addr, Listener.FreedEvents[0].Code);
 }
 

unittests/ExecutionEngine/JIT/JITTest.cpp

@@ -9,6 +9,7 @@
 
 #include "gtest/gtest.h"
 #include "llvm/ADT/OwningPtr.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Assembly/Parser.h"
 #include "llvm/BasicBlock.h"
 #include "llvm/Constant.h"
@@ -28,6 +29,8 @@
 #include "llvm/Target/TargetSelect.h"
 #include "llvm/Type.h"
 
+#include <vector>
+
 using namespace llvm;
 
 namespace {
@@ -47,13 +50,141 @@ Function *makeReturnGlobal(std::string Name, GlobalVariable *G, Module *M) {
   return F;
 }
 
+std::string DumpFunction(const Function *F) {
+  std::string Result;
+  raw_string_ostream(Result) << "" << *F;
+  return Result;
+}
+
+class RecordingJITMemoryManager : public JITMemoryManager {
+  const OwningPtr<JITMemoryManager> Base;
+public:
+  RecordingJITMemoryManager()
+    : Base(JITMemoryManager::CreateDefaultMemManager()) {
+  }
+
+  virtual void setMemoryWritable() { Base->setMemoryWritable(); }
+  virtual void setMemoryExecutable() { Base->setMemoryExecutable(); }
+  virtual void setPoisonMemory(bool poison) { Base->setPoisonMemory(poison); }
+  virtual void AllocateGOT() { Base->AllocateGOT(); }
+  virtual uint8_t *getGOTBase() const { return Base->getGOTBase(); }
+  virtual void SetDlsymTable(void *ptr) { Base->SetDlsymTable(ptr); }
+  virtual void *getDlsymTable() const { return Base->getDlsymTable(); }
+  struct StartFunctionBodyCall {
+    StartFunctionBodyCall(uint8_t *Result, const Function *F,
+                          uintptr_t ActualSize, uintptr_t ActualSizeResult)
+      : Result(Result), F(F), F_dump(DumpFunction(F)),
+        ActualSize(ActualSize), ActualSizeResult(ActualSizeResult) {}
+    uint8_t *Result;
+    const Function *F;
+    std::string F_dump;
+    uintptr_t ActualSize;
+    uintptr_t ActualSizeResult;
+  };
+  std::vector<StartFunctionBodyCall> startFunctionBodyCalls;
+  virtual uint8_t *startFunctionBody(const Function *F,
+                                     uintptr_t &ActualSize) {
+    uintptr_t InitialActualSize = ActualSize;
+    uint8_t *Result = Base->startFunctionBody(F, ActualSize);
+    startFunctionBodyCalls.push_back(
+      StartFunctionBodyCall(Result, F, InitialActualSize, ActualSize));
+    return Result;
+  }
+  virtual uint8_t *allocateStub(const GlobalValue* F, unsigned StubSize,
+                                unsigned Alignment) {
+    return Base->allocateStub(F, StubSize, Alignment);
+  }
+  struct EndFunctionBodyCall {
+    EndFunctionBodyCall(const Function *F, uint8_t *FunctionStart,
+                        uint8_t *FunctionEnd)
+      : F(F), F_dump(DumpFunction(F)),
+        FunctionStart(FunctionStart), FunctionEnd(FunctionEnd) {}
+    const Function *F;
+    std::string F_dump;
+    uint8_t *FunctionStart;
+    uint8_t *FunctionEnd;
+  };
+  std::vector<EndFunctionBodyCall> endFunctionBodyCalls;
+  virtual void endFunctionBody(const Function *F, uint8_t *FunctionStart,
+                               uint8_t *FunctionEnd) {
+    endFunctionBodyCalls.push_back(
+      EndFunctionBodyCall(F, FunctionStart, FunctionEnd));
+    Base->endFunctionBody(F, FunctionStart, FunctionEnd);
+  }
+  virtual uint8_t *allocateSpace(intptr_t Size, unsigned Alignment) {
+    return Base->allocateSpace(Size, Alignment);
+  }
+  virtual uint8_t *allocateGlobal(uintptr_t Size, unsigned Alignment) {
+    return Base->allocateGlobal(Size, Alignment);
+  }
+  struct DeallocateFunctionBodyCall {
+    DeallocateFunctionBodyCall(const void *Body) : Body(Body) {}
+    const void *Body;
+  };
+  std::vector<DeallocateFunctionBodyCall> deallocateFunctionBodyCalls;
+  virtual void deallocateFunctionBody(void *Body) {
+    deallocateFunctionBodyCalls.push_back(DeallocateFunctionBodyCall(Body));
+    Base->deallocateFunctionBody(Body);
+  }
+  struct DeallocateExceptionTableCall {
+    DeallocateExceptionTableCall(const void *ET) : ET(ET) {}
+    const void *ET;
+  };
+  std::vector<DeallocateExceptionTableCall> deallocateExceptionTableCalls;
+  virtual void deallocateExceptionTable(void *ET) {
+    deallocateExceptionTableCalls.push_back(DeallocateExceptionTableCall(ET));
+    Base->deallocateExceptionTable(ET);
+  }
+  struct StartExceptionTableCall {
+    StartExceptionTableCall(uint8_t *Result, const Function *F,
+                            uintptr_t ActualSize, uintptr_t ActualSizeResult)
+      : Result(Result), F(F), F_dump(DumpFunction(F)),
+        ActualSize(ActualSize), ActualSizeResult(ActualSizeResult) {}
+    uint8_t *Result;
+    const Function *F;
+    std::string F_dump;
+    uintptr_t ActualSize;
+    uintptr_t ActualSizeResult;
+  };
+  std::vector<StartExceptionTableCall> startExceptionTableCalls;
+  virtual uint8_t* startExceptionTable(const Function* F,
+                                       uintptr_t &ActualSize) {
+    uintptr_t InitialActualSize = ActualSize;
+    uint8_t *Result = Base->startExceptionTable(F, ActualSize);
+    startExceptionTableCalls.push_back(
+      StartExceptionTableCall(Result, F, InitialActualSize, ActualSize));
+    return Result;
+  }
+  struct EndExceptionTableCall {
+    EndExceptionTableCall(const Function *F, uint8_t *TableStart,
+                          uint8_t *TableEnd, uint8_t* FrameRegister)
+      : F(F), F_dump(DumpFunction(F)),
+        TableStart(TableStart), TableEnd(TableEnd),
+        FrameRegister(FrameRegister) {}
+    const Function *F;
+    std::string F_dump;
+    uint8_t *TableStart;
+    uint8_t *TableEnd;
+    uint8_t *FrameRegister;
+  };
+  std::vector<EndExceptionTableCall> endExceptionTableCalls;
+  virtual void endExceptionTable(const Function *F, uint8_t *TableStart,
+                                 uint8_t *TableEnd, uint8_t* FrameRegister) {
+      endExceptionTableCalls.push_back(
+          EndExceptionTableCall(F, TableStart, TableEnd, FrameRegister));
+    return Base->endExceptionTable(F, TableStart, TableEnd, FrameRegister);
+  }
+};
+
 class JITTest : public testing::Test {
  protected:
   virtual void SetUp() {
     M = new Module("<main>", Context);
     MP = new ExistingModuleProvider(M);
+    RJMM = new RecordingJITMemoryManager;
     std::string Error;
     TheJIT.reset(EngineBuilder(MP).setEngineKind(EngineKind::JIT)
+                 .setJITMemoryManager(RJMM)
                  .setErrorStr(&Error).create());
     ASSERT_TRUE(TheJIT.get() != NULL) << Error;
   }
@@ -70,6 +201,7 @@ class JITTest : public testing::Test {
   LLVMContext Context;
   Module *M;  // Owned by MP.
   ModuleProvider *MP;  // Owned by ExecutionEngine.
+  RecordingJITMemoryManager *RJMM;
   OwningPtr<ExecutionEngine> TheJIT;
 };
 
@@ -289,6 +421,34 @@ TEST_F(JITTest, ModuleDeletion) {
   Function *func = M->getFunction("main");
   TheJIT->getPointerToFunction(func);
   TheJIT->deleteModuleProvider(MP);
+
+  SmallPtrSet<const void*, 2> FunctionsDeallocated;
+  for (unsigned i = 0, e = RJMM->deallocateFunctionBodyCalls.size();
+       i != e; ++i) {
+    FunctionsDeallocated.insert(RJMM->deallocateFunctionBodyCalls[i].Body);
+  }
+  for (unsigned i = 0, e = RJMM->startFunctionBodyCalls.size(); i != e; ++i) {
+    EXPECT_TRUE(FunctionsDeallocated.count(
+                  RJMM->startFunctionBodyCalls[i].Result))
+      << "Function leaked: \n" << RJMM->startFunctionBodyCalls[i].F_dump;
+  }
+  EXPECT_EQ(RJMM->startFunctionBodyCalls.size(),
+            RJMM->deallocateFunctionBodyCalls.size());
+
+  SmallPtrSet<const void*, 2> ExceptionTablesDeallocated;
+  for (unsigned i = 0, e = RJMM->deallocateExceptionTableCalls.size();
+       i != e; ++i) {
+    ExceptionTablesDeallocated.insert(
+        RJMM->deallocateExceptionTableCalls[i].ET);
+  }
+  for (unsigned i = 0, e = RJMM->startExceptionTableCalls.size(); i != e; ++i) {
+    EXPECT_TRUE(ExceptionTablesDeallocated.count(
+                  RJMM->startExceptionTableCalls[i].Result))
+      << "Function's exception table leaked: \n"
+      << RJMM->startExceptionTableCalls[i].F_dump;
+  }
+  EXPECT_EQ(RJMM->startExceptionTableCalls.size(),
+            RJMM->deallocateExceptionTableCalls.size());
 }
 
 // This code is copied from JITEventListenerTest, but it only runs once for all