Index: llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h =================================================================== --- llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h +++ llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h @@ -26,6 +26,7 @@ #include "llvm/ExecutionEngine/Orc/LazyReexports.h" #include "llvm/ExecutionEngine/Orc/Legacy.h" #include "llvm/ExecutionEngine/Orc/OrcError.h" +#include "llvm/ExecutionEngine/Orc/Speculation.h" #include "llvm/ExecutionEngine/RuntimeDyld.h" #include "llvm/IR/Attributes.h" #include "llvm/IR/Constant.h" @@ -91,6 +92,9 @@ /// Sets the partition function. void setPartitionFunction(PartitionFunction Partition); + /// Sets the Imap Pointer + void setImap(Imap *Imp); + /// 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(MaterializationResponsibility R, ThreadSafeModule TSM) override; @@ -128,6 +132,7 @@ PerDylibResourcesMap DylibResources; PartitionFunction Partition = compileRequested; SymbolLinkagePromoter PromoteSymbols; + Imap *AliaseeImpls = nullptr; }; /// Compile-on-demand layer. Index: llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h =================================================================== --- llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h +++ llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h @@ -18,6 +18,7 @@ #include "llvm/ExecutionEngine/Orc/Core.h" #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h" +#include "llvm/ExecutionEngine/Orc/Speculation.h" namespace llvm { @@ -159,7 +160,7 @@ IndirectStubsManager &ISManager, JITDylib &SourceJD, SymbolAliasMap CallableAliases, - VModuleKey K); + Imap *SrcJDLoc, VModuleKey K); StringRef getName() const override; @@ -174,6 +175,7 @@ SymbolAliasMap CallableAliases; std::shared_ptr NotifyResolved; + Imap *AliaseeTable; }; /// Define lazy-reexports based on the given SymbolAliasMap. Each lazy re-export @@ -182,9 +184,10 @@ inline std::unique_ptr lazyReexports(LazyCallThroughManager &LCTManager, IndirectStubsManager &ISManager, JITDylib &SourceJD, - SymbolAliasMap CallableAliases, VModuleKey K = VModuleKey()) { + SymbolAliasMap CallableAliases, Imap *SrcJDLoc, + VModuleKey K = VModuleKey()) { return llvm::make_unique( - LCTManager, ISManager, SourceJD, std::move(CallableAliases), + LCTManager, ISManager, SourceJD, std::move(CallableAliases), SrcJDLoc, std::move(K)); } Index: llvm/include/llvm/ExecutionEngine/Orc/Speculation.h =================================================================== --- /dev/null +++ llvm/include/llvm/ExecutionEngine/Orc/Speculation.h @@ -0,0 +1,136 @@ +//===-- Speculation.h - Speculative Compilation --*- 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 +// +//===----------------------------------------------------------------------===// +// +// Contains the definition to support speculative compilation when laziness is +// enabled. +//===----------------------------------------------------------------------===// + +#ifndef LLVM_EXECUTIONENGINE_ORC_SPECULATION_H +#define LLVM_EXECUTIONENGINE_ORC_SPECULATION_H + +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ExecutionEngine/Orc/Core.h" +#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/Layer.h" + +#include +#include +#include + +namespace llvm { +namespace orc { + +class Speculator; + +// Track the Impls (JITDylib,Symbols) of Symbols while lazy call through +// trampolines are created. Operations are guarded by locks tp ensure that Imap +// stays in consistent state after read/write + +class Imap { + friend class Speculator; + +public: + using ImplPair = std::pair; + using ImapTy = DenseMap; + + void trackImpls(SymbolAliasMap ImplMaps, JITDylib *SrcJD); + +private: + ImplPair getImplFor(const SymbolStringPtr &StubAddr) { + std::lock_guard Lockit(ConcurrentAccess); + auto Position = Maps.find(StubAddr); + assert(Position != Maps.end() && + "ImplSymbols are not tracked for this Symbol?"); + return Position->getSecond(); + } + + std::mutex ConcurrentAccess; + ImapTy Maps; +}; + +class Speculator { +public: + using TargetFAddr = JITTargetAddress; + using SpeculationMap = DenseMap; + + explicit Speculator(Imap &Impl, ExecutionSession &ref) + : AliaseeImplTable(Impl), ES(ref) {} + + // Speculation Layer registers likely function through this method. + void registerSymbolsWithAddr(TargetFAddr, SymbolNameSet); + + // Speculatively compile likely functions for the given Stub Address. + // destination of __orc_speculate_for jump + void speculateFor(JITTargetAddress); + + Imap &getImapRef() { return AliaseeImplTable; } + +private: + void launchCompile(JITTargetAddress FAddr) { + auto It = GlobalSpecMap.find(FAddr); + assert(It != GlobalSpecMap.end() && + "launching speculative compiles for unexpected function address?"); + for (auto &Calle : It->getSecond()) { + auto ImplSymbol = AliaseeImplTable.getImplFor(Calle); + const auto &ImplSymbolName = ImplSymbol.first; + auto *ImplJD = ImplSymbol.second; + ES.lookup(JITDylibSearchList({{ImplJD, true}}), + SymbolNameSet({ImplSymbolName}), SymbolState::Ready, + [this](Expected Result) { + if (auto Err = Result.takeError()) + ES.reportError(std::move(Err)); + }, + NoDependenciesToRegister); + } + } + + std::mutex ConcurrentAccess; + Imap &AliaseeImplTable; + SpeculationMap GlobalSpecMap; + ExecutionSession &ES; +}; + +// Walks the LLVM Module and collect likely functions for each LLVM Function if +// any, and instruments the IR with runtime call named __orc_speculate_for. +// Construct IR level function mapping. + +class SpeculationLayer : public IRLayer { + +public: + using CalleSet = DenseSet; + using WalkerTy = std::function; + + SpeculationLayer(ExecutionSession &ES, IRCompileLayer &BaseLayer, + Speculator &Spec) + : IRLayer(ES), NextLayer(BaseLayer), S(Spec) {} + + // For each function F + // Walk F to find likely functions. + // If F has callers, then instrument F + void emit(MaterializationResponsibility R, ThreadSafeModule TSM); + + void setModuleWalker(WalkerTy W) { Walker = std::move(W); } + + Speculator &getSpeculator() const { return S; } + + IRCompileLayer &getCompileLayer() const { return NextLayer; } + + static CalleSet Walk(const Function &); + + DenseSet internAllFns(CalleSet &&AR); + +private: + IRCompileLayer &NextLayer; + Speculator &S; + WalkerTy Walker = Walk; +}; + +} // namespace orc +} // namespace llvm +#endif \ No newline at end of file Index: llvm/include/llvm/ExecutionEngine/Orc/SpeculationRT.h =================================================================== --- /dev/null +++ llvm/include/llvm/ExecutionEngine/Orc/SpeculationRT.h @@ -0,0 +1,22 @@ +//===-- Speculation-rt.h - Runtime Calls --*- 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 +// +//===----------------------------------------------------------------------===// +// Interfaces for Runtime Calls +//===----------------------------------------------------------------------===// + +#ifndef LLVM_EXECUTIONENGINE_ORC_SPECULATIONRT_H +#define LLVM_EXECUTIONENGINE_ORC_SPECULATIONRT_H + +#include +namespace llvm { +namespace orc { +extern "C" { +void __orc_speculate_for(uint64_t stub_id); +} +} // namespace orc +} // namespace llvm +#endif \ No newline at end of file Index: llvm/lib/ExecutionEngine/Orc/CMakeLists.txt =================================================================== --- llvm/lib/ExecutionEngine/Orc/CMakeLists.txt +++ llvm/lib/ExecutionEngine/Orc/CMakeLists.txt @@ -21,7 +21,8 @@ RPCUtils.cpp RTDyldObjectLinkingLayer.cpp ThreadSafeModule.cpp - + Speculation.cpp + SpeculationRT.cpp ADDITIONAL_HEADER_DIRS ${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc Index: llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp =================================================================== --- llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp +++ llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp @@ -117,6 +117,8 @@ this->Partition = std::move(Partition); } +void CompileOnDemandLayer::setImap(Imap *Imp) { this->AliaseeImpls = Imp; } + void CompileOnDemandLayer::emit(MaterializationResponsibility R, ThreadSafeModule TSM) { assert(TSM.getModule() && "Null module"); @@ -158,7 +160,7 @@ R.replace(reexports(PDR.getImplDylib(), std::move(NonCallables), true)); R.replace(lazyReexports(LCTMgr, PDR.getISManager(), PDR.getImplDylib(), - std::move(Callables))); + std::move(Callables), AliaseeImpls)); } CompileOnDemandLayer::PerDylibResources & Index: llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp =================================================================== --- llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp +++ llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp @@ -121,7 +121,8 @@ LazyReexportsMaterializationUnit::LazyReexportsMaterializationUnit( LazyCallThroughManager &LCTManager, IndirectStubsManager &ISManager, - JITDylib &SourceJD, SymbolAliasMap CallableAliases, VModuleKey K) + JITDylib &SourceJD, SymbolAliasMap CallableAliases, Imap *SrcJDLoc, + VModuleKey K) : MaterializationUnit(extractFlags(CallableAliases), std::move(K)), LCTManager(LCTManager), ISManager(ISManager), SourceJD(SourceJD), CallableAliases(std::move(CallableAliases)), @@ -129,7 +130,8 @@ [&ISManager](JITDylib &JD, const SymbolStringPtr &SymbolName, JITTargetAddress ResolvedAddr) { return ISManager.updatePointer(*SymbolName, ResolvedAddr); - })) {} + })), + AliaseeTable(SrcJDLoc) {} StringRef LazyReexportsMaterializationUnit::getName() const { return ""; @@ -149,7 +151,7 @@ if (!CallableAliases.empty()) R.replace(lazyReexports(LCTManager, ISManager, SourceJD, - std::move(CallableAliases))); + std::move(CallableAliases), AliaseeTable)); IndirectStubsManager::StubInitsMap StubInits; for (auto &Alias : RequestedAliases) { @@ -167,6 +169,9 @@ StubInits[*Alias.first] = std::make_pair(*CallThroughTrampoline, Alias.second.AliasFlags); } + // enable speculative compilation + if (AliaseeTable != nullptr) + AliaseeTable->trackImpls(RequestedAliases, &SourceJD); if (auto Err = ISManager.createStubs(StubInits)) { SourceJD.getExecutionSession().reportError(std::move(Err)); Index: llvm/lib/ExecutionEngine/Orc/Speculation.cpp =================================================================== --- /dev/null +++ llvm/lib/ExecutionEngine/Orc/Speculation.cpp @@ -0,0 +1,127 @@ +//===---------- speculation.cpp - Utilities for Speculation ----------===// +// +// 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/Speculation.h" + +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Instruction.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Type.h" +#include "llvm/IR/Verifier.h" + +#include +namespace llvm { + +namespace orc { + +// Imap methods +void Imap::trackImpls(SymbolAliasMap ImplMaps, JITDylib *SrcJD) { + std::lock_guard lockit(ConcurrentAccess); + for (auto &I : ImplMaps) { + auto It = Maps.insert({I.first, {I.second.Aliasee, SrcJD}}); + assert(It.second && "ImplSymbols are already tracked for this Symbol?"); + } +} + +// Speculator methods +// FIX ME: Register with Unified Stub Address, after JITLink Fix. +void Speculator::registerSymbolsWithAddr(TargetFAddr ImplAddr, + SymbolNameSet likelySymbols) { + std::lock_guard lockit(ConcurrentAccess); + GlobalSpecMap.insert({ImplAddr, std::move(likelySymbols)}); +} + +void Speculator::speculateFor(JITTargetAddress FAddr) { launchCompile(FAddr); } + +void SpeculationLayer::emit(MaterializationResponsibility R, + ThreadSafeModule TSM) { + auto Module = TSM.getModule(); + assert(Module && "Speculation Layer received Null Module"); + auto &ESession = this->getExecutionSession(); + auto &InContext = Module->getContext(); + + // reinterpret_cast of Stub Address to i64 + auto RTFTy = FunctionType::get(Type::getVoidTy(InContext), + Type::getInt64Ty(InContext), false); + auto RTFn = Function::Create(RTFTy, Function::LinkageTypes::ExternalLinkage, + "__orc_speculate_for", *Module); + IRBuilder<> Mutator(InContext); + + DenseMap> SymbolsToMap; + + for (auto &Fn : Module->getFunctionList()) { + auto Candidates = Walker(Fn); + // callees + if (!Candidates.empty()) { + Mutator.SetInsertPoint(&(Fn.getEntryBlock().front())); + auto ImplAddrToUint = + Mutator.CreatePtrToInt(&Fn, Type::getInt64Ty(InContext)); + Mutator.CreateCall(RTFTy, RTFn, {ImplAddrToUint}); + auto FuncName = ESession.intern(Fn.getName()); + SymbolsToMap[FuncName] = internAllFns(std::move(Candidates)); + } + } + + assert(!verifyModule(*Module) && "Speculation Instrumentation breaks IR?"); + auto &SpecMap = S; + for (auto &Symbol : SymbolsToMap) { + auto Target = Symbol.first; + auto Likely = Symbol.second; + // Appending Queries on the same symbol but with different callback + // Callback should be OnReady,So Imap can track Impl Symbols. + ESession.lookup( + JITDylibSearchList({{&R.getTargetJITDylib(), false}}), + SymbolNameSet({Target}), SymbolState::Ready, + // FIX ME: Move Capture - Likely, Target, when we have C++14 + [&SpecMap, Likely, Target, &ESession](Expected ResSymMap) { + if (ResSymMap) { + auto RAddr = (*ResSymMap)[Target].getAddress(); + SpecMap.registerSymbolsWithAddr(RAddr, Likely); + SpecMap.speculateFor(RAddr); + } else { + ESession.reportError(ResSymMap.takeError()); + return; + } + }, + NoDependenciesToRegister); + } + NextLayer.emit(std::move(R), std::move(TSM)); +} + +SpeculationLayer::CalleSet SpeculationLayer::Walk(const Function &F) { + CalleSet Candidates; + for (auto &BB : F) { + for (auto &I : BB) { + if (auto Call = (dyn_cast(&I))) { + auto Callee = Call->getCalledFunction(); + if (Callee) + Candidates.insert(Callee); + } else if (auto Call = (dyn_cast(&I))) { + auto Callee = Call->getCalledFunction(); + if (Callee) + Candidates.insert(Callee); + } else { + } + } + } + return Candidates; +} +DenseSet SpeculationLayer::internAllFns(CalleSet &&AR) { + DenseSet Symbols; + auto &ESession = this->getExecutionSession(); + for (auto Name : AR) + Symbols.insert(ESession.intern(Name->getName())); + return Symbols; +} + +} // namespace orc +} // namespace llvm \ No newline at end of file Index: llvm/lib/ExecutionEngine/Orc/SpeculationRT.cpp =================================================================== --- /dev/null +++ llvm/lib/ExecutionEngine/Orc/SpeculationRT.cpp @@ -0,0 +1,22 @@ +//===---------- Speculation-rt.cpp - Utilities for lazy reexports +//----------===// +// +// 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/SpeculationRT.h" + +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "orc" + +namespace llvm { +namespace orc { +extern "C" void __orc_speculate_for(uint64_t StubId) { + // trigger compilation +} +} // namespace orc +} // namespace llvm \ No newline at end of file