https://github.com/sunho created https://github.com/llvm/llvm-project/pull/67050
None >From 9d759677e725c244a7a3eb06864d6474253b6d86 Mon Sep 17 00:00:00 2001 From: Sunho Kim <ksunhokim...@gmail.com> Date: Mon, 21 Aug 2023 15:48:26 +0900 Subject: [PATCH 1/4] [ORC] Introduce RedirectionManager interface and implementation using JITLink. --- .../Orc/JITLinkRedirectableSymbolManager.h | 106 +++++++++++ .../ExecutionEngine/Orc/RedirectionManager.h | 102 ++++++++++ llvm/lib/ExecutionEngine/Orc/CMakeLists.txt | 2 + .../Orc/JITLinkRedirectableSymbolManager.cpp | 179 ++++++++++++++++++ .../Orc/RedirectionManager.cpp | 28 +++ .../ExecutionEngine/Orc/CMakeLists.txt | 1 + .../Orc/JITLinkRedirectionManagerTest.cpp | 100 ++++++++++ 7 files changed, 518 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 000000000000000..5de0da1f52d0db6 --- /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 000000000000000..6116102f4fa4d4a --- /dev/null +++ b/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h @@ -0,0 +1,102 @@ +//===- 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 3256ed8b7362c66..17e611d5795cec3 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 000000000000000..6d76d50271b9f2d --- /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 000000000000000..61d765066c8e4cd --- /dev/null +++ b/llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp @@ -0,0 +1,28 @@ +//===---- 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(); + + if (auto Err = JD.define(std::make_unique<RedirectableMaterializationUnit>( + *this, InitialDests), + RT)) + return Err; + + return Error::success(); +} diff --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt index 37768e91fd44729..e932ff94361a7b7 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 000000000000000..0f87c1b7433238b --- /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 8e96f8ce2b54fbf404e5df2064da851515920aef Mon Sep 17 00:00:00 2001 From: Zequan Wu <zequa...@google.com> Date: Tue, 19 Sep 2023 14:59:34 -0400 Subject: [PATCH 2/4] [Coverage] Skip visiting ctor member initializers with invalid source locations. --- clang/lib/CodeGen/CoverageMappingGen.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/clang/lib/CodeGen/CoverageMappingGen.cpp b/clang/lib/CodeGen/CoverageMappingGen.cpp index bb814dfbfd580f1..32c67739915380c 100644 --- a/clang/lib/CodeGen/CoverageMappingGen.cpp +++ b/clang/lib/CodeGen/CoverageMappingGen.cpp @@ -1044,7 +1044,8 @@ struct CounterCoverageMappingBuilder for (auto *Initializer : Ctor->inits()) { if (Initializer->isWritten()) { auto *Init = Initializer->getInit(); - propagateCounts(BodyCounter, Init); + if (getStart(Init).isValid() && getEnd(Init).isValid()) + propagateCounts(BodyCounter, Init); } } } >From e576d28381104c5cd8243c6e6475ac300be2702d Mon Sep 17 00:00:00 2001 From: Sunho Kim <ksunhokim...@gmail.com> Date: Wed, 20 Sep 2023 05:11:15 +0900 Subject: [PATCH 3/4] [ORC] Introduce IRPartitionLayer for common partition functionality. --- .../BuildingAJIT/Chapter3/KaleidoscopeJIT.h | 6 +- .../SpeculativeJIT/SpeculativeJIT.cpp | 9 +- .../Orc/CompileOnDemandLayer.h | 36 +-- .../ExecutionEngine/Orc/IRPartitionLayer.h | 83 +++++ llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h | 7 +- llvm/lib/ExecutionEngine/Orc/CMakeLists.txt | 1 + .../Orc/CompileOnDemandLayer.cpp | 282 +--------------- .../ExecutionEngine/Orc/IRPartitionLayer.cpp | 303 ++++++++++++++++++ llvm/lib/ExecutionEngine/Orc/LLJIT.cpp | 7 +- llvm/tools/lli/lli.cpp | 3 +- 10 files changed, 415 insertions(+), 322 deletions(-) create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h create mode 100644 llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp diff --git a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h index fd0e081ff2b4bd0..fee2d26e5d9254f 100644 --- a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h +++ b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h @@ -21,6 +21,7 @@ #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h" #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h" #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" @@ -48,6 +49,7 @@ class KaleidoscopeJIT { RTDyldObjectLinkingLayer ObjectLayer; IRCompileLayer CompileLayer; IRTransformLayer OptimizeLayer; + IRPartitionLayer IPLayer; CompileOnDemandLayer CODLayer; JITDylib &MainJD; @@ -68,8 +70,8 @@ class KaleidoscopeJIT { CompileLayer(*this->ES, ObjectLayer, std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))), OptimizeLayer(*this->ES, CompileLayer, optimizeModule), - CODLayer(*this->ES, OptimizeLayer, - this->EPCIU->getLazyCallThroughManager(), + IPLayer(*this->ES, OptimizeLayer), + CODLayer(*this->ES, IPLayer, this->EPCIU->getLazyCallThroughManager(), [this] { return this->EPCIU->createIndirectStubsManager(); }), MainJD(this->ES->createBareJITDylib("<main>")) { MainJD.addGenerator( diff --git a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp b/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp index fdd376d82da5d89..a1722748c6ce069 100644 --- a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp +++ b/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp @@ -3,6 +3,7 @@ #include "llvm/ExecutionEngine/Orc/Core.h" #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h" #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" @@ -107,13 +108,14 @@ class SpeculativeJIT { IndirectStubsManagerBuilderFunction ISMBuilder, std::unique_ptr<DynamicLibrarySearchGenerator> ProcessSymbolsGenerator) : ES(std::move(ES)), DL(std::move(DL)), - MainJD(this->ES->createBareJITDylib("<main>")), LCTMgr(std::move(LCTMgr)), + MainJD(this->ES->createBareJITDylib("<main>")), + LCTMgr(std::move(LCTMgr)), CompileLayer(*this->ES, ObjLayer, std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))), S(Imps, *this->ES), SpeculateLayer(*this->ES, CompileLayer, S, Mangle, BlockFreqQuery()), - CODLayer(*this->ES, SpeculateLayer, *this->LCTMgr, - std::move(ISMBuilder)) { + IPLayer(*this->ES, SpeculateLayer), + CODLayer(*this->ES, IPLayer, *this->LCTMgr, std::move(ISMBuilder)) { MainJD.addGenerator(std::move(ProcessSymbolsGenerator)); this->CODLayer.setImplMap(&Imps); this->ES->setDispatchTask( @@ -147,6 +149,7 @@ class SpeculativeJIT { Speculator S; RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr}; IRSpeculationLayer SpeculateLayer; + IRPartitionLayer IPLayer; CompileOnDemandLayer CODLayer; }; diff --git a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h index 23a64b5836db85f..ee556d11b708352 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h @@ -54,37 +54,15 @@ namespace llvm { namespace orc { class CompileOnDemandLayer : public IRLayer { - friend class PartitioningIRMaterializationUnit; - public: /// Builder for IndirectStubsManagers. using IndirectStubsManagerBuilder = std::function<std::unique_ptr<IndirectStubsManager>()>; - using GlobalValueSet = std::set<const GlobalValue *>; - - /// Partitioning function. - using PartitionFunction = - std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>; - - /// Off-the-shelf partitioning which compiles all requested symbols (usually - /// a single function at a time). - static std::optional<GlobalValueSet> - compileRequested(GlobalValueSet Requested); - - /// Off-the-shelf partitioning which compiles whole modules whenever any - /// symbol in them is requested. - static std::optional<GlobalValueSet> - compileWholeModule(GlobalValueSet Requested); - /// Construct a CompileOnDemandLayer. CompileOnDemandLayer(ExecutionSession &ES, IRLayer &BaseLayer, - LazyCallThroughManager &LCTMgr, - IndirectStubsManagerBuilder BuildIndirectStubsManager); - - /// Sets the partition function. - void setPartitionFunction(PartitionFunction Partition); - + LazyCallThroughManager &LCTMgr, + IndirectStubsManagerBuilder BuildIndirectStubsManager); /// Sets the ImplSymbolMap void setImplMap(ImplSymbolMap *Imp); @@ -111,22 +89,12 @@ class CompileOnDemandLayer : public IRLayer { PerDylibResources &getPerDylibResources(JITDylib &TargetD); - void cleanUpModule(Module &M); - - void expandPartition(GlobalValueSet &Partition); - - void emitPartition(std::unique_ptr<MaterializationResponsibility> R, - ThreadSafeModule TSM, - IRMaterializationUnit::SymbolNameToDefinitionMap Defs); - mutable std::mutex CODLayerMutex; IRLayer &BaseLayer; LazyCallThroughManager &LCTMgr; IndirectStubsManagerBuilder BuildIndirectStubsManager; PerDylibResourcesMap DylibResources; - PartitionFunction Partition = compileRequested; - SymbolLinkagePromoter PromoteSymbols; ImplSymbolMap *AliaseeImpls = nullptr; }; diff --git a/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h new file mode 100644 index 000000000000000..0535ce473837474 --- /dev/null +++ b/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h @@ -0,0 +1,83 @@ +//===- IRPartitionLayer.h - Partition IR module on lookup -------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// JIT layer for breaking up modules into smaller submodules that only contains +// looked up symbols. +// +//===----------------------------------------------------------------------===// + +#ifndef LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H +#define LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H + +#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h" +#include "llvm/ExecutionEngine/Orc/Layer.h" +#include "llvm/IR/Attributes.h" +#include "llvm/IR/Constant.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/DataLayout.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/GlobalAlias.h" +#include "llvm/IR/GlobalValue.h" +#include "llvm/IR/GlobalVariable.h" +#include "llvm/IR/Instruction.h" +#include "llvm/IR/Mangler.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Type.h" + +namespace llvm { +namespace orc { + +class IRPartitionLayer : public IRLayer { + friend class PartitioningIRMaterializationUnit; + +public: + using GlobalValueSet = std::set<const GlobalValue *>; + + /// Partitioning function. + using PartitionFunction = + std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>; + + /// Construct a IRPartitionLayer. + IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer); + + /// Off-the-shelf partitioning which compiles all requested symbols (usually + /// a single function at a time). + static std::optional<GlobalValueSet> + compileRequested(GlobalValueSet Requested); + + /// Off-the-shelf partitioning which compiles whole modules whenever any + /// symbol in them is requested. + static std::optional<GlobalValueSet> + compileWholeModule(GlobalValueSet Requested); + + /// Sets the partition function. + void setPartitionFunction(PartitionFunction Partition); + + /// Emits the given module. This should not be called by clients: it will be + /// called by the JIT when a definition added via the add method is requested. + void emit(std::unique_ptr<MaterializationResponsibility> R, + ThreadSafeModule TSM) override; + +private: + void cleanUpModule(Module &M); + + void expandPartition(GlobalValueSet &Partition); + + void emitPartition(std::unique_ptr<MaterializationResponsibility> R, + ThreadSafeModule TSM, + IRMaterializationUnit::SymbolNameToDefinitionMap Defs); + + IRLayer &BaseLayer; + PartitionFunction Partition = compileRequested; + SymbolLinkagePromoter PromoteSymbols; +}; + +} // namespace orc +} // namespace llvm + +#endif diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h index ab54e0c2c288205..896de521f13af80 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h @@ -17,6 +17,7 @@ #include "llvm/ExecutionEngine/Orc/CompileUtils.h" #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h" #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h" @@ -271,9 +272,8 @@ class LLLazyJIT : public LLJIT { public: /// Sets the partition function. - void - setPartitionFunction(CompileOnDemandLayer::PartitionFunction Partition) { - CODLayer->setPartitionFunction(std::move(Partition)); + void setPartitionFunction(IRPartitionLayer::PartitionFunction Partition) { + IPLayer->setPartitionFunction(std::move(Partition)); } /// Returns a reference to the on-demand layer. @@ -293,6 +293,7 @@ class LLLazyJIT : public LLJIT { LLLazyJIT(LLLazyJITBuilderState &S, Error &Err); std::unique_ptr<LazyCallThroughManager> LCTMgr; + std::unique_ptr<IRPartitionLayer> IPLayer; std::unique_ptr<CompileOnDemandLayer> CODLayer; }; diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt index 17e611d5795cec3..379c3cc240e8c7b 100644 --- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt +++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt @@ -27,6 +27,7 @@ add_llvm_component_library(LLVMOrcJIT IndirectionUtils.cpp IRCompileLayer.cpp IRTransformLayer.cpp + IRPartitionLayer.cpp JITTargetMachineBuilder.cpp LazyReexports.cpp Layer.cpp diff --git a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp index 6448adaa0ceb36f..9296bc2b389ab24 100644 --- a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp +++ b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp @@ -9,6 +9,7 @@ #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h" #include "llvm/ADT/Hashing.h" #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" +#include "llvm/ExecutionEngine/Orc/Layer.h" #include "llvm/IR/Mangler.h" #include "llvm/IR/Module.h" #include "llvm/Support/FormatVariadic.h" @@ -17,101 +18,6 @@ using namespace llvm; using namespace llvm::orc; -static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM, - StringRef Suffix, - GVPredicate ShouldExtract) { - - auto DeleteExtractedDefs = [](GlobalValue &GV) { - // Bump the linkage: this global will be provided by the external module. - GV.setLinkage(GlobalValue::ExternalLinkage); - - // Delete the definition in the source module. - if (isa<Function>(GV)) { - auto &F = cast<Function>(GV); - F.deleteBody(); - F.setPersonalityFn(nullptr); - } else if (isa<GlobalVariable>(GV)) { - cast<GlobalVariable>(GV).setInitializer(nullptr); - } else if (isa<GlobalAlias>(GV)) { - // We need to turn deleted aliases into function or variable decls based - // on the type of their aliasee. - auto &A = cast<GlobalAlias>(GV); - Constant *Aliasee = A.getAliasee(); - assert(A.hasName() && "Anonymous alias?"); - assert(Aliasee->hasName() && "Anonymous aliasee"); - std::string AliasName = std::string(A.getName()); - - if (isa<Function>(Aliasee)) { - auto *F = cloneFunctionDecl(*A.getParent(), *cast<Function>(Aliasee)); - A.replaceAllUsesWith(F); - A.eraseFromParent(); - F->setName(AliasName); - } else if (isa<GlobalVariable>(Aliasee)) { - auto *G = cloneGlobalVariableDecl(*A.getParent(), - *cast<GlobalVariable>(Aliasee)); - A.replaceAllUsesWith(G); - A.eraseFromParent(); - G->setName(AliasName); - } else - llvm_unreachable("Alias to unsupported type"); - } else - llvm_unreachable("Unsupported global type"); - }; - - auto NewTSM = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs); - NewTSM.withModuleDo([&](Module &M) { - M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str()); - }); - - return NewTSM; -} - -namespace llvm { -namespace orc { - -class PartitioningIRMaterializationUnit : public IRMaterializationUnit { -public: - PartitioningIRMaterializationUnit(ExecutionSession &ES, - const IRSymbolMapper::ManglingOptions &MO, - ThreadSafeModule TSM, - CompileOnDemandLayer &Parent) - : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {} - - PartitioningIRMaterializationUnit( - ThreadSafeModule TSM, Interface I, - SymbolNameToDefinitionMap SymbolToDefinition, - CompileOnDemandLayer &Parent) - : IRMaterializationUnit(std::move(TSM), std::move(I), - std::move(SymbolToDefinition)), - Parent(Parent) {} - -private: - void materialize(std::unique_ptr<MaterializationResponsibility> R) override { - Parent.emitPartition(std::move(R), std::move(TSM), - std::move(SymbolToDefinition)); - } - - void discard(const JITDylib &V, const SymbolStringPtr &Name) override { - // All original symbols were materialized by the CODLayer and should be - // final. The function bodies provided by M should never be overridden. - llvm_unreachable("Discard should never be called on an " - "ExtractingIRMaterializationUnit"); - } - - mutable std::mutex SourceModuleMutex; - CompileOnDemandLayer &Parent; -}; - -std::optional<CompileOnDemandLayer::GlobalValueSet> -CompileOnDemandLayer::compileRequested(GlobalValueSet Requested) { - return std::move(Requested); -} - -std::optional<CompileOnDemandLayer::GlobalValueSet> -CompileOnDemandLayer::compileWholeModule(GlobalValueSet Requested) { - return std::nullopt; -} - CompileOnDemandLayer::CompileOnDemandLayer( ExecutionSession &ES, IRLayer &BaseLayer, LazyCallThroughManager &LCTMgr, IndirectStubsManagerBuilder BuildIndirectStubsManager) @@ -119,13 +25,10 @@ CompileOnDemandLayer::CompileOnDemandLayer( LCTMgr(LCTMgr), BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)) {} -void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) { - this->Partition = std::move(Partition); -} - void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) { this->AliaseeImpls = Imp; } + void CompileOnDemandLayer::emit( std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM) { assert(TSM && "Null module"); @@ -138,10 +41,6 @@ void CompileOnDemandLayer::emit( SymbolAliasMap NonCallables; SymbolAliasMap Callables; - TSM.withModuleDo([&](Module &M) { - // First, do some cleanup on the module: - cleanUpModule(M); - }); for (auto &KV : R->getSymbols()) { auto &Name = KV.first; @@ -152,11 +51,10 @@ void CompileOnDemandLayer::emit( NonCallables[Name] = SymbolAliasMapEntry(Name, Flags); } - // Create a partitioning materialization unit and lodge it with the - // implementation dylib. + // Lodge symbols with the implementation dylib. if (auto Err = PDR.getImplDylib().define( - std::make_unique<PartitioningIRMaterializationUnit>( - ES, *getManglingOptions(), std::move(TSM), *this))) { + std::make_unique<BasicIRLayerMaterializationUnit>( + BaseLayer, *getManglingOptions(), std::move(TSM)))) { ES.reportError(std::move(Err)); R->failMaterialization(); return; @@ -210,173 +108,3 @@ CompileOnDemandLayer::getPerDylibResources(JITDylib &TargetD) { return I->second; } - -void CompileOnDemandLayer::cleanUpModule(Module &M) { - for (auto &F : M.functions()) { - if (F.isDeclaration()) - continue; - - if (F.hasAvailableExternallyLinkage()) { - F.deleteBody(); - F.setPersonalityFn(nullptr); - continue; - } - } -} - -void CompileOnDemandLayer::expandPartition(GlobalValueSet &Partition) { - // Expands the partition to ensure the following rules hold: - // (1) If any alias is in the partition, its aliasee is also in the partition. - // (2) If any aliasee is in the partition, its aliases are also in the - // partiton. - // (3) If any global variable is in the partition then all global variables - // are in the partition. - assert(!Partition.empty() && "Unexpected empty partition"); - - const Module &M = *(*Partition.begin())->getParent(); - bool ContainsGlobalVariables = false; - std::vector<const GlobalValue *> GVsToAdd; - - for (const auto *GV : Partition) - if (isa<GlobalAlias>(GV)) - GVsToAdd.push_back( - cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee())); - else if (isa<GlobalVariable>(GV)) - ContainsGlobalVariables = true; - - for (auto &A : M.aliases()) - if (Partition.count(cast<GlobalValue>(A.getAliasee()))) - GVsToAdd.push_back(&A); - - if (ContainsGlobalVariables) - for (auto &G : M.globals()) - GVsToAdd.push_back(&G); - - for (const auto *GV : GVsToAdd) - Partition.insert(GV); -} - -void CompileOnDemandLayer::emitPartition( - std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM, - IRMaterializationUnit::SymbolNameToDefinitionMap Defs) { - - // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the - // extracted module key, extracted module, and source module key - // together. This could be used, for example, to provide a specific - // memory manager instance to the linking layer. - - auto &ES = getExecutionSession(); - GlobalValueSet RequestedGVs; - for (auto &Name : R->getRequestedSymbols()) { - if (Name == R->getInitializerSymbol()) - TSM.withModuleDo([&](Module &M) { - for (auto &GV : getStaticInitGVs(M)) - RequestedGVs.insert(&GV); - }); - else { - assert(Defs.count(Name) && "No definition for symbol"); - RequestedGVs.insert(Defs[Name]); - } - } - - /// Perform partitioning with the context lock held, since the partition - /// function is allowed to access the globals to compute the partition. - auto GVsToExtract = - TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); }); - - // Take a 'None' partition to mean the whole module (as opposed to an empty - // partition, which means "materialize nothing"). Emit the whole module - // unmodified to the base layer. - if (GVsToExtract == std::nullopt) { - Defs.clear(); - BaseLayer.emit(std::move(R), std::move(TSM)); - return; - } - - // If the partition is empty, return the whole module to the symbol table. - if (GVsToExtract->empty()) { - if (auto Err = - R->replace(std::make_unique<PartitioningIRMaterializationUnit>( - std::move(TSM), - MaterializationUnit::Interface(R->getSymbols(), - R->getInitializerSymbol()), - std::move(Defs), *this))) { - getExecutionSession().reportError(std::move(Err)); - R->failMaterialization(); - return; - } - return; - } - - // Ok -- we actually need to partition the symbols. Promote the symbol - // linkages/names, expand the partition to include any required symbols - // (i.e. symbols that can't be separated from our partition), and - // then extract the partition. - // - // FIXME: We apply this promotion once per partitioning. It's safe, but - // overkill. - auto ExtractedTSM = - TSM.withModuleDo([&](Module &M) -> Expected<ThreadSafeModule> { - auto PromotedGlobals = PromoteSymbols(M); - if (!PromotedGlobals.empty()) { - - MangleAndInterner Mangle(ES, M.getDataLayout()); - SymbolFlagsMap SymbolFlags; - IRSymbolMapper::add(ES, *getManglingOptions(), - PromotedGlobals, SymbolFlags); - - if (auto Err = R->defineMaterializing(SymbolFlags)) - return std::move(Err); - } - - expandPartition(*GVsToExtract); - - // Submodule name is given by hashing the names of the globals. - std::string SubModuleName; - { - std::vector<const GlobalValue*> HashGVs; - HashGVs.reserve(GVsToExtract->size()); - for (const auto *GV : *GVsToExtract) - HashGVs.push_back(GV); - llvm::sort(HashGVs, [](const GlobalValue *LHS, const GlobalValue *RHS) { - return LHS->getName() < RHS->getName(); - }); - hash_code HC(0); - for (const auto *GV : HashGVs) { - assert(GV->hasName() && "All GVs to extract should be named by now"); - auto GVName = GV->getName(); - HC = hash_combine(HC, hash_combine_range(GVName.begin(), GVName.end())); - } - raw_string_ostream(SubModuleName) - << ".submodule." - << formatv(sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}", - static_cast<size_t>(HC)) - << ".ll"; - } - - // Extract the requested partiton (plus any necessary aliases) and - // put the rest back into the impl dylib. - auto ShouldExtract = [&](const GlobalValue &GV) -> bool { - return GVsToExtract->count(&GV); - }; - - return extractSubModule(TSM, SubModuleName , ShouldExtract); - }); - - if (!ExtractedTSM) { - ES.reportError(ExtractedTSM.takeError()); - R->failMaterialization(); - return; - } - - if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>( - ES, *getManglingOptions(), std::move(TSM), *this))) { - ES.reportError(std::move(Err)); - R->failMaterialization(); - return; - } - BaseLayer.emit(std::move(R), std::move(*ExtractedTSM)); -} - -} // end namespace orc -} // end namespace llvm diff --git a/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp b/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp new file mode 100644 index 000000000000000..9ad171beac7fe2e --- /dev/null +++ b/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp @@ -0,0 +1,303 @@ +//===----- IRPartitionLayer.cpp - Partition IR module into submodules -----===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" +#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" +#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h" + +using namespace llvm; +using namespace llvm::orc; + +static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM, + StringRef Suffix, + GVPredicate ShouldExtract) { + + auto DeleteExtractedDefs = [](GlobalValue &GV) { + // Bump the linkage: this global will be provided by the external module. + GV.setLinkage(GlobalValue::ExternalLinkage); + + // Delete the definition in the source module. + if (isa<Function>(GV)) { + auto &F = cast<Function>(GV); + F.deleteBody(); + F.setPersonalityFn(nullptr); + } else if (isa<GlobalVariable>(GV)) { + cast<GlobalVariable>(GV).setInitializer(nullptr); + } else if (isa<GlobalAlias>(GV)) { + // We need to turn deleted aliases into function or variable decls based + // on the type of their aliasee. + auto &A = cast<GlobalAlias>(GV); + Constant *Aliasee = A.getAliasee(); + assert(A.hasName() && "Anonymous alias?"); + assert(Aliasee->hasName() && "Anonymous aliasee"); + std::string AliasName = std::string(A.getName()); + + if (isa<Function>(Aliasee)) { + auto *F = cloneFunctionDecl(*A.getParent(), *cast<Function>(Aliasee)); + A.replaceAllUsesWith(F); + A.eraseFromParent(); + F->setName(AliasName); + } else if (isa<GlobalVariable>(Aliasee)) { + auto *G = cloneGlobalVariableDecl(*A.getParent(), + *cast<GlobalVariable>(Aliasee)); + A.replaceAllUsesWith(G); + A.eraseFromParent(); + G->setName(AliasName); + } else + llvm_unreachable("Alias to unsupported type"); + } else + llvm_unreachable("Unsupported global type"); + }; + + auto NewTSM = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs); + NewTSM.withModuleDo([&](Module &M) { + M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str()); + }); + + return NewTSM; +} + +namespace llvm { +namespace orc { + +class PartitioningIRMaterializationUnit : public IRMaterializationUnit { +public: + PartitioningIRMaterializationUnit(ExecutionSession &ES, + const IRSymbolMapper::ManglingOptions &MO, + ThreadSafeModule TSM, + IRPartitionLayer &Parent) + : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {} + + PartitioningIRMaterializationUnit( + ThreadSafeModule TSM, Interface I, + SymbolNameToDefinitionMap SymbolToDefinition, IRPartitionLayer &Parent) + : IRMaterializationUnit(std::move(TSM), std::move(I), + std::move(SymbolToDefinition)), + Parent(Parent) {} + +private: + void materialize(std::unique_ptr<MaterializationResponsibility> R) override { + Parent.emitPartition(std::move(R), std::move(TSM), + std::move(SymbolToDefinition)); + } + + void discard(const JITDylib &V, const SymbolStringPtr &Name) override { + // All original symbols were materialized by the CODLayer and should be + // final. The function bodies provided by M should never be overridden. + llvm_unreachable("Discard should never be called on an " + "ExtractingIRMaterializationUnit"); + } + + IRPartitionLayer &Parent; +}; + +} // namespace orc +} // namespace llvm + +IRPartitionLayer::IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer) + : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer) {} + +void IRPartitionLayer::setPartitionFunction(PartitionFunction Partition) { + this->Partition = Partition; +} + +std::optional<IRPartitionLayer::GlobalValueSet> +IRPartitionLayer::compileRequested(GlobalValueSet Requested) { + return std::move(Requested); +} + +std::optional<IRPartitionLayer::GlobalValueSet> +IRPartitionLayer::compileWholeModule(GlobalValueSet Requested) { + return std::nullopt; +} + +void IRPartitionLayer::emit(std::unique_ptr<MaterializationResponsibility> R, + ThreadSafeModule TSM) { + assert(TSM && "Null module"); + + auto &ES = getExecutionSession(); + TSM.withModuleDo([&](Module &M) { + // First, do some cleanup on the module: + cleanUpModule(M); + }); + + // Create a partitioning materialization unit and pass the responsibility. + if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>( + ES, *getManglingOptions(), std::move(TSM), *this))) { + ES.reportError(std::move(Err)); + R->failMaterialization(); + return; + } +} + +void IRPartitionLayer::cleanUpModule(Module &M) { + for (auto &F : M.functions()) { + if (F.isDeclaration()) + continue; + + if (F.hasAvailableExternallyLinkage()) { + F.deleteBody(); + F.setPersonalityFn(nullptr); + continue; + } + } +} + +void IRPartitionLayer::expandPartition(GlobalValueSet &Partition) { + // Expands the partition to ensure the following rules hold: + // (1) If any alias is in the partition, its aliasee is also in the partition. + // (2) If any aliasee is in the partition, its aliases are also in the + // partiton. + // (3) If any global variable is in the partition then all global variables + // are in the partition. + assert(!Partition.empty() && "Unexpected empty partition"); + + const Module &M = *(*Partition.begin())->getParent(); + bool ContainsGlobalVariables = false; + std::vector<const GlobalValue *> GVsToAdd; + + for (const auto *GV : Partition) + if (isa<GlobalAlias>(GV)) + GVsToAdd.push_back( + cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee())); + else if (isa<GlobalVariable>(GV)) + ContainsGlobalVariables = true; + + for (auto &A : M.aliases()) + if (Partition.count(cast<GlobalValue>(A.getAliasee()))) + GVsToAdd.push_back(&A); + + if (ContainsGlobalVariables) + for (auto &G : M.globals()) + GVsToAdd.push_back(&G); + + for (const auto *GV : GVsToAdd) + Partition.insert(GV); +} + +void IRPartitionLayer::emitPartition( + std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM, + IRMaterializationUnit::SymbolNameToDefinitionMap Defs) { + + // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the + // extracted module key, extracted module, and source module key + // together. This could be used, for example, to provide a specific + // memory manager instance to the linking layer. + + auto &ES = getExecutionSession(); + GlobalValueSet RequestedGVs; + for (auto &Name : R->getRequestedSymbols()) { + if (Name == R->getInitializerSymbol()) + TSM.withModuleDo([&](Module &M) { + for (auto &GV : getStaticInitGVs(M)) + RequestedGVs.insert(&GV); + }); + else { + assert(Defs.count(Name) && "No definition for symbol"); + RequestedGVs.insert(Defs[Name]); + } + } + + /// Perform partitioning with the context lock held, since the partition + /// function is allowed to access the globals to compute the partition. + auto GVsToExtract = + TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); }); + + // Take a 'None' partition to mean the whole module (as opposed to an empty + // partition, which means "materialize nothing"). Emit the whole module + // unmodified to the base layer. + if (GVsToExtract == std::nullopt) { + Defs.clear(); + BaseLayer.emit(std::move(R), std::move(TSM)); + return; + } + + // If the partition is empty, return the whole module to the symbol table. + if (GVsToExtract->empty()) { + if (auto Err = + R->replace(std::make_unique<PartitioningIRMaterializationUnit>( + std::move(TSM), + MaterializationUnit::Interface(R->getSymbols(), + R->getInitializerSymbol()), + std::move(Defs), *this))) { + getExecutionSession().reportError(std::move(Err)); + R->failMaterialization(); + return; + } + return; + } + + // Ok -- we actually need to partition the symbols. Promote the symbol + // linkages/names, expand the partition to include any required symbols + // (i.e. symbols that can't be separated from our partition), and + // then extract the partition. + // + // FIXME: We apply this promotion once per partitioning. It's safe, but + // overkill. + auto ExtractedTSM = TSM.withModuleDo([&](Module &M) + -> Expected<ThreadSafeModule> { + auto PromotedGlobals = PromoteSymbols(M); + if (!PromotedGlobals.empty()) { + + MangleAndInterner Mangle(ES, M.getDataLayout()); + SymbolFlagsMap SymbolFlags; + IRSymbolMapper::add(ES, *getManglingOptions(), PromotedGlobals, + SymbolFlags); + + if (auto Err = R->defineMaterializing(SymbolFlags)) + return std::move(Err); + } + + expandPartition(*GVsToExtract); + + // Submodule name is given by hashing the names of the globals. + std::string SubModuleName; + { + std::vector<const GlobalValue *> HashGVs; + HashGVs.reserve(GVsToExtract->size()); + for (const auto *GV : *GVsToExtract) + HashGVs.push_back(GV); + llvm::sort(HashGVs, [](const GlobalValue *LHS, const GlobalValue *RHS) { + return LHS->getName() < RHS->getName(); + }); + hash_code HC(0); + for (const auto *GV : HashGVs) { + assert(GV->hasName() && "All GVs to extract should be named by now"); + auto GVName = GV->getName(); + HC = hash_combine(HC, hash_combine_range(GVName.begin(), GVName.end())); + } + raw_string_ostream(SubModuleName) + << ".submodule." + << formatv(sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}", + static_cast<size_t>(HC)) + << ".ll"; + } + + // Extract the requested partiton (plus any necessary aliases) and + // put the rest back into the impl dylib. + auto ShouldExtract = [&](const GlobalValue &GV) -> bool { + return GVsToExtract->count(&GV); + }; + + return extractSubModule(TSM, SubModuleName, ShouldExtract); + }); + + if (!ExtractedTSM) { + ES.reportError(ExtractedTSM.takeError()); + R->failMaterialization(); + return; + } + + if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>( + ES, *getManglingOptions(), std::move(TSM), *this))) { + ES.reportError(std::move(Err)); + R->failMaterialization(); + return; + } + BaseLayer.emit(std::move(R), std::move(*ExtractedTSM)); +} diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp index 18ab9e6b940fefe..15fc50bad94aabf 100644 --- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp +++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp @@ -1288,9 +1288,12 @@ LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) { return; } + // Create the IP Layer. + IPLayer = std::make_unique<IRPartitionLayer>(*ES, *InitHelperTransformLayer); + // Create the COD layer. - CODLayer = std::make_unique<CompileOnDemandLayer>( - *ES, *InitHelperTransformLayer, *LCTMgr, std::move(ISMBuilder)); + CODLayer = std::make_unique<CompileOnDemandLayer>(*ES, *IPLayer, *LCTMgr, + std::move(ISMBuilder)); if (S.NumCompileThreads > 0) CODLayer->setCloneToNewContextOnEmit(true); diff --git a/llvm/tools/lli/lli.cpp b/llvm/tools/lli/lli.cpp index 360ffa628a47b3f..1378622cb71511d 100644 --- a/llvm/tools/lli/lli.cpp +++ b/llvm/tools/lli/lli.cpp @@ -30,6 +30,7 @@ #include "llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h" #include "llvm/ExecutionEngine/Orc/EPCGenericRTDyldMemoryManager.h" #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" +#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h" #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" #include "llvm/ExecutionEngine/Orc/LLJIT.h" #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" @@ -984,7 +985,7 @@ int runOrcJIT(const char *ProgName) { } if (PerModuleLazy) - J->setPartitionFunction(orc::CompileOnDemandLayer::compileWholeModule); + J->setPartitionFunction(orc::IRPartitionLayer::compileWholeModule); auto Dump = createDebugDumper(); >From 342adf162fc54a20e77b18ad5580622871657f8c Mon Sep 17 00:00:00 2001 From: Sunho Kim <ksunhokim...@gmail.com> Date: Wed, 20 Sep 2023 05:22:43 +0900 Subject: [PATCH 4/4] [ORC] Implement basic reoptimization. --- compiler-rt/lib/orc/common.h | 6 +- compiler-rt/lib/orc/elfnix_platform.cpp | 1 + .../ExecutionEngine/Orc/ReOptimizeLayer.h | 179 +++++++++++ llvm/lib/ExecutionEngine/Orc/CMakeLists.txt | 1 + .../ExecutionEngine/Orc/ReOptimizeLayer.cpp | 278 ++++++++++++++++++ .../ExecutionEngine/Orc/CMakeLists.txt | 1 + .../Orc/ReOptimizeLayerTest.cpp | 150 ++++++++++ 7 files changed, 613 insertions(+), 3 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 73c5c4a2bd8d474..f34229a615341df 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 c087e71038f9504..f76a070240315fe 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/ReOptimizeLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h new file mode 100644 index 000000000000000..df2f005825b5a7d --- /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/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt index 379c3cc240e8c7b..550a36e7afc4771 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/ReOptimizeLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp new file mode 100644 index 000000000000000..2ec82f210e6f625 --- /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 e932ff94361a7b7..310cb8276d894a6 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/ReOptimizeLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp new file mode 100644 index 000000000000000..c7e795c749ad000 --- /dev/null +++ b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp @@ -0,0 +1,150 @@ +#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(*ES, *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); +} _______________________________________________ cfe-commits mailing list cfe-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits