[compiler-rt] [llvm] [ORC] Implement basic reoptimization. (PR #67050)
Sunho Kim via llvm-commits
llvm-commits at lists.llvm.org
Wed Apr 10 21:42:08 PDT 2024
https://github.com/sunho updated https://github.com/llvm/llvm-project/pull/67050
>From 7f4db328ad13c37cb68e10a759cecc3851bef89b Mon Sep 17 00:00:00 2001
From: Sunho Kim <ksunhokim123 at gmail.com>
Date: Mon, 21 Aug 2023 15:48:26 +0900
Subject: [PATCH 1/3] [ORC] Introduce RedirectionManager interface and
implementation using JITLink.
---
.../Orc/JITLinkRedirectableSymbolManager.h | 106 +++++++++++
.../ExecutionEngine/Orc/RedirectionManager.h | 101 ++++++++++
llvm/lib/ExecutionEngine/Orc/CMakeLists.txt | 2 +
.../Orc/JITLinkRedirectableSymbolManager.cpp | 179 ++++++++++++++++++
.../Orc/RedirectionManager.cpp | 24 +++
.../ExecutionEngine/Orc/CMakeLists.txt | 1 +
.../Orc/JITLinkRedirectionManagerTest.cpp | 100 ++++++++++
7 files changed, 513 insertions(+)
create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
create mode 100644 llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
create mode 100644 llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp
create mode 100644 llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
new file mode 100644
index 00000000000000..5de0da1f52d0db
--- /dev/null
+++ b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
@@ -0,0 +1,106 @@
+//===- JITLinkRedirectableSymbolManager.h - JITLink redirection -*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Redirectable Symbol Manager implementation using JITLink
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_JITLINKREDIRECABLEMANAGER_H
+#define LLVM_EXECUTIONENGINE_ORC_JITLINKREDIRECABLEMANAGER_H
+
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
+#include "llvm/Support/StringSaver.h"
+
+namespace llvm {
+namespace orc {
+
+class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
+ public ResourceManager {
+public:
+ /// Create redirection manager that uses JITLink based implementaion.
+ static Expected<std::unique_ptr<RedirectableSymbolManager>>
+ Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
+ JITDylib &JD) {
+ Error Err = Error::success();
+ auto RM = std::unique_ptr<RedirectableSymbolManager>(
+ new JITLinkRedirectableSymbolManager(ES, ObjLinkingLayer, JD, Err));
+ if (Err)
+ return Err;
+ return std::move(RM);
+ }
+
+ void emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> R,
+ const SymbolAddrMap &InitialDests) override;
+
+ Error redirect(JITDylib &TargetJD, const SymbolAddrMap &NewDests) override;
+
+ Error handleRemoveResources(JITDylib &TargetJD, ResourceKey K) override;
+
+ void handleTransferResources(JITDylib &TargetJD, ResourceKey DstK,
+ ResourceKey SrcK) override;
+
+private:
+ using StubHandle = unsigned;
+ constexpr static unsigned StubBlockSize = 256;
+ constexpr static StringRef JumpStubPrefix = "$__IND_JUMP_STUBS";
+ constexpr static StringRef StubPtrPrefix = "$IND_JUMP_PTR_";
+ constexpr static StringRef JumpStubTableName = "$IND_JUMP_";
+ constexpr static StringRef StubPtrTableName = "$__IND_JUMP_PTRS";
+
+ JITLinkRedirectableSymbolManager(ExecutionSession &ES,
+ ObjectLinkingLayer &ObjLinkingLayer,
+ JITDylib &JD, Error &Err)
+ : ES(ES), ObjLinkingLayer(ObjLinkingLayer), JD(JD),
+ AnonymousPtrCreator(
+ jitlink::getAnonymousPointerCreator(ES.getTargetTriple())),
+ PtrJumpStubCreator(
+ jitlink::getPointerJumpStubCreator(ES.getTargetTriple())) {
+ if (!AnonymousPtrCreator || !PtrJumpStubCreator)
+ Err = make_error<StringError>("Architecture not supported",
+ inconvertibleErrorCode());
+ if (Err)
+ return;
+ ES.registerResourceManager(*this);
+ }
+
+ ~JITLinkRedirectableSymbolManager() { ES.deregisterResourceManager(*this); }
+
+ StringRef JumpStubSymbolName(unsigned I) {
+ return *ES.intern((JumpStubPrefix + Twine(I)).str());
+ }
+
+ StringRef StubPtrSymbolName(unsigned I) {
+ return *ES.intern((StubPtrPrefix + Twine(I)).str());
+ }
+
+ unsigned GetNumAvailableStubs() const { return AvailableStubs.size(); }
+
+ Error redirectInner(JITDylib &TargetJD, const SymbolAddrMap &NewDests);
+ Error grow(unsigned Need);
+
+ ExecutionSession &ES;
+ ObjectLinkingLayer &ObjLinkingLayer;
+ JITDylib &JD;
+ jitlink::AnonymousPointerCreator AnonymousPtrCreator;
+ jitlink::PointerJumpStubCreator PtrJumpStubCreator;
+
+ std::vector<StubHandle> AvailableStubs;
+ using SymbolToStubMap = DenseMap<SymbolStringPtr, StubHandle>;
+ DenseMap<JITDylib *, SymbolToStubMap> SymbolToStubs;
+ std::vector<ExecutorSymbolDef> JumpStubs;
+ std::vector<ExecutorSymbolDef> StubPointers;
+ DenseMap<ResourceKey, std::vector<SymbolStringPtr>> TrackedResources;
+
+ std::mutex Mutex;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h b/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
new file mode 100644
index 00000000000000..87a81b0e529ccd
--- /dev/null
+++ b/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
@@ -0,0 +1,101 @@
+//===- RedirectionManager.h - Redirection manager interface -----*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Redirection manager interface that redirects a call to symbol to another.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_REDIRECTIONMANAGER_H
+#define LLVM_EXECUTIONENGINE_ORC_REDIRECTIONMANAGER_H
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
+
+namespace llvm {
+namespace orc {
+
+/// Base class for performing redirection of call to symbol to another symbol in
+/// runtime.
+class RedirectionManager {
+public:
+ /// Symbol name to symbol definition map.
+ using SymbolAddrMap = DenseMap<SymbolStringPtr, ExecutorSymbolDef>;
+
+ virtual ~RedirectionManager() = default;
+ /// Change the redirection destination of given symbols to new destination
+ /// symbols.
+ virtual Error redirect(JITDylib &JD, const SymbolAddrMap &NewDests) = 0;
+
+ /// Change the redirection destination of given symbol to new destination
+ /// symbol.
+ virtual Error redirect(JITDylib &JD, SymbolStringPtr Symbol,
+ ExecutorSymbolDef NewDest) {
+ return redirect(JD, {{Symbol, NewDest}});
+ }
+
+private:
+ virtual void anchor();
+};
+
+/// Base class for managing redirectable symbols in which a call
+/// gets redirected to another symbol in runtime.
+class RedirectableSymbolManager : public RedirectionManager {
+public:
+ /// Create redirectable symbols with given symbol names and initial
+ /// desitnation symbol addresses.
+ Error createRedirectableSymbols(ResourceTrackerSP RT,
+ const SymbolMap &InitialDests);
+
+ /// Create a single redirectable symbol with given symbol name and initial
+ /// desitnation symbol address.
+ Error createRedirectableSymbol(ResourceTrackerSP RT, SymbolStringPtr Symbol,
+ ExecutorSymbolDef InitialDest) {
+ return createRedirectableSymbols(RT, {{Symbol, InitialDest}});
+ }
+
+ /// Emit redirectable symbol
+ virtual void
+ emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> MR,
+ const SymbolMap &InitialDests) = 0;
+};
+
+class RedirectableMaterializationUnit : public MaterializationUnit {
+public:
+ RedirectableMaterializationUnit(RedirectableSymbolManager &RM,
+ const SymbolMap &InitialDests)
+ : MaterializationUnit(convertToFlags(InitialDests)), RM(RM),
+ InitialDests(InitialDests) {}
+
+ StringRef getName() const override {
+ return "RedirectableSymbolMaterializationUnit";
+ }
+
+ void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
+ RM.emitRedirectableSymbols(std::move(R), std::move(InitialDests));
+ }
+
+ void discard(const JITDylib &JD, const SymbolStringPtr &Name) override {
+ InitialDests.erase(Name);
+ }
+
+private:
+ static MaterializationUnit::Interface
+ convertToFlags(const SymbolMap &InitialDests) {
+ SymbolFlagsMap Flags;
+ for (auto [K, V] : InitialDests)
+ Flags[K] = V.getFlags();
+ return MaterializationUnit::Interface(Flags, {});
+ }
+
+ RedirectableSymbolManager &RM;
+ SymbolMap InitialDests;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif
diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index 3256ed8b7362c6..17e611d5795cec 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
@@ -49,6 +49,8 @@ add_llvm_component_library(LLVMOrcJIT
ExecutorProcessControl.cpp
TaskDispatch.cpp
ThreadSafeModule.cpp
+ RedirectionManager.cpp
+ JITLinkRedirectableSymbolManager.cpp
ADDITIONAL_HEADER_DIRS
${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc
diff --git a/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
new file mode 100644
index 00000000000000..6d76d50271b9f2
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
@@ -0,0 +1,179 @@
+//===-- JITLinkRedirectableSymbolManager.cpp - JITLink redirection in Orc -===//
+//
+// 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/JITLinkRedirectableSymbolManager.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+
+#define DEBUG_TYPE "orc"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
+ std::unique_ptr<MaterializationResponsibility> R,
+ const SymbolAddrMap &InitialDests) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ if (GetNumAvailableStubs() < InitialDests.size())
+ if (auto Err = grow(InitialDests.size() - GetNumAvailableStubs())) {
+ ES.reportError(std::move(Err));
+ R->failMaterialization();
+ return;
+ }
+
+ JITDylib &TargetJD = R->getTargetJITDylib();
+ SymbolMap NewSymbolDefs;
+ std::vector<SymbolStringPtr> Symbols;
+ for (auto &[K, V] : InitialDests) {
+ StubHandle StubID = AvailableStubs.back();
+ if (SymbolToStubs[&TargetJD].count(K)) {
+ ES.reportError(make_error<StringError>(
+ "Tried to create duplicate redirectable symbols",
+ inconvertibleErrorCode()));
+ R->failMaterialization();
+ return;
+ }
+ dbgs() << *K << "\n";
+ SymbolToStubs[&TargetJD][K] = StubID;
+ NewSymbolDefs[K] = JumpStubs[StubID];
+ NewSymbolDefs[K].setFlags(V.getFlags());
+ Symbols.push_back(K);
+ AvailableStubs.pop_back();
+ }
+
+ if (auto Err = R->replace(absoluteSymbols(NewSymbolDefs))) {
+ ES.reportError(std::move(Err));
+ R->failMaterialization();
+ return;
+ }
+
+ if (auto Err = redirectInner(TargetJD, InitialDests)) {
+ ES.reportError(std::move(Err));
+ R->failMaterialization();
+ return;
+ }
+
+ auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
+ TrackedResources[Key].insert(TrackedResources[Key].end(), Symbols.begin(),
+ Symbols.end());
+ });
+ if (Err) {
+ ES.reportError(std::move(Err));
+ R->failMaterialization();
+ return;
+ }
+}
+
+Error JITLinkRedirectableSymbolManager::redirect(
+ JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ return redirectInner(TargetJD, NewDests);
+}
+
+Error JITLinkRedirectableSymbolManager::redirectInner(
+ JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
+ std::vector<tpctypes::PointerWrite> PtrWrites;
+ for (auto &[K, V] : NewDests) {
+ if (!SymbolToStubs[&TargetJD].count(K))
+ return make_error<StringError>(
+ "Tried to redirect non-existent redirectalbe symbol",
+ inconvertibleErrorCode());
+ StubHandle StubID = SymbolToStubs[&TargetJD].at(K);
+ PtrWrites.push_back({StubPointers[StubID].getAddress(), V.getAddress()});
+ }
+ if (auto Err = ES.getExecutorProcessControl().getMemoryAccess().writePointers(
+ PtrWrites))
+ return Err;
+ return Error::success();
+}
+
+Error JITLinkRedirectableSymbolManager::grow(unsigned Need) {
+ unsigned OldSize = JumpStubs.size();
+ unsigned NumNewStubs = alignTo(Need, StubBlockSize);
+ unsigned NewSize = OldSize + NumNewStubs;
+
+ JumpStubs.resize(NewSize);
+ StubPointers.resize(NewSize);
+ AvailableStubs.reserve(NewSize);
+
+ SymbolLookupSet LookupSymbols;
+ DenseMap<SymbolStringPtr, ExecutorSymbolDef *> NewDefsMap;
+
+ Triple TT = ES.getTargetTriple();
+ auto G = std::make_unique<jitlink::LinkGraph>(
+ "<INDIRECT STUBS>", TT, TT.isArch64Bit() ? 8 : 4,
+ TT.isLittleEndian() ? support::little : support::big,
+ jitlink::getGenericEdgeKindName);
+ auto &PointerSection =
+ G->createSection(StubPtrTableName, MemProt::Write | MemProt::Read);
+ auto &StubsSection =
+ G->createSection(JumpStubTableName, MemProt::Exec | MemProt::Read);
+
+ for (size_t I = OldSize; I < NewSize; I++) {
+ auto Pointer = AnonymousPtrCreator(*G, PointerSection, nullptr, 0);
+ if (auto Err = Pointer.takeError())
+ return Err;
+
+ StringRef PtrSymName = StubPtrSymbolName(I);
+ Pointer->setName(PtrSymName);
+ Pointer->setScope(jitlink::Scope::Default);
+ LookupSymbols.add(ES.intern(PtrSymName));
+ NewDefsMap[ES.intern(PtrSymName)] = &StubPointers[I];
+
+ auto Stub = PtrJumpStubCreator(*G, StubsSection, *Pointer);
+ if (auto Err = Stub.takeError())
+ return Err;
+
+ StringRef JumpStubSymName = JumpStubSymbolName(I);
+ Stub->setName(JumpStubSymName);
+ Stub->setScope(jitlink::Scope::Default);
+ LookupSymbols.add(ES.intern(JumpStubSymName));
+ NewDefsMap[ES.intern(JumpStubSymName)] = &JumpStubs[I];
+ }
+
+ if (auto Err = ObjLinkingLayer.add(JD, std::move(G)))
+ return Err;
+
+ auto LookupResult = ES.lookup(makeJITDylibSearchOrder(&JD), LookupSymbols);
+ if (auto Err = LookupResult.takeError())
+ return Err;
+
+ for (auto &[K, V] : *LookupResult)
+ *NewDefsMap.at(K) = V;
+
+ for (size_t I = OldSize; I < NewSize; I++)
+ AvailableStubs.push_back(I);
+
+ return Error::success();
+}
+
+Error JITLinkRedirectableSymbolManager::handleRemoveResources(
+ JITDylib &TargetJD, ResourceKey K) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ for (auto &Symbol : TrackedResources[K]) {
+ if (!SymbolToStubs[&TargetJD].count(Symbol))
+ return make_error<StringError>(
+ "Tried to remove non-existent redirectable symbol",
+ inconvertibleErrorCode());
+ AvailableStubs.push_back(SymbolToStubs[&TargetJD].at(Symbol));
+ SymbolToStubs[&TargetJD].erase(Symbol);
+ if (SymbolToStubs[&TargetJD].empty())
+ SymbolToStubs.erase(&TargetJD);
+ }
+ TrackedResources.erase(K);
+
+ return Error::success();
+}
+
+void JITLinkRedirectableSymbolManager::handleTransferResources(
+ JITDylib &TargetJD, ResourceKey DstK, ResourceKey SrcK) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ TrackedResources[DstK].insert(TrackedResources[DstK].end(),
+ TrackedResources[SrcK].begin(),
+ TrackedResources[SrcK].end());
+ TrackedResources.erase(SrcK);
+}
diff --git a/llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp b/llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp
new file mode 100644
index 00000000000000..cbc77c5034303a
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp
@@ -0,0 +1,24 @@
+//===---- RedirectionManager.cpp - Redirection manager interface in Orc ---===//
+//
+// 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/RedirectionManager.h"
+
+#define DEBUG_TYPE "orc"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+void RedirectionManager::anchor() {}
+
+Error RedirectableSymbolManager::createRedirectableSymbols(
+ ResourceTrackerSP RT, const SymbolMap &InitialDests) {
+ auto &JD = RT->getJITDylib();
+ return JD.define(
+ std::make_unique<RedirectableMaterializationUnit>(*this, InitialDests),
+ RT);
+}
diff --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
index 37768e91fd4472..e932ff94361a7b 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
@@ -40,6 +40,7 @@ add_llvm_unittest(OrcJITTests
TaskDispatchTest.cpp
ThreadSafeModuleTest.cpp
WrapperFunctionUtilsTest.cpp
+ JITLinkRedirectionManagerTest.cpp
)
target_link_libraries(OrcJITTests PRIVATE
diff --git a/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
new file mode 100644
index 00000000000000..0f87c1b7433238
--- /dev/null
+++ b/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
@@ -0,0 +1,100 @@
+#include "OrcTestCommon.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
+#include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::jitlink;
+
+static int initialTarget() { return 42; }
+static int middleTarget() { return 13; }
+static int finalTarget() { return 53; }
+
+class JITLinkRedirectionManagerTest : public testing::Test {
+public:
+ ~JITLinkRedirectionManagerTest() {
+ if (ES)
+ if (auto Err = ES->endSession())
+ ES->reportError(std::move(Err));
+ }
+
+protected:
+ void SetUp() override {
+ auto JTMB = JITTargetMachineBuilder::detectHost();
+ // Bail out if we can not detect the host.
+ if (!JTMB) {
+ consumeError(JTMB.takeError());
+ GTEST_SKIP();
+ }
+
+ ES = std::make_unique<ExecutionSession>(
+ std::make_unique<UnsupportedExecutorProcessControl>(
+ nullptr, nullptr, JTMB->getTargetTriple().getTriple()));
+ JD = &ES->createBareJITDylib("main");
+ ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(
+ *ES, std::make_unique<InProcessMemoryManager>(4096));
+ DL = std::make_unique<DataLayout>(
+ cantFail(JTMB->getDefaultDataLayoutForTarget()));
+ }
+ JITDylib *JD{nullptr};
+ std::unique_ptr<ExecutionSession> ES;
+ std::unique_ptr<ObjectLinkingLayer> ObjLinkingLayer;
+ std::unique_ptr<DataLayout> DL;
+};
+
+TEST_F(JITLinkRedirectionManagerTest, BasicRedirectionOperation) {
+ auto RM =
+ JITLinkRedirectableSymbolManager::Create(*ES, *ObjLinkingLayer, *JD);
+ // Bail out if we can not create
+ if (!RM) {
+ consumeError(RM.takeError());
+ GTEST_SKIP();
+ }
+
+ auto DefineTarget = [&](StringRef TargetName, ExecutorAddr Addr) {
+ SymbolStringPtr Target = ES->intern(TargetName);
+ cantFail(JD->define(std::make_unique<SimpleMaterializationUnit>(
+ SymbolFlagsMap({{Target, JITSymbolFlags::Exported}}),
+ [&](std::unique_ptr<MaterializationResponsibility> R) -> void {
+ // No dependencies registered, can't fail.
+ cantFail(
+ R->notifyResolved({{Target, {Addr, JITSymbolFlags::Exported}}}));
+ cantFail(R->notifyEmitted());
+ })));
+ return cantFail(ES->lookup({JD}, TargetName));
+ };
+
+ auto InitialTarget =
+ DefineTarget("InitialTarget", ExecutorAddr::fromPtr(&initialTarget));
+ auto MiddleTarget =
+ DefineTarget("MiddleTarget", ExecutorAddr::fromPtr(&middleTarget));
+ auto FinalTarget =
+ DefineTarget("FinalTarget", ExecutorAddr::fromPtr(&finalTarget));
+
+ auto RedirectableSymbol = ES->intern("RedirectableTarget");
+ EXPECT_THAT_ERROR(
+ (*RM)->createRedirectableSymbols(JD->getDefaultResourceTracker(),
+ {{RedirectableSymbol, InitialTarget}}),
+ Succeeded());
+ auto RTDef = cantFail(ES->lookup({JD}, RedirectableSymbol));
+
+ auto RTPtr = RTDef.getAddress().toPtr<int (*)()>();
+ auto Result = RTPtr();
+ EXPECT_EQ(Result, 42) << "Failed to call initial target";
+
+ EXPECT_THAT_ERROR((*RM)->redirect(*JD, {{RedirectableSymbol, MiddleTarget}}),
+ Succeeded());
+ Result = RTPtr();
+ EXPECT_EQ(Result, 13) << "Failed to call middle redirected target";
+
+ EXPECT_THAT_ERROR((*RM)->redirect(*JD, {{RedirectableSymbol, FinalTarget}}),
+ Succeeded());
+ Result = RTPtr();
+ EXPECT_EQ(Result, 53) << "Failed to call redirected target";
+}
>From 6059562cb908e27675b2505cc03bad4bab9566a1 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 2/3] [ORC] Introduce IRPartitionLayer for common partition
functionality.
---
.../BuildingAJIT/Chapter3/KaleidoscopeJIT.h | 6 +-
.../SpeculativeJIT/SpeculativeJIT.cpp | 9 +-
.../Orc/CompileOnDemandLayer.h | 36 +--
.../ExecutionEngine/Orc/IRPartitionLayer.h | 85 +++++
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, 417 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 fd0e081ff2b4bd..fee2d26e5d9254 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 fdd376d82da5d8..a1722748c6ce06 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 23a64b5836db85..ee556d11b70835 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 00000000000000..a2981fea61d3ea
--- /dev/null
+++ b/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h
@@ -0,0 +1,85 @@
+//===- 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 {
+
+/// A layer that breaks up IR modules into smaller submodules that only contains
+/// looked up symbols.
+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 ab54e0c2c28820..896de521f13af8 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 17e611d5795cec..379c3cc240e8c7 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 6448adaa0ceb36..9296bc2b389ab2 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 00000000000000..9ad171beac7fe2
--- /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 18ab9e6b940fef..15fc50bad94aab 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 360ffa628a47b3..1378622cb71511 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();
>From 6b5eac29215b7a6d70917377be10e873946db9dd Mon Sep 17 00:00:00 2001
From: Sunho Kim <ksunhokim123 at gmail.com>
Date: Wed, 10 Apr 2024 21:41:54 -0700
Subject: [PATCH 3/3] [ORC] Implement basic reoptimization.
---
compiler-rt/lib/orc/common.h | 6 +-
compiler-rt/lib/orc/elfnix_platform.cpp | 1 +
.../Orc/JITLinkRedirectableSymbolManager.h | 31 +-
.../ExecutionEngine/Orc/ReOptimizeLayer.h | 179 +++++++++++
.../ExecutionEngine/Orc/RedirectionManager.h | 2 +
llvm/lib/ExecutionEngine/Orc/CMakeLists.txt | 1 +
.../Orc/JITLinkRedirectableSymbolManager.cpp | 17 +-
.../ExecutionEngine/Orc/ReOptimizeLayer.cpp | 278 ++++++++++++++++++
.../ExecutionEngine/Orc/CMakeLists.txt | 1 +
.../Orc/JITLinkRedirectionManagerTest.cpp | 3 +-
.../Orc/ReOptimizeLayerTest.cpp | 149 ++++++++++
11 files changed, 641 insertions(+), 27 deletions(-)
create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
create mode 100644 llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
create mode 100644 llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
diff --git a/compiler-rt/lib/orc/common.h b/compiler-rt/lib/orc/common.h
index 73c5c4a2bd8d47..f34229a615341d 100644
--- a/compiler-rt/lib/orc/common.h
+++ b/compiler-rt/lib/orc/common.h
@@ -19,9 +19,9 @@
/// This macro should be used to define tags that will be associated with
/// handlers in the JIT process, and call can be used to define tags f
-#define ORC_RT_JIT_DISPATCH_TAG(X) \
-extern "C" char X; \
-char X = 0;
+#define ORC_RT_JIT_DISPATCH_TAG(X) \
+ ORC_RT_INTERFACE char X; \
+ char X = 0;
/// Opaque struct for external symbols.
struct __orc_rt_Opaque {};
diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
index c087e71038f950..f76a070240315f 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -28,6 +28,7 @@ using namespace __orc_rt;
using namespace __orc_rt::elfnix;
// Declare function tags for functions in the JIT process.
+ORC_RT_JIT_DISPATCH_TAG(__orc_rt_reoptimize_tag)
ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_get_initializers_tag)
ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_get_deinitializers_tag)
ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_symbol_lookup_tag)
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
index 5de0da1f52d0db..52f284c89bdade 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
@@ -25,11 +25,10 @@ class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
public:
/// Create redirection manager that uses JITLink based implementaion.
static Expected<std::unique_ptr<RedirectableSymbolManager>>
- Create(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
- JITDylib &JD) {
+ Create(ObjectLinkingLayer &ObjLinkingLayer, JITDylib &JD) {
Error Err = Error::success();
auto RM = std::unique_ptr<RedirectableSymbolManager>(
- new JITLinkRedirectableSymbolManager(ES, ObjLinkingLayer, JD, Err));
+ new JITLinkRedirectableSymbolManager(ObjLinkingLayer, JD, Err));
if (Err)
return Err;
return std::move(RM);
@@ -53,30 +52,33 @@ class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
constexpr static StringRef JumpStubTableName = "$IND_JUMP_";
constexpr static StringRef StubPtrTableName = "$__IND_JUMP_PTRS";
- JITLinkRedirectableSymbolManager(ExecutionSession &ES,
- ObjectLinkingLayer &ObjLinkingLayer,
+ JITLinkRedirectableSymbolManager(ObjectLinkingLayer &ObjLinkingLayer,
JITDylib &JD, Error &Err)
- : ES(ES), ObjLinkingLayer(ObjLinkingLayer), JD(JD),
- AnonymousPtrCreator(
- jitlink::getAnonymousPointerCreator(ES.getTargetTriple())),
- PtrJumpStubCreator(
- jitlink::getPointerJumpStubCreator(ES.getTargetTriple())) {
+ : ObjLinkingLayer(ObjLinkingLayer), JD(JD),
+ AnonymousPtrCreator(jitlink::getAnonymousPointerCreator(
+ ObjLinkingLayer.getExecutionSession().getTargetTriple())),
+ PtrJumpStubCreator(jitlink::getPointerJumpStubCreator(
+ ObjLinkingLayer.getExecutionSession().getTargetTriple())) {
if (!AnonymousPtrCreator || !PtrJumpStubCreator)
Err = make_error<StringError>("Architecture not supported",
inconvertibleErrorCode());
if (Err)
return;
- ES.registerResourceManager(*this);
+ ObjLinkingLayer.getExecutionSession().registerResourceManager(*this);
}
- ~JITLinkRedirectableSymbolManager() { ES.deregisterResourceManager(*this); }
+ ~JITLinkRedirectableSymbolManager() {
+ ObjLinkingLayer.getExecutionSession().deregisterResourceManager(*this);
+ }
StringRef JumpStubSymbolName(unsigned I) {
- return *ES.intern((JumpStubPrefix + Twine(I)).str());
+ return *ObjLinkingLayer.getExecutionSession().intern(
+ (JumpStubPrefix + Twine(I)).str());
}
StringRef StubPtrSymbolName(unsigned I) {
- return *ES.intern((StubPtrPrefix + Twine(I)).str());
+ return *ObjLinkingLayer.getExecutionSession().intern(
+ (StubPtrPrefix + Twine(I)).str());
}
unsigned GetNumAvailableStubs() const { return AvailableStubs.size(); }
@@ -84,7 +86,6 @@ class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
Error redirectInner(JITDylib &TargetJD, const SymbolAddrMap &NewDests);
Error grow(unsigned Need);
- ExecutionSession &ES;
ObjectLinkingLayer &ObjLinkingLayer;
JITDylib &JD;
jitlink::AnonymousPointerCreator AnonymousPtrCreator;
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
new file mode 100644
index 00000000000000..df2f005825b5a7
--- /dev/null
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
@@ -0,0 +1,179 @@
+//===- ReOptimizeLayer.h - Re-optimization layer interface ------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Re-optimization layer interface.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_EXECUTIONENGINE_ORC_REOPTLAYER_H
+#define LLVM_EXECUTIONENGINE_ORC_REOPTLAYER_H
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/Layer.h"
+#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
+#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+
+namespace llvm {
+namespace orc {
+
+using ReOptMaterializationUnitID = uint64_t;
+
+class ReOptimizeLayer : public IRLayer, public ResourceManager {
+public:
+ /// AddProfilerFunc will be called when ReOptimizeLayer emits the first
+ /// version of a materialization unit in order to inject profiling code and
+ /// reoptimization request code.
+ using AddProfilerFunc = unique_function<Error(
+ ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
+ unsigned CurVersion, ThreadSafeModule &TSM)>;
+
+ /// ReOptimizeFunc will be called when ReOptimizeLayer reoptimization of a
+ /// materialization unit was requested in order to reoptimize the IR module
+ /// based on profile data. OldRT is the ResourceTracker that tracks the old
+ /// function definitions. The OldRT must be kept alive until it can be
+ /// guaranteed that every invocation of the old function definitions has been
+ /// terminated.
+ using ReOptimizeFunc = unique_function<Error(
+ ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
+ unsigned CurVersion, ResourceTrackerSP OldRT, ThreadSafeModule &TSM)>;
+
+ ReOptimizeLayer(ExecutionSession &ES, IRLayer &BaseLayer,
+ RedirectableSymbolManager &RM)
+ : IRLayer(ES, BaseLayer.getManglingOptions()), ES(ES),
+ BaseLayer(BaseLayer), RSManager(RM), ReOptFunc(identity),
+ ProfilerFunc(reoptimizeIfCallFrequent) {}
+
+ void setReoptimizeFunc(ReOptimizeFunc ReOptFunc) {
+ this->ReOptFunc = std::move(ReOptFunc);
+ }
+
+ void setAddProfilerFunc(AddProfilerFunc ProfilerFunc) {
+ this->ProfilerFunc = std::move(ProfilerFunc);
+ }
+
+ /// Registers reoptimize runtime dispatch handlers to given PlatformJD. The
+ /// reoptimization request will not be handled if dispatch handler is not
+ /// registered by using this function.
+ Error reigsterRuntimeFunctions(JITDylib &PlatformJD);
+
+ /// 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;
+
+ static const uint64_t CallCountThreshold = 10;
+
+ /// Basic AddProfilerFunc that reoptimizes the function when the call count
+ /// exceeds CallCountThreshold.
+ static Error reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
+ ReOptMaterializationUnitID MUID,
+ unsigned CurVersion,
+ ThreadSafeModule &TSM);
+
+ static Error identity(ReOptimizeLayer &Parent,
+ ReOptMaterializationUnitID MUID, unsigned CurVersion,
+ ResourceTrackerSP OldRT, ThreadSafeModule &TSM) {
+ return Error::success();
+ }
+
+ // Create IR reoptimize request fucntion call.
+ static void createReoptimizeCall(Module &M, Instruction &IP,
+ GlobalVariable *ArgBuffer);
+
+ Error handleRemoveResources(JITDylib &JD, ResourceKey K) override;
+ void handleTransferResources(JITDylib &JD, ResourceKey DstK,
+ ResourceKey SrcK) override;
+
+private:
+ class ReOptMaterializationUnitState {
+ public:
+ ReOptMaterializationUnitState() = default;
+ ReOptMaterializationUnitState(ReOptMaterializationUnitID ID,
+ ThreadSafeModule TSM)
+ : ID(ID), TSM(std::move(TSM)) {}
+ ReOptMaterializationUnitState(ReOptMaterializationUnitState &&Other)
+ : ID(Other.ID), TSM(std::move(Other.TSM)), RT(std::move(Other.RT)),
+ Reoptimizing(std::move(Other.Reoptimizing)),
+ CurVersion(Other.CurVersion) {}
+
+ ReOptMaterializationUnitID getID() { return ID; }
+
+ const ThreadSafeModule &getThreadSafeModule() { return TSM; }
+
+ ResourceTrackerSP getResourceTracker() {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ return RT;
+ }
+
+ void setResourceTracker(ResourceTrackerSP RT) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ this->RT = RT;
+ }
+
+ uint32_t getCurVersion() {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ return CurVersion;
+ }
+
+ bool tryStartReoptimize();
+ void reoptimizeSucceeded();
+ void reoptimizeFailed();
+
+ private:
+ std::mutex Mutex;
+ ReOptMaterializationUnitID ID;
+ ThreadSafeModule TSM;
+ ResourceTrackerSP RT;
+ bool Reoptimizing = false;
+ uint32_t CurVersion = 0;
+ };
+
+ using SPSReoptimizeArgList =
+ shared::SPSArgList<ReOptMaterializationUnitID, uint32_t>;
+ using SendErrorFn = unique_function<void(Error)>;
+
+ Expected<SymbolMap> emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
+ uint32_t Version, JITDylib &JD,
+ ThreadSafeModule TSM);
+
+ void rt_reoptimize(SendErrorFn SendResult, ReOptMaterializationUnitID MUID,
+ uint32_t CurVersion);
+
+ static Expected<Constant *>
+ createReoptimizeArgBuffer(Module &M, ReOptMaterializationUnitID MUID,
+ uint32_t CurVersion);
+
+ ReOptMaterializationUnitState &
+ createMaterializationUnitState(const ThreadSafeModule &TSM);
+
+ void
+ registerMaterializationUnitResource(ResourceKey Key,
+ ReOptMaterializationUnitState &State);
+
+ ReOptMaterializationUnitState &
+ getMaterializationUnitState(ReOptMaterializationUnitID MUID);
+
+ ExecutionSession &ES;
+ IRLayer &BaseLayer;
+ RedirectableSymbolManager &RSManager;
+
+ ReOptimizeFunc ReOptFunc;
+ AddProfilerFunc ProfilerFunc;
+
+ std::mutex Mutex;
+ std::map<ReOptMaterializationUnitID, ReOptMaterializationUnitState> MUStates;
+ DenseMap<ResourceKey, DenseSet<ReOptMaterializationUnitID>> MUResources;
+ ReOptMaterializationUnitID NextID = 1;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h b/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
index 87a81b0e529ccd..4004c42d914684 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
@@ -63,6 +63,8 @@ class RedirectableSymbolManager : public RedirectionManager {
const SymbolMap &InitialDests) = 0;
};
+/// RedirectableMaterializationUnit materializes redirectable symbol
+/// by invoking RedirectableSymbolManager::emitRedirectableSymbols
class RedirectableMaterializationUnit : public MaterializationUnit {
public:
RedirectableMaterializationUnit(RedirectableSymbolManager &RM,
diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index 379c3cc240e8c7..550a36e7afc477 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
@@ -52,6 +52,7 @@ add_llvm_component_library(LLVMOrcJIT
ThreadSafeModule.cpp
RedirectionManager.cpp
JITLinkRedirectableSymbolManager.cpp
+ ReOptimizeLayer.cpp
ADDITIONAL_HEADER_DIRS
${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc
diff --git a/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
index 6d76d50271b9f2..0e123b98201b28 100644
--- a/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
@@ -17,6 +17,7 @@ using namespace llvm::orc;
void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
std::unique_ptr<MaterializationResponsibility> R,
const SymbolAddrMap &InitialDests) {
+ auto &ES = ObjLinkingLayer.getExecutionSession();
std::unique_lock<std::mutex> Lock(Mutex);
if (GetNumAvailableStubs() < InitialDests.size())
if (auto Err = grow(InitialDests.size() - GetNumAvailableStubs())) {
@@ -37,7 +38,6 @@ void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
R->failMaterialization();
return;
}
- dbgs() << *K << "\n";
SymbolToStubs[&TargetJD][K] = StubID;
NewSymbolDefs[K] = JumpStubs[StubID];
NewSymbolDefs[K].setFlags(V.getFlags());
@@ -45,13 +45,14 @@ void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
AvailableStubs.pop_back();
}
- if (auto Err = R->replace(absoluteSymbols(NewSymbolDefs))) {
+ // FIXME: when this fails we can return stubs to the pool
+ if (auto Err = redirectInner(TargetJD, InitialDests)) {
ES.reportError(std::move(Err));
R->failMaterialization();
return;
}
- if (auto Err = redirectInner(TargetJD, InitialDests)) {
+ if (auto Err = R->replace(absoluteSymbols(NewSymbolDefs))) {
ES.reportError(std::move(Err));
R->failMaterialization();
return;
@@ -85,10 +86,10 @@ Error JITLinkRedirectableSymbolManager::redirectInner(
StubHandle StubID = SymbolToStubs[&TargetJD].at(K);
PtrWrites.push_back({StubPointers[StubID].getAddress(), V.getAddress()});
}
- if (auto Err = ES.getExecutorProcessControl().getMemoryAccess().writePointers(
- PtrWrites))
- return Err;
- return Error::success();
+ return ObjLinkingLayer.getExecutionSession()
+ .getExecutorProcessControl()
+ .getMemoryAccess()
+ .writePointers(PtrWrites);
}
Error JITLinkRedirectableSymbolManager::grow(unsigned Need) {
@@ -103,6 +104,7 @@ Error JITLinkRedirectableSymbolManager::grow(unsigned Need) {
SymbolLookupSet LookupSymbols;
DenseMap<SymbolStringPtr, ExecutorSymbolDef *> NewDefsMap;
+ auto &ES = ObjLinkingLayer.getExecutionSession();
Triple TT = ES.getTargetTriple();
auto G = std::make_unique<jitlink::LinkGraph>(
"<INDIRECT STUBS>", TT, TT.isArch64Bit() ? 8 : 4,
@@ -113,6 +115,7 @@ Error JITLinkRedirectableSymbolManager::grow(unsigned Need) {
auto &StubsSection =
G->createSection(JumpStubTableName, MemProt::Exec | MemProt::Read);
+ // FIXME: We can batch the stubs into one block and use address to access them
for (size_t I = OldSize; I < NewSize; I++) {
auto Pointer = AnonymousPtrCreator(*G, PointerSection, nullptr, 0);
if (auto Err = Pointer.takeError())
diff --git a/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
new file mode 100644
index 00000000000000..2ec82f210e6f62
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
@@ -0,0 +1,278 @@
+#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
+
+using namespace llvm;
+using namespace orc;
+
+bool ReOptimizeLayer::ReOptMaterializationUnitState::tryStartReoptimize() {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ if (Reoptimizing)
+ return false;
+
+ Reoptimizing = true;
+ return true;
+}
+
+void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeSucceeded() {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
+ Reoptimizing = false;
+ CurVersion++;
+}
+
+void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeFailed() {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
+ Reoptimizing = false;
+}
+
+Error ReOptimizeLayer::reigsterRuntimeFunctions(JITDylib &PlatformJD) {
+ ExecutionSession::JITDispatchHandlerAssociationMap WFs;
+ using ReoptimizeSPSSig = shared::SPSError(uint64_t, uint32_t);
+ WFs[ES.intern("__orc_rt_reoptimize_tag")] =
+ ES.wrapAsyncWithSPS<ReoptimizeSPSSig>(this,
+ &ReOptimizeLayer::rt_reoptimize);
+ return ES.registerJITDispatchHandlers(PlatformJD, std::move(WFs));
+}
+
+void ReOptimizeLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
+ ThreadSafeModule TSM) {
+ auto &JD = R->getTargetJITDylib();
+
+ bool HasNonCallable = false;
+ for (auto &KV : R->getSymbols()) {
+ auto &Flags = KV.second;
+ if (!Flags.isCallable())
+ HasNonCallable = true;
+ }
+
+ if (HasNonCallable) {
+ BaseLayer.emit(std::move(R), std::move(TSM));
+ return;
+ }
+
+ auto &MUState = createMaterializationUnitState(TSM);
+
+ if (auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
+ registerMaterializationUnitResource(Key, MUState);
+ })) {
+ ES.reportError(std::move(Err));
+ R->failMaterialization();
+ return;
+ }
+
+ if (auto Err =
+ ProfilerFunc(*this, MUState.getID(), MUState.getCurVersion(), TSM)) {
+ ES.reportError(std::move(Err));
+ R->failMaterialization();
+ return;
+ }
+
+ auto InitialDests =
+ emitMUImplSymbols(MUState, MUState.getCurVersion(), JD, std::move(TSM));
+ if (!InitialDests) {
+ ES.reportError(InitialDests.takeError());
+ R->failMaterialization();
+ return;
+ }
+
+ RSManager.emitRedirectableSymbols(std::move(R), std::move(*InitialDests));
+}
+
+Error ReOptimizeLayer::reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
+ ReOptMaterializationUnitID MUID,
+ unsigned CurVersion,
+ ThreadSafeModule &TSM) {
+ return TSM.withModuleDo([&](Module &M) -> Error {
+ Type *I64Ty = Type::getInt64Ty(M.getContext());
+ GlobalVariable *Counter = new GlobalVariable(
+ M, I64Ty, false, GlobalValue::InternalLinkage,
+ Constant::getNullValue(I64Ty), "__orc_reopt_counter");
+ auto ArgBufferConst = createReoptimizeArgBuffer(M, MUID, CurVersion);
+ if (auto Err = ArgBufferConst.takeError())
+ return Err;
+ GlobalVariable *ArgBuffer =
+ new GlobalVariable(M, (*ArgBufferConst)->getType(), true,
+ GlobalValue::InternalLinkage, (*ArgBufferConst));
+ for (auto &F : M) {
+ if (F.isDeclaration())
+ continue;
+ auto &BB = F.getEntryBlock();
+ auto *IP = &*BB.getFirstInsertionPt();
+ IRBuilder<> IRB(IP);
+ Value *Threshold = ConstantInt::get(I64Ty, CallCountThreshold, true);
+ Value *Cnt = IRB.CreateLoad(I64Ty, Counter);
+ // Use EQ to prevent further reoptimize calls.
+ Value *Cmp = IRB.CreateICmpEQ(Cnt, Threshold);
+ Value *Added = IRB.CreateAdd(Cnt, ConstantInt::get(I64Ty, 1));
+ (void)IRB.CreateStore(Added, Counter);
+ Instruction *SplitTerminator = SplitBlockAndInsertIfThen(Cmp, IP, false);
+ createReoptimizeCall(M, *SplitTerminator, ArgBuffer);
+ }
+ return Error::success();
+ });
+}
+
+Expected<SymbolMap>
+ReOptimizeLayer::emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
+ uint32_t Version, JITDylib &JD,
+ ThreadSafeModule TSM) {
+ DenseMap<SymbolStringPtr, SymbolStringPtr> RenamedMap;
+ cantFail(TSM.withModuleDo([&](Module &M) -> Error {
+ MangleAndInterner Mangle(ES, M.getDataLayout());
+ for (auto &F : M)
+ if (!F.isDeclaration()) {
+ std::string NewName =
+ (F.getName() + ".__def__." + Twine(Version)).str();
+ RenamedMap[Mangle(F.getName())] = Mangle(NewName);
+ F.setName(NewName);
+ }
+ return Error::success();
+ }));
+
+ auto RT = JD.createResourceTracker();
+ if (auto Err =
+ JD.define(std::make_unique<BasicIRLayerMaterializationUnit>(
+ BaseLayer, *getManglingOptions(), std::move(TSM)),
+ RT))
+ return Err;
+ MUState.setResourceTracker(RT);
+
+ SymbolLookupSet LookupSymbols;
+ for (auto [K, V] : RenamedMap)
+ LookupSymbols.add(V);
+
+ auto ImplSymbols =
+ ES.lookup({{&JD, JITDylibLookupFlags::MatchAllSymbols}}, LookupSymbols,
+ LookupKind::Static, SymbolState::Resolved);
+ if (auto Err = ImplSymbols.takeError())
+ return Err;
+
+ SymbolMap Result;
+ for (auto [K, V] : RenamedMap)
+ Result[K] = (*ImplSymbols)[V];
+
+ return Result;
+}
+
+void ReOptimizeLayer::rt_reoptimize(SendErrorFn SendResult,
+ ReOptMaterializationUnitID MUID,
+ uint32_t CurVersion) {
+ auto &MUState = getMaterializationUnitState(MUID);
+ if (CurVersion < MUState.getCurVersion() || !MUState.tryStartReoptimize()) {
+ SendResult(Error::success());
+ return;
+ }
+
+ ThreadSafeModule TSM = cloneToNewContext(MUState.getThreadSafeModule());
+ auto OldRT = MUState.getResourceTracker();
+ auto &JD = OldRT->getJITDylib();
+
+ if (auto Err = ReOptFunc(*this, MUID, CurVersion + 1, OldRT, TSM)) {
+ ES.reportError(std::move(Err));
+ MUState.reoptimizeFailed();
+ SendResult(Error::success());
+ return;
+ }
+
+ auto SymbolDests =
+ emitMUImplSymbols(MUState, CurVersion + 1, JD, std::move(TSM));
+ if (!SymbolDests) {
+ ES.reportError(SymbolDests.takeError());
+ MUState.reoptimizeFailed();
+ SendResult(Error::success());
+ return;
+ }
+
+ if (auto Err = RSManager.redirect(JD, std::move(*SymbolDests))) {
+ ES.reportError(std::move(Err));
+ MUState.reoptimizeFailed();
+ SendResult(Error::success());
+ return;
+ }
+
+ MUState.reoptimizeSucceeded();
+ SendResult(Error::success());
+}
+
+Expected<Constant *> ReOptimizeLayer::createReoptimizeArgBuffer(
+ Module &M, ReOptMaterializationUnitID MUID, uint32_t CurVersion) {
+ size_t ArgBufferSize = SPSReoptimizeArgList::size(MUID, CurVersion);
+ std::vector<char> ArgBuffer(ArgBufferSize);
+ shared::SPSOutputBuffer OB(ArgBuffer.data(), ArgBuffer.size());
+ if (!SPSReoptimizeArgList::serialize(OB, MUID, CurVersion))
+ return make_error<StringError>("Could not serealize args list",
+ inconvertibleErrorCode());
+ return ConstantDataArray::get(M.getContext(), ArrayRef(ArgBuffer));
+}
+
+void ReOptimizeLayer::createReoptimizeCall(Module &M, Instruction &IP,
+ GlobalVariable *ArgBuffer) {
+ GlobalVariable *DispatchCtx =
+ M.getGlobalVariable("__orc_rt_jit_dispatch_ctx");
+ if (!DispatchCtx)
+ DispatchCtx = new GlobalVariable(M, Type::getInt8PtrTy(M.getContext()),
+ false, GlobalValue::ExternalLinkage,
+ nullptr, "__orc_rt_jit_dispatch_ctx");
+ GlobalVariable *ReoptimizeTag =
+ M.getGlobalVariable("__orc_rt_reoptimize_tag");
+ if (!ReoptimizeTag)
+ ReoptimizeTag = new GlobalVariable(M, Type::getInt8PtrTy(M.getContext()),
+ false, GlobalValue::ExternalLinkage,
+ nullptr, "__orc_rt_reoptimize_tag");
+ Function *DispatchFunc = M.getFunction("__orc_rt_jit_dispatch");
+ if (!DispatchFunc) {
+ std::vector<Type *> Args = {Type::getInt8PtrTy(M.getContext()),
+ Type::getInt8PtrTy(M.getContext()),
+ Type::getInt8PtrTy(M.getContext()),
+ IntegerType::get(M.getContext(), 64)};
+ FunctionType *FuncTy =
+ FunctionType::get(Type::getVoidTy(M.getContext()), Args, false);
+ DispatchFunc = Function::Create(FuncTy, GlobalValue::ExternalLinkage,
+ "__orc_rt_jit_dispatch", &M);
+ }
+ size_t ArgBufferSizeConst =
+ SPSReoptimizeArgList::size(ReOptMaterializationUnitID{}, uint32_t{});
+ Constant *ArgBufferSize = ConstantInt::get(
+ IntegerType::get(M.getContext(), 64), ArgBufferSizeConst, false);
+ IRBuilder<> IRB(&IP);
+ (void)IRB.CreateCall(DispatchFunc,
+ {DispatchCtx, ReoptimizeTag, ArgBuffer, ArgBufferSize});
+}
+
+ReOptimizeLayer::ReOptMaterializationUnitState &
+ReOptimizeLayer::createMaterializationUnitState(const ThreadSafeModule &TSM) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ ReOptMaterializationUnitID MUID = NextID;
+ MUStates.emplace(MUID,
+ ReOptMaterializationUnitState(MUID, cloneToNewContext(TSM)));
+ ++NextID;
+ return MUStates.at(MUID);
+}
+
+ReOptimizeLayer::ReOptMaterializationUnitState &
+ReOptimizeLayer::getMaterializationUnitState(ReOptMaterializationUnitID MUID) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ return MUStates.at(MUID);
+}
+
+void ReOptimizeLayer::registerMaterializationUnitResource(
+ ResourceKey Key, ReOptMaterializationUnitState &State) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ MUResources[Key].insert(State.getID());
+}
+
+Error ReOptimizeLayer::handleRemoveResources(JITDylib &JD, ResourceKey K) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ for (auto MUID : MUResources[K])
+ MUStates.erase(MUID);
+
+ MUResources.erase(K);
+ return Error::success();
+}
+
+void ReOptimizeLayer::handleTransferResources(JITDylib &JD, ResourceKey DstK,
+ ResourceKey SrcK) {
+ std::unique_lock<std::mutex> Lock(Mutex);
+ MUResources[DstK].insert(MUResources[SrcK].begin(), MUResources[SrcK].end());
+ MUResources.erase(SrcK);
+}
diff --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
index e932ff94361a7b..310cb8276d894a 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
@@ -41,6 +41,7 @@ add_llvm_unittest(OrcJITTests
ThreadSafeModuleTest.cpp
WrapperFunctionUtilsTest.cpp
JITLinkRedirectionManagerTest.cpp
+ ReOptimizeLayerTest.cpp
)
target_link_libraries(OrcJITTests PRIVATE
diff --git a/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
index 0f87c1b7433238..f063111354ac90 100644
--- a/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
@@ -49,8 +49,7 @@ class JITLinkRedirectionManagerTest : public testing::Test {
};
TEST_F(JITLinkRedirectionManagerTest, BasicRedirectionOperation) {
- auto RM =
- JITLinkRedirectableSymbolManager::Create(*ES, *ObjLinkingLayer, *JD);
+ auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer, *JD);
// Bail out if we can not create
if (!RM) {
consumeError(RM.takeError());
diff --git a/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
new file mode 100644
index 00000000000000..f97468576081c2
--- /dev/null
+++ b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
@@ -0,0 +1,149 @@
+#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
+#include "OrcTestCommon.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/CompileUtils.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/JITLinkRedirectableSymbolManager.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/Support/CodeGen.h"
+#include "llvm/TargetParser/Host.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::jitlink;
+
+class ReOptimizeLayerTest : public testing::Test {
+public:
+ ~ReOptimizeLayerTest() {
+ if (ES)
+ if (auto Err = ES->endSession())
+ ES->reportError(std::move(Err));
+ }
+
+protected:
+ void SetUp() override {
+ auto JTMB = JITTargetMachineBuilder::detectHost();
+ // Bail out if we can not detect the host.
+ if (!JTMB) {
+ consumeError(JTMB.takeError());
+ GTEST_SKIP();
+ }
+
+ auto EPC = SelfExecutorProcessControl::Create();
+ if (!EPC) {
+ consumeError(EPC.takeError());
+ GTEST_SKIP();
+ }
+ ES = std::make_unique<ExecutionSession>(std::move(*EPC));
+ JD = &ES->createBareJITDylib("main");
+ ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(
+ *ES, std::make_unique<InProcessMemoryManager>(4096));
+ DL = std::make_unique<DataLayout>(
+ cantFail(JTMB->getDefaultDataLayoutForTarget()));
+
+ auto TM = JTMB->createTargetMachine();
+ if (!TM) {
+ consumeError(TM.takeError());
+ GTEST_SKIP();
+ }
+ auto CompileFunction =
+ std::make_unique<TMOwningSimpleCompiler>(std::move(*TM));
+ CompileLayer = std::make_unique<IRCompileLayer>(*ES, *ObjLinkingLayer,
+ std::move(CompileFunction));
+ }
+
+ Error addIRModule(ResourceTrackerSP RT, ThreadSafeModule TSM) {
+ assert(TSM && "Can not add null module");
+
+ TSM.withModuleDo([&](Module &M) { M.setDataLayout(*DL); });
+
+ return ROLayer->add(std::move(RT), std::move(TSM));
+ }
+
+ JITDylib *JD{nullptr};
+ std::unique_ptr<ExecutionSession> ES;
+ std::unique_ptr<ObjectLinkingLayer> ObjLinkingLayer;
+ std::unique_ptr<IRCompileLayer> CompileLayer;
+ std::unique_ptr<ReOptimizeLayer> ROLayer;
+ std::unique_ptr<DataLayout> DL;
+};
+
+static Function *createRetFunction(Module *M, StringRef Name,
+ uint32_t ReturnCode) {
+ Function *Result = Function::Create(
+ FunctionType::get(Type::getInt32Ty(M->getContext()), {}, false),
+ GlobalValue::ExternalLinkage, Name, M);
+
+ BasicBlock *BB = BasicBlock::Create(M->getContext(), Name, Result);
+ IRBuilder<> Builder(M->getContext());
+ Builder.SetInsertPoint(BB);
+
+ Value *RetValue = ConstantInt::get(M->getContext(), APInt(32, ReturnCode));
+ Builder.CreateRet(RetValue);
+ return Result;
+}
+
+TEST_F(ReOptimizeLayerTest, BasicReOptimization) {
+ auto &EPC = ES->getExecutorProcessControl();
+ EXPECT_THAT_ERROR(JD->define(absoluteSymbols(
+ {{ES->intern("__orc_rt_jit_dispatch"),
+ {EPC.getJITDispatchInfo().JITDispatchFunction,
+ JITSymbolFlags::Exported}},
+ {ES->intern("__orc_rt_jit_dispatch_ctx"),
+ {EPC.getJITDispatchInfo().JITDispatchContext,
+ JITSymbolFlags::Exported}},
+ {ES->intern("__orc_rt_reoptimize_tag"),
+ {ExecutorAddr(), JITSymbolFlags::Exported}}})),
+ Succeeded());
+
+ auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer, *JD);
+ EXPECT_THAT_ERROR(RM.takeError(), Succeeded());
+
+ ROLayer = std::make_unique<ReOptimizeLayer>(*ES, *CompileLayer, **RM);
+ ROLayer->setReoptimizeFunc(
+ [&](ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
+ unsigned CurVerison, ResourceTrackerSP OldRT, ThreadSafeModule &TSM) {
+ TSM.withModuleDo([&](Module &M) {
+ for (auto &F : M) {
+ if (F.isDeclaration())
+ continue;
+ for (auto &B : F) {
+ for (auto &I : B) {
+ if (ReturnInst *Ret = dyn_cast<ReturnInst>(&I)) {
+ Value *RetValue =
+ ConstantInt::get(M.getContext(), APInt(32, 53));
+ Ret->setOperand(0, RetValue);
+ }
+ }
+ }
+ }
+ });
+ return Error::success();
+ });
+ EXPECT_THAT_ERROR(ROLayer->reigsterRuntimeFunctions(*JD), Succeeded());
+
+ ThreadSafeContext Ctx(std::make_unique<LLVMContext>());
+ auto M = std::make_unique<Module>("<main>", *Ctx.getContext());
+ M->setTargetTriple(sys::getProcessTriple());
+
+ (void)createRetFunction(M.get(), "main", 42);
+
+ EXPECT_THAT_ERROR(addIRModule(JD->getDefaultResourceTracker(),
+ ThreadSafeModule(std::move(M), std::move(Ctx))),
+ Succeeded());
+
+ auto Result = cantFail(ES->lookup({JD}, "main"));
+ auto FuncPtr = Result.getAddress().toPtr<int (*)()>();
+ for (size_t I = 0; I <= ReOptimizeLayer::CallCountThreshold; I++)
+ EXPECT_EQ(FuncPtr(), 42);
+ EXPECT_EQ(FuncPtr(), 53);
+}
More information about the llvm-commits
mailing list