[llvm] [ORC] Introduce IRPartitionLayer for common partition functionality. (PR #66812)
Sunho Kim via llvm-commits
llvm-commits at lists.llvm.org
Tue Sep 19 13:12:50 PDT 2023
https://github.com/sunho created https://github.com/llvm/llvm-project/pull/66812
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 "CompileOnDemandLayer -> IRPartitionLayer -> ReOptLayer" as well as ordinary reoptimization like "IRPartitionLayer -> ReOptLayer"
>From adac015517504aa40eafb15fded4f7ab4b4b780d Mon Sep 17 00:00:00 2001
From: Sunho Kim <ksunhokim123 at gmail.com>
Date: Wed, 20 Sep 2023 05:11:15 +0900
Subject: [PATCH] [ORC] Introduce IRPartitionLayer for common partition
functionality.
---
.../BuildingAJIT/Chapter3/KaleidoscopeJIT.h | 6 +-
.../SpeculativeJIT/SpeculativeJIT.cpp | 9 +-
.../Orc/CompileOnDemandLayer.h | 36 +--
.../ExecutionEngine/Orc/IRPartitionLayer.h | 83 +++++
llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h | 7 +-
llvm/lib/ExecutionEngine/Orc/CMakeLists.txt | 1 +
.../Orc/CompileOnDemandLayer.cpp | 282 +---------------
.../ExecutionEngine/Orc/IRPartitionLayer.cpp | 303 ++++++++++++++++++
llvm/lib/ExecutionEngine/Orc/LLJIT.cpp | 7 +-
llvm/tools/lli/lli.cpp | 3 +-
10 files changed, 415 insertions(+), 322 deletions(-)
create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h
create mode 100644 llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp
diff --git a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h
index fd0e081ff2b4bd0..fee2d26e5d9254f 100644
--- a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h
+++ b/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(
diff --git a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp b/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp
index fdd376d82da5d89..a1722748c6ce069 100644
--- a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp
+++ b/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;
};
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
index 23a64b5836db85f..ee556d11b708352 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
+++ b/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;
};
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h
new file mode 100644
index 000000000000000..0535ce473837474
--- /dev/null
+++ b/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
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
index ab54e0c2c288205..896de521f13af80 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
+++ b/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;
};
diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index 3256ed8b7362c66..51d1ff1a86f2dbc 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/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
diff --git a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
index 6448adaa0ceb36f..9296bc2b389ab24 100644
--- a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
@@ -9,6 +9,7 @@
#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"
@@ -17,101 +18,6 @@
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)
@@ -119,13 +25,10 @@ CompileOnDemandLayer::CompileOnDemandLayer(
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");
@@ -138,10 +41,6 @@ void CompileOnDemandLayer::emit(
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;
@@ -152,11 +51,10 @@ void CompileOnDemandLayer::emit(
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;
@@ -210,173 +108,3 @@ CompileOnDemandLayer::getPerDylibResources(JITDylib &TargetD) {
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
diff --git a/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp b/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp
new file mode 100644
index 000000000000000..9ad171beac7fe2e
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp
@@ -0,0 +1,303 @@
+//===----- 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/IRPartitionLayer.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+
+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,
+ IRPartitionLayer &Parent)
+ : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {}
+
+ PartitioningIRMaterializationUnit(
+ ThreadSafeModule TSM, Interface I,
+ 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");
+ }
+
+ IRPartitionLayer &Parent;
+};
+
+} // namespace orc
+} // namespace llvm
+
+IRPartitionLayer::IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer)
+ : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer) {}
+
+void IRPartitionLayer::setPartitionFunction(PartitionFunction Partition) {
+ this->Partition = Partition;
+}
+
+std::optional<IRPartitionLayer::GlobalValueSet>
+IRPartitionLayer::compileRequested(GlobalValueSet Requested) {
+ return std::move(Requested);
+}
+
+std::optional<IRPartitionLayer::GlobalValueSet>
+IRPartitionLayer::compileWholeModule(GlobalValueSet Requested) {
+ return std::nullopt;
+}
+
+void IRPartitionLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
+ ThreadSafeModule TSM) {
+ assert(TSM && "Null module");
+
+ auto &ES = getExecutionSession();
+ TSM.withModuleDo([&](Module &M) {
+ // First, do some cleanup on the module:
+ cleanUpModule(M);
+ });
+
+ // Create a partitioning materialization unit and pass the responsibility.
+ if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
+ ES, *getManglingOptions(), std::move(TSM), *this))) {
+ ES.reportError(std::move(Err));
+ R->failMaterialization();
+ return;
+ }
+}
+
+void IRPartitionLayer::cleanUpModule(Module &M) {
+ for (auto &F : M.functions()) {
+ if (F.isDeclaration())
+ continue;
+
+ if (F.hasAvailableExternallyLinkage()) {
+ F.deleteBody();
+ F.setPersonalityFn(nullptr);
+ continue;
+ }
+ }
+}
+
+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");
+
+ 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 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.
+
+ 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));
+}
diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index 18ab9e6b940fefe..15fc50bad94aabf 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
@@ -1288,9 +1288,12 @@ LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) {
return;
}
+ // Create the IP Layer.
+ IPLayer = std::make_unique<IRPartitionLayer>(*ES, *InitHelperTransformLayer);
+
// Create the COD layer.
- CODLayer = std::make_unique<CompileOnDemandLayer>(
- *ES, *InitHelperTransformLayer, *LCTMgr, std::move(ISMBuilder));
+ CODLayer = std::make_unique<CompileOnDemandLayer>(*ES, *IPLayer, *LCTMgr,
+ std::move(ISMBuilder));
if (S.NumCompileThreads > 0)
CODLayer->setCloneToNewContextOnEmit(true);
diff --git a/llvm/tools/lli/lli.cpp b/llvm/tools/lli/lli.cpp
index 360ffa628a47b3f..1378622cb71511d 100644
--- a/llvm/tools/lli/lli.cpp
+++ b/llvm/tools/lli/lli.cpp
@@ -30,6 +30,7 @@
#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"
@@ -984,7 +985,7 @@ int runOrcJIT(const char *ProgName) {
}
if (PerModuleLazy)
- J->setPartitionFunction(orc::CompileOnDemandLayer::compileWholeModule);
+ J->setPartitionFunction(orc::IRPartitionLayer::compileWholeModule);
auto Dump = createDebugDumper();
More information about the llvm-commits
mailing list