Swift stable jit updates

llvm/docs/tutorial/BuildingAJIT4.rst

@@ -1,6 +1,6 @@
-===========================================================================
-Building a JIT: Extreme Laziness - Using Compile Callbacks to JIT from ASTs
-===========================================================================
+=======================================================================
+Building a JIT: Extreme Laziness - Using LazyReexports to JIT from ASTs
+=======================================================================
 
 .. contents::
    :local:
@@ -13,9 +13,9 @@ Chapter 4 Introduction
 ======================
 
 Welcome to Chapter 4 of the "Building an ORC-based JIT in LLVM" tutorial. This
-chapter introduces the Compile Callbacks and Indirect Stubs APIs and shows how
-they can be used to replace the CompileOnDemand layer from
-`Chapter 3 <BuildingAJIT3.html>`_ with a custom lazy-JITing scheme that JITs
+chapter introduces custom MaterializationUnits and Layers, and the lazy
+reexports API. Together these will be used to replace the CompileOnDemandLayer
+from `Chapter 3 <BuildingAJIT3.html>`_ with a custom lazy-JITing scheme that JITs
 directly from Kaleidoscope ASTs.
 
 **To be done:**

llvm/examples/Kaleidoscope/BuildingAJIT/CMakeLists.txt

@@ -3,6 +3,6 @@ add_subdirectory(Chapter2)
 add_subdirectory(Chapter3)
 add_subdirectory(Chapter4)
 
-if (NOT WIN32)
-  add_subdirectory(Chapter5)
-endif()
+# if (NOT WIN32)
+#  add_subdirectory(Chapter5)
+# endif()

llvm/examples/Kaleidoscope/BuildingAJIT/Chapter1/KaleidoscopeJIT.h

@@ -21,6 +21,7 @@
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/TargetProcessControl.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/LLVMContext.h"
@@ -31,53 +32,67 @@ namespace orc {
 
 class KaleidoscopeJIT {
 private:
-  ExecutionSession ES;
-  RTDyldObjectLinkingLayer ObjectLayer;
-  IRCompileLayer CompileLayer;
+  std::unique_ptr<TargetProcessControl> TPC;
+  std::unique_ptr<ExecutionSession> ES;
 
   DataLayout DL;
   MangleAndInterner Mangle;
-  ThreadSafeContext Ctx;
+
+  RTDyldObjectLinkingLayer ObjectLayer;
+  IRCompileLayer CompileLayer;
 
   JITDylib &MainJD;
 
 public:
-  KaleidoscopeJIT(JITTargetMachineBuilder JTMB, DataLayout DL)
-      : ObjectLayer(ES,
+  KaleidoscopeJIT(std::unique_ptr<TargetProcessControl> TPC,
+                  std::unique_ptr<ExecutionSession> ES,
+                  JITTargetMachineBuilder JTMB, DataLayout DL)
+      : TPC(std::move(TPC)), ES(std::move(ES)), DL(std::move(DL)),
+        Mangle(*this->ES, this->DL),
+        ObjectLayer(*this->ES,
                     []() { return std::make_unique<SectionMemoryManager>(); }),
-        CompileLayer(ES, ObjectLayer,
+        CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
-        DL(std::move(DL)), Mangle(ES, this->DL),
-        Ctx(std::make_unique<LLVMContext>()),
-        MainJD(ES.createBareJITDylib("<main>")) {
+        MainJD(this->ES->createBareJITDylib("<main>")) {
     MainJD.addGenerator(
         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
             DL.getGlobalPrefix())));
   }
 
+  ~KaleidoscopeJIT() {
+    if (auto Err = ES->endSession())
+      ES->reportError(std::move(Err));
+  }
+
   static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
-    auto JTMB = JITTargetMachineBuilder::detectHost();
+    auto TPC = SelfTargetProcessControl::Create();
+    if (!TPC)
+      return TPC.takeError();
+
+    auto ES = std::make_unique<ExecutionSession>();
 
-    if (!JTMB)
-      return JTMB.takeError();
+    JITTargetMachineBuilder JTMB((*TPC)->getTargetTriple());
 
-    auto DL = JTMB->getDefaultDataLayoutForTarget();
+    auto DL = JTMB.getDefaultDataLayoutForTarget();
     if (!DL)
       return DL.takeError();
 
-    return std::make_unique<KaleidoscopeJIT>(std::move(*JTMB), std::move(*DL));
+    return std::make_unique<KaleidoscopeJIT>(std::move(*TPC), std::move(ES),
+                                             std::move(JTMB), std::move(*DL));
   }
 
   const DataLayout &getDataLayout() const { return DL; }
 
-  LLVMContext &getContext() { return *Ctx.getContext(); }
+  JITDylib &getMainJITDylib() { return MainJD; }
 
-  Error addModule(std::unique_ptr<Module> M) {
-    return CompileLayer.add(MainJD, ThreadSafeModule(std::move(M), Ctx));
+  Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
+    if (!RT)
+      RT = MainJD.getDefaultResourceTracker();
+    return CompileLayer.add(RT, std::move(TSM));
   }
 
   Expected<JITEvaluatedSymbol> lookup(StringRef Name) {
-    return ES.lookup({&MainJD}, Mangle(Name.str()));
+    return ES->lookup({&MainJD}, Mangle(Name.str()));
   }
 };
 

llvm/examples/Kaleidoscope/BuildingAJIT/Chapter1/toy.cpp

@@ -676,11 +676,11 @@ static std::unique_ptr<FunctionAST> ParseDefinition() {
 }
 
 /// toplevelexpr ::= expression
-static std::unique_ptr<FunctionAST> ParseTopLevelExpr(unsigned ExprCount) {
+static std::unique_ptr<FunctionAST> ParseTopLevelExpr() {
   if (auto E = ParseExpression()) {
     // Make an anonymous proto.
-    auto Proto = std::make_unique<PrototypeAST>
-        (("__anon_expr" + Twine(ExprCount)).str(), std::vector<std::string>());
+    auto Proto = std::make_unique<PrototypeAST>("__anon_expr",
+                                                std::vector<std::string>());
     return std::make_unique<FunctionAST>(std::move(Proto), std::move(E));
   }
   return nullptr;
@@ -697,7 +697,7 @@ static std::unique_ptr<PrototypeAST> ParseExtern() {
 //===----------------------------------------------------------------------===//
 
 static std::unique_ptr<KaleidoscopeJIT> TheJIT;
-static LLVMContext *TheContext;
+static std::unique_ptr<LLVMContext> TheContext;
 static std::unique_ptr<IRBuilder<>> Builder;
 static std::unique_ptr<Module> TheModule;
 static std::map<std::string, AllocaInst *> NamedValues;
@@ -1102,7 +1102,8 @@ Function *FunctionAST::codegen() {
 //===----------------------------------------------------------------------===//
 
 static void InitializeModule() {
-  // Open a new module.
+  // Open a new context and module.
+  TheContext = std::make_unique<LLVMContext>();
   TheModule = std::make_unique<Module>("my cool jit", *TheContext);
   TheModule->setDataLayout(TheJIT->getDataLayout());
 
@@ -1116,7 +1117,8 @@ static void HandleDefinition() {
       fprintf(stderr, "Read function definition:");
       FnIR->print(errs());
       fprintf(stderr, "\n");
-      ExitOnErr(TheJIT->addModule(std::move(TheModule)));
+      auto TSM = ThreadSafeModule(std::move(TheModule), std::move(TheContext));
+      ExitOnErr(TheJIT->addModule(std::move(TSM)));
       InitializeModule();
     }
   } else {
@@ -1140,27 +1142,27 @@ static void HandleExtern() {
 }
 
 static void HandleTopLevelExpression() {
-  static unsigned ExprCount = 0;
-
-  // Update ExprCount. This number will be added to anonymous expressions to
-  // prevent them from clashing.
-  ++ExprCount;
-
   // Evaluate a top-level expression into an anonymous function.
-  if (auto FnAST = ParseTopLevelExpr(ExprCount)) {
+  if (auto FnAST = ParseTopLevelExpr()) {
     if (FnAST->codegen()) {
-      // JIT the module containing the anonymous expression, keeping a handle so
-      // we can free it later.
-      ExitOnErr(TheJIT->addModule(std::move(TheModule)));
+      // Create a ResourceTracker to track JIT'd memory allocated to our
+      // anonymous expression -- that way we can free it after executing.
+      auto RT = TheJIT->getMainJITDylib().createResourceTracker();
+
+      auto TSM = ThreadSafeModule(std::move(TheModule), std::move(TheContext));
+      ExitOnErr(TheJIT->addModule(std::move(TSM), RT));
       InitializeModule();
 
       // Get the anonymous expression's JITSymbol.
-      auto Sym =
-          ExitOnErr(TheJIT->lookup(("__anon_expr" + Twine(ExprCount)).str()));
+      auto Sym = ExitOnErr(TheJIT->lookup("__anon_expr"));
 
+      // Get the symbol's address and cast it to the right type (takes no
+      // arguments, returns a double) so we can call it as a native function.
       auto *FP = (double (*)())(intptr_t)Sym.getAddress();
-      assert(FP && "Failed to codegen function");
       fprintf(stderr, "Evaluated to %f\n", FP());
+
+      // Delete the anonymous expression module from the JIT.
+      ExitOnErr(RT->remove());
     }
   } else {
     // Skip token for error recovery.
@@ -1229,8 +1231,6 @@ int main() {
   getNextToken();
 
   TheJIT = ExitOnErr(KaleidoscopeJIT::Create());
-  TheContext = &TheJIT->getContext();
-
   InitializeModule();
 
   // Run the main "interpreter loop" now.

llvm/examples/Kaleidoscope/BuildingAJIT/Chapter2/KaleidoscopeJIT.h

@@ -22,6 +22,7 @@
 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/TargetProcessControl.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/LLVMContext.h"
@@ -36,54 +37,70 @@ namespace orc {
 
 class KaleidoscopeJIT {
 private:
-  ExecutionSession ES;
-  RTDyldObjectLinkingLayer ObjectLayer;
-  IRCompileLayer CompileLayer;
-  IRTransformLayer OptimizeLayer;
+  std::unique_ptr<TargetProcessControl> TPC;
+  std::unique_ptr<ExecutionSession> ES;
 
   DataLayout DL;
   MangleAndInterner Mangle;
-  ThreadSafeContext Ctx;
+
+  RTDyldObjectLinkingLayer ObjectLayer;
+  IRCompileLayer CompileLayer;
+  IRTransformLayer OptimizeLayer;
 
   JITDylib &MainJD;
 
 public:
-  KaleidoscopeJIT(JITTargetMachineBuilder JTMB, DataLayout DL)
-      : ObjectLayer(ES,
+  KaleidoscopeJIT(std::unique_ptr<TargetProcessControl> TPC,
+                  std::unique_ptr<ExecutionSession> ES,
+                  JITTargetMachineBuilder JTMB, DataLayout DL)
+      : TPC(std::move(TPC)), ES(std::move(ES)), DL(std::move(DL)),
+        Mangle(*this->ES, this->DL),
+        ObjectLayer(*this->ES,
                     []() { return std::make_unique<SectionMemoryManager>(); }),
-        CompileLayer(ES, ObjectLayer,
+        CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
-        OptimizeLayer(ES, CompileLayer, optimizeModule), DL(std::move(DL)),
-        Mangle(ES, this->DL), Ctx(std::make_unique<LLVMContext>()),
-        MainJD(ES.createBareJITDylib("<main>")) {
+        OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
+        MainJD(this->ES->createBareJITDylib("<main>")) {
     MainJD.addGenerator(
         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
             DL.getGlobalPrefix())));
   }
 
-  const DataLayout &getDataLayout() const { return DL; }
-
-  LLVMContext &getContext() { return *Ctx.getContext(); }
+  ~KaleidoscopeJIT() {
+    if (auto Err = ES->endSession())
+      ES->reportError(std::move(Err));
+  }
 
   static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
-    auto JTMB = JITTargetMachineBuilder::detectHost();
+    auto TPC = SelfTargetProcessControl::Create();
+    if (!TPC)
+      return TPC.takeError();
 
-    if (!JTMB)
-      return JTMB.takeError();
+    auto ES = std::make_unique<ExecutionSession>();
 
-    auto DL = JTMB->getDefaultDataLayoutForTarget();
+    JITTargetMachineBuilder JTMB((*TPC)->getTargetTriple());
+
+    auto DL = JTMB.getDefaultDataLayoutForTarget();
     if (!DL)
       return DL.takeError();
 
-    return std::make_unique<KaleidoscopeJIT>(std::move(*JTMB), std::move(*DL));
+    return std::make_unique<KaleidoscopeJIT>(std::move(*TPC), std::move(ES),
+                                             std::move(JTMB), std::move(*DL));
   }
 
-  Error addModule(std::unique_ptr<Module> M) {
-    return OptimizeLayer.add(MainJD, ThreadSafeModule(std::move(M), Ctx));
+  const DataLayout &getDataLayout() const { return DL; }
+
+  JITDylib &getMainJITDylib() { return MainJD; }
+
+  Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
+    if (!RT)
+      RT = MainJD.getDefaultResourceTracker();
+
+    return OptimizeLayer.add(RT, std::move(TSM));
   }
 
   Expected<JITEvaluatedSymbol> lookup(StringRef Name) {
-    return ES.lookup({&MainJD}, Mangle(Name.str()));
+    return ES->lookup({&MainJD}, Mangle(Name.str()));
   }
 
 private: