[ORC] Update LLJIT to automatically run specially named initializer functions.

The GenericLLVMIRPlatformSupport class runs a transform on all LLVM IR added to
the LLJIT instance to replace instances of llvm.global_ctors with a specially
named function that runs the corresponing static initializers (See
(GlobalCtorDtorScraper from lib/ExecutionEngine/Orc/LLJIT.cpp). This patch
updates the GenericIRPlatform class to check for this specially named function
in other materialization units that are added to the JIT and, if found, add
the function to the initializer work queue. Doing this allows object files
that were compiled from IR and cached to be reloaded in subsequent JIT sessions
without their initializers being skipped.

To enable testing this patch also updates the lli tool's -jit-kind=orc-lazy mode
to respect the -enable-cache-manager and -object-cache-dir options, and modifies
the CompileOnDemandLayer to rename extracted submodules to include a hash of the
names of their symbol definitions. This allows a simple object caching scheme
based on module names (which was already implemented in lli) to work with the
lazy JIT.

lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp

@@ -7,9 +7,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
+
+#include "llvm/ADT/Hashing.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Support/FormatVariadic.h"
 
 using namespace llvm;
 using namespace llvm::orc;
@@ -294,29 +297,52 @@ void CompileOnDemandLayer::emitPartition(
   //
   // FIXME: We apply this promotion once per partitioning. It's safe, but
   // overkill.
-
   auto ExtractedTSM =
       TSM.withModuleDo([&](Module &M) -> Expected<ThreadSafeModule> {
         auto PromotedGlobals = PromoteSymbols(M);
         if (!PromotedGlobals.empty()) {
+
           MangleAndInterner Mangle(ES, M.getDataLayout());
           SymbolFlagsMap SymbolFlags;
-          for (auto &GV : PromotedGlobals)
-            SymbolFlags[Mangle(GV->getName())] =
-                JITSymbolFlags::fromGlobalValue(*GV);
+          IRSymbolMapper::add(ES, *getManglingOptions(),
+                              PromotedGlobals, SymbolFlags);
+
           if (auto Err = R.defineMaterializing(SymbolFlags))
             return std::move(Err);
         }
 
         expandPartition(*GVsToExtract);
 
+        // Submodule name is given by hashing the names of the globals.
+        std::string SubModuleName;
+        {
+          std::vector<const GlobalValue*> HashGVs;
+          HashGVs.reserve(GVsToExtract->size());
+          for (auto *GV : *GVsToExtract)
+            HashGVs.push_back(GV);
+          llvm::sort(HashGVs, [](const GlobalValue *LHS, const GlobalValue *RHS) {
+              return LHS->getName() < RHS->getName();
+            });
+          hash_code HC(0);
+          for (auto *GV : HashGVs) {
+            assert(GV->hasName() && "All GVs to extract should be named by now");
+            auto GVName = GV->getName();
+            HC = hash_combine(HC, hash_combine_range(GVName.begin(), GVName.end()));
+          }
+          raw_string_ostream(SubModuleName)
+            << ".submodule."
+            << formatv(sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}",
+                       static_cast<size_t>(HC))
+            << ".ll";
+        }
+
         // Extract the requested partiton (plus any necessary aliases) and
         // put the rest back into the impl dylib.
         auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
           return GVsToExtract->count(&GV);
         };
 
-        return extractSubModule(TSM, ".submodule", ShouldExtract);
+        return extractSubModule(TSM, SubModuleName , ShouldExtract);
       });
 
   if (!ExtractedTSM) {

lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -107,12 +107,16 @@ private:
 /// llvm.global_ctors.
 class GlobalCtorDtorScraper {
 public:
-  GlobalCtorDtorScraper(GenericLLVMIRPlatformSupport &PS) : PS(PS) {}
+
+  GlobalCtorDtorScraper(GenericLLVMIRPlatformSupport &PS,
+                        StringRef InitFunctionPrefix)
+    : PS(PS), InitFunctionPrefix(InitFunctionPrefix) {}
   Expected<ThreadSafeModule> operator()(ThreadSafeModule TSM,
                                         MaterializationResponsibility &R);
 
 private:
   GenericLLVMIRPlatformSupport &PS;
+  StringRef InitFunctionPrefix;
 };
 
 /// Generic IR Platform Support
@@ -125,12 +129,14 @@ public:
   // GenericLLVMIRPlatform &P) : P(P) {
   GenericLLVMIRPlatformSupport(LLJIT &J) : J(J) {
 
+    MangleAndInterner Mangle(getExecutionSession(), J.getDataLayout());
+    InitFunctionPrefix = Mangle("__orc_init_func.");
+
     getExecutionSession().setPlatform(
         std::make_unique<GenericLLVMIRPlatform>(*this));
 
-    setInitTransform(J, GlobalCtorDtorScraper(*this));
+    setInitTransform(J, GlobalCtorDtorScraper(*this, *InitFunctionPrefix));
 
-    MangleAndInterner Mangle(getExecutionSession(), J.getDataLayout());
     SymbolMap StdInterposes;
 
     StdInterposes[Mangle("__lljit.platform_support_instance")] =
@@ -169,6 +175,18 @@ public:
     std::lock_guard<std::mutex> Lock(PlatformSupportMutex);
     if (auto &InitSym = MU.getInitializerSymbol())
       InitSymbols[&JD].add(InitSym);
+    else {
+      // If there's no identified init symbol attached, but there is a symbol
+      // with the GenericIRPlatform::InitFunctionPrefix, then treat that as
+      // an init function. Add the symbol to both the InitSymbols map (which
+      // will trigger a lookup to materialize the module) and the InitFunctions
+      // map (which holds the names of the symbols to execute).
+      for (auto &KV : MU.getSymbols())
+        if ((*KV.first).startswith(*InitFunctionPrefix)) {
+          InitSymbols[&JD].add(KV.first);
+          InitFunctions[&JD].add(KV.first);
+        }
+    }
     return Error::success();
   }
 
@@ -387,6 +405,7 @@ private:
 
   std::mutex PlatformSupportMutex;
   LLJIT &J;
+  SymbolStringPtr InitFunctionPrefix;
   DenseMap<JITDylib *, SymbolLookupSet> InitSymbols;
   DenseMap<JITDylib *, SymbolLookupSet> InitFunctions;
   DenseMap<JITDylib *, SymbolLookupSet> DeInitFunctions;
@@ -415,7 +434,7 @@ GlobalCtorDtorScraper::operator()(ThreadSafeModule TSM,
 
     std::string InitFunctionName;
     raw_string_ostream(InitFunctionName)
-        << "__orc_init." << M.getModuleIdentifier();
+        << InitFunctionPrefix << M.getModuleIdentifier();
 
     MangleAndInterner Mangle(PS.getExecutionSession(), M.getDataLayout());
     auto InternedName = Mangle(InitFunctionName);

test/ExecutionEngine/OrcLazy/static-initializers-in-objectfiles.ll

@@ -0,0 +1,28 @@
+; RUN: rm -rf %t
+; RUN: mkdir -p %t
+; RUN: lli -jit-kind=orc-lazy -enable-cache-manager -object-cache-dir=%t %s
+; RUN: lli -jit-kind=orc-lazy -enable-cache-manager -object-cache-dir=%t %s
+;
+; Verify that LLJIT Platforms respect static initializers in cached objects.
+; This IR file contains a static initializer that must execute for main to exit
+; with value zero. The first execution will populate an object cache for the
+; second. The initializer in the cached objects must also be successfully run
+; for the test to pass.
+
+@HasError = global i8 1, align 1
+@llvm.global_ctors = appending global [1 x { i32, void ()*, i8* }] [{ i32, void ()*, i8* } { i32 65535, void ()* @resetHasError, i8* null }]
+
+define void @resetHasError() {
+entry:
+  store i8 0, i8* @HasError, align 1
+  ret void
+}
+
+define i32 @main(i32 %argc, i8** %argv) #2 {
+entry:
+  %0 = load i8, i8* @HasError, align 1
+  %tobool = trunc i8 %0 to i1
+  %conv = zext i1 %tobool to i32
+  ret i32 %conv
+}
+

tools/lli/lli.cpp

@@ -274,6 +274,7 @@ public:
       SmallString<128> dir(sys::path::parent_path(CacheName));
       sys::fs::create_directories(Twine(dir));
     }
+
     std::error_code EC;
     raw_fd_ostream outfile(CacheName, EC, sys::fs::OF_None);
     outfile.write(Obj.getBufferStart(), Obj.getBufferSize());
@@ -306,14 +307,16 @@ private:
     size_t PrefixLength = Prefix.length();
     if (ModID.substr(0, PrefixLength) != Prefix)
       return false;
-        std::string CacheSubdir = ModID.substr(PrefixLength);
+
+    std::string CacheSubdir = ModID.substr(PrefixLength);
 #if defined(_WIN32)
-        // Transform "X:\foo" => "/X\foo" for convenience.
-        if (isalpha(CacheSubdir[0]) && CacheSubdir[1] == ':') {
-          CacheSubdir[1] = CacheSubdir[0];
-          CacheSubdir[0] = '/';
-        }
+    // Transform "X:\foo" => "/X\foo" for convenience.
+    if (isalpha(CacheSubdir[0]) && CacheSubdir[1] == ':') {
+      CacheSubdir[1] = CacheSubdir[0];
+      CacheSubdir[0] = '/';
+    }
 #endif
+
     CacheName = CacheDir + CacheSubdir;
     size_t pos = CacheName.rfind('.');
     CacheName.replace(pos, CacheName.length() - pos, ".o");
@@ -777,30 +780,56 @@ Error loadDylibs() {
 
 static void exitOnLazyCallThroughFailure() { exit(1); }
 
+Expected<orc::ThreadSafeModule>
+loadModule(StringRef Path, orc::ThreadSafeContext TSCtx) {
+  SMDiagnostic Err;
+  auto M = parseIRFile(Path, Err, *TSCtx.getContext());
+  if (!M) {
+    std::string ErrMsg;
+    {
+      raw_string_ostream ErrMsgStream(ErrMsg);
+      Err.print("lli", ErrMsgStream);
+    }
+    return make_error<StringError>(std::move(ErrMsg), inconvertibleErrorCode());
+  }
+
+  if (EnableCacheManager)
+    M->setModuleIdentifier("file:" + M->getModuleIdentifier());
+
+  return orc::ThreadSafeModule(std::move(M), std::move(TSCtx));
+}
+
 int runOrcLazyJIT(const char *ProgName) {
   // Start setting up the JIT environment.
 
   // Parse the main module.
   orc::ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
-  SMDiagnostic Err;
-  auto MainModule = parseIRFile(InputFile, Err, *TSCtx.getContext());
-  if (!MainModule)
-    reportError(Err, ProgName);
+  auto MainModule = ExitOnErr(loadModule(InputFile, TSCtx));
+
+  // Get TargetTriple and DataLayout from the main module if they're explicitly
+  // set.
+  Optional<Triple> TT;
+  Optional<DataLayout> DL;
+  MainModule.withModuleDo([&](Module &M) {
+      if (!M.getTargetTriple().empty())
+        TT = Triple(M.getTargetTriple());
+      if (!M.getDataLayout().isDefault())
+        DL = M.getDataLayout();
+    });
 
-  Triple TT(MainModule->getTargetTriple());
   orc::LLLazyJITBuilder Builder;
 
   Builder.setJITTargetMachineBuilder(
-      MainModule->getTargetTriple().empty()
-          ? ExitOnErr(orc::JITTargetMachineBuilder::detectHost())
-          : orc::JITTargetMachineBuilder(TT));
+      TT ? orc::JITTargetMachineBuilder(*TT)
+         : ExitOnErr(orc::JITTargetMachineBuilder::detectHost()));
+
+  TT = Builder.getJITTargetMachineBuilder()->getTargetTriple();
+  if (DL)
+    Builder.setDataLayout(DL);
 
   if (!MArch.empty())
     Builder.getJITTargetMachineBuilder()->getTargetTriple().setArchName(MArch);
 
-  if (!MainModule->getDataLayout().isDefault())
-    Builder.setDataLayout(MainModule->getDataLayout());
-
   Builder.getJITTargetMachineBuilder()
       ->setCPU(getCPUStr())
       .addFeatures(getFeatureList())
@@ -815,11 +844,34 @@ int runOrcLazyJIT(const char *ProgName) {
       pointerToJITTargetAddress(exitOnLazyCallThroughFailure));
   Builder.setNumCompileThreads(LazyJITCompileThreads);
 
+  // If the object cache is enabled then set a custom compile function
+  // creator to use the cache.
+  std::unique_ptr<LLIObjectCache> CacheManager;
+  if (EnableCacheManager) {
+
+    CacheManager = std::make_unique<LLIObjectCache>(ObjectCacheDir);
+
+    Builder.setCompileFunctionCreator(
+      [&](orc::JITTargetMachineBuilder JTMB)
+            -> Expected<std::unique_ptr<orc::IRCompileLayer::IRCompiler>> {
+        if (LazyJITCompileThreads > 0)
+          return std::make_unique<orc::ConcurrentIRCompiler>(std::move(JTMB),
+                                                        CacheManager.get());
+
+        auto TM = JTMB.createTargetMachine();
+        if (!TM)
+          return TM.takeError();
+
+        return std::make_unique<orc::TMOwningSimpleCompiler>(std::move(*TM),
+                                                        CacheManager.get());
+      });
+  }
+
   // Set up LLJIT platform.
   {
     LLJITPlatform P = Platform;
     if (P == LLJITPlatform::DetectHost) {
-      if (TT.isOSBinFormatMachO())
+      if (TT->isOSBinFormatMachO())
         P = LLJITPlatform::MachO;
       else
         P = LLJITPlatform::GenericIR;
@@ -871,8 +923,7 @@ int runOrcLazyJIT(const char *ProgName) {
             })));
 
   // Add the main module.
-  ExitOnErr(
-      J->addLazyIRModule(orc::ThreadSafeModule(std::move(MainModule), TSCtx)));
+  ExitOnErr(J->addLazyIRModule(std::move(MainModule)));
 
   // Create JITDylibs and add any extra modules.
   {
@@ -894,16 +945,13 @@ int runOrcLazyJIT(const char *ProgName) {
 
     for (auto EMItr = ExtraModules.begin(), EMEnd = ExtraModules.end();
          EMItr != EMEnd; ++EMItr) {
-      auto M = parseIRFile(*EMItr, Err, *TSCtx.getContext());
-      if (!M)
-        reportError(Err, ProgName);
+      auto M = ExitOnErr(loadModule(*EMItr, TSCtx));
 
       auto EMIdx = ExtraModules.getPosition(EMItr - ExtraModules.begin());
       assert(EMIdx != 0 && "ExtraModule should have index > 0");
       auto JDItr = std::prev(IdxToDylib.lower_bound(EMIdx));
       auto &JD = *JDItr->second;
-      ExitOnErr(
-          J->addLazyIRModule(JD, orc::ThreadSafeModule(std::move(M), TSCtx)));
+      ExitOnErr(J->addLazyIRModule(JD, std::move(M)));
     }
 
     for (auto EAItr = ExtraArchives.begin(), EAEnd = ExtraArchives.end();