[ORC] LLJITTest: API test and usage examples with actual codegen

[sgraenitz]

Summary

4 tests to illustrate the LLJIT API in combinations of sequential/concurrent and single/multi dylib.
Requires: https://reviews.llvm.org/D51126

I am using thread-local scope-exit guards to join worker threads back to their correct origins. This works surprisingly well! It's definitely dangerous with long-living threads as resources can get stuck for long. OTOH it's quite simple compared to tracking the instances in containers that require locks etc...

What do you think?

[sgraenitz]

Use LLJIT::CreateMultiThreaded()

[sgraenitz]

I will do some experiments with LLLazyJIT next.

[sgraenitz]

[ORC] LLJITTest: API test and usage examples with actual codegen

[lhames]

This looks great! Thanks Stefan!

The thread / LLVMContext creation code is awkward, but I think that's because ORC isn't really helping yet. I was thinking I might add an ExecutionSession::createCompileThread(...) utility to spawn a thread, provide a portable id that could be used for registering LLVMContexts, provide portable thread teardown, etc. This seems like a good example to apply the idea to to see if it simplifies things.

[lhames]

I do not think this landed in the end (it does not look like it). It may be able to be rewritten more naturally now that LLJIT has native support for multiple threads. Also ThreadSafeModule and ThreadSafeContext will help with the context management issues.

unittests/ExecutionEngine/Orc/CMakeLists.txt

 
 set(LLVM_LINK_COMPONENTS
   Core
   ExecutionEngine
+  IRReader
   Object
   OrcJIT
   RuntimeDyld
   Support
   native
   )
 
 add_llvm_unittest(OrcJITTests
   CompileOnDemandLayerTest.cpp
   CoreAPIsTest.cpp
   IndirectionUtilsTest.cpp
   GlobalMappingLayerTest.cpp
   LazyEmittingLayerTest.cpp
   LegacyAPIInteropTest.cpp
+  LLJITTest.cpp
   ObjectTransformLayerTest.cpp
   OrcCAPITest.cpp
   OrcTestCommon.cpp
   QueueChannel.cpp
   RemoteObjectLayerTest.cpp
   RPCUtilsTest.cpp
   RTDyldObjectLinkingLayerTest.cpp
   RTDyldObjectLinkingLayer2Test.cpp
   SymbolStringPoolTest.cpp
   )
 
 target_link_libraries(OrcJITTests PRIVATE ${ORC_JIT_TEST_LIBS})

unittests/ExecutionEngine/Orc/LLJITTest.cpp

+//===- LLJITTest.cpp - Unit tests for LLIJIT ------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
+#include "OrcTestCommon.h"
+
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Support/SourceMgr.h"
+
+#include "gtest/gtest.h"
+
+#include <memory>
+#include <thread>
+
+using namespace llvm;
+using namespace llvm::orc;
+
+//===----------------------------------------------------------------------===//
+//    Some utils shared by all tests
+//===----------------------------------------------------------------------===//
+
+#define LLVM_IR(...) #__VA_ARGS__
+
+const char *FooIR = LLVM_IR(
+  define i32 @foo() {
+    ret i32 1
+  }
+);
+
+const char *BarIR = LLVM_IR(
+  define i32 @bar() {
+    ret i32 1
+  }
+);
+
+const char *BuzIR = LLVM_IR(
+  declare i32 @foo()
+  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(bool Multithreaded) {
+  InitializeNativeTarget();
+  InitializeNativeTargetAsmParser();
+  InitializeNativeTargetAsmPrinter();
+
+  auto ES = llvm::make_unique<orc::ExecutionSession>();
+  auto JTMB = cantFail(JITTargetMachineBuilder::detectHost());
+  auto DL = cantFail(JTMB.createTargetMachine())->createDataLayout();
+  return cantFail(LLJIT::Create(std::move(ES), std::move(JTMB), std::move(DL),
+                                Multithreaded));
+}
+
+static int Exec(JITEvaluatedSymbol FuncSym) {
+  auto FuncPtr = (int (*)())FuncSym.getAddress();
+  return FuncPtr();
+}
+
+static void JoinAll(std::vector<std::thread> &Ts) {
+  for (std::thread &T : Ts)
+    if (T.joinable())
+      T.join();
+}
+
+//===----------------------------------------------------------------------===//
+
+TEST(LLJIT, UnthreadedSingleDylib) {
+  std::unique_ptr<LLJIT> Jit = Init(false);
+
+  SMDiagnostic E;
+  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])));
+
+  JITEvaluatedSymbol FooSym = cantFail(Jit->lookup("foo"));
+  EXPECT_EQ(1, Exec(FooSym));
+
+  JITEvaluatedSymbol BarSym = cantFail(Jit->lookup("bar"));
+  EXPECT_EQ(1, Exec(BarSym));
+
+  JITEvaluatedSymbol BuzSym = cantFail(Jit->lookup("buz"));
+  EXPECT_EQ(2, Exec(BuzSym));
+}
+
+//===----------------------------------------------------------------------===//
+
+TEST(LLJIT, UnthreadedMultiDylib) {
+  std::unique_ptr<LLJIT> Jit = Init(false);
+  ExecutionSession &ES = Jit->getExecutionSession();
+
+  JITDylib &Main = Jit->getMainJITDylib();
+  JITDylib &Extra1 = ES.createJITDylib("1");
+  JITDylib &Extra2 = ES.createJITDylib("2");
+
+  SMDiagnostic E;
+  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])));
+
+  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.
+  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(true);
+  ExecutionSession &ES = Jit->getExecutionSession();
+
+  SMDiagnostic E;
+  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 jobs to separate threads. Below we invoke lookups from
+  // runner threads. Eventually execution arrives here and spawns additional
+  // materialization threads. These must be joined back to the original runner
+  // threads, which is quite hard to track. Thus, we use thread-local scope-exit
+  // guards here. It's not a good idea when calling from the main thread though!
+  std::thread::id MainThreadId = std::this_thread::get_id();
+  ES.setDispatchMaterialization([&](JITDylib &JD,
+                                    std::unique_ptr<MaterializationUnit> MU) {
+    assert(std::this_thread::get_id() != MainThreadId &&
+           "Don't join via thread-local scope-exit "
+           "when calling from main thread!!!");
+
+    auto SMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
+    auto Worker =
+        std::make_shared<std::thread>([SMU, &JD]() { SMU->doMaterialize(JD); });
+
+    thread_local auto JoinWorker =
+        llvm::make_scope_exit([Worker]() { Worker->join(); });
+  });
+
+  // Run functions in different threads.
+  std::vector<std::thread> Ts;
+
+  Ts.emplace_back([&Jit]() {
+    JITEvaluatedSymbol BuzSym = cantFail(Jit->lookup("buz"));
+    EXPECT_EQ(2, Exec(BuzSym));
+  });
+
+  Ts.emplace_back([&Jit]() {
+    JITEvaluatedSymbol FooSym = cantFail(Jit->lookup("foo"));
+    EXPECT_EQ(1, Exec(FooSym));
+  });
+
+  Ts.emplace_back([&Jit]() {
+    JITEvaluatedSymbol BarSym = cantFail(Jit->lookup("bar"));
+    EXPECT_EQ(1, Exec(BarSym));
+  });
+
+  JoinAll(Ts);
+}
+
+//===----------------------------------------------------------------------===//
+
+TEST(LLJIT, ThreadedMultiDylib) {
+  std::unique_ptr<LLJIT> Jit = Init(true);
+  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<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 jobs to separate threads. Below we invoke lookups from
+  // runner threads. Eventually execution arrives here and spawns additional
+  // materialization threads. These must be joined back to the original runner
+  // threads, which is quite hard to track. Thus, we use thread-local scope-exit
+  // guards here. It's not a good idea when calling from the main thread though!
+  std::thread::id MainThreadId = std::this_thread::get_id();
+  ES.setDispatchMaterialization([&](JITDylib &JD,
+                                    std::unique_ptr<MaterializationUnit> MU) {
+    assert(std::this_thread::get_id() != MainThreadId &&
+           "Don't join via thread-local scope-exit "
+           "when calling from main thread!!!");
+
+    auto SMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
+    auto Worker =
+        std::make_shared<std::thread>([SMU, &JD]() { SMU->doMaterialize(JD); });
+
+    thread_local auto JoinWorker =
+        llvm::make_scope_exit([Worker]() { Worker->join(); });
+  });
+
+  // Allow buz to resolve foo and bar from the other dylibs.
+  Extra2.addToSearchOrder(Main);
+  Extra2.addToSearchOrder(Extra1);
+
+  // Run functions in different threads.
+  std::vector<std::thread> Ts;
+
+  Ts.emplace_back([&]() {
+    JITEvaluatedSymbol BuzSym = cantFail(Jit->lookup(Extra2, "buz"));
+    EXPECT_EQ(2, Exec(BuzSym));
+  });
+
+  Ts.emplace_back([&]() {
+    JITEvaluatedSymbol FooSym = cantFail(Jit->lookup("foo"));
+    EXPECT_EQ(1, Exec(FooSym));
+  });
+
+  Ts.emplace_back([&]() {
+    JITEvaluatedSymbol BarSym = cantFail(Jit->lookup(Extra1, "bar"));
+    EXPECT_EQ(1, Exec(BarSym));
+  });
+
+  JoinAll(Ts);
+}