[ORC] Introduce IRPartitionLayer for common partition functionality.

[sunho]

Summary

The IR module partitioning logic has been embeded in CompileOnDemandLayer but this functionality is also required for other purpose such as re-optimization.
This patch splits that IR partitioning logic out of CompileOnDemandLayer into a new IR layer IRPartitionLayer.

A motivation behind making it into new IR layer is for the possibility of "lazy re-optimizing JIT" where we'd like to have a composition like "IRPartitionLayer -> CompileOnDemandLayer -> ReOptLayer"

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

 #include "llvm/ADT/StringRef.h"
 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/EPCIndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassManager.h"
private:
 
   DataLayout DL;
   MangleAndInterner Mangle;
 
   RTDyldObjectLinkingLayer ObjectLayer;
   IRCompileLayer CompileLayer;
   IRTransformLayer OptimizeLayer;
   CompileOnDemandLayer CODLayer;
+  IRPartitionLayer IPLayer;
 
   JITDylib &MainJD;
 
   static void handleLazyCallThroughError() {
     errs() << "LazyCallThrough error: Could not find function body";
     exit(1);
   }
 
 public:
   KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
                   std::unique_ptr<EPCIndirectionUtils> EPCIU,
                   JITTargetMachineBuilder JTMB, DataLayout DL)
       : ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),
         Mangle(*this->ES, this->DL),
         ObjectLayer(*this->ES,
                     []() { return std::make_unique<SectionMemoryManager>(); }),
         CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
         CODLayer(*this->ES, OptimizeLayer,
                  this->EPCIU->getLazyCallThroughManager(),
                  [this] { return this->EPCIU->createIndirectStubsManager(); }),
+        IPLayer(*this->ES, CODLayer),
         MainJD(this->ES->createBareJITDylib("<main>")) {
     MainJD.addGenerator(
         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
             DL.getGlobalPrefix())));
   }
 
   ~KaleidoscopeJIT() {
     if (auto Err = ES->endSession())
public:
   const DataLayout &getDataLayout() const { return DL; }
 
   JITDylib &getMainJITDylib() { return MainJD; }
 
   Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
     if (!RT)
       RT = MainJD.getDefaultResourceTracker();
 
-    return CODLayer.add(RT, std::move(TSM));
+    return IPLayer.add(RT, std::move(TSM));
   }
 
   Expected<ExecutorSymbolDef> lookup(StringRef Name) {
     return ES->lookup({&MainJD}, Mangle(Name.str()));
   }
 
 private:
   static Expected<ThreadSafeModule>

llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp

 #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/IRPartitionLayer.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/Orc/SpeculateAnalyses.h"
 #include "llvm/ExecutionEngine/Orc/Speculation.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/TargetExecutionUtils.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IRReader/IRReader.h"
std::unique_ptr<SpeculativeJIT> SJ(new SpeculativeJIT(
         std::move(ES), std::move(*DL), std::move(*JTMB), std::move(*LCTMgr),
         std::move(ISMBuilder), std::move(*ProcessSymbolsSearchGenerator)));
     return std::move(SJ);
   }
 
   ExecutionSession &getES() { return *ES; }
 
   Error addModule(ThreadSafeModule TSM) {
-    return CODLayer.add(MainJD, std::move(TSM));
+    return IPLayer.add(MainJD, std::move(TSM));
   }
 
   Expected<ExecutorSymbolDef> lookup(StringRef UnmangledName) {
     return ES->lookup({&MainJD}, Mangle(UnmangledName));
   }
 
   ~SpeculativeJIT() { CompileThreads.wait(); }
 
 private:
   using IndirectStubsManagerBuilderFunction =
       std::function<std::unique_ptr<IndirectStubsManager>()>;
 
   static void explodeOnLazyCompileFailure() {
     errs() << "Lazy compilation failed, Symbol Implmentation not found!\n";
     exit(1);
   }
 
   SpeculativeJIT(
       std::unique_ptr<ExecutionSession> ES, DataLayout DL,
       orc::JITTargetMachineBuilder JTMB,
       std::unique_ptr<LazyCallThroughManager> LCTMgr,
       IndirectStubsManagerBuilderFunction ISMBuilder,
       std::unique_ptr<DynamicLibrarySearchGenerator> ProcessSymbolsGenerator)
       : ES(std::move(ES)), DL(std::move(DL)),
-        MainJD(this->ES->createBareJITDylib("<main>")), LCTMgr(std::move(LCTMgr)),
+        MainJD(this->ES->createBareJITDylib("<main>")),
+        LCTMgr(std::move(LCTMgr)),
         CompileLayer(*this->ES, ObjLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         S(Imps, *this->ES),
         SpeculateLayer(*this->ES, CompileLayer, S, Mangle, BlockFreqQuery()),
         CODLayer(*this->ES, SpeculateLayer, *this->LCTMgr,
-                 std::move(ISMBuilder)) {
+                 std::move(ISMBuilder)),
+        IPLayer(*this->ES, CODLayer) {
     MainJD.addGenerator(std::move(ProcessSymbolsGenerator));
     this->CODLayer.setImplMap(&Imps);
     this->ES->setDispatchTask(
         [this](std::unique_ptr<Task> T) {
           CompileThreads.async(
               [UnownedT = T.release()]() {
                 std::unique_ptr<Task> T(UnownedT);
                 T->run();
private:
   Triple TT;
   std::unique_ptr<LazyCallThroughManager> LCTMgr;
   IRCompileLayer CompileLayer;
   ImplSymbolMap Imps;
   Speculator S;
   RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr};
   IRSpeculationLayer SpeculateLayer;
   CompileOnDemandLayer CODLayer;
+  IRPartitionLayer IPLayer;
 };
 
 int main(int argc, char *argv[]) {
   // Initialize LLVM.
   InitLLVM X(argc, argv);
 
   InitializeNativeTarget();
   InitializeNativeTargetAsmPrinter();

llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

 #include <set>
 #include <utility>
 #include <vector>
 
 namespace llvm {
 namespace orc {
 
 class CompileOnDemandLayer : public IRLayer {
-  friend class PartitioningIRMaterializationUnit;
-
 public:
   /// Builder for IndirectStubsManagers.
   using IndirectStubsManagerBuilder =
       std::function<std::unique_ptr<IndirectStubsManager>()>;
 
-  using GlobalValueSet = std::set<const GlobalValue *>;
-
-  /// Partitioning function.
-  using PartitionFunction =
-      std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>;
-
-  /// Off-the-shelf partitioning which compiles all requested symbols (usually
-  /// a single function at a time).
-  static std::optional<GlobalValueSet>
-  compileRequested(GlobalValueSet Requested);
-
-  /// Off-the-shelf partitioning which compiles whole modules whenever any
-  /// symbol in them is requested.
-  static std::optional<GlobalValueSet>
-  compileWholeModule(GlobalValueSet Requested);
-
   /// Construct a CompileOnDemandLayer.
   CompileOnDemandLayer(ExecutionSession &ES, IRLayer &BaseLayer,
-                        LazyCallThroughManager &LCTMgr,
-                        IndirectStubsManagerBuilder BuildIndirectStubsManager);
-
-  /// Sets the partition function.
-  void setPartitionFunction(PartitionFunction Partition);
-
   /// Sets the ImplSymbolMap
   void setImplMap(ImplSymbolMap *Imp);
 
   /// Emits the given module. This should not be called by clients: it will be
   /// called by the JIT when a definition added via the add method is requested.
   void emit(std::unique_ptr<MaterializationResponsibility> R,
             ThreadSafeModule TSM) override;
 
private:
     JITDylib &ImplD;
     std::unique_ptr<IndirectStubsManager> ISMgr;
   };
 
   using PerDylibResourcesMap = std::map<const JITDylib *, PerDylibResources>;
 
   PerDylibResources &getPerDylibResources(JITDylib &TargetD);
 
-  void cleanUpModule(Module &M);
-
-  void expandPartition(GlobalValueSet &Partition);
-
-  void emitPartition(std::unique_ptr<MaterializationResponsibility> R,
-                     ThreadSafeModule TSM,
-                     IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
-
   mutable std::mutex CODLayerMutex;
 
   IRLayer &BaseLayer;
   LazyCallThroughManager &LCTMgr;
   IndirectStubsManagerBuilder BuildIndirectStubsManager;
   PerDylibResourcesMap DylibResources;
-  PartitionFunction Partition = compileRequested;
-  SymbolLinkagePromoter PromoteSymbols;
   ImplSymbolMap *AliaseeImpls = nullptr;
 };
 
 } // end namespace orc
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H

llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h

-//===- CompileOnDemandLayer.h - Compile each function on demand -*- C++ -*-===//
+//===- IRPartitionLayer.h - Partition IR module on lookup -------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
-// JIT layer for breaking up modules and inserting callbacks to allow
-// individual functions to be compiled on demand.
+// JIT layer for breaking up modules into smaller submodules that only contains
+// looked up symbols.
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
-#define LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
+#ifndef LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
+#define LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
 
-#include "llvm/ADT/APInt.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/StringRef.h"
-#include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/Layer.h"
-#include "llvm/ExecutionEngine/Orc/LazyReexports.h"
-#include "llvm/ExecutionEngine/Orc/Shared/OrcError.h"
-#include "llvm/ExecutionEngine/Orc/Speculation.h"
-#include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
-#include "llvm/Support/Casting.h"
-#include "llvm/Support/raw_ostream.h"
-#include "llvm/Transforms/Utils/ValueMapper.h"
-#include <algorithm>
-#include <cassert>
-#include <functional>
-#include <iterator>
-#include <list>
-#include <memory>
-#include <optional>
-#include <set>
-#include <utility>
-#include <vector>
 
 namespace llvm {
 namespace orc {
 
-class CompileOnDemandLayer : public IRLayer {
+class IRPartitionLayer : public IRLayer {
   friend class PartitioningIRMaterializationUnit;
 
 public:
-  /// Builder for IndirectStubsManagers.
-  using IndirectStubsManagerBuilder =
-      std::function<std::unique_ptr<IndirectStubsManager>()>;
-
   using GlobalValueSet = std::set<const GlobalValue *>;
 
   /// Partitioning function.
   using PartitionFunction =
       std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>;
 
+  /// Construct a IRPartitionLayer.
+  IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer);
+
   /// Off-the-shelf partitioning which compiles all requested symbols (usually
   /// a single function at a time).
   static std::optional<GlobalValueSet>
   compileRequested(GlobalValueSet Requested);
 
   /// Off-the-shelf partitioning which compiles whole modules whenever any
   /// symbol in them is requested.
   static std::optional<GlobalValueSet>
   compileWholeModule(GlobalValueSet Requested);
 
-  /// Construct a CompileOnDemandLayer.
-  CompileOnDemandLayer(ExecutionSession &ES, IRLayer &BaseLayer,
-                        LazyCallThroughManager &LCTMgr,
-                        IndirectStubsManagerBuilder BuildIndirectStubsManager);
-
   /// Sets the partition function.
   void setPartitionFunction(PartitionFunction Partition);
 
-  /// Sets the ImplSymbolMap
-  void setImplMap(ImplSymbolMap *Imp);
-
   /// Emits the given module. This should not be called by clients: it will be
   /// called by the JIT when a definition added via the add method is requested.
   void emit(std::unique_ptr<MaterializationResponsibility> R,
             ThreadSafeModule TSM) override;
 
 private:
-  struct PerDylibResources {
-  public:
-    PerDylibResources(JITDylib &ImplD,
-                      std::unique_ptr<IndirectStubsManager> ISMgr)
-        : ImplD(ImplD), ISMgr(std::move(ISMgr)) {}
-    JITDylib &getImplDylib() { return ImplD; }
-    IndirectStubsManager &getISManager() { return *ISMgr; }
-
-  private:
-    JITDylib &ImplD;
-    std::unique_ptr<IndirectStubsManager> ISMgr;
-  };
-
-  using PerDylibResourcesMap = std::map<const JITDylib *, PerDylibResources>;
-
-  PerDylibResources &getPerDylibResources(JITDylib &TargetD);
-
   void cleanUpModule(Module &M);
 
   void expandPartition(GlobalValueSet &Partition);
 
   void emitPartition(std::unique_ptr<MaterializationResponsibility> R,
                      ThreadSafeModule TSM,
                      IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
 
-  mutable std::mutex CODLayerMutex;
-
   IRLayer &BaseLayer;
-  LazyCallThroughManager &LCTMgr;
-  IndirectStubsManagerBuilder BuildIndirectStubsManager;
-  PerDylibResourcesMap DylibResources;
   PartitionFunction Partition = compileRequested;
   SymbolLinkagePromoter PromoteSymbols;
-  ImplSymbolMap *AliaseeImpls = nullptr;
 };
 
-} // end namespace orc
-} // end namespace llvm
+} // namespace orc
+} // namespace llvm
 
-#endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
+#endif

llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h

 
 #ifndef LLVM_EXECUTIONENGINE_ORC_LLJIT_H
 #define LLVM_EXECUTIONENGINE_ORC_LLJIT_H
 
 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ThreadPool.h"
 #include <variant>
 
 namespace llvm {

 /// An extended version of LLJIT that supports lazy function-at-a-time
 /// compilation of LLVM IR.
 class LLLazyJIT : public LLJIT {
   template <typename, typename, typename> friend class LLJITBuilderSetters;
 
 public:
 
   /// Sets the partition function.
-  void
-  setPartitionFunction(CompileOnDemandLayer::PartitionFunction Partition) {
-    CODLayer->setPartitionFunction(std::move(Partition));
+  void setPartitionFunction(IRPartitionLayer::PartitionFunction Partition) {
+    IPLayer->setPartitionFunction(std::move(Partition));
   }
 
   /// Returns a reference to the on-demand layer.
   CompileOnDemandLayer &getCompileOnDemandLayer() { return *CODLayer; }
 
   /// Add a module to be lazily compiled to JITDylib JD.
   Error addLazyIRModule(JITDylib &JD, ThreadSafeModule M);
 
   /// Add a module to be lazily compiled to the main JITDylib.
   Error addLazyIRModule(ThreadSafeModule M) {
     return addLazyIRModule(*Main, std::move(M));
   }
 
 private:
 
   // Create a single-threaded LLLazyJIT instance.
   LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
 
   std::unique_ptr<LazyCallThroughManager> LCTMgr;
+  std::unique_ptr<IRPartitionLayer> IPLayer;
   std::unique_ptr<CompileOnDemandLayer> CODLayer;
 };
 
 class LLJITBuilderState {
 public:
   using ObjectLinkingLayerCreator =
       std::function<Expected<std::unique_ptr<ObjectLayer>>(ExecutionSession &,
                                                            const Triple &)>;

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

add_llvm_component_library(LLVMOrcJIT
   EPCGenericJITLinkMemoryManager.cpp
   EPCGenericRTDyldMemoryManager.cpp
   EPCIndirectionUtils.cpp
   ExecutionUtils.cpp
   ObjectFileInterface.cpp
   IndirectionUtils.cpp
   IRCompileLayer.cpp
   IRTransformLayer.cpp
+  IRPartitionLayer.cpp
   JITTargetMachineBuilder.cpp
   LazyReexports.cpp
   Layer.cpp
   LookupAndRecordAddrs.cpp
   LLJIT.cpp
   MachOPlatform.cpp
   MapperJITLinkMemoryManager.cpp
   MemoryMapper.cpp

llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp

 //===----- CompileOnDemandLayer.cpp - Lazily emit IR on first call --------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/Layer.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/FormatVariadic.h"
 #include <string>
 
 using namespace llvm;
 using namespace llvm::orc;
 
-static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM,
-                                         StringRef Suffix,
-                                         GVPredicate ShouldExtract) {
-
-  auto DeleteExtractedDefs = [](GlobalValue &GV) {
-    // Bump the linkage: this global will be provided by the external module.
-    GV.setLinkage(GlobalValue::ExternalLinkage);
-
-    // Delete the definition in the source module.
-    if (isa<Function>(GV)) {
-      auto &F = cast<Function>(GV);
-      F.deleteBody();
-      F.setPersonalityFn(nullptr);
-    } else if (isa<GlobalVariable>(GV)) {
-      cast<GlobalVariable>(GV).setInitializer(nullptr);
-    } else if (isa<GlobalAlias>(GV)) {
-      // We need to turn deleted aliases into function or variable decls based
-      // on the type of their aliasee.
-      auto &A = cast<GlobalAlias>(GV);
-      Constant *Aliasee = A.getAliasee();
-      assert(A.hasName() && "Anonymous alias?");
-      assert(Aliasee->hasName() && "Anonymous aliasee");
-      std::string AliasName = std::string(A.getName());
-
-      if (isa<Function>(Aliasee)) {
-        auto *F = cloneFunctionDecl(*A.getParent(), *cast<Function>(Aliasee));
-        A.replaceAllUsesWith(F);
-        A.eraseFromParent();
-        F->setName(AliasName);
-      } else if (isa<GlobalVariable>(Aliasee)) {
-        auto *G = cloneGlobalVariableDecl(*A.getParent(),
-                                          *cast<GlobalVariable>(Aliasee));
-        A.replaceAllUsesWith(G);
-        A.eraseFromParent();
-        G->setName(AliasName);
-      } else
-        llvm_unreachable("Alias to unsupported type");
-    } else
-      llvm_unreachable("Unsupported global type");
-  };
-
-  auto NewTSM = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs);
-  NewTSM.withModuleDo([&](Module &M) {
-    M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
-  });
-
-  return NewTSM;
-}
-
-namespace llvm {
-namespace orc {
-
-class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
-public:
-  PartitioningIRMaterializationUnit(ExecutionSession &ES,
-                                    const IRSymbolMapper::ManglingOptions &MO,
-                                    ThreadSafeModule TSM,
-                                    CompileOnDemandLayer &Parent)
-      : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {}
-
-  PartitioningIRMaterializationUnit(
-      ThreadSafeModule TSM, Interface I,
-      SymbolNameToDefinitionMap SymbolToDefinition,
-      CompileOnDemandLayer &Parent)
-      : IRMaterializationUnit(std::move(TSM), std::move(I),
-                              std::move(SymbolToDefinition)),
-        Parent(Parent) {}
-
-private:
-  void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
-    Parent.emitPartition(std::move(R), std::move(TSM),
-                         std::move(SymbolToDefinition));
-  }
-
-  void discard(const JITDylib &V, const SymbolStringPtr &Name) override {
-    // All original symbols were materialized by the CODLayer and should be
-    // final. The function bodies provided by M should never be overridden.
-    llvm_unreachable("Discard should never be called on an "
-                     "ExtractingIRMaterializationUnit");
-  }
-
-  mutable std::mutex SourceModuleMutex;
-  CompileOnDemandLayer &Parent;
-};
-
-std::optional<CompileOnDemandLayer::GlobalValueSet>
-CompileOnDemandLayer::compileRequested(GlobalValueSet Requested) {
-  return std::move(Requested);
-}
-
-std::optional<CompileOnDemandLayer::GlobalValueSet>
-CompileOnDemandLayer::compileWholeModule(GlobalValueSet Requested) {
-  return std::nullopt;
-}
-
 CompileOnDemandLayer::CompileOnDemandLayer(
     ExecutionSession &ES, IRLayer &BaseLayer, LazyCallThroughManager &LCTMgr,
     IndirectStubsManagerBuilder BuildIndirectStubsManager)
     : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer),
       LCTMgr(LCTMgr),
       BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)) {}
 
-void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) {
-  this->Partition = std::move(Partition);
-}
-
 void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) {
   this->AliaseeImpls = Imp;
 }
+
 void CompileOnDemandLayer::emit(
     std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM) {
   assert(TSM && "Null module");
 
   auto &ES = getExecutionSession();
 
   // Sort the callables and non-callables, build re-exports and lodge the
   // actual module with the implementation dylib.
   auto &PDR = getPerDylibResources(R->getTargetJITDylib());
 
   SymbolAliasMap NonCallables;
   SymbolAliasMap Callables;
-  TSM.withModuleDo([&](Module &M) {
-    // First, do some cleanup on the module:
-    cleanUpModule(M);
-  });
 
   for (auto &KV : R->getSymbols()) {
     auto &Name = KV.first;
     auto &Flags = KV.second;
     if (Flags.isCallable())
       Callables[Name] = SymbolAliasMapEntry(Name, Flags);
     else
       NonCallables[Name] = SymbolAliasMapEntry(Name, Flags);
   }
 
-  // Create a partitioning materialization unit and lodge it with the
-  // implementation dylib.
+  // Lodge symbols with the implementation dylib.
   if (auto Err = PDR.getImplDylib().define(
-          std::make_unique<PartitioningIRMaterializationUnit>(
-              ES, *getManglingOptions(), std::move(TSM), *this))) {
+          std::make_unique<BasicIRLayerMaterializationUnit>(
+              BaseLayer, *getManglingOptions(), std::move(TSM)))) {
     ES.reportError(std::move(Err));
     R->failMaterialization();
     return;
   }
 
   if (!NonCallables.empty())
     if (auto Err =
             R->replace(reexports(PDR.getImplDylib(), std::move(NonCallables),
if (I == DylibResources.end()) {
     TargetD.setLinkOrder(std::move(NewLinkOrder), false);
 
     PerDylibResources PDR(ImplD, BuildIndirectStubsManager());
     I = DylibResources.insert(std::make_pair(&TargetD, std::move(PDR))).first;
   }
 
   return I->second;
 }
-
-void CompileOnDemandLayer::cleanUpModule(Module &M) {
-  for (auto &F : M.functions()) {
-    if (F.isDeclaration())
-      continue;
-
-    if (F.hasAvailableExternallyLinkage()) {
-      F.deleteBody();
-      F.setPersonalityFn(nullptr);
-      continue;
-    }
-  }
-}
-
-void CompileOnDemandLayer::expandPartition(GlobalValueSet &Partition) {
-  // Expands the partition to ensure the following rules hold:
-  // (1) If any alias is in the partition, its aliasee is also in the partition.
-  // (2) If any aliasee is in the partition, its aliases are also in the
-  //     partiton.
-  // (3) If any global variable is in the partition then all global variables
-  //     are in the partition.
-  assert(!Partition.empty() && "Unexpected empty partition");
-
-  const Module &M = *(*Partition.begin())->getParent();
-  bool ContainsGlobalVariables = false;
-  std::vector<const GlobalValue *> GVsToAdd;
-
-  for (const auto *GV : Partition)
-    if (isa<GlobalAlias>(GV))
-      GVsToAdd.push_back(
-          cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee()));
-    else if (isa<GlobalVariable>(GV))
-      ContainsGlobalVariables = true;
-
-  for (auto &A : M.aliases())
-    if (Partition.count(cast<GlobalValue>(A.getAliasee())))
-      GVsToAdd.push_back(&A);
-
-  if (ContainsGlobalVariables)
-    for (auto &G : M.globals())
-      GVsToAdd.push_back(&G);
-
-  for (const auto *GV : GVsToAdd)
-    Partition.insert(GV);
-}
-
-void CompileOnDemandLayer::emitPartition(
-    std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM,
-    IRMaterializationUnit::SymbolNameToDefinitionMap Defs) {
-
-  // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
-  //        extracted module key, extracted module, and source module key
-  //        together. This could be used, for example, to provide a specific
-  //        memory manager instance to the linking layer.
-
-  auto &ES = getExecutionSession();
-  GlobalValueSet RequestedGVs;
-  for (auto &Name : R->getRequestedSymbols()) {
-    if (Name == R->getInitializerSymbol())
-      TSM.withModuleDo([&](Module &M) {
-        for (auto &GV : getStaticInitGVs(M))
-          RequestedGVs.insert(&GV);
-      });
-    else {
-      assert(Defs.count(Name) && "No definition for symbol");
-      RequestedGVs.insert(Defs[Name]);
-    }
-  }
-
-  /// Perform partitioning with the context lock held, since the partition
-  /// function is allowed to access the globals to compute the partition.
-  auto GVsToExtract =
-      TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); });
-
-  // Take a 'None' partition to mean the whole module (as opposed to an empty
-  // partition, which means "materialize nothing"). Emit the whole module
-  // unmodified to the base layer.
-  if (GVsToExtract == std::nullopt) {
-    Defs.clear();
-    BaseLayer.emit(std::move(R), std::move(TSM));
-    return;
-  }
-
-  // If the partition is empty, return the whole module to the symbol table.
-  if (GVsToExtract->empty()) {
-    if (auto Err =
-            R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
-                std::move(TSM),
-                MaterializationUnit::Interface(R->getSymbols(),
-                                               R->getInitializerSymbol()),
-                std::move(Defs), *this))) {
-      getExecutionSession().reportError(std::move(Err));
-      R->failMaterialization();
-      return;
-    }
-    return;
-  }
-
-  // Ok -- we actually need to partition the symbols. Promote the symbol
-  // linkages/names, expand the partition to include any required symbols
-  // (i.e. symbols that can't be separated from our partition), and
-  // then extract the partition.
-  //
-  // 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;
-          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 (const 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 (const 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, SubModuleName , ShouldExtract);
-      });
-
-  if (!ExtractedTSM) {
-    ES.reportError(ExtractedTSM.takeError());
-    R->failMaterialization();
-    return;
-  }
-
-  if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
-          ES, *getManglingOptions(), std::move(TSM), *this))) {
-    ES.reportError(std::move(Err));
-    R->failMaterialization();
-    return;
-  }
-  BaseLayer.emit(std::move(R), std::move(*ExtractedTSM));
-}
-
-} // end namespace orc
-} // end namespace llvm

llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp

-//===----- CompileOnDemandLayer.cpp - Lazily emit IR on first call --------===//
+//===----- IRPartitionLayer.cpp - Partition IR module into submodules -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
-#include "llvm/ADT/Hashing.h"
+#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
-#include "llvm/IR/Mangler.h"
-#include "llvm/IR/Module.h"
-#include "llvm/Support/FormatVariadic.h"
-#include <string>
+#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 
 using namespace llvm;
 using namespace llvm::orc;
 
 static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM,
                                          StringRef Suffix,
                                          GVPredicate ShouldExtract) {
 

 namespace llvm {
 namespace orc {
 
 class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
 public:
   PartitioningIRMaterializationUnit(ExecutionSession &ES,
                                     const IRSymbolMapper::ManglingOptions &MO,
                                     ThreadSafeModule TSM,
-                                    CompileOnDemandLayer &Parent)
+                                    IRPartitionLayer &Parent)
       : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {}
 
   PartitioningIRMaterializationUnit(
       ThreadSafeModule TSM, Interface I,
-      SymbolNameToDefinitionMap SymbolToDefinition,
-      CompileOnDemandLayer &Parent)
+      SymbolNameToDefinitionMap SymbolToDefinition, IRPartitionLayer &Parent)
       : IRMaterializationUnit(std::move(TSM), std::move(I),
                               std::move(SymbolToDefinition)),
         Parent(Parent) {}
 
 private:
   void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
     Parent.emitPartition(std::move(R), std::move(TSM),
                          std::move(SymbolToDefinition));
   }
 
   void discard(const JITDylib &V, const SymbolStringPtr &Name) override {
     // All original symbols were materialized by the CODLayer and should be
     // final. The function bodies provided by M should never be overridden.
     llvm_unreachable("Discard should never be called on an "
                      "ExtractingIRMaterializationUnit");
   }
 
-  mutable std::mutex SourceModuleMutex;
-  CompileOnDemandLayer &Parent;
+  IRPartitionLayer &Parent;
 };
 
-std::optional<CompileOnDemandLayer::GlobalValueSet>
-CompileOnDemandLayer::compileRequested(GlobalValueSet Requested) {
-  return std::move(Requested);
-}
+} // namespace orc
+} // namespace llvm
 
-std::optional<CompileOnDemandLayer::GlobalValueSet>
-CompileOnDemandLayer::compileWholeModule(GlobalValueSet Requested) {
-  return std::nullopt;
+IRPartitionLayer::IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer)
+    : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer) {}
+
+void IRPartitionLayer::setPartitionFunction(PartitionFunction Partition) {
+  this->Partition = Partition;
 }
 
-CompileOnDemandLayer::CompileOnDemandLayer(
-    ExecutionSession &ES, IRLayer &BaseLayer, LazyCallThroughManager &LCTMgr,
-    IndirectStubsManagerBuilder BuildIndirectStubsManager)
-    : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer),
-      LCTMgr(LCTMgr),
-      BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)) {}
-
-void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) {
-  this->Partition = std::move(Partition);
+std::optional<IRPartitionLayer::GlobalValueSet>
+IRPartitionLayer::compileRequested(GlobalValueSet Requested) {
+  return std::move(Requested);
 }
 
-void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) {
-  this->AliaseeImpls = Imp;
+std::optional<IRPartitionLayer::GlobalValueSet>
+IRPartitionLayer::compileWholeModule(GlobalValueSet Requested) {
+  return std::nullopt;
 }
-void CompileOnDemandLayer::emit(
-    std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM) {
+
+void IRPartitionLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
+                            ThreadSafeModule TSM) {
   assert(TSM && "Null module");
 
   auto &ES = getExecutionSession();
-
-  // Sort the callables and non-callables, build re-exports and lodge the
-  // actual module with the implementation dylib.
-  auto &PDR = getPerDylibResources(R->getTargetJITDylib());
-
-  SymbolAliasMap NonCallables;
-  SymbolAliasMap Callables;
   TSM.withModuleDo([&](Module &M) {
     // First, do some cleanup on the module:
     cleanUpModule(M);
   });
 
-  for (auto &KV : R->getSymbols()) {
-    auto &Name = KV.first;
-    auto &Flags = KV.second;
-    if (Flags.isCallable())
-      Callables[Name] = SymbolAliasMapEntry(Name, Flags);
-    else
-      NonCallables[Name] = SymbolAliasMapEntry(Name, Flags);
-  }
-
-  // Create a partitioning materialization unit and lodge it with the
-  // implementation dylib.
-  if (auto Err = PDR.getImplDylib().define(
-          std::make_unique<PartitioningIRMaterializationUnit>(
-              ES, *getManglingOptions(), std::move(TSM), *this))) {
+  // Create a partitioning materialization unit and pass the responsibility.
+  if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
     ES.reportError(std::move(Err));
     R->failMaterialization();
     return;
   }
-
-  if (!NonCallables.empty())
-    if (auto Err =
-            R->replace(reexports(PDR.getImplDylib(), std::move(NonCallables),
-                                 JITDylibLookupFlags::MatchAllSymbols))) {
-      getExecutionSession().reportError(std::move(Err));
-      R->failMaterialization();
-      return;
-    }
-  if (!Callables.empty()) {
-    if (auto Err = R->replace(
-            lazyReexports(LCTMgr, PDR.getISManager(), PDR.getImplDylib(),
-                          std::move(Callables), AliaseeImpls))) {
-      getExecutionSession().reportError(std::move(Err));
-      R->failMaterialization();
-      return;
-    }
-  }
 }
 
-CompileOnDemandLayer::PerDylibResources &
-CompileOnDemandLayer::getPerDylibResources(JITDylib &TargetD) {
-  std::lock_guard<std::mutex> Lock(CODLayerMutex);
-
-  auto I = DylibResources.find(&TargetD);
-  if (I == DylibResources.end()) {
-    auto &ImplD =
-        getExecutionSession().createBareJITDylib(TargetD.getName() + ".impl");
-    JITDylibSearchOrder NewLinkOrder;
-    TargetD.withLinkOrderDo([&](const JITDylibSearchOrder &TargetLinkOrder) {
-      NewLinkOrder = TargetLinkOrder;
-    });
-
-    assert(!NewLinkOrder.empty() && NewLinkOrder.front().first == &TargetD &&
-           NewLinkOrder.front().second ==
-               JITDylibLookupFlags::MatchAllSymbols &&
-           "TargetD must be at the front of its own search order and match "
-           "non-exported symbol");
-    NewLinkOrder.insert(std::next(NewLinkOrder.begin()),
-                        {&ImplD, JITDylibLookupFlags::MatchAllSymbols});
-    ImplD.setLinkOrder(NewLinkOrder, false);
-    TargetD.setLinkOrder(std::move(NewLinkOrder), false);
-
-    PerDylibResources PDR(ImplD, BuildIndirectStubsManager());
-    I = DylibResources.insert(std::make_pair(&TargetD, std::move(PDR))).first;
-  }
-
-  return I->second;
-}
-
-void CompileOnDemandLayer::cleanUpModule(Module &M) {
+void IRPartitionLayer::cleanUpModule(Module &M) {
   for (auto &F : M.functions()) {
     if (F.isDeclaration())
       continue;
 
     if (F.hasAvailableExternallyLinkage()) {
       F.deleteBody();
       F.setPersonalityFn(nullptr);
       continue;
     }
   }
 }
 
-void CompileOnDemandLayer::expandPartition(GlobalValueSet &Partition) {
+void IRPartitionLayer::expandPartition(GlobalValueSet &Partition) {
   // Expands the partition to ensure the following rules hold:
   // (1) If any alias is in the partition, its aliasee is also in the partition.
   // (2) If any aliasee is in the partition, its aliases are also in the
   //     partiton.
   // (3) If any global variable is in the partition then all global variables
   //     are in the partition.
   assert(!Partition.empty() && "Unexpected empty partition");
 
void IRPartitionLayer::expandPartition(GlobalValueSet &Partition) {
   if (ContainsGlobalVariables)
     for (auto &G : M.globals())
       GVsToAdd.push_back(&G);
 
   for (const auto *GV : GVsToAdd)
     Partition.insert(GV);
 }
 
-void CompileOnDemandLayer::emitPartition(
+void IRPartitionLayer::emitPartition(
     std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM,
     IRMaterializationUnit::SymbolNameToDefinitionMap Defs) {
 
   // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
   //        extracted module key, extracted module, and source module key
   //        together. This could be used, for example, to provide a specific
   //        memory manager instance to the linking layer.
 
void IRPartitionLayer::emitPartition(
 
   // Ok -- we actually need to partition the symbols. Promote the symbol
   // linkages/names, expand the partition to include any required symbols
   // (i.e. symbols that can't be separated from our partition), and
   // then extract the partition.
   //
   // 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()) {
+  auto ExtractedTSM = TSM.withModuleDo([&](Module &M)
+                                           -> Expected<ThreadSafeModule> {
 
-          MangleAndInterner Mangle(ES, M.getDataLayout());
-          SymbolFlagsMap SymbolFlags;
-          IRSymbolMapper::add(ES, *getManglingOptions(),
-                              PromotedGlobals, SymbolFlags);
+      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 (const 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 (const 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";
-        }
+      std::vector<const GlobalValue *> HashGVs;
 
-        // 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, SubModuleName , ShouldExtract);
-      });
+    return extractSubModule(TSM, SubModuleName, ShouldExtract);
 
   if (!ExtractedTSM) {
     ES.reportError(ExtractedTSM.takeError());
     R->failMaterialization();
     return;
   }
 
   if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
           ES, *getManglingOptions(), std::move(TSM), *this))) {
     ES.reportError(std::move(Err));
     R->failMaterialization();
     return;
   }
   BaseLayer.emit(std::move(R), std::move(*ExtractedTSM));
 }
-
-} // end namespace orc
-} // end namespace llvm

llvm/tools/lli/lli.cpp

 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/MCJIT.h"
 #include "llvm/ExecutionEngine/ObjectCache.h"
 #include "llvm/ExecutionEngine/Orc/DebugUtils.h"
 #include "llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h"
 #include "llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h"
 #include "llvm/ExecutionEngine/Orc/EPCGenericRTDyldMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h"
 #include "llvm/ExecutionEngine/Orc/SymbolStringPool.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/RegisterEHFrames.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/TargetExecutionUtils.h"

 #endif
 #if LLVM_USE_PERF
     RTDyldObjLayer->registerJITEventListener(
         *JITEventListener::createPerfJITEventListener());
 #endif
   }
 
   if (PerModuleLazy)
-    J->setPartitionFunction(orc::CompileOnDemandLayer::compileWholeModule);
+    J->setPartitionFunction(orc::IRPartitionLayer::compileWholeModule);
 
   auto Dump = createDebugDumper();
 
   J->getIRTransformLayer().setTransform(
       [&](orc::ThreadSafeModule TSM,
           const orc::MaterializationResponsibility &R) {
         TSM.withModuleDo([&](Module &M) {
           if (verifyModule(M, &dbgs())) {