Updating Kaleidoscope tutorial to Synchronize with ORC v2 APIs

[pree-jackie]

Summary

Adds functionality to lazily compile LLVM IR to the Kaleidoscope JIT.

[pree-jackie]

toy.cpp is already updated for chapter 1&2.

[pree-jackie]

Adding IRTransformLayer to Optimize IR *Modules* Lazily

[pree-jackie]

User Documentation - Part 1

[pree-jackie]

Added user documentation part 2, It is ready to go now.

[pree-jackie]

ping,
Whether docs is okay?

[lhames]

This looks good to me. Thanks Praveen!

[pree-jackie]

Always Welcome

llvm/docs/tutorial/BuildingAJIT3.rst

 .. contents::
    :local:
 
-**This tutorial is under active development. It is incomplete and details may
-change frequently.** Nonetheless we invite you to try it out as it stands, and
-we welcome any feedback.
-
 Chapter 3 Introduction
 ======================
 
-**Warning: This text is currently out of date due to ORC API updates.**
-
-**The example code has been updated and can be used. The text will be updated
-once the API churn dies down.**
-
 Welcome to Chapter 3 of the "Building an ORC-based JIT in LLVM" tutorial. This
 chapter discusses lazy JITing and shows you how to enable it by adding an ORC
-CompileOnDemand layer the JIT from `Chapter 2 <BuildingAJIT2.html>`_.
+CompileOnDemand layer on the top of JIT from `Chapter 2 <BuildingAJIT2.html>`_.
 
 Lazy Compilation
 ================
 
 When we add a module to the KaleidoscopeJIT class from Chapter 2 it is
 immediately optimized, compiled and linked for us by the IRTransformLayer,
-IRCompileLayer and RTDyldObjectLinkingLayer respectively. This scheme, where all the
-work to make a Module executable is done up front, is simple to understand and
+IRCompileLayer and RTDyldObjectLinkingLayer respectively when a Symbol from the Materialization Unit 
+is looked up. This scheme, where all the work to make a Module executable is done up front, is simple to understand and
 its performance characteristics are easy to reason about. However, it will lead
 to very high startup times if the amount of code to be compiled is large, and
 may also do a lot of unnecessary compilation if only a few compiled functions

 the ORC APIs provide us with a layer to lazily compile LLVM IR:
 *CompileOnDemandLayer*.
 
-The CompileOnDemandLayer class conforms to the layer interface described in
-Chapter 2, but its addModule method behaves quite differently from the layers
-we have seen so far: rather than doing any work up front, it just scans the
-Modules being added and arranges for each function in them to be compiled the
-first time it is called. To do this, the CompileOnDemandLayer creates two small
-utilities for each function that it scans: a *stub* and a *compile
-callback*. The stub is a pair of a function pointer (which will be pointed at
-the function's implementation once the function has been compiled) and an
-indirect jump through the pointer. By fixing the address of the indirect jump
-for the lifetime of the program we can give the function a permanent "effective
-address", one that can be safely used for indirection and function pointer
-comparison even if the function's implementation is never compiled, or if it is
-compiled more than once (due to, for example, recompiling the function at a
-higher optimization level) and changes address. The second utility, the compile
-callback, represents a re-entry point from the program into the compiler that
-will trigger compilation and then execution of a function. By initializing the
-function's stub to point at the function's compile callback, we enable lazy
-compilation: The first attempted call to the function will follow the function
-pointer and trigger the compile callback instead. The compile callback will
-compile the function, update the function pointer for the stub, then execute
-the function. On all subsequent calls to the function, the function pointer
-will point at the already-compiled function, so there is no further overhead
-from the compiler. We will look at this process in more detail in the next
-chapter of this tutorial, but for now we'll trust the CompileOnDemandLayer to
-set all the stubs and callbacks up for us. All we need to do is to add the
-CompileOnDemandLayer to the top of our stack and we'll get the benefits of
-lazy compilation. We just need a few changes to the source:
+The CompileOnDemandLayer class conforms to the Layer interface described in chapter 2,
+When we lookup symbols through ExecutionSession, associated materialize method create stub definitions 
+in memory and return immediately without compiling. When they Stubs are called at runtime, stub will re-enter 
+to the JIT and trigger the actual compilation of the functions. So to compile a function that it added
+to compile on demand layer require two lookups : 1) Create a Stub 2) Trigger compilation
+
+On the first lookup of a symbol in a Materialization Unit, COD emit method gets called, it partitions the Callable and NonCallable 
+Symbols and creates a new JITDylib with name suffix ".impl",wraps the module into a PartitioningIRMaterializationUnit and move to ".impl" JITDylib.
+It then creates two facade materialization units "ReexportsMaterializationUnit" for NonCallable Symbols and "LazyReexportsMaterializationUnit" 
+for Callable Symbols to re export data symbols and lazy reexport function symbols respectively. These Materialization Unit takes the responsibilty 
+for materializing the Symbols it contains.
+Whenever a Symbol is lookedup, ExecutionSession finds which Materialization unit is responsible for materializing that symbol and trigger the 
+corresponding materialize method.
+
+LazyReexportsMaterializationUnit::Materialize method create stub definitions for the requested symbols. Whenever a Stub is called at runtime, ORC
+find out the which Symbol definition the Stubs exports and find the corresponding materialization unit (in this case, PartitioningIRMaterializationUnit)
+and call the PartitioningIRMaterializationUnit::materialize method, ORC will pack the requested function symbols into a new module and move then to the 
+layer below.
+
+
+To faciliate laziness, ORC uses two utilities: 1) IndirectStubManager 2) Lazy Call Through manager
+IndirectStubManager creates a Stubs for each symbol. 
 
 .. code-block:: c++
 

   #include "llvm/ExecutionEngine/SectionMemoryManager.h"
   #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
   #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
+  #include "llvm/ExecutionEngine/Orc/LazyReexports.h"
+  #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
   ...
 
   ...
   class KaleidoscopeJIT {
   private:
-    std::unique_ptr<TargetMachine> TM;
-    const DataLayout DL;
+    ExecutionSession ES;
+    std::unique_ptr<LazyCallThroughManager> LCTM;
     RTDyldObjectLinkingLayer ObjectLayer;
-    IRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer;
-
-    using OptimizeFunction =
-        std::function<std::shared_ptr<Module>(std::shared_ptr<Module>)>;
-
-    IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;
-
-    std::unique_ptr<JITCompileCallbackManager> CompileCallbackManager;
-    CompileOnDemandLayer<decltype(OptimizeLayer)> CODLayer;
-
-  public:
-    using ModuleHandle = decltype(CODLayer)::ModuleHandleT;
+    IRCompileLayer CompileLayer;
+    IRTransformLayer OptimizeIRLayer;
+    CompileOnDemandLayer CODLayer;
+    DataLayout DL;
+    MangleAndInterner Mangle;
+    ThreadSafeContext Ctx;
 
 First we need to include the CompileOnDemandLayer.h header, then add two new
-members: a std::unique_ptr<JITCompileCallbackManager> and a CompileOnDemandLayer,
-to our class. The CompileCallbackManager member is used by the CompileOnDemandLayer
-to create the compile callback needed for each function.
+members: a std::unique_ptr<LazyCallThroughManager> and a CompileOnDemandLayer,
+to our class. The LazyCallThroughManager member is used by the CompileOnDemandLayer
+to create the call through trampoline needed for each function.
 
 .. code-block:: c++
 
-  KaleidoscopeJIT()
-      : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
-        ObjectLayer([]() { return std::make_shared<SectionMemoryManager>(); }),
-        CompileLayer(ObjectLayer, SimpleCompiler(*TM)),
-        OptimizeLayer(CompileLayer,
-                      [this](std::shared_ptr<Module> M) {
-                        return optimizeModule(std::move(M));
-                      }),
-        CompileCallbackManager(
-            orc::createLocalCompileCallbackManager(TM->getTargetTriple(), 0)),
-        CODLayer(OptimizeLayer,
-                 [this](Function &F) { return std::set<Function*>({&F}); },
-                 *CompileCallbackManager,
-                 orc::createLocalIndirectStubsManagerBuilder(
-                   TM->getTargetTriple())) {
-    llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
+  KaleidoscopeJIT(): LCTM(cantFail(createLocalLazyCallThroughManager(JTMB.getTargetTriple(),
+                                                        this->ES, 0))),
+        ObjectLayer(ES,
+                    []() { return llvm::make_unique<SectionMemoryManager>(); }),
+        CompileLayer(ES, ObjectLayer, ConcurrentIRCompiler(std::move(JTMB))),
+        OptimizeIRLayer(ES, CompileLayer, optimizeModule),
+        CODLayer(this->ES, OptimizeIRLayer, *LCTM,
+                 createLocalIndirectStubsManagerBuilder(T)),
+        DL(std::move(DL)), Mangle(ES, this->DL),
+        Ctx(llvm::make_unique<LLVMContext>()) {
+    ES.getMainJITDylib().setGenerator(
+        cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
+            DL.getGlobalPrefix())));
   }
 
 Next we have to update our constructor to initialize the new members. To create
 an appropriate compile callback manager we use the
-createLocalCompileCallbackManager function, which takes a TargetMachine and a
-JITTargetAddress to call if it receives a request to compile an unknown
-function.  In our simple JIT this situation is unlikely to come up, so we'll
-cheat and just pass '0' here. In a production quality JIT you could give the
-address of a function that throws an exception in order to unwind the JIT'd
-code's stack.
+createLocalLazyCallThroughManager function, and we use createLocalIndirectStubsManagerBuilder
+to create IndirectStub Builder.
 
 Now we can construct our CompileOnDemandLayer. Following the pattern from
 previous layers we start by passing a reference to the next layer down in our

 memory allocated to the stubs. We supply this function by using the
 createLocalIndirectStubsManagerBuilder utility.
 
-.. code-block:: c++
-
-  // ...
-          if (auto Sym = CODLayer.findSymbol(Name, false))
-  // ...
-  return cantFail(CODLayer.addModule(std::move(Ms),
-                                     std::move(Resolver)));
-  // ...
-
-  // ...
-  return CODLayer.findSymbol(MangledNameStream.str(), true);
-  // ...
-
-  // ...
-  CODLayer.removeModule(H);
-  // ...
-
-Finally, we need to replace the references to OptimizeLayer in our addModule,
-findSymbol, and removeModule methods. With that, we're up and running.
-
 **To be done:**
 
 ** Chapter conclusion.**

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

 #ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
 #define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
 
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ExecutionEngine/ExecutionEngine.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
-#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
+#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/LazyReexports.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
-#include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
-#include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassManager.h"
-#include "llvm/IR/Mangler.h"
-#include "llvm/Support/DynamicLibrary.h"
-#include "llvm/Support/raw_ostream.h"
-#include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Scalar/GVN.h"
-#include <algorithm>
-#include <map>
 #include <memory>
-#include <set>
-#include <string>
-#include <vector>
 
 namespace llvm {
 namespace orc {

 class KaleidoscopeJIT {
 private:
   ExecutionSession ES;
-  std::map<VModuleKey, std::shared_ptr<SymbolResolver>> Resolvers;
-  std::unique_ptr<TargetMachine> TM;
-  const DataLayout DL;
-  LegacyRTDyldObjectLinkingLayer ObjectLayer;
-  LegacyIRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer;
-
-  using OptimizeFunction =
-      std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)>;
-
-  LegacyIRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;
-
-  std::unique_ptr<JITCompileCallbackManager> CompileCallbackManager;
-  LegacyCompileOnDemandLayer<decltype(OptimizeLayer)> CODLayer;
+  std::unique_ptr<LazyCallThroughManager> LCTM;
+  RTDyldObjectLinkingLayer ObjectLayer;
+  IRCompileLayer CompileLayer;
+  IRTransformLayer OptimizeIRLayer;
+  CompileOnDemandLayer CODLayer;
+  DataLayout DL;
+  MangleAndInterner Mangle;
+  ThreadSafeContext Ctx;
 
 public:
-  KaleidoscopeJIT()
-      : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
+  KaleidoscopeJIT(JITTargetMachineBuilder JTMB, DataLayout DL, Triple T)
+      : LCTM(cantFail(createLocalLazyCallThroughManager(JTMB.getTargetTriple(),
+                                                        this->ES, 0))),
         ObjectLayer(ES,
-                    [this](VModuleKey K) {
-                      return LegacyRTDyldObjectLinkingLayer::Resources{
-                          std::make_shared<SectionMemoryManager>(),
-                          Resolvers[K]};
-                    }),
-        CompileLayer(ObjectLayer, SimpleCompiler(*TM)),
-        OptimizeLayer(CompileLayer,
-                      [this](std::unique_ptr<Module> M) {
-                        return optimizeModule(std::move(M));
-                      }),
-        CompileCallbackManager(cantFail(orc::createLocalCompileCallbackManager(
-            TM->getTargetTriple(), ES, 0))),
-        CODLayer(ES, OptimizeLayer,
-                 [&](orc::VModuleKey K) { return Resolvers[K]; },
-                 [&](orc::VModuleKey K, std::shared_ptr<SymbolResolver> R) {
-                   Resolvers[K] = std::move(R);
-                 },
-                 [](Function &F) { return std::set<Function *>({&F}); },
-                 *CompileCallbackManager,
-                 orc::createLocalIndirectStubsManagerBuilder(
-                     TM->getTargetTriple())) {
-    llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
+                    []() { return llvm::make_unique<SectionMemoryManager>(); }),
+        CompileLayer(ES, ObjectLayer, ConcurrentIRCompiler(std::move(JTMB))),
+        OptimizeIRLayer(ES,CompileLayer,optimizeModule),
+        CODLayer(this->ES, OptimizeIRLayer, *LCTM,
+                 createLocalIndirectStubsManagerBuilder(T)),
+        DL(std::move(DL)), Mangle(ES, this->DL),
+        Ctx(llvm::make_unique<LLVMContext>()) {
+    ES.getMainJITDylib().setGenerator(
+        cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
+            DL.getGlobalPrefix())));
   }
 
-  TargetMachine &getTargetMachine() { return *TM; }
-
-  VModuleKey addModule(std::unique_ptr<Module> M) {
-    // Create a new VModuleKey.
-    VModuleKey K = ES.allocateVModule();
-
-    // Build a resolver and associate it with the new key.
-    Resolvers[K] = createLegacyLookupResolver(
-        ES,
-        [this](const std::string &Name) -> JITSymbol {
-          if (auto Sym = CompileLayer.findSymbol(Name, false))
-            return Sym;
-          else if (auto Err = Sym.takeError())
-            return std::move(Err);
-          if (auto SymAddr =
-                  RTDyldMemoryManager::getSymbolAddressInProcess(Name))
-            return JITSymbol(SymAddr, JITSymbolFlags::Exported);
-          return nullptr;
-        },
-        [](Error Err) { cantFail(std::move(Err), "lookupFlags failed"); });
-
-    // Add the module to the JIT with the new key.
-    cantFail(CODLayer.addModule(K, std::move(M)));
-    return K;
+  static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
+    auto JTMB = JITTargetMachineBuilder::detectHost();
+
+    if (!JTMB)
+      return JTMB.takeError();
+
+    auto DL = JTMB->getDefaultDataLayoutForTarget();
+    if (!DL)
+      return DL.takeError();
+
+    auto T = JTMB->getTargetTriple();
+    return llvm::make_unique<KaleidoscopeJIT>(std::move(*JTMB), std::move(*DL),
+                                              std::move(T));
   }
 
-  JITSymbol findSymbol(const std::string Name) {
-    std::string MangledName;
-    raw_string_ostream MangledNameStream(MangledName);
-    Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
-    return CODLayer.findSymbol(MangledNameStream.str(), true);
+  const DataLayout &getDataLayout() const { return DL; }
+
+  LLVMContext &getContext() { return *Ctx.getContext(); }
+
+  Error addModule(std::unique_ptr<Module> M) {
+    return CODLayer.add(ES.getMainJITDylib(),
+                        ThreadSafeModule(std::move(M), Ctx));
   }
 
-  void removeModule(VModuleKey K) {
-    cantFail(CODLayer.removeModule(K));
+  Expected<JITEvaluatedSymbol> lookup(StringRef Name) {
+    return ES.lookup({&ES.getMainJITDylib()}, Mangle(Name.str()));
   }
 
-private:
-  std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) {
-    // Create a function pass manager.
-    auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get());
+  void dumpState() { ES.dump(llvm::errs()); }
+  
+  private:
+  static Expected<ThreadSafeModule>
+  optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
+    // Create a Legacy function pass manager.
+    auto FPM = llvm::make_unique<legacy::FunctionPassManager>(TSM.getModule());
 
     // Add some optimizations.
     FPM->add(createInstructionCombiningPass());

     FPM->add(createCFGSimplificationPass());
     FPM->doInitialization();
 
-    // Run the optimizations over all functions in the module being added to
-    // the JIT.
-    for (auto &F : *M)
+    // Run the optimizations over functions that are packed up by COD and added to OptimizeIRLayer
+    for (auto &F : *TSM.getModule())
       FPM->run(F);
-
-    return M;
+    llvm::errs() << *TSM.getModule();
+    return TSM;
   }
 };
 

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

 }
 
 /// toplevelexpr ::= expression
-static std::unique_ptr<FunctionAST> ParseTopLevelExpr() {
+static std::unique_ptr<FunctionAST> ParseTopLevelExpr(unsigned ExprCount) {
   if (auto E = ParseExpression()) {
     // Make an anonymous proto.
-    auto Proto = llvm::make_unique<PrototypeAST>("__anon_expr",
-                                                 std::vector<std::string>());
+    auto Proto = llvm::make_unique<PrototypeAST>(
+        ("__anon_expr" + Twine(ExprCount)).str(), std::vector<std::string>());
     return llvm::make_unique<FunctionAST>(std::move(Proto), std::move(E));
   }
   return nullptr;

 // Code Generation
 //===----------------------------------------------------------------------===//
 
-static LLVMContext TheContext;
-static IRBuilder<> Builder(TheContext);
+static std::unique_ptr<KaleidoscopeJIT> TheJIT;
+static LLVMContext *TheContext;
+static std::unique_ptr<IRBuilder<>> Builder;
 static std::unique_ptr<Module> TheModule;
 static std::map<std::string, AllocaInst *> NamedValues;
-static std::unique_ptr<KaleidoscopeJIT> TheJIT;
 static std::map<std::string, std::unique_ptr<PrototypeAST>> FunctionProtos;
+static ExitOnError ExitOnErr;
 
 Value *LogErrorV(const char *Str) {
   LogError(Str);

                                           const std::string &VarName) {
   IRBuilder<> TmpB(&TheFunction->getEntryBlock(),
                    TheFunction->getEntryBlock().begin());
-  return TmpB.CreateAlloca(Type::getDoubleTy(TheContext), nullptr, VarName);
+  return TmpB.CreateAlloca(Type::getDoubleTy(*TheContext), nullptr, VarName);
 }
 
 Value *NumberExprAST::codegen() {
-  return ConstantFP::get(TheContext, APFloat(Val));
+  return ConstantFP::get(*TheContext, APFloat(Val));
 }
 
 Value *VariableExprAST::codegen() {

     return LogErrorV("Unknown variable name");
 
   // Load the value.
-  return Builder.CreateLoad(V, Name.c_str());
+  return Builder->CreateLoad(V, Name.c_str());
 }
 
 Value *UnaryExprAST::codegen() {

   if (!F)
     return LogErrorV("Unknown unary operator");
 
-  return Builder.CreateCall(F, OperandV, "unop");
+  return Builder->CreateCall(F, OperandV, "unop");
 }
 
 Value *BinaryExprAST::codegen() {

     if (!Variable)
       return LogErrorV("Unknown variable name");
 
-    Builder.CreateStore(Val, Variable);
+    Builder->CreateStore(Val, Variable);
     return Val;
   }
 

 
   switch (Op) {
   case '+':
-    return Builder.CreateFAdd(L, R, "addtmp");
+    return Builder->CreateFAdd(L, R, "addtmp");
   case '-':
-    return Builder.CreateFSub(L, R, "subtmp");
+    return Builder->CreateFSub(L, R, "subtmp");
   case '*':
-    return Builder.CreateFMul(L, R, "multmp");
+    return Builder->CreateFMul(L, R, "multmp");
   case '<':
-    L = Builder.CreateFCmpULT(L, R, "cmptmp");
+    L = Builder->CreateFCmpULT(L, R, "cmptmp");
     // Convert bool 0/1 to double 0.0 or 1.0
-    return Builder.CreateUIToFP(L, Type::getDoubleTy(TheContext), "booltmp");
+    return Builder->CreateUIToFP(L, Type::getDoubleTy(*TheContext), "booltmp");
   default:
     break;
   }

   assert(F && "binary operator not found!");
 
   Value *Ops[] = {L, R};
-  return Builder.CreateCall(F, Ops, "binop");
+  return Builder->CreateCall(F, Ops, "binop");
 }
 
 Value *CallExprAST::codegen() {

       return nullptr;
   }
 
-  return Builder.CreateCall(CalleeF, ArgsV, "calltmp");
+  return Builder->CreateCall(CalleeF, ArgsV, "calltmp");
 }
 
 Value *IfExprAST::codegen() {

     return nullptr;
 
   // Convert condition to a bool by comparing equal to 0.0.
-  CondV = Builder.CreateFCmpONE(
-      CondV, ConstantFP::get(TheContext, APFloat(0.0)), "ifcond");
+  CondV = Builder->CreateFCmpONE(
+      CondV, ConstantFP::get(*TheContext, APFloat(0.0)), "ifcond");
 
-  Function *TheFunction = Builder.GetInsertBlock()->getParent();
+  Function *TheFunction = Builder->GetInsertBlock()->getParent();
 
   // Create blocks for the then and else cases.  Insert the 'then' block at the
   // end of the function.
-  BasicBlock *ThenBB = BasicBlock::Create(TheContext, "then", TheFunction);
-  BasicBlock *ElseBB = BasicBlock::Create(TheContext, "else");
-  BasicBlock *MergeBB = BasicBlock::Create(TheContext, "ifcont");
+  BasicBlock *ThenBB = BasicBlock::Create(*TheContext, "then", TheFunction);
+  BasicBlock *ElseBB = BasicBlock::Create(*TheContext, "else");
+  BasicBlock *MergeBB = BasicBlock::Create(*TheContext, "ifcont");
 
-  Builder.CreateCondBr(CondV, ThenBB, ElseBB);
+  Builder->CreateCondBr(CondV, ThenBB, ElseBB);
 
   // Emit then value.
-  Builder.SetInsertPoint(ThenBB);
+  Builder->SetInsertPoint(ThenBB);
 
   Value *ThenV = Then->codegen();
   if (!ThenV)
     return nullptr;
 
-  Builder.CreateBr(MergeBB);
+  Builder->CreateBr(MergeBB);
   // Codegen of 'Then' can change the current block, update ThenBB for the PHI.
-  ThenBB = Builder.GetInsertBlock();
+  ThenBB = Builder->GetInsertBlock();
 
   // Emit else block.
   TheFunction->getBasicBlockList().push_back(ElseBB);
-  Builder.SetInsertPoint(ElseBB);
+  Builder->SetInsertPoint(ElseBB);
 
   Value *ElseV = Else->codegen();
   if (!ElseV)
     return nullptr;
 
-  Builder.CreateBr(MergeBB);
+  Builder->CreateBr(MergeBB);
   // Codegen of 'Else' can change the current block, update ElseBB for the PHI.
-  ElseBB = Builder.GetInsertBlock();
+  ElseBB = Builder->GetInsertBlock();
 
   // Emit merge block.
   TheFunction->getBasicBlockList().push_back(MergeBB);
-  Builder.SetInsertPoint(MergeBB);
-  PHINode *PN = Builder.CreatePHI(Type::getDoubleTy(TheContext), 2, "iftmp");
+  Builder->SetInsertPoint(MergeBB);
+  PHINode *PN = Builder->CreatePHI(Type::getDoubleTy(*TheContext), 2, "iftmp");
 
   PN->addIncoming(ThenV, ThenBB);
   PN->addIncoming(ElseV, ElseBB);

 //   br endcond, loop, endloop
 // outloop:
 Value *ForExprAST::codegen() {
-  Function *TheFunction = Builder.GetInsertBlock()->getParent();
+  Function *TheFunction = Builder->GetInsertBlock()->getParent();
 
   // Create an alloca for the variable in the entry block.
   AllocaInst *Alloca = CreateEntryBlockAlloca(TheFunction, VarName);

     return nullptr;
 
   // Store the value into the alloca.
-  Builder.CreateStore(StartVal, Alloca);
+  Builder->CreateStore(StartVal, Alloca);
 
   // Make the new basic block for the loop header, inserting after current
   // block.
-  BasicBlock *LoopBB = BasicBlock::Create(TheContext, "loop", TheFunction);
+  BasicBlock *LoopBB = BasicBlock::Create(*TheContext, "loop", TheFunction);
 
   // Insert an explicit fall through from the current block to the LoopBB.
-  Builder.CreateBr(LoopBB);
+  Builder->CreateBr(LoopBB);
 
   // Start insertion in LoopBB.
-  Builder.SetInsertPoint(LoopBB);
+  Builder->SetInsertPoint(LoopBB);
 
   // Within the loop, the variable is defined equal to the PHI node.  If it
   // shadows an existing variable, we have to restore it, so save it now.

       return nullptr;
   } else {
     // If not specified, use 1.0.
-    StepVal = ConstantFP::get(TheContext, APFloat(1.0));
+    StepVal = ConstantFP::get(*TheContext, APFloat(1.0));
   }
 
   // Compute the end condition.

 
   // Reload, increment, and restore the alloca.  This handles the case where
   // the body of the loop mutates the variable.
-  Value *CurVar = Builder.CreateLoad(Alloca, VarName.c_str());
-  Value *NextVar = Builder.CreateFAdd(CurVar, StepVal, "nextvar");
-  Builder.CreateStore(NextVar, Alloca);
+  Value *CurVar = Builder->CreateLoad(Alloca, VarName.c_str());
+  Value *NextVar = Builder->CreateFAdd(CurVar, StepVal, "nextvar");
+  Builder->CreateStore(NextVar, Alloca);
 
   // Convert condition to a bool by comparing equal to 0.0.
-  EndCond = Builder.CreateFCmpONE(
-      EndCond, ConstantFP::get(TheContext, APFloat(0.0)), "loopcond");
+  EndCond = Builder->CreateFCmpONE(
+      EndCond, ConstantFP::get(*TheContext, APFloat(0.0)), "loopcond");
 
   // Create the "after loop" block and insert it.
   BasicBlock *AfterBB =
-      BasicBlock::Create(TheContext, "afterloop", TheFunction);
+      BasicBlock::Create(*TheContext, "afterloop", TheFunction);
 
   // Insert the conditional branch into the end of LoopEndBB.
-  Builder.CreateCondBr(EndCond, LoopBB, AfterBB);
+  Builder->CreateCondBr(EndCond, LoopBB, AfterBB);
 
   // Any new code will be inserted in AfterBB.
-  Builder.SetInsertPoint(AfterBB);
+  Builder->SetInsertPoint(AfterBB);
 
   // Restore the unshadowed variable.
   if (OldVal)

     NamedValues.erase(VarName);
 
   // for expr always returns 0.0.
-  return Constant::getNullValue(Type::getDoubleTy(TheContext));
+  return Constant::getNullValue(Type::getDoubleTy(*TheContext));
 }
 
 Value *VarExprAST::codegen() {
   std::vector<AllocaInst *> OldBindings;
 
-  Function *TheFunction = Builder.GetInsertBlock()->getParent();
+  Function *TheFunction = Builder->GetInsertBlock()->getParent();
 
   // Register all variables and emit their initializer.
   for (unsigned i = 0, e = VarNames.size(); i != e; ++i) {

       if (!InitVal)
         return nullptr;
     } else { // If not specified, use 0.0.
-      InitVal = ConstantFP::get(TheContext, APFloat(0.0));
+      InitVal = ConstantFP::get(*TheContext, APFloat(0.0));
     }
 
     AllocaInst *Alloca = CreateEntryBlockAlloca(TheFunction, VarName);
-    Builder.CreateStore(InitVal, Alloca);
+    Builder->CreateStore(InitVal, Alloca);
 
     // Remember the old variable binding so that we can restore the binding when
     // we unrecurse.

 
 Function *PrototypeAST::codegen() {
   // Make the function type:  double(double,double) etc.
-  std::vector<Type *> Doubles(Args.size(), Type::getDoubleTy(TheContext));
+  std::vector<Type *> Doubles(Args.size(), Type::getDoubleTy(*TheContext));
   FunctionType *FT =
-      FunctionType::get(Type::getDoubleTy(TheContext), Doubles, false);
+      FunctionType::get(Type::getDoubleTy(*TheContext), Doubles, false);
 
   Function *F =
       Function::Create(FT, Function::ExternalLinkage, Name, TheModule.get());

     BinopPrecedence[P.getOperatorName()] = P.getBinaryPrecedence();
 
   // Create a new basic block to start insertion into.
-  BasicBlock *BB = BasicBlock::Create(TheContext, "entry", TheFunction);
-  Builder.SetInsertPoint(BB);
+  BasicBlock *BB = BasicBlock::Create(*TheContext, "entry", TheFunction);
+  Builder->SetInsertPoint(BB);
 
   // Record the function arguments in the NamedValues map.
   NamedValues.clear();

     AllocaInst *Alloca = CreateEntryBlockAlloca(TheFunction, Arg.getName());
 
     // Store the initial value into the alloca.
-    Builder.CreateStore(&Arg, Alloca);
+    Builder->CreateStore(&Arg, Alloca);
 
     // Add arguments to variable symbol table.
     NamedValues[Arg.getName()] = Alloca;

 
   if (Value *RetVal = Body->codegen()) {
     // Finish off the function.
-    Builder.CreateRet(RetVal);
+    Builder->CreateRet(RetVal);
 
     // Validate the generated code, checking for consistency.
     verifyFunction(*TheFunction);

 
 static void InitializeModule() {
   // Open a new module.
-  TheModule = llvm::make_unique<Module>("my cool jit", TheContext);
-  TheModule->setDataLayout(TheJIT->getTargetMachine().createDataLayout());
+  TheModule = llvm::make_unique<Module>("my cool jit", *TheContext);
+  TheModule->setDataLayout(TheJIT->getDataLayout());
+
+  // Create a new builder for the module.
+  Builder = llvm::make_unique<IRBuilder<>>(*TheContext);
 }
 
 static void HandleDefinition() {

       fprintf(stderr, "Read function definition:");
       FnIR->print(errs());
       fprintf(stderr, "\n");
-      TheJIT->addModule(std::move(TheModule));
+      ExitOnErr(TheJIT->addModule(std::move(TheModule)));
       InitializeModule();
     }
   } else {

 }
 
 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()) {
+  if (auto FnAST = ParseTopLevelExpr(ExprCount)) {
     if (FnAST->codegen()) {
       // JIT the module containing the anonymous expression, keeping a handle so
       // we can free it later.
-      auto H = TheJIT->addModule(std::move(TheModule));
+      ExitOnErr(TheJIT->addModule(std::move(TheModule)));
       InitializeModule();
-
-      // Search the JIT for the __anon_expr symbol.
-      auto ExprSymbol = TheJIT->findSymbol("__anon_expr");
-      assert(ExprSymbol && "Function not found");
-
-      // 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.
-      double (*FP)() = (double (*)())(intptr_t)cantFail(ExprSymbol.getAddress());
+      // Get the anonymous expression's JITSymbol.
+      auto Sym =
+          ExitOnErr(TheJIT->lookup(("__anon_expr" + Twine(ExprCount)).str()));
+      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.
-      TheJIT->removeModule(H);
     }
   } else {
     // Skip token for error recovery.

   fprintf(stderr, "ready> ");
   getNextToken();
 
-  TheJIT = llvm::make_unique<KaleidoscopeJIT>();
+  TheJIT = ExitOnErr(KaleidoscopeJIT::Create());
+  TheContext = &TheJIT->getContext();
 
   InitializeModule();