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,8 @@ /// 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 +131,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,125 @@ +//===-- 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 Impl JITDylibs of Symbols while lazy call through trampolines +// are created and send the impl details to GlobalSpeculationObserver + +class Imap { + friend class Speculator; + +public: + using ImplPair = std::pair; + using ImapTy = DenseMap; + void saveImpls(SymbolAliasMap ImplMaps, JITDylib *SrcJD); + +private: + void recordImpl(SymbolStringPtr FaceSymbol, ImplPair Implementations) { + auto It = Maps.insert({FaceSymbol, Implementations}); + if (It.second == false) + assert(0 && "Source Entities are already tracked for this Symbol?"); + } + + ImplPair getImplFor(SymbolStringPtr StubAddr) { + auto Position = Maps.find(StubAddr); + if (Position != Maps.end()) { + return Position->getSecond(); + } + assert(0 && "Source Entities are not tracked for a Symbol?"); + } + + std::mutex ConcurrentAccess; + ImapTy Maps; +}; + +class Speculator { + +public: + using TargetFAddr = JITTargetAddress; + using SpeculationMap = DenseMap; + + Speculator(Imap &Impl) : AliaseeImplTable(Impl) {} + + // Speculation Layer registers likely function through this method. + void registerSymbolsWithAddr(TargetFAddr, SymbolNameSet); + + // Speculatively compile likely functions for the given Stub Address. + void speculateFor(JITTargetAddress); + +private: + void launchCompile(JITTargetAddress FAddr) { + auto It = GlobalSpecMap.find(FAddr); + if (It != GlobalSpecMap.end()) { + for (auto &Pos : It->getSecond()) { + auto SourceEntities = AliaseeImplTable.getImplFor(Pos); + } + } else + assert(0 && "launching compiles for Unexpected Function Address?"); + } + + std::mutex ConcurrentAccess; + Imap &AliaseeImplTable; + SpeculationMap GlobalSpecMap; +}; + +// 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 WalkerResultTy = 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; } + + static WalkerResultTy Walk(const Function &); + + DenseSet internAllFns(WalkerResultTy &&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,6 +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->saveImpls(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,131 @@ +//===---------- 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::saveImpls(SymbolAliasMap ImplMaps, JITDylib *SrcJD) { + std::lock_guard lockit(ConcurrentAccess); + for (auto &I : ImplMaps) { + recordImpl(I.first, {I.second.Aliasee, SrcJD}); + } +} + +// Speculator methods + +void Speculator::registerSymbolsWithAddr(TargetFAddr ImplAddr, + SymbolNameSet likelySymbols) { + std::lock_guard lockit(ConcurrentAccess); + GlobalSpecMap.insert({ImplAddr, std::move(likelySymbols)}); +} + +void Speculator::speculateFor(JITTargetAddress FAddr) { + std::lock_guard lockit(ConcurrentAccess); + launchCompile(FAddr); +} + +// SpeculationLayer methods + +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 = getSpeculator(); + for (auto &Symbol : SymbolsToMap) { + auto Target = Symbol.first; + auto likely = Symbol.second; + // Appending Queries on the same symbol but with different callback action + ESession.lookup( + JITDylibSearchList({{&R.getTargetJITDylib(), false}}), + SymbolNameSet({Symbol.first}), SymbolState::Resolved, + [&SpecMap, likely, Target, &ESession](Expected ResSymMap) { + if (ResSymMap) { + auto RAddr = (*ResSymMap)[Target].getAddress(); + SpecMap.registerSymbolsWithAddr(std::move(RAddr), likely); + } else { + ESession.reportError(ResSymMap.takeError()); + return; // fail MaterializationResponsibility is called by first + // queued query. + } + }, + NoDependenciesToRegister); + } + NextLayer.emit(std::move(R), std::move(TSM)); +} + +DenseSet SpeculationLayer::internAllFns(WalkerResultTy &&AR) { + DenseSet Symbols; + auto &ESession = this->getExecutionSession(); + for (auto Name : AR) { + Symbols.insert(ESession.intern(Name->getName())); + } + return Symbols; +} + +SpeculationLayer::WalkerResultTy SpeculationLayer::Walk(const Function &F) { + WalkerResultTy 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; +} + +} // 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