[ORC] Introduce IRPartitionLayer for common partition functionality.

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

@@ -21,6 +21,7 @@
 #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"
@@ -48,6 +49,7 @@ class KaleidoscopeJIT {
   RTDyldObjectLinkingLayer ObjectLayer;
   IRCompileLayer CompileLayer;
   IRTransformLayer OptimizeLayer;
+  IRPartitionLayer IPLayer;
   CompileOnDemandLayer CODLayer;
 
   JITDylib &MainJD;
@@ -68,8 +70,8 @@ class KaleidoscopeJIT {
         CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
-        CODLayer(*this->ES, OptimizeLayer,
-                 this->EPCIU->getLazyCallThroughManager(),
+        IPLayer(*this->ES, OptimizeLayer),
+        CODLayer(*this->ES, IPLayer, this->EPCIU->getLazyCallThroughManager(),
                  [this] { return this->EPCIU->createIndirectStubsManager(); }),
         MainJD(this->ES->createBareJITDylib("<main>")) {
     MainJD.addGenerator(

llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp

@@ -3,6 +3,7 @@
 #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"
@@ -107,13 +108,14 @@ class SpeculativeJIT {
       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)) {
+        IPLayer(*this->ES, SpeculateLayer),
+        CODLayer(*this->ES, IPLayer, *this->LCTMgr, std::move(ISMBuilder)) {
     MainJD.addGenerator(std::move(ProcessSymbolsGenerator));
     this->CODLayer.setImplMap(&Imps);
     this->ES->setDispatchTask(
@@ -147,6 +149,7 @@ class SpeculativeJIT {
   Speculator S;
   RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr};
   IRSpeculationLayer SpeculateLayer;
+  IRPartitionLayer IPLayer;
   CompileOnDemandLayer CODLayer;
 };
 

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

@@ -54,37 +54,15 @@ 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);
-
+                       LazyCallThroughManager &LCTMgr,
+                       IndirectStubsManagerBuilder BuildIndirectStubsManager);
   /// Sets the ImplSymbolMap
   void setImplMap(ImplSymbolMap *Imp);
 
@@ -111,22 +89,12 @@ class CompileOnDemandLayer : public IRLayer {
 
   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;
 };
 

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

@@ -0,0 +1,83 @@
+//===- 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 into smaller submodules that only contains
+// looked up symbols.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
+#define LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
+
+#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+#include "llvm/ExecutionEngine/Orc/Layer.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"
+
+namespace llvm {
+namespace orc {
+
+class IRPartitionLayer : public IRLayer {
+  friend class PartitioningIRMaterializationUnit;
+
+public:
+  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);
+
+  /// Sets the partition function.
+  void setPartitionFunction(PartitionFunction Partition);
+
+  /// 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:
+  void cleanUpModule(Module &M);
+
+  void expandPartition(GlobalValueSet &Partition);
+
+  void emitPartition(std::unique_ptr<MaterializationResponsibility> R,
+                     ThreadSafeModule TSM,
+                     IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
+
+  IRLayer &BaseLayer;
+  PartitionFunction Partition = compileRequested;
+  SymbolLinkagePromoter PromoteSymbols;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif

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

@@ -17,6 +17,7 @@
 #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"
@@ -271,9 +272,8 @@ class LLLazyJIT : public LLJIT {
 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.
@@ -293,6 +293,7 @@ class LLLazyJIT : public LLJIT {
   LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
 
   std::unique_ptr<LazyCallThroughManager> LCTMgr;
+  std::unique_ptr<IRPartitionLayer> IPLayer;
   std::unique_ptr<CompileOnDemandLayer> CODLayer;
 };
 

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -27,6 +27,7 @@ add_llvm_component_library(LLVMOrcJIT
   IndirectionUtils.cpp
   IRCompileLayer.cpp
   IRTransformLayer.cpp
+  IRPartitionLayer.cpp
   JITTargetMachineBuilder.cpp
   LazyReexports.cpp
   Layer.cpp