1 //===- ReOptimizeLayer.h - Re-optimization layer interface ------*- C++ -*-===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // Re-optimization layer interface. 10 // 11 //===----------------------------------------------------------------------===// 12 #ifndef LLVM_EXECUTIONENGINE_ORC_REOPTIMIZELAYER_H 13 #define LLVM_EXECUTIONENGINE_ORC_REOPTIMIZELAYER_H 14 15 #include "llvm/ExecutionEngine/Orc/Core.h" 16 #include "llvm/ExecutionEngine/Orc/Layer.h" 17 #include "llvm/ExecutionEngine/Orc/Mangling.h" 18 #include "llvm/ExecutionEngine/Orc/RedirectionManager.h" 19 #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h" 20 #include "llvm/IR/IRBuilder.h" 21 #include "llvm/Support/Compiler.h" 22 #include "llvm/Transforms/Utils/BasicBlockUtils.h" 23 #include "llvm/Transforms/Utils/Cloning.h" 24 25 namespace llvm { 26 namespace orc { 27 28 class LLVM_ABI ReOptimizeLayer : public IRLayer, public ResourceManager { 29 public: 30 using ReOptMaterializationUnitID = uint64_t; 31 32 /// AddProfilerFunc will be called when ReOptimizeLayer emits the first 33 /// version of a materialization unit in order to inject profiling code and 34 /// reoptimization request code. 35 using AddProfilerFunc = unique_function<Error( 36 ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID, 37 unsigned CurVersion, ThreadSafeModule &TSM)>; 38 39 /// ReOptimizeFunc will be called when ReOptimizeLayer reoptimization of a 40 /// materialization unit was requested in order to reoptimize the IR module 41 /// based on profile data. OldRT is the ResourceTracker that tracks the old 42 /// function definitions. The OldRT must be kept alive until it can be 43 /// guaranteed that every invocation of the old function definitions has been 44 /// terminated. 45 using ReOptimizeFunc = unique_function<Error( 46 ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID, 47 unsigned CurVersion, ResourceTrackerSP OldRT, ThreadSafeModule &TSM)>; 48 ReOptimizeLayer(ExecutionSession & ES,DataLayout & DL,IRLayer & BaseLayer,RedirectableSymbolManager & RM)49 ReOptimizeLayer(ExecutionSession &ES, DataLayout &DL, IRLayer &BaseLayer, 50 RedirectableSymbolManager &RM) 51 : IRLayer(ES, BaseLayer.getManglingOptions()), ES(ES), Mangle(ES, DL), 52 BaseLayer(BaseLayer), RSManager(RM), ReOptFunc(identity), 53 ProfilerFunc(reoptimizeIfCallFrequent) {} 54 setReoptimizeFunc(ReOptimizeFunc ReOptFunc)55 void setReoptimizeFunc(ReOptimizeFunc ReOptFunc) { 56 this->ReOptFunc = std::move(ReOptFunc); 57 } 58 setAddProfilerFunc(AddProfilerFunc ProfilerFunc)59 void setAddProfilerFunc(AddProfilerFunc ProfilerFunc) { 60 this->ProfilerFunc = std::move(ProfilerFunc); 61 } 62 63 /// Registers reoptimize runtime dispatch handlers to given PlatformJD. The 64 /// reoptimization request will not be handled if dispatch handler is not 65 /// registered by using this function. 66 Error reigsterRuntimeFunctions(JITDylib &PlatformJD); 67 68 /// Emits the given module. This should not be called by clients: it will be 69 /// called by the JIT when a definition added via the add method is requested. 70 void emit(std::unique_ptr<MaterializationResponsibility> R, 71 ThreadSafeModule TSM) override; 72 73 static const uint64_t CallCountThreshold = 10; 74 75 /// Basic AddProfilerFunc that reoptimizes the function when the call count 76 /// exceeds CallCountThreshold. 77 static Error reoptimizeIfCallFrequent(ReOptimizeLayer &Parent, 78 ReOptMaterializationUnitID MUID, 79 unsigned CurVersion, 80 ThreadSafeModule &TSM); 81 identity(ReOptimizeLayer & Parent,ReOptMaterializationUnitID MUID,unsigned CurVersion,ResourceTrackerSP OldRT,ThreadSafeModule & TSM)82 static Error identity(ReOptimizeLayer &Parent, 83 ReOptMaterializationUnitID MUID, unsigned CurVersion, 84 ResourceTrackerSP OldRT, ThreadSafeModule &TSM) { 85 return Error::success(); 86 } 87 88 // Create IR reoptimize request fucntion call. 89 static void createReoptimizeCall(Module &M, Instruction &IP, 90 GlobalVariable *ArgBuffer); 91 92 Error handleRemoveResources(JITDylib &JD, ResourceKey K) override; 93 void handleTransferResources(JITDylib &JD, ResourceKey DstK, 94 ResourceKey SrcK) override; 95 96 private: 97 class ReOptMaterializationUnitState { 98 public: 99 ReOptMaterializationUnitState() = default; ReOptMaterializationUnitState(ReOptMaterializationUnitID ID,ThreadSafeModule TSM)100 ReOptMaterializationUnitState(ReOptMaterializationUnitID ID, 101 ThreadSafeModule TSM) 102 : ID(ID), TSM(std::move(TSM)) {} ReOptMaterializationUnitState(ReOptMaterializationUnitState && Other)103 ReOptMaterializationUnitState(ReOptMaterializationUnitState &&Other) 104 : ID(Other.ID), TSM(std::move(Other.TSM)), RT(std::move(Other.RT)), 105 Reoptimizing(std::move(Other.Reoptimizing)), 106 CurVersion(Other.CurVersion) {} 107 getID()108 ReOptMaterializationUnitID getID() { return ID; } 109 getThreadSafeModule()110 const ThreadSafeModule &getThreadSafeModule() { return TSM; } 111 getResourceTracker()112 ResourceTrackerSP getResourceTracker() { 113 std::unique_lock<std::mutex> Lock(Mutex); 114 return RT; 115 } 116 setResourceTracker(ResourceTrackerSP RT)117 void setResourceTracker(ResourceTrackerSP RT) { 118 std::unique_lock<std::mutex> Lock(Mutex); 119 this->RT = RT; 120 } 121 getCurVersion()122 uint32_t getCurVersion() { 123 std::unique_lock<std::mutex> Lock(Mutex); 124 return CurVersion; 125 } 126 127 LLVM_ABI bool tryStartReoptimize(); 128 LLVM_ABI void reoptimizeSucceeded(); 129 LLVM_ABI void reoptimizeFailed(); 130 131 private: 132 std::mutex Mutex; 133 ReOptMaterializationUnitID ID; 134 ThreadSafeModule TSM; 135 ResourceTrackerSP RT; 136 bool Reoptimizing = false; 137 uint32_t CurVersion = 0; 138 }; 139 140 using SPSReoptimizeArgList = 141 shared::SPSArgList<ReOptMaterializationUnitID, uint32_t>; 142 using SendErrorFn = unique_function<void(Error)>; 143 144 Expected<SymbolMap> emitMUImplSymbols(ReOptMaterializationUnitState &MUState, 145 uint32_t Version, JITDylib &JD, 146 ThreadSafeModule TSM); 147 148 void rt_reoptimize(SendErrorFn SendResult, ReOptMaterializationUnitID MUID, 149 uint32_t CurVersion); 150 151 static Expected<Constant *> 152 createReoptimizeArgBuffer(Module &M, ReOptMaterializationUnitID MUID, 153 uint32_t CurVersion); 154 155 ReOptMaterializationUnitState & 156 createMaterializationUnitState(const ThreadSafeModule &TSM); 157 158 void 159 registerMaterializationUnitResource(ResourceKey Key, 160 ReOptMaterializationUnitState &State); 161 162 ReOptMaterializationUnitState & 163 getMaterializationUnitState(ReOptMaterializationUnitID MUID); 164 165 ExecutionSession &ES; 166 MangleAndInterner Mangle; 167 IRLayer &BaseLayer; 168 RedirectableSymbolManager &RSManager; 169 170 ReOptimizeFunc ReOptFunc; 171 AddProfilerFunc ProfilerFunc; 172 173 std::mutex Mutex; 174 std::map<ReOptMaterializationUnitID, ReOptMaterializationUnitState> MUStates; 175 DenseMap<ResourceKey, DenseSet<ReOptMaterializationUnitID>> MUResources; 176 ReOptMaterializationUnitID NextID = 1; 177 }; 178 179 } // namespace orc 180 } // namespace llvm 181 182 #endif // LLVM_EXECUTIONENGINE_ORC_REOPTIMIZELAYER_H 183