[WIP] LLJITTest

[sgraenitz]

Summary

Getting familiar with the new API

[sgraenitz]

Split up into 4 usage examples. This goes from a very simple use case to more advanced ones.

It illustrates 2 current issues:

• PassManager::run() seems to be not thread-safe, thus the extra mutex in IRCompileLayer2::emit().
• JITing from different threads causes deadlock during symbol lookup.

[sgraenitz]

Add more polishing, explanations, fixmes

[sgraenitz]

Maybe it makes sense in oder to illustrate the usage of your API, to give people a very simple example (like UnthreadedSingleDylib) and extended it step by step to more advanced ones (like ThreadedMultiDylib).

I don't see how to avoid the PassManagerRunMutex in IRCompileLayer2::emit():

• I tried to add mutexes in MCAssembler, TargetLoweringObjectFile, AtomicExpandPass, etc., but it got very complicated very quickly. Everything goes through PassManager::run(), but a mutex there would affect a lot of other code. IRCompileLayer2 appeared to be a good point.
• I always thought that support for multithreaded codegen is implicitly given when using different LLVMContext instances. It doesn't look like that here. Maybe there's more requirements to it?

unittests/ExecutionEngine/Orc/LLJITTest.cpp

55

Changed example functions: foo and bar are disjunct now, buz is the one that depends on both of them.

161

The thread that runs DispatchMaterialization can't be blocked, right? So it must spawn a new thread, which must be joined back to this thread later. How do we know the correct one? Track it in MaterializationUnit and join in destructor? Also see my below comments.

184

Can't join the DispatchMaterialization thread here, because I don't know whether it's M[0] or M[1].

190

Concurrent lookup for foo and bar works due to the mutex in IRCompileLayer2::emit().

194

This randomly causes failures. I think threads should always be joined in the thread that they were forked from. Not simple, even with a fixed-size thread pool.

269

No barrier with JoinAll necessary here. Maybe because foo and bar are in other dylibs?

[sgraenitz]

lib/ExecutionEngine/Orc/IRCompileLayer.cpp

26 ↗

(On Diff #161950)

PassManagerRunMutex is here. It's not included in the reviews main diff. Not sure why.

[sgraenitz]

unittests/ExecutionEngine/Orc/LLJITTest.cpp

161

Track it in MaterializationUnit and join in destructor?

The last shared ref to SMU below is in the newly spawned thread, so this won't work.

[sgraenitz]

I collected quite some unnecessary noise here. Will close this and open dedicated ones for proposals.

unittests/ExecutionEngine/Orc/LLJITTest.cpp

const char *BuzIR = LLVM_IR(
   declare i32 @bar()
 
   define i32 @buz() {
     %1 = call i32 @foo()
     %2 = call i32 @bar()
     %3 = add nsw i32 %1, %2
     ret i32 %3
   }
 );
 
 static const MemoryBufferRef Foo{FooIR, "FooIR"};
 static const MemoryBufferRef Bar{BarIR, "BarIR"};
 static const MemoryBufferRef Buz{BuzIR, "BuzIR"};
 
 static std::unique_ptr<LLJIT> Init() {
   InitializeNativeTarget();
   InitializeNativeTargetAsmParser();
for (size_t i = 0; i < 3; i++)
     cantFail(Jit->addIRModule(*Ds[i], parseIR(Bs[i], E, Cs[i])));
   
   JITEvaluatedSymbol FooSym = cantFail(Jit->lookup("foo"));
   EXPECT_EQ(1, Exec(FooSym));
   
   JITEvaluatedSymbol BarSym = cantFail(Jit->lookup(Extra1, "bar"));
   EXPECT_EQ(1, Exec(BarSym));
 
-  // Allow buz to resolve foo and bar from the other dylibs:
+  // Allow buz to resolve foo and bar from the other dylibs.
   Extra2.addToSearchOrder(Main);
   Extra2.addToSearchOrder(Extra1);
 
   JITEvaluatedSymbol BuzSym = cantFail(Jit->lookup(Extra2, "buz"));
   EXPECT_EQ(2, Exec(BuzSym));
 }
 
 //===----------------------------------------------------------------------===//
 
 TEST(LLJIT, ThreadedSingleDylib) {
   std::unique_ptr<LLJIT> Jit = Init();
   ExecutionSession &ES = Jit->getExecutionSession();
   
   SMDiagnostic E;
   std::vector<std::thread> Ts;
   std::vector<LLVMContext> Cs(3);
   std::vector<MemoryBufferRef> Bs{Foo, Bar, Buz};
 
   for (size_t i = 0; i < 3; i++)
     cantFail(Jit->addIRModule(parseIR(Bs[i], E, Cs[i])));
 
-  // Dispatch compile job to different threads.
+  // Dispatch compile jobs to different threads.
   std::vector<std::thread> Ms;
-  static constexpr size_t MaxMaterializerThreads = 3;
+  static constexpr size_t MaxMaterializerThreads = 4;
   Ms.reserve(MaxMaterializerThreads);
 
   ES.setDispatchMaterialization(
       [&](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
         assert(Ms.size() < MaxMaterializerThreads &&
                "More slots used than preallocated");
 
         auto SMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
         Ms.emplace_back([SMU, &JD]() {
           SMU->doMaterialize(JD);
         });
       });
   
+  /*
+  // FIXME: Trying to run this before the other functions always causes a
+  // deadlock.
+  Ts.emplace_back([&Jit]() {
+    JITEvaluatedSymbol BuzSym = cantFail(Jit->lookup("buz"));
+    EXPECT_EQ(2, Exec(BuzSym)); // Automatically waits for materializer.
+  });
+  
+  JoinAll(Ts);
+  JoinAll(Ms); // FIXME: Ideally this happened in the individual runner threads.
+  */
+
   // Run functions in different threads.
   Ts.emplace_back([&Jit]() {
     JITEvaluatedSymbol FooSym = cantFail(Jit->lookup("foo"));
     EXPECT_EQ(1, Exec(FooSym)); // Automatically waits for materializer.
   });
 
   Ts.emplace_back([&Jit]() {
     JITEvaluatedSymbol BarSym = cantFail(Jit->lookup("bar"));
     EXPECT_EQ(1, Exec(BarSym)); // Automatically waits for materializer.
   });
   
   // FIXME: Lookup for "buz" can't run in parallel.
   JoinAll(Ts);
-  JoinAll(Ms);
+  JoinAll(Ms); // This should happen in the above runner threads.
 
   Ts.emplace_back([&Jit]() {
     JITEvaluatedSymbol BuzSym = cantFail(Jit->lookup("buz"));
     EXPECT_EQ(2, Exec(BuzSym)); // Automatically waits for materializer.
   });
   
   JoinAll(Ts);
-  JoinAll(Ms);
+  JoinAll(Ms); // FIXME: Ideally this happened in the individual runner threads.
 }
 
 //===----------------------------------------------------------------------===//
 
 TEST(LLJIT, ThreadedMultiDylib) {
   std::unique_ptr<LLJIT> Jit = Init();
   ExecutionSession &ES = Jit->getExecutionSession();
   
   JITDylib &Main = Jit->getMainJITDylib();
   JITDylib &Extra1 = ES.createJITDylib("1");
   JITDylib &Extra2 = ES.createJITDylib("2");
   Main.addToSearchOrder(Extra1);
   Main.addToSearchOrder(Extra2);
   
   SMDiagnostic E;
   std::vector<std::thread> Ts;
   std::vector<LLVMContext> Cs(3);
   std::vector<MemoryBufferRef> Bs{Foo, Bar, Buz};
   std::vector<JITDylib *> Ds{&Main, &Extra1, &Extra2};
   
   for (size_t i = 0; i < 3; i++)
     cantFail(Jit->addIRModule(*Ds[i], parseIR(Bs[i], E, Cs[i])));
 
-  // Dispatch compile job to different threads.
+  // Dispatch compile jobs to different threads.
   std::vector<std::thread> Ms;
-  static constexpr size_t MaxMaterializerThreads = 3;
+  static constexpr size_t MaxMaterializerThreads = 4; // one more than necessary
   Ms.reserve(MaxMaterializerThreads);
 
   ES.setDispatchMaterialization(
       [&](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
         assert(Ms.size() < MaxMaterializerThreads &&
                "More slots used than preallocated");
 
         auto SMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
         Ms.emplace_back([SMU, &JD]() {
           SMU->doMaterialize(JD);
         });
       });
 
+  // Allow buz to resolve foo and bar from the other dylibs.
+  Extra2.addToSearchOrder(Main);
+  Extra2.addToSearchOrder(Extra1);
+
+  /*
+  // FIXME: Trying to run this before the other functions always causes a
+  // deadlock.
+  Ts.emplace_back([&]() {
+    JITEvaluatedSymbol BuzSym = cantFail(Jit->lookup(Extra2, "buz"));
+    EXPECT_EQ(2, Exec(BuzSym)); // Automatically waits for materializer.
+  });
+  
+  JoinAll(Ts);
+  JoinAll(Ms); // FIXME: Ideally this happened in the individual runner threads.
+  */
+
   // Run functions in different threads.
   Ts.emplace_back([&]() {
     JITEvaluatedSymbol FooSym = cantFail(Jit->lookup("foo"));
     EXPECT_EQ(1, Exec(FooSym)); // Automatically waits for materializer.
   });
 
   Ts.emplace_back([&]() {
     JITEvaluatedSymbol BarSym = cantFail(Jit->lookup(Extra1, "bar"));
     EXPECT_EQ(1, Exec(BarSym)); // Automatically waits for materializer.
   });
   
-  // FIXME: Lookup for "buz" can't run in parallel.
-  JoinAll(Ts);
-  JoinAll(Ms);
-
-  // Allow buz to resolve foo and bar from the other dylibs.
-  Extra2.addToSearchOrder(Main);
-  Extra2.addToSearchOrder(Extra1);
-
+  // Lookup for "buz" works in parallel here even without a barrier.
   Ts.emplace_back([&]() {
     JITEvaluatedSymbol BuzSym = cantFail(Jit->lookup(Extra2, "buz"));
     EXPECT_EQ(2, Exec(BuzSym)); // Automatically waits for materializer.
   });
 
   JoinAll(Ts);
-  JoinAll(Ms);
+  JoinAll(Ms); // FIXME: Ideally this happened in the individual runner threads.
 }