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