Index: include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h =================================================================== --- include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h +++ 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 ImplSymbolMap + void setImplMap(ImplSymbolMap *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; + ImplSymbolMap *AliaseeImpls = nullptr; }; /// Compile-on-demand layer. Index: include/llvm/ExecutionEngine/Orc/LazyReexports.h =================================================================== --- include/llvm/ExecutionEngine/Orc/LazyReexports.h +++ 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); + ImplSymbolMap *SrcJDLoc, VModuleKey K); StringRef getName() const override; @@ -174,6 +175,7 @@ SymbolAliasMap CallableAliases; std::shared_ptr NotifyResolved; + ImplSymbolMap *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, ImplSymbolMap *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: include/llvm/ExecutionEngine/Orc/Speculation.h =================================================================== --- /dev/null +++ include/llvm/ExecutionEngine/Orc/Speculation.h @@ -0,0 +1,242 @@ +//===-- 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/IR/PassManager.h" +#include "llvm/Passes/PassBuilder.h" +#include +#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 ImplSymbolMap { + 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; +}; + +// Defines Speculator Concept, clients can re-use this common functionality +// for different program representation + +class Speculator { +public: + using TargetFAddr = JITTargetAddress; + using TargetAndLikelies = DenseMap; + using StubAddrLikelies = DenseMap; + +private: + void registerSymbolsWithAddr(TargetFAddr ImplAddr, + SymbolNameSet likelySymbols) { + std::lock_guard lockit(ConcurrentAccess); + GlobalSpecMap.insert({ImplAddr, std::move(likelySymbols)}); + } + + void launchCompile(JITTargetAddress FAddr) { + auto It = GlobalSpecMap.find(FAddr); + assert(It != GlobalSpecMap.end() && + "launching speculative compiles for unexpected function address?"); + for (auto &Callee : It->getSecond()) { + auto ImplSymbol = AliaseeImplTable.getImplFor(Callee); + 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); + } + } + +public: + Speculator(ImplSymbolMap &Impl, ExecutionSession &ref) + : AliaseeImplTable(Impl), ES(ref), GlobalSpecMap(0) {} + Speculator(const Speculator &) = delete; + Speculator(Speculator &&) = delete; + Speculator &operator=(const Speculator &) = delete; + Speculator &operator=(Speculator &&) = delete; + virtual ~Speculator() {} + + // Speculatively compile likely functions for the given Stub Address. + // destination of __orc_speculate_for jump + void speculateFor(TargetFAddr StubAddr) { launchCompile(StubAddr); } + + // FIXME : Register with Stub Address, after JITLink Fix. + void registerSymbols(TargetAndLikelies Candidates, JITDylib *JD) { + for (auto &SymPair : Candidates) { + auto Target = SymPair.first; + auto Likely = SymPair.second; + + auto OnReadyFixUp = [Likely, Target, + this](Expected ReadySymbol) { + if (ReadySymbol) { + auto RAddr = (*ReadySymbol)[Target].getAddress(); + registerSymbolsWithAddr(RAddr, std::move(Likely)); + } else + this->getES().reportError(ReadySymbol.takeError()); + }; + ES.lookup(JITDylibSearchList({{JD, false}}), SymbolNameSet({Target}), + SymbolState::Ready, OnReadyFixUp, NoDependenciesToRegister); + } + } + + ExecutionSession &getES() { return ES; } + +private: + std::mutex ConcurrentAccess; + ImplSymbolMap &AliaseeImplTable; + ExecutionSession &ES; + StubAddrLikelies GlobalSpecMap; +}; + +// Speculator Interface for LLVM IR representation +class IRSpeculatorImpl : private Speculator { + +public: + using CompilationUnitTy = Module; + using PassIRUnit = Function; + using IRlikiesStrRef = DenseMap>; + +private: + // register all function analysis pass in LLVM + void setUpAnalysesPass() { PB.registerFunctionAnalyses(FAM); } + +public: + IRSpeculatorImpl(ImplSymbolMap &impl, ExecutionSession &ES) + : Speculator(impl, ES) { + setUpAnalysesPass(); + } + + template < + typename AnalysisTy, + typename std::enable_if< + std::is_base_of, AnalysisTy>::value, + bool>::type = true> + void registerAnalysis() { + FAM.registerPass([]() { return AnalysisTy(); }); + } + + void runOn(const CompilationUnitTy &ProgramRepr, JITDylib &LookupDylib) { + for (auto &Fn : ProgramRepr.getFunctionList()) { + if (!Fn.isDeclaration()) { + auto InternedNames = internToJITSymbols(collect(Fn)); + registerSymbols(InternedNames, &LookupDylib); + } + } + } + + FunctionAnalysisManager &getAnalysisManager() { return FAM; } + +private: + virtual IRlikiesStrRef collect(const PassIRUnit &) = 0; + + TargetAndLikelies internToJITSymbols(IRlikiesStrRef IRNames) { + TargetAndLikelies InternedNames; + DenseSet TargetJITNames; + ExecutionSession &Es = getES(); + for (auto &NamePair : IRNames) { + for (auto &TargetNames : NamePair.second) + TargetJITNames.insert(Es.intern(TargetNames)); + + InternedNames.insert( + {Es.intern(NamePair.first), std::move(TargetJITNames)}); + } + return InternedNames; + } + + PassBuilder PB; + FunctionAnalysisManager FAM; +}; + +// Create Object of this class templated with top-level analysis pass, type is +// used to query the Function Analysis Manager, Pass a function object to help +// speculator to retrieve likely next executable based on the Analysis Result + +template < + typename AnalysisType, + typename std::enable_if< + std::is_base_of, AnalysisType>::value, + bool>::type = true> +class IRSpeculator : public IRSpeculatorImpl { +public: + using EvalResultTy = std::function; + IRSpeculator(ImplSymbolMap &Impl, ExecutionSession &ES, + EvalResultTy &&Interpreter) + : IRSpeculatorImpl(Impl, ES) { + ResultResolver = Interpreter; + } + +private: + IRlikiesStrRef collect(const PassIRUnit &Fn) override { + auto &V = getAnalysisManager().template getResult( + const_cast(Fn)); + auto IRNames = ResultResolver(V, Fn); + return IRNames; + } + EvalResultTy ResultResolver; +}; + +// Helps to Speculate and Instrument LLVM IR +class IRSpeculationLayer : public IRLayer { +public: + IRSpeculationLayer(ExecutionSession &ES, IRCompileLayer &BaseLayer, + IRSpeculatorImpl &Spec) + : IRLayer(ES), NextLayer(BaseLayer), S(Spec) {} + + void emit(MaterializationResponsibility R, ThreadSafeModule TSM); + +private: + IRCompileLayer &NextLayer; + IRSpeculatorImpl &S; +}; + +// Runtime Function Interface +extern "C" { +void __orc_speculate_for(Speculator *, uint64_t stub_id); +} + +} // namespace orc +} // namespace llvm +#endif \ No newline at end of file Index: lib/ExecutionEngine/Orc/CMakeLists.txt =================================================================== --- lib/ExecutionEngine/Orc/CMakeLists.txt +++ lib/ExecutionEngine/Orc/CMakeLists.txt @@ -21,7 +21,7 @@ RPCUtils.cpp RTDyldObjectLinkingLayer.cpp ThreadSafeModule.cpp - + Speculation.cpp ADDITIONAL_HEADER_DIRS ${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc @@ -33,4 +33,5 @@ PRIVATE LLVMBitReader LLVMBitWriter + LLVMAnalysis ) Index: lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp =================================================================== --- lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp +++ lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp @@ -117,6 +117,9 @@ this->Partition = std::move(Partition); } +void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) { + this->AliaseeImpls = Imp; +} void CompileOnDemandLayer::emit(MaterializationResponsibility R, ThreadSafeModule TSM) { assert(TSM.getModule() && "Null module"); @@ -158,7 +161,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: lib/ExecutionEngine/Orc/LazyReexports.cpp =================================================================== --- lib/ExecutionEngine/Orc/LazyReexports.cpp +++ 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, ImplSymbolMap *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) { @@ -168,6 +170,9 @@ std::make_pair(*CallThroughTrampoline, Alias.second.AliasFlags); } + if (AliaseeTable != nullptr) + AliaseeTable->trackImpls(RequestedAliases, &SourceJD); + if (auto Err = ISManager.createStubs(StubInits)) { SourceJD.getExecutionSession().reportError(std::move(Err)); R.failMaterialization(); Index: lib/ExecutionEngine/Orc/Speculation.cpp =================================================================== --- /dev/null +++ lib/ExecutionEngine/Orc/Speculation.cpp @@ -0,0 +1,83 @@ +//===---------- 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 { + +// ImplSymbolMap methods +void ImplSymbolMap::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?"); + } +} + +void IRSpeculationLayer::emit(MaterializationResponsibility R, + ThreadSafeModule TSM) { + auto M = TSM.getModule(); + + // Do not spawn new thread. + S.runOn(*M, R.getTargetJITDylib()); + + assert(M && "Speculation Layer received Null Module ?"); + + // Instrumentation of runtime calls + auto &InContext = M->getContext(); + auto SpeculatorVTy = StructType::create(InContext, "Class.Speculator"); + auto RuntimeCallTy = FunctionType::get( + Type::getVoidTy(InContext), + {SpeculatorVTy->getPointerTo(), Type::getInt64Ty(InContext)}, false); + + auto RuntimeCall = + Function::Create(RuntimeCallTy, Function::LinkageTypes::ExternalLinkage, + "__orc_speculate_for", *M); + + auto SpeclAddr = new GlobalVariable( + *M, SpeculatorVTy, false, GlobalValue::LinkageTypes::ExternalLinkage, + nullptr, "__orc_speculator"); + + IRBuilder<> Mutator(InContext); + for (auto &Fn : M->getFunctionList()) { + if (!Fn.isDeclaration()) { + Mutator.SetInsertPoint(&(Fn.getEntryBlock().front())); + auto ImplAddrToUint = + Mutator.CreatePtrToInt(&Fn, Type::getInt64Ty(InContext)); + Mutator.CreateCall(RuntimeCallTy, RuntimeCall, + {SpeclAddr, ImplAddrToUint}); + } + } + + assert(!verifyModule(*M) && "Speculation Instrumentation breaks IR?"); + + NextLayer.emit(std::move(R), std::move(TSM)); +} + +// Runtime Function Implementation +extern "C" void __orc_speculate_for(Speculator *Ptr, uint64_t StubId) { + assert(Ptr && " Null Address Received in orc_speculate_for "); + Ptr->speculateFor(StubId); +} + +} // namespace orc +} // namespace llvm \ No newline at end of file