[ORC] Simplify use of lazyReexports with LLJIT.

This patch makes the target triple available via the LLJIT interface, and moves
the IRTransformLayer from LLLazyJIT down into LLJIT. Together these changes make
it easier to use the lazyReexports utility with LLJIT, and to apply IR
transforms to code as it is compiled in LLJIT (rather than requiring transforms
to be applied manually before code is added). An code example is added in
llvm/examples/LLJITExamples/LLJITWithLazyReexports

examples/LLJITExamples/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(LLJITDumpObjects)
 add_subdirectory(LLJITWithObjectCache)
 add_subdirectory(LLJITWithCustomObjectLinkingLayer)
+add_subdirectory(LLJITWithLazyReexports)

examples/LLJITExamples/LLJITWithLazyReexports/CMakeLists.txt

@@ -0,0 +1,12 @@
+set(LLVM_LINK_COMPONENTS
+  Core
+  ExecutionEngine
+  IRReader
+  OrcJIT
+  Support
+  nativecodegen
+  )
+
+add_llvm_example(LLJITWithLazyReexports
+  LLJITWithLazyReexports.cpp
+  )

examples/LLJITExamples/LLJITWithLazyReexports/LLJITWithLazyReexports.cpp

@@ -0,0 +1,163 @@
+//===--- LLJITWithLazyReexports.cpp - LLJIT example with custom laziness --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// In this example we will use the lazy re-exports utility to lazily compile
+// IR modules. We will do this in seven steps:
+//
+// 1. Create an LLJIT instance.
+// 2. Install a transform so that we is being compiled.
+// 3. Create an indirect stubs manager and lazy call-through manager.
+// 4. Add two modules that will be conditionally compiled, plus a main module.
+// 5. Add lazy-rexports of the symbols in the conditionally compiled modules.
+// 6. Dump the ExecutionSession state to see the symbol table prior to
+//    executing any code.
+// 7. Verify that only modules containing executed code are compiled.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "../ExampleModules.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+ExitOnError ExitOnErr;
+
+// Example IR modules.
+//
+// Note that in the conditionally compiled modules, FooMod and BarMod, functions
+// have been given an _body suffix. This is to ensure that their names do not
+// clash with their lazy-reexports.
+// For clients who do not wish to rename function bodies (e.g. because they want
+// to re-use cached objects between static and JIT compiles) techniques exist to
+// avoid renaming. See the lazy-reexports section of the ORCv2 design doc.
+
+const llvm::StringRef FooMod =
+    R"(
+  define i32 @foo_body() {
+  entry:
+    ret i32 1
+  }
+)";
+
+const llvm::StringRef BarMod =
+    R"(
+  define i32 @bar_body(i32 %x) {
+  entry:
+    ret i32 2
+  }
+)";
+
+const llvm::StringRef MainMod =
+    R"(
+
+  define i32 @entry(i32 %argc) {
+  entry:
+    %and = and i32 %argc, 1
+    %tobool = icmp eq i32 %and, 0
+    br i1 %tobool, label %if.end, label %if.then
+
+  if.then:                                          ; preds = %entry
+    %call = tail call i32 @foo() #2
+    br label %return
+
+  if.end:                                           ; preds = %entry
+    %call1 = tail call i32 @bar() #2
+    br label %return
+
+  return:                                           ; preds = %if.end, %if.then
+    %retval.0 = phi i32 [ %call, %if.then ], [ %call1, %if.end ]
+    ret i32 %retval.0
+  }
+
+  declare i32 @foo()
+  declare i32 @bar()
+)";
+
+cl::list<std::string> InputArgv(cl::Positional,
+                                cl::desc("<program arguments>..."));
+
+int main(int argc, char *argv[]) {
+  // Initialize LLVM.
+  InitLLVM X(argc, argv);
+
+  InitializeNativeTarget();
+  InitializeNativeTargetAsmPrinter();
+
+  cl::ParseCommandLineOptions(argc, argv, "LLJITWithLazyReexports");
+  ExitOnErr.setBanner(std::string(argv[0]) + ": ");
+
+  // (1) Create LLJIT instance.
+  auto J = ExitOnErr(LLJITBuilder().create());
+
+  // (2) Install transform to print modules as they are compiled:
+  J->getIRTransformLayer().setTransform(
+      [](ThreadSafeModule TSM,
+         const MaterializationResponsibility &R) -> Expected<ThreadSafeModule> {
+        TSM.withModuleDo([](Module &M) { dbgs() << "---Compiling---\n" << M; });
+        return TSM;
+      });
+
+  // (3) Create stubs and call-through managers:
+  std::unique_ptr<IndirectStubsManager> ISM;
+  {
+    auto ISMBuilder =
+        createLocalIndirectStubsManagerBuilder(J->getTargetTriple());
+    if (!ISMBuilder())
+      ExitOnErr(make_error<StringError>("Could not create stubs manager for " +
+                                            J->getTargetTriple().str(),
+                                        inconvertibleErrorCode()));
+    ISM = ISMBuilder();
+  }
+  auto LCTM = ExitOnErr(createLocalLazyCallThroughManager(
+      J->getTargetTriple(), J->getExecutionSession(), 0));
+
+  // (4) Add modules.
+  ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(FooMod, "foo-mod"))));
+  ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(BarMod, "bar-mod"))));
+  ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(MainMod, "main-mod"))));
+
+  // (5) Add lazy reexports.
+  MangleAndInterner Mangle(J->getExecutionSession(), J->getDataLayout());
+  SymbolAliasMap ReExports(
+      {{Mangle("foo"),
+        {Mangle("foo_body"),
+         JITSymbolFlags::Exported | JITSymbolFlags::Callable}},
+       {Mangle("bar"),
+        {Mangle("bar_body"),
+         JITSymbolFlags::Exported | JITSymbolFlags::Callable}}});
+  ExitOnErr(J->getMainJITDylib().define(
+      lazyReexports(*LCTM, *ISM, J->getMainJITDylib(), std::move(ReExports))));
+
+  // (6) Dump the ExecutionSession state.
+  dbgs() << "---Session state---\n";
+  J->getExecutionSession().dump(dbgs());
+  dbgs() << "\n";
+
+  // (7) Execute the JIT'd main function and pass the example's command line
+  // arguments unmodified. This should cause either ExampleMod1 or ExampleMod2
+  // to be compiled, and either "1" or "2" returned depending on the number of
+  // arguments passed.
+
+  // Look up the JIT'd function, cast it to a function pointer, then call it.
+  auto EntrySym = ExitOnErr(J->lookup("entry"));
+  auto *Entry = (int (*)(int))EntrySym.getAddress();
+
+  int Result = Entry(argc);
+  outs() << "---Result---\n"
+         << "entry(" << argc << ") = " << Result << "\n";
+
+  return 0;
+}

include/llvm/ExecutionEngine/Orc/LLJIT.h

@@ -45,6 +45,9 @@ public:
   /// Returns the ExecutionSession for this instance.
   ExecutionSession &getExecutionSession() { return *ES; }
 
+  /// Returns a reference to the triple for this instance.
+  const Triple &getTargetTriple() const { return TT; }
+
   /// Returns a reference to the DataLayout for this instance.
   const DataLayout &getDataLayout() const { return DL; }
 
@@ -120,6 +123,9 @@ public:
   /// Returns a reference to the object transform layer.
   ObjectTransformLayer &getObjTransformLayer() { return ObjTransformLayer; }
 
+  /// Returns a reference to the IR transform layer.
+  IRTransformLayer &getIRTransformLayer() { return *TransformLayer; }
+
 protected:
   static std::unique_ptr<ObjectLayer>
   createObjectLinkingLayer(LLJITBuilderState &S, ExecutionSession &ES);
@@ -140,11 +146,13 @@ protected:
   JITDylib &Main;
 
   DataLayout DL;
+  Triple TT;
   std::unique_ptr<ThreadPool> CompileThreads;
 
   std::unique_ptr<ObjectLayer> ObjLinkingLayer;
   ObjectTransformLayer ObjTransformLayer;
   std::unique_ptr<IRCompileLayer> CompileLayer;
+  std::unique_ptr<IRTransformLayer> TransformLayer;
 
   CtorDtorRunner CtorRunner, DtorRunner;
 };
@@ -156,12 +164,6 @@ class LLLazyJIT : public LLJIT {
 
 public:
 
-  /// Set an IR transform (e.g. pass manager pipeline) to run on each function
-  /// when it is compiled.
-  void setLazyCompileTransform(IRTransformLayer::TransformFunction Transform) {
-    TransformLayer->setTransform(std::move(Transform));
-  }
-
   /// Sets the partition function.
   void
   setPartitionFunction(CompileOnDemandLayer::PartitionFunction Partition) {
@@ -182,7 +184,6 @@ private:
   LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
 
   std::unique_ptr<LazyCallThroughManager> LCTMgr;
-  std::unique_ptr<IRTransformLayer> TransformLayer;
   std::unique_ptr<CompileOnDemandLayer> CODLayer;
 };
 

lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -67,7 +67,7 @@ Error LLJIT::addIRModule(JITDylib &JD, ThreadSafeModule TSM) {
           TSM.withModuleDo([&](Module &M) { return applyDataLayout(M); }))
     return Err;
 
-  return CompileLayer->add(JD, std::move(TSM), ES->allocateVModule());
+  return TransformLayer->add(JD, std::move(TSM), ES->allocateVModule());
 }
 
 Error LLJIT::addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj) {
@@ -128,6 +128,7 @@ LLJIT::createCompileFunction(LLJITBuilderState &S,
 LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
     : ES(S.ES ? std::move(S.ES) : std::make_unique<ExecutionSession>()),
       Main(this->ES->createJITDylib("<main>")), DL(""),
+      TT(S.JTMB->getTargetTriple()),
       ObjLinkingLayer(createObjectLinkingLayer(S, *ES)),
       ObjTransformLayer(*this->ES, *ObjLinkingLayer), CtorRunner(Main),
       DtorRunner(Main) {
@@ -162,6 +163,8 @@ LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
           CompileThreads->async(std::move(Work));
         });
   }
+
+  TransformLayer = std::make_unique<IRTransformLayer>(*ES, *CompileLayer);
 }
 
 std::string LLJIT::mangle(StringRef UnmangledName) {
@@ -249,9 +252,6 @@ LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) {
     return;
   }
 
-  // Create the transform layer.
-  TransformLayer = std::make_unique<IRTransformLayer>(*ES, *CompileLayer);
-
   // Create the COD layer.
   CODLayer = std::make_unique<CompileOnDemandLayer>(
       *ES, *TransformLayer, *LCTMgr, std::move(ISMBuilder));

tools/lli/lli.cpp

@@ -781,17 +781,18 @@ int runOrcLazyJIT(const char *ProgName) {
 
   auto Dump = createDebugDumper();
 
-  J->setLazyCompileTransform([&](orc::ThreadSafeModule TSM,
-                                 const orc::MaterializationResponsibility &R) {
-    TSM.withModuleDo([&](Module &M) {
-      if (verifyModule(M, &dbgs())) {
-        dbgs() << "Bad module: " << &M << "\n";
-        exit(1);
-      }
-      Dump(M);
-    });
-    return TSM;
-  });
+  J->getIRTransformLayer().setTransform(
+      [&](orc::ThreadSafeModule TSM,
+          const orc::MaterializationResponsibility &R) {
+        TSM.withModuleDo([&](Module &M) {
+          if (verifyModule(M, &dbgs())) {
+            dbgs() << "Bad module: " << &M << "\n";
+            exit(1);
+          }
+          Dump(M);
+        });
+        return TSM;
+      });
 
   orc::MangleAndInterner Mangle(J->getExecutionSession(), J->getDataLayout());
   J->getMainJITDylib().addGenerator(