Make it possible to create multiple JIT instances at the same time, by removing

the global TheJIT and TheJITResolver variables.  Lazy compilation is supported
by a global map from a stub address to the JITResolver that knows how to
compile it.

Patch by Olivier Meurant!


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

docs/ReleaseNotes.html

@@ -467,6 +467,10 @@ href="http://llvm.org/viewvc/llvm-project?view=rev&revision=85295">defaults
 to compiling eagerly</a> to avoid a race condition in the lazy JIT.
 Clients that still want the lazy JIT can switch it on by calling
 <tt>ExecutionEngine::DisableLazyCompilation(false)</tt>.</li>
+<li>It is now possible to create more than one JIT instance in the same process.
+These JITs can generate machine code in parallel,
+although <a href="http://llvm.org/docs/ProgrammersManual.html#jitthreading">you
+still have to obey the other threading restrictions</a>.</li>
 </ul>
 
 </div>

lib/ExecutionEngine/JIT/JIT.cpp

@@ -18,6 +18,7 @@
 #include "llvm/Function.h"
 #include "llvm/GlobalVariable.h"
 #include "llvm/Instructions.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/CodeGen/JITCodeEmitter.h"
 #include "llvm/CodeGen/MachineCodeInfo.h"
 #include "llvm/ExecutionEngine/GenericValue.h"
@@ -27,6 +28,7 @@
 #include "llvm/Target/TargetJITInfo.h"
 #include "llvm/Support/Dwarf.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/MutexGuard.h"
 #include "llvm/System/DynamicLibrary.h"
 #include "llvm/Config/config.h"
@@ -237,9 +239,53 @@ ExecutionEngine *JIT::createJIT(Module *M,
   }
 }
 
+namespace {
+/// This class supports the global getPointerToNamedFunction(), which allows
+/// bugpoint or gdb users to search for a function by name without any context.
+class JitPool {
+  SmallPtrSet<JIT*, 1> JITs;  // Optimize for process containing just 1 JIT.
+  mutable sys::Mutex Lock;
+public:
+  void Add(JIT *jit) {
+    MutexGuard guard(Lock);
+    JITs.insert(jit);
+  }
+  void Remove(JIT *jit) {
+    MutexGuard guard(Lock);
+    JITs.erase(jit);
+  }
+  void *getPointerToNamedFunction(const char *Name) const {
+    MutexGuard guard(Lock);
+    assert(JITs.size() != 0 && "No Jit registered");
+    //search function in every instance of JIT
+    for (SmallPtrSet<JIT*, 1>::const_iterator Jit = JITs.begin(),
+           end = JITs.end();
+         Jit != end; ++Jit) {
+      if (Function *F = (*Jit)->FindFunctionNamed(Name))
+        return (*Jit)->getPointerToFunction(F);
+    }
+    // The function is not available : fallback on the first created (will
+    // search in symbol of the current program/library)
+    return (*JITs.begin())->getPointerToNamedFunction(Name);
+  }
+};
+ManagedStatic<JitPool> AllJits;
+}
+extern "C" {
+  // getPointerToNamedFunction - This function is used as a global wrapper to
+  // JIT::getPointerToNamedFunction for the purpose of resolving symbols when
+  // bugpoint is debugging the JIT. In that scenario, we are loading an .so and
+  // need to resolve function(s) that are being mis-codegenerated, so we need to
+  // resolve their addresses at runtime, and this is the way to do it.
+  void *getPointerToNamedFunction(const char *Name) {
+    return AllJits->getPointerToNamedFunction(Name);
+  }
+}
+
 JIT::JIT(Module *M, TargetMachine &tm, TargetJITInfo &tji,
          JITMemoryManager *JMM, CodeGenOpt::Level OptLevel, bool GVsWithCode)
-  : ExecutionEngine(M), TM(tm), TJI(tji), AllocateGVsWithCode(GVsWithCode) {
+  : ExecutionEngine(M), TM(tm), TJI(tji), AllocateGVsWithCode(GVsWithCode),
+    isAlreadyCodeGenerating(false) {
   setTargetData(TM.getTargetData());
 
   jitstate = new JITState(M);
@@ -247,6 +293,9 @@ JIT::JIT(Module *M, TargetMachine &tm, TargetJITInfo &tji,
   // Initialize JCE
   JCE = createEmitter(*this, JMM, TM);
 
+  // Register in global list of all JITs.
+  AllJits->Add(this);
+
   // Add target data
   MutexGuard locked(lock);
   FunctionPassManager &PM = jitstate->getPM(locked);
@@ -281,6 +330,7 @@ JIT::JIT(Module *M, TargetMachine &tm, TargetJITInfo &tji,
 }
 
 JIT::~JIT() {
+  AllJits->Remove(this);
   delete jitstate;
   delete JCE;
   delete &TM;
@@ -570,7 +620,6 @@ void JIT::runJITOnFunction(Function *F, MachineCodeInfo *MCI) {
 }
 
 void JIT::runJITOnFunctionUnlocked(Function *F, const MutexGuard &locked) {
-  static bool isAlreadyCodeGenerating = false;
   assert(!isAlreadyCodeGenerating && "Error: Recursive compilation detected!");
 
   // JIT the function

lib/ExecutionEngine/JIT/JIT.h

@@ -61,6 +61,10 @@ class JIT : public ExecutionEngine {
   /// should be set to true.  Doing so breaks freeMachineCodeForFunction.
   bool AllocateGVsWithCode;
 
+  /// True while the JIT is generating code.  Used to assert against recursive
+  /// entry.
+  bool isAlreadyCodeGenerating;
+
   JITState *jitstate;
 
   JIT(Module *M, TargetMachine &tm, TargetJITInfo &tji,

lib/ExecutionEngine/JIT/JITEmitter.cpp

@@ -37,6 +37,7 @@
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/MutexGuard.h"
 #include "llvm/Support/ValueHandle.h"
 #include "llvm/Support/raw_ostream.h"
@@ -57,7 +58,6 @@ using namespace llvm;
 STATISTIC(NumBytes, "Number of bytes of machine code compiled");
 STATISTIC(NumRelos, "Number of relocations applied");
 STATISTIC(NumRetries, "Number of retries with more memory");
-static JIT *TheJIT = 0;
 
 
 // A declaration may stop being a declaration once it's fully read from bitcode.
@@ -109,9 +109,13 @@ namespace {
     /// particular GlobalVariable so that we can reuse them if necessary.
     GlobalToIndirectSymMapTy GlobalToIndirectSymMap;
 
+    /// Instance of the JIT this ResolverState serves.
+    JIT *TheJIT;
+
   public:
-    JITResolverState() : FunctionToLazyStubMap(this),
+    JITResolverState(JIT *jit) : FunctionToLazyStubMap(this),
-                         FunctionToCallSitesMap(this) {}
+                                 FunctionToCallSitesMap(this),
+                                 TheJIT(jit) {}
 
     FunctionToLazyStubMapTy& getFunctionToLazyStubMap(
       const MutexGuard& locked) {
@@ -227,18 +231,13 @@ namespace {
 
     JITEmitter &JE;
 
-    static JITResolver *TheJITResolver;
+    /// Instance of JIT corresponding to this Resolver.
+    JIT *TheJIT;
+
   public:
-    explicit JITResolver(JIT &jit, JITEmitter &je) : nextGOTIndex(0), JE(je) {
+    explicit JITResolver(JIT &jit, JITEmitter &je)
-      TheJIT = &jit;
+      : state(&jit), nextGOTIndex(0), JE(je), TheJIT(&jit) {
-
       LazyResolverFn = jit.getJITInfo().getLazyResolverFunction(JITCompilerFn);
-      assert(TheJITResolver == 0 && "Multiple JIT resolvers?");
-      TheJITResolver = this;
-    }
-
-    ~JITResolver() {
-      TheJITResolver = 0;
     }
 
     /// getLazyFunctionStubIfAvailable - This returns a pointer to a function's
@@ -273,6 +272,44 @@ namespace {
     static void *JITCompilerFn(void *Stub);
   };
 
+  class StubToResolverMapTy {
+    /// Map a stub address to a specific instance of a JITResolver so that
+    /// lazily-compiled functions can find the right resolver to use.
+    ///
+    /// Guarded by Lock.
+    std::map<void*, JITResolver*> Map;
+
+    /// Guards Map from concurrent accesses.
+    mutable sys::Mutex Lock;
+
+  public:
+    /// Registers a Stub to be resolved by Resolver.
+    void RegisterStubResolver(void *Stub, JITResolver *Resolver) {
+      MutexGuard guard(Lock);
+      Map.insert(std::make_pair(Stub, Resolver));
+    }
+    /// Unregisters the Stub when it's invalidated.
+    void UnregisterStubResolver(void *Stub) {
+      MutexGuard guard(Lock);
+      Map.erase(Stub);
+    }
+    /// Returns the JITResolver instance that owns the Stub.
+    JITResolver *getResolverFromStub(void *Stub) const {
+      MutexGuard guard(Lock);
+      // The address given to us for the stub may not be exactly right, it might
+      // be a little bit after the stub.  As such, use upper_bound to find it.
+      // This is the same trick as in LookupFunctionFromCallSite from
+      // JITResolverState.
+      std::map<void*, JITResolver*>::const_iterator I = Map.upper_bound(Stub);
+      assert(I != Map.begin() && "This is not a known stub!");
+      --I;
+      return I->second;
+    }
+  };
+  /// This needs to be static so that a lazy call stub can access it with no
+  /// context except the address of the stub.
+  ManagedStatic<StubToResolverMapTy> StubToResolverMap;
+
   /// JITEmitter - The JIT implementation of the MachineCodeEmitter, which is
   /// used to output functions to memory for execution.
   class JITEmitter : public JITCodeEmitter {
@@ -371,10 +408,13 @@ namespace {
 
     DILocation PrevDLT;
 
+    /// Instance of the JIT
+    JIT *TheJIT;
+
   public:
     JITEmitter(JIT &jit, JITMemoryManager *JMM, TargetMachine &TM)
       : SizeEstimate(0), Resolver(jit, *this), MMI(0), CurFn(0),
-        EmittedFunctions(this), PrevDLT(NULL) {
+        EmittedFunctions(this), PrevDLT(NULL), TheJIT(&jit) {
       MemMgr = JMM ? JMM : JITMemoryManager::CreateDefaultMemManager();
       if (jit.getJITInfo().needsGOT()) {
         MemMgr->AllocateGOT();
@@ -495,8 +535,6 @@ namespace {
   };
 }
 
-JITResolver *JITResolver::TheJITResolver = 0;
-
 void CallSiteValueMapConfig::onDelete(JITResolverState *JRS, Function *F) {
   JRS->EraseAllCallSitesPrelocked(F);
 }
@@ -551,6 +589,10 @@ void *JITResolver::getLazyFunctionStub(Function *F) {
   DEBUG(dbgs() << "JIT: Lazy stub emitted at [" << Stub << "] for function '"
         << F->getName() << "'\n");
 
+  // Register this JITResolver as the one corresponding to this call site so
+  // JITCompilerFn will be able to find it.
+  StubToResolverMap->RegisterStubResolver(Stub, this);
+
   // Finally, keep track of the stub-to-Function mapping so that the
   // JITCompilerFn knows which function to compile!
   state.AddCallSite(locked, Stub, F);
@@ -637,6 +679,9 @@ void JITResolver::getRelocatableGVs(SmallVectorImpl<GlobalValue*> &GVs,
 GlobalValue *JITResolver::invalidateStub(void *Stub) {
   MutexGuard locked(TheJIT->lock);
 
+  // Remove the stub from the StubToResolverMap.
+  StubToResolverMap->UnregisterStubResolver(Stub);
+
   GlobalToIndirectSymMapTy &GM = state.getGlobalToIndirectSymMap(locked);
 
   // Look up the cheap way first, to see if it's a function stub we are
@@ -671,7 +716,8 @@ GlobalValue *JITResolver::invalidateStub(void *Stub) {
 /// been entered.  It looks up which function this stub corresponds to, compiles
 /// it if necessary, then returns the resultant function pointer.
 void *JITResolver::JITCompilerFn(void *Stub) {
-  JITResolver &JR = *TheJITResolver;
+  JITResolver *JR = StubToResolverMap->getResolverFromStub(Stub);
+  assert(JR && "Unable to find the corresponding JITResolver to the call site");
 
   Function* F = 0;
   void* ActualPtr = 0;
@@ -680,24 +726,24 @@ void *JITResolver::JITCompilerFn(void *Stub) {
     // Only lock for getting the Function. The call getPointerToFunction made
     // in this function might trigger function materializing, which requires
     // JIT lock to be unlocked.
-    MutexGuard locked(TheJIT->lock);
+    MutexGuard locked(JR->TheJIT->lock);
 
     // The address given to us for the stub may not be exactly right, it might
     // be a little bit after the stub.  As such, use upper_bound to find it.
     pair<void*, Function*> I =
-      JR.state.LookupFunctionFromCallSite(locked, Stub);
+      JR->state.LookupFunctionFromCallSite(locked, Stub);
     F = I.second;
     ActualPtr = I.first;
   }
 
   // If we have already code generated the function, just return the address.
-  void *Result = TheJIT->getPointerToGlobalIfAvailable(F);
+  void *Result = JR->TheJIT->getPointerToGlobalIfAvailable(F);
 
   if (!Result) {
     // Otherwise we don't have it, do lazy compilation now.
 
     // If lazy compilation is disabled, emit a useful error message and abort.
-    if (!TheJIT->isCompilingLazily()) {
+    if (!JR->TheJIT->isCompilingLazily()) {
       llvm_report_error("LLVM JIT requested to do lazy compilation of function '"
                         + F->getName() + "' when lazy compiles are disabled!");
     }
@@ -706,11 +752,11 @@ void *JITResolver::JITCompilerFn(void *Stub) {
           << "' In stub ptr = " << Stub << " actual ptr = "
           << ActualPtr << "\n");
 
-    Result = TheJIT->getPointerToFunction(F);
+    Result = JR->TheJIT->getPointerToFunction(F);
   }
 
   // Reacquire the lock to update the GOT map.
-  MutexGuard locked(TheJIT->lock);
+  MutexGuard locked(JR->TheJIT->lock);
 
   // We might like to remove the call site from the CallSiteToFunction map, but
   // we can't do that! Multiple threads could be stuck, waiting to acquire the
@@ -725,8 +771,8 @@ void *JITResolver::JITCompilerFn(void *Stub) {
   // if they see it still using the stub address.
   // Note: this is done so the Resolver doesn't have to manage GOT memory
   // Do this without allocating map space if the target isn't using a GOT
-  if(JR.revGOTMap.find(Stub) != JR.revGOTMap.end())
+  if(JR->revGOTMap.find(Stub) != JR->revGOTMap.end())
-    JR.revGOTMap[Result] = JR.revGOTMap[Stub];
+    JR->revGOTMap[Result] = JR->revGOTMap[Stub];
 
   return Result;
 }
@@ -839,7 +885,7 @@ static unsigned GetConstantPoolSizeInBytes(MachineConstantPool *MCP,
   return Size;
 }
 
-static unsigned GetJumpTableSizeInBytes(MachineJumpTableInfo *MJTI) {
+static unsigned GetJumpTableSizeInBytes(MachineJumpTableInfo *MJTI, JIT *jit) {
   const std::vector<MachineJumpTableEntry> &JT = MJTI->getJumpTables();
   if (JT.empty()) return 0;
 
@@ -847,7 +893,7 @@ static unsigned GetJumpTableSizeInBytes(MachineJumpTableInfo *MJTI) {
   for (unsigned i = 0, e = JT.size(); i != e; ++i)
     NumEntries += JT[i].MBBs.size();
 
-  return NumEntries * MJTI->getEntrySize(*TheJIT->getTargetData());
+  return NumEntries * MJTI->getEntrySize(*jit->getTargetData());
 }
 
 static uintptr_t RoundUpToAlign(uintptr_t Size, unsigned Alignment) {
@@ -1032,7 +1078,7 @@ void JITEmitter::startFunction(MachineFunction &F) {
                              MJTI->getEntryAlignment(*TheJIT->getTargetData()));
 
       // Add the jump table size
-      ActualSize += GetJumpTableSizeInBytes(MJTI);
+      ActualSize += GetJumpTableSizeInBytes(MJTI, TheJIT);
     }
 
     // Add the alignment for the function
@@ -1552,19 +1598,6 @@ JITCodeEmitter *JIT::createEmitter(JIT &jit, JITMemoryManager *JMM,
   return new JITEmitter(jit, JMM, tm);
 }
 
-// getPointerToNamedFunction - This function is used as a global wrapper to
-// JIT::getPointerToNamedFunction for the purpose of resolving symbols when
-// bugpoint is debugging the JIT. In that scenario, we are loading an .so and
-// need to resolve function(s) that are being mis-codegenerated, so we need to
-// resolve their addresses at runtime, and this is the way to do it.
-extern "C" {
-  void *getPointerToNamedFunction(const char *Name) {
-    if (Function *F = TheJIT->FindFunctionNamed(Name))
-      return TheJIT->getPointerToFunction(F);
-    return TheJIT->getPointerToNamedFunction(Name);
-  }
-}
-
 // getPointerToFunctionOrStub - If the specified function has been
 // code-gen'd, return a pointer to the function.  If not, compile it, or use
 // a stub to implement lazy compilation if available.

unittests/ExecutionEngine/JIT/MultiJITTest.cpp

@@ -0,0 +1,164 @@
+//===- MultiJITTest.cpp - Unit tests for instantiating multiple JITs ------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "gtest/gtest.h"
+#include "llvm/LLVMContext.h"
+#include "llvm/Module.h"
+#include "llvm/Assembly/Parser.h"
+#include "llvm/ExecutionEngine/GenericValue.h"
+#include "llvm/ExecutionEngine/JIT.h"
+#include "llvm/Support/SourceMgr.h"
+#include <vector>
+
+using namespace llvm;
+
+namespace {
+
+bool LoadAssemblyInto(Module *M, const char *assembly) {
+  SMDiagnostic Error;
+  bool success =
+    NULL != ParseAssemblyString(assembly, M, Error, M->getContext());
+  std::string errMsg;
+  raw_string_ostream os(errMsg);
+  Error.Print("", os);
+  EXPECT_TRUE(success) << os.str();
+  return success;
+}
+
+void createModule1(LLVMContext &Context1, Module *&M1, Function *&FooF1) {
+  M1 = new Module("test1", Context1);
+  LoadAssemblyInto(M1,
+                   "define i32 @add1(i32 %ArgX1) { "
+                   "entry: "
+                   "  %addresult = add i32 1, %ArgX1 "
+                   "  ret i32 %addresult "
+                   "} "
+                   " "
+                   "define i32 @foo1() { "
+                   "entry: "
+                   "  %add1 = call i32 @add1(i32 10) "
+                   "  ret i32 %add1 "
+                   "} ");
+  FooF1 = M1->getFunction("foo1");
+}
+
+void createModule2(LLVMContext &Context2, Module *&M2, Function *&FooF2) {
+  M2 = new Module("test2", Context2);
+  LoadAssemblyInto(M2,
+                   "define i32 @add2(i32 %ArgX2) { "
+                   "entry: "
+                   "  %addresult = add i32 2, %ArgX2 "
+                   "  ret i32 %addresult "
+                   "} "
+                   " "
+                   "define i32 @foo2() { "
+                   "entry: "
+                   "  %add2 = call i32 @add2(i32 10) "
+                   "  ret i32 %add2 "
+                   "} ");
+  FooF2 = M2->getFunction("foo2");
+}
+
+TEST(MultiJitTest, EagerMode) {
+  LLVMContext Context1;
+  Module *M1 = 0;
+  Function *FooF1 = 0;
+  createModule1(Context1, M1, FooF1);
+
+  LLVMContext Context2;
+  Module *M2 = 0;
+  Function *FooF2 = 0;
+  createModule2(Context2, M2, FooF2);
+
+  // Now we create the JIT in eager mode
+  OwningPtr<ExecutionEngine> EE1(EngineBuilder(M1).create());
+  EE1->DisableLazyCompilation(true);
+  OwningPtr<ExecutionEngine> EE2(EngineBuilder(M2).create());
+  EE2->DisableLazyCompilation(true);
+
+  // Call the `foo' function with no arguments:
+  std::vector<GenericValue> noargs;
+  GenericValue gv1 = EE1->runFunction(FooF1, noargs);
+  GenericValue gv2 = EE2->runFunction(FooF2, noargs);
+
+  // Import result of execution:
+  EXPECT_EQ(gv1.IntVal, 11);
+  EXPECT_EQ(gv2.IntVal, 12);
+
+  EE1->freeMachineCodeForFunction(FooF1);
+  EE2->freeMachineCodeForFunction(FooF2);
+}
+
+TEST(MultiJitTest, LazyMode) {
+  LLVMContext Context1;
+  Module *M1 = 0;
+  Function *FooF1 = 0;
+  createModule1(Context1, M1, FooF1);
+
+  LLVMContext Context2;
+  Module *M2 = 0;
+  Function *FooF2 = 0;
+  createModule2(Context2, M2, FooF2);
+
+  // Now we create the JIT in lazy mode
+  OwningPtr<ExecutionEngine> EE1(EngineBuilder(M1).create());
+  EE1->DisableLazyCompilation(false);
+  OwningPtr<ExecutionEngine> EE2(EngineBuilder(M2).create());
+  EE2->DisableLazyCompilation(false);
+
+  // Call the `foo' function with no arguments:
+  std::vector<GenericValue> noargs;
+  GenericValue gv1 = EE1->runFunction(FooF1, noargs);
+  GenericValue gv2 = EE2->runFunction(FooF2, noargs);
+
+  // Import result of execution:
+  EXPECT_EQ(gv1.IntVal, 11);
+  EXPECT_EQ(gv2.IntVal, 12);
+
+  EE1->freeMachineCodeForFunction(FooF1);
+  EE2->freeMachineCodeForFunction(FooF2);
+}
+
+extern "C" {
+  extern void *getPointerToNamedFunction(const char *Name);
+}
+
+TEST(MultiJitTest, JitPool) {
+  LLVMContext Context1;
+  Module *M1 = 0;
+  Function *FooF1 = 0;
+  createModule1(Context1, M1, FooF1);
+
+  LLVMContext Context2;
+  Module *M2 = 0;
+  Function *FooF2 = 0;
+  createModule2(Context2, M2, FooF2);
+
+  // Now we create two JITs
+  OwningPtr<ExecutionEngine> EE1(EngineBuilder(M1).create());
+  OwningPtr<ExecutionEngine> EE2(EngineBuilder(M2).create());
+
+  Function *F1 = EE1->FindFunctionNamed("foo1");
+  void *foo1 = EE1->getPointerToFunction(F1);
+
+  Function *F2 = EE2->FindFunctionNamed("foo2");
+  void *foo2 = EE2->getPointerToFunction(F2);
+
+  // Function in M1
+  EXPECT_EQ(getPointerToNamedFunction("foo1"), foo1);
+
+  // Function in M2
+  EXPECT_EQ(getPointerToNamedFunction("foo2"), foo2);
+
+  // Symbol search
+  EXPECT_EQ((intptr_t)getPointerToNamedFunction("getPointerToNamedFunction"),
+            (intptr_t)&getPointerToNamedFunction);
+}
+
+}  // anonymous namespace