Skip to content

Commit 2c01b27

Browse files
sunhosunho
committed
[ORC] Introduce RedirectionManager interface and implementation using JITLink.
1 parent 50866e8 commit 2c01b27

File tree

7 files changed

+513
-0
lines changed

7 files changed

+513
-0
lines changed

‎llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
//===- JITLinkRedirectableSymbolManager.h - JITLink redirection -*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// Redirectable Symbol Manager implementation using JITLink
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#ifndef LLVM_EXECUTIONENGINE_ORC_JITLINKREDIRECABLEMANAGER_H
14+
#define LLVM_EXECUTIONENGINE_ORC_JITLINKREDIRECABLEMANAGER_H
15+
16+
#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
17+
#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
18+
#include "llvm/Support/StringSaver.h"
19+
20+
namespace llvm {
21+
namespace orc {
22+
23+
class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
24+
public ResourceManager {
25+
public:
26+
/// Create redirection manager that uses JITLink based implementaion.
27+
static Expected<std::unique_ptr<RedirectableSymbolManager>>
28+
Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
29+
JITDylib &JD) {
30+
Error Err = Error::success();
31+
auto RM = std::unique_ptr<RedirectableSymbolManager>(
32+
new JITLinkRedirectableSymbolManager(ES, ObjLinkingLayer, JD, Err));
33+
if (Err)
34+
return Err;
35+
return std::move(RM);
36+
}
37+
38+
void emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> R,
39+
const SymbolAddrMap &InitialDests) override;
40+
41+
Error redirect(JITDylib &TargetJD, const SymbolAddrMap &NewDests) override;
42+
43+
Error handleRemoveResources(JITDylib &TargetJD, ResourceKey K) override;
44+
45+
void handleTransferResources(JITDylib &TargetJD, ResourceKey DstK,
46+
ResourceKey SrcK) override;
47+
48+
private:
49+
using StubHandle = unsigned;
50+
constexpr static unsigned StubBlockSize = 256;
51+
constexpr static StringRef JumpStubPrefix = "$__IND_JUMP_STUBS";
52+
constexpr static StringRef StubPtrPrefix = "$IND_JUMP_PTR_";
53+
constexpr static StringRef JumpStubTableName = "$IND_JUMP_";
54+
constexpr static StringRef StubPtrTableName = "$__IND_JUMP_PTRS";
55+
56+
JITLinkRedirectableSymbolManager(ExecutionSession &ES,
57+
ObjectLinkingLayer &ObjLinkingLayer,
58+
JITDylib &JD, Error &Err)
59+
: ES(ES), ObjLinkingLayer(ObjLinkingLayer), JD(JD),
60+
AnonymousPtrCreator(
61+
jitlink::getAnonymousPointerCreator(ES.getTargetTriple())),
62+
PtrJumpStubCreator(
63+
jitlink::getPointerJumpStubCreator(ES.getTargetTriple())) {
64+
if (!AnonymousPtrCreator || !PtrJumpStubCreator)
65+
Err = make_error<StringError>("Architecture not supported",
66+
inconvertibleErrorCode());
67+
if (Err)
68+
return;
69+
ES.registerResourceManager(*this);
70+
}
71+
72+
~JITLinkRedirectableSymbolManager() { ES.deregisterResourceManager(*this); }
73+
74+
StringRef JumpStubSymbolName(unsigned I) {
75+
return *ES.intern((JumpStubPrefix + Twine(I)).str());
76+
}
77+
78+
StringRef StubPtrSymbolName(unsigned I) {
79+
return *ES.intern((StubPtrPrefix + Twine(I)).str());
80+
}
81+
82+
unsigned GetNumAvailableStubs() const { return AvailableStubs.size(); }
83+
84+
Error redirectInner(JITDylib &TargetJD, const SymbolAddrMap &NewDests);
85+
Error grow(unsigned Need);
86+
87+
ExecutionSession &ES;
88+
ObjectLinkingLayer &ObjLinkingLayer;
89+
JITDylib &JD;
90+
jitlink::AnonymousPointerCreator AnonymousPtrCreator;
91+
jitlink::PointerJumpStubCreator PtrJumpStubCreator;
92+
93+
std::vector<StubHandle> AvailableStubs;
94+
using SymbolToStubMap = DenseMap<SymbolStringPtr, StubHandle>;
95+
DenseMap<JITDylib *, SymbolToStubMap> SymbolToStubs;
96+
std::vector<ExecutorSymbolDef> JumpStubs;
97+
std::vector<ExecutorSymbolDef> StubPointers;
98+
DenseMap<ResourceKey, std::vector<SymbolStringPtr>> TrackedResources;
99+
100+
std::mutex Mutex;
101+
};
102+
103+
} // namespace orc
104+
} // namespace llvm
105+
106+
#endif

‎llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
//===- RedirectionManager.h - Redirection manager interface -----*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// Redirection manager interface that redirects a call to symbol to another.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#ifndef LLVM_EXECUTIONENGINE_ORC_REDIRECTIONMANAGER_H
14+
#define LLVM_EXECUTIONENGINE_ORC_REDIRECTIONMANAGER_H
15+
16+
#include "llvm/ExecutionEngine/Orc/Core.h"
17+
18+
namespace llvm {
19+
namespace orc {
20+
21+
/// Base class for performing redirection of call to symbol to another symbol in
22+
/// runtime.
23+
class RedirectionManager {
24+
public:
25+
/// Symbol name to symbol definition map.
26+
using SymbolAddrMap = DenseMap<SymbolStringPtr, ExecutorSymbolDef>;
27+
28+
virtual ~RedirectionManager() = default;
29+
/// Change the redirection destination of given symbols to new destination
30+
/// symbols.
31+
virtual Error redirect(JITDylib &JD, const SymbolAddrMap &NewDests) = 0;
32+
33+
/// Change the redirection destination of given symbol to new destination
34+
/// symbol.
35+
virtual Error redirect(JITDylib &JD, SymbolStringPtr Symbol,
36+
ExecutorSymbolDef NewDest) {
37+
return redirect(JD, {{Symbol, NewDest}});
38+
}
39+
40+
private:
41+
virtual void anchor();
42+
};
43+
44+
/// Base class for managing redirectable symbols in which a call
45+
/// gets redirected to another symbol in runtime.
46+
class RedirectableSymbolManager : public RedirectionManager {
47+
public:
48+
/// Create redirectable symbols with given symbol names and initial
49+
/// desitnation symbol addresses.
50+
Error createRedirectableSymbols(ResourceTrackerSP RT,
51+
const SymbolMap &InitialDests);
52+
53+
/// Create a single redirectable symbol with given symbol name and initial
54+
/// desitnation symbol address.
55+
Error createRedirectableSymbol(ResourceTrackerSP RT, SymbolStringPtr Symbol,
56+
ExecutorSymbolDef InitialDest) {
57+
return createRedirectableSymbols(RT, {{Symbol, InitialDest}});
58+
}
59+
60+
/// Emit redirectable symbol
61+
virtual void
62+
emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> MR,
63+
const SymbolMap &InitialDests) = 0;
64+
};
65+
66+
class RedirectableMaterializationUnit : public MaterializationUnit {
67+
public:
68+
RedirectableMaterializationUnit(RedirectableSymbolManager &RM,
69+
const SymbolMap &InitialDests)
70+
: MaterializationUnit(convertToFlags(InitialDests)), RM(RM),
71+
InitialDests(InitialDests) {}
72+
73+
StringRef getName() const override {
74+
return "RedirectableSymbolMaterializationUnit";
75+
}
76+
77+
void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
78+
RM.emitRedirectableSymbols(std::move(R), std::move(InitialDests));
79+
}
80+
81+
void discard(const JITDylib &JD, const SymbolStringPtr &Name) override {
82+
InitialDests.erase(Name);
83+
}
84+
85+
private:
86+
static MaterializationUnit::Interface
87+
convertToFlags(const SymbolMap &InitialDests) {
88+
SymbolFlagsMap Flags;
89+
for (auto [K, V] : InitialDests)
90+
Flags[K] = V.getFlags();
91+
return MaterializationUnit::Interface(Flags, {});
92+
}
93+
94+
RedirectableSymbolManager &RM;
95+
SymbolMap InitialDests;
96+
};
97+
98+
} // namespace orc
99+
} // namespace llvm
100+
101+
#endif

‎llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ add_llvm_component_library(LLVMOrcJIT
5050
ExecutorProcessControl.cpp
5151
TaskDispatch.cpp
5252
ThreadSafeModule.cpp
53+
RedirectionManager.cpp
54+
JITLinkRedirectableSymbolManager.cpp
5355
ADDITIONAL_HEADER_DIRS
5456
${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc
5557

‎llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
//===-- JITLinkRedirectableSymbolManager.cpp - JITLink redirection in Orc -===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"
10+
#include "llvm/ExecutionEngine/Orc/Core.h"
11+
12+
#define DEBUG_TYPE "orc"
13+
14+
using namespace llvm;
15+
using namespace llvm::orc;
16+
17+
void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
18+
std::unique_ptr<MaterializationResponsibility> R,
19+
const SymbolAddrMap &InitialDests) {
20+
std::unique_lock<std::mutex> Lock(Mutex);
21+
if (GetNumAvailableStubs() < InitialDests.size())
22+
if (auto Err = grow(InitialDests.size() - GetNumAvailableStubs())) {
23+
ES.reportError(std::move(Err));
24+
R->failMaterialization();
25+
return;
26+
}
27+
28+
JITDylib &TargetJD = R->getTargetJITDylib();
29+
SymbolMap NewSymbolDefs;
30+
std::vector<SymbolStringPtr> Symbols;
31+
for (auto &[K, V] : InitialDests) {
32+
StubHandle StubID = AvailableStubs.back();
33+
if (SymbolToStubs[&TargetJD].count(K)) {
34+
ES.reportError(make_error<StringError>(
35+
"Tried to create duplicate redirectable symbols",
36+
inconvertibleErrorCode()));
37+
R->failMaterialization();
38+
return;
39+
}
40+
dbgs() << *K << "\n";
41+
SymbolToStubs[&TargetJD][K] = StubID;
42+
NewSymbolDefs[K] = JumpStubs[StubID];
43+
NewSymbolDefs[K].setFlags(V.getFlags());
44+
Symbols.push_back(K);
45+
AvailableStubs.pop_back();
46+
}
47+
48+
if (auto Err = R->replace(absoluteSymbols(NewSymbolDefs))) {
49+
ES.reportError(std::move(Err));
50+
R->failMaterialization();
51+
return;
52+
}
53+
54+
if (auto Err = redirectInner(TargetJD, InitialDests)) {
55+
ES.reportError(std::move(Err));
56+
R->failMaterialization();
57+
return;
58+
}
59+
60+
auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
61+
TrackedResources[Key].insert(TrackedResources[Key].end(), Symbols.begin(),
62+
Symbols.end());
63+
});
64+
if (Err) {
65+
ES.reportError(std::move(Err));
66+
R->failMaterialization();
67+
return;
68+
}
69+
}
70+
71+
Error JITLinkRedirectableSymbolManager::redirect(
72+
JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
73+
std::unique_lock<std::mutex> Lock(Mutex);
74+
return redirectInner(TargetJD, NewDests);
75+
}
76+
77+
Error JITLinkRedirectableSymbolManager::redirectInner(
78+
JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
79+
std::vector<tpctypes::PointerWrite> PtrWrites;
80+
for (auto &[K, V] : NewDests) {
81+
if (!SymbolToStubs[&TargetJD].count(K))
82+
return make_error<StringError>(
83+
"Tried to redirect non-existent redirectalbe symbol",
84+
inconvertibleErrorCode());
85+
StubHandle StubID = SymbolToStubs[&TargetJD].at(K);
86+
PtrWrites.push_back({StubPointers[StubID].getAddress(), V.getAddress()});
87+
}
88+
if (auto Err = ES.getExecutorProcessControl().getMemoryAccess().writePointers(
89+
PtrWrites))
90+
return Err;
91+
return Error::success();
92+
}
93+
94+
Error JITLinkRedirectableSymbolManager::grow(unsigned Need) {
95+
unsigned OldSize = JumpStubs.size();
96+
unsigned NumNewStubs = alignTo(Need, StubBlockSize);
97+
unsigned NewSize = OldSize + NumNewStubs;
98+
99+
JumpStubs.resize(NewSize);
100+
StubPointers.resize(NewSize);
101+
AvailableStubs.reserve(NewSize);
102+
103+
SymbolLookupSet LookupSymbols;
104+
DenseMap<SymbolStringPtr, ExecutorSymbolDef *> NewDefsMap;
105+
106+
Triple TT = ES.getTargetTriple();
107+
auto G = std::make_unique<jitlink::LinkGraph>(
108+
"<INDIRECT STUBS>", TT, TT.isArch64Bit() ? 8 : 4,
109+
TT.isLittleEndian() ? support::little : support::big,
110+
jitlink::getGenericEdgeKindName);
111+
auto &PointerSection =
112+
G->createSection(StubPtrTableName, MemProt::Write | MemProt::Read);
113+
auto &StubsSection =
114+
G->createSection(JumpStubTableName, MemProt::Exec | MemProt::Read);
115+
116+
for (size_t I = OldSize; I < NewSize; I++) {
117+
auto Pointer = AnonymousPtrCreator(*G, PointerSection, nullptr, 0);
118+
if (auto Err = Pointer.takeError())
119+
return Err;
120+
121+
StringRef PtrSymName = StubPtrSymbolName(I);
122+
Pointer->setName(PtrSymName);
123+
Pointer->setScope(jitlink::Scope::Default);
124+
LookupSymbols.add(ES.intern(PtrSymName));
125+
NewDefsMap[ES.intern(PtrSymName)] = &StubPointers[I];
126+
127+
auto Stub = PtrJumpStubCreator(*G, StubsSection, *Pointer);
128+
if (auto Err = Stub.takeError())
129+
return Err;
130+
131+
StringRef JumpStubSymName = JumpStubSymbolName(I);
132+
Stub->setName(JumpStubSymName);
133+
Stub->setScope(jitlink::Scope::Default);
134+
LookupSymbols.add(ES.intern(JumpStubSymName));
135+
NewDefsMap[ES.intern(JumpStubSymName)] = &JumpStubs[I];
136+
}
137+
138+
if (auto Err = ObjLinkingLayer.add(JD, std::move(G)))
139+
return Err;
140+
141+
auto LookupResult = ES.lookup(makeJITDylibSearchOrder(&JD), LookupSymbols);
142+
if (auto Err = LookupResult.takeError())
143+
return Err;
144+
145+
for (auto &[K, V] : *LookupResult)
146+
*NewDefsMap.at(K) = V;
147+
148+
for (size_t I = OldSize; I < NewSize; I++)
149+
AvailableStubs.push_back(I);
150+
151+
return Error::success();
152+
}
153+
154+
Error JITLinkRedirectableSymbolManager::handleRemoveResources(
155+
JITDylib &TargetJD, ResourceKey K) {
156+
std::unique_lock<std::mutex> Lock(Mutex);
157+
for (auto &Symbol : TrackedResources[K]) {
158+
if (!SymbolToStubs[&TargetJD].count(Symbol))
159+
return make_error<StringError>(
160+
"Tried to remove non-existent redirectable symbol",
161+
inconvertibleErrorCode());
162+
AvailableStubs.push_back(SymbolToStubs[&TargetJD].at(Symbol));
163+
SymbolToStubs[&TargetJD].erase(Symbol);
164+
if (SymbolToStubs[&TargetJD].empty())
165+
SymbolToStubs.erase(&TargetJD);
166+
}
167+
TrackedResources.erase(K);
168+
169+
return Error::success();
170+
}
171+
172+
void JITLinkRedirectableSymbolManager::handleTransferResources(
173+
JITDylib &TargetJD, ResourceKey DstK, ResourceKey SrcK) {
174+
std::unique_lock<std::mutex> Lock(Mutex);
175+
TrackedResources[DstK].insert(TrackedResources[DstK].end(),
176+
TrackedResources[SrcK].begin(),
177+
TrackedResources[SrcK].end());
178+
TrackedResources.erase(SrcK);
179+
}

0 commit comments

Comments
 (0)