ExecutionEngine/Orc/ReOptimizeLayer.h

0001 //===- ReOptimizeLayer.h - Re-optimization layer interface ------*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 //
0009 // Re-optimization layer interface.
0010 //
0011 //===----------------------------------------------------------------------===//
0012 #ifndef LLVM_EXECUTIONENGINE_ORC_REOPTIMIZELAYER_H
0013 #define LLVM_EXECUTIONENGINE_ORC_REOPTIMIZELAYER_H
0014
0015 #include "llvm/ExecutionEngine/Orc/Core.h"
0016 #include "llvm/ExecutionEngine/Orc/Layer.h"
0017 #include "llvm/ExecutionEngine/Orc/Mangling.h"
0018 #include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
0019 #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
0020 #include "llvm/IR/IRBuilder.h"
0021 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
0022 #include "llvm/Transforms/Utils/Cloning.h"
0023
0024 namespace llvm {
0025 namespace orc {
0026
0027 class ReOptimizeLayer : public IRLayer, public ResourceManager {
0028 public:
0029   using ReOptMaterializationUnitID = uint64_t;
0030
0031   /// AddProfilerFunc will be called when ReOptimizeLayer emits the first
0032   /// version of a materialization unit in order to inject profiling code and
0033   /// reoptimization request code.
0034   using AddProfilerFunc = unique_function<Error(
0035       ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
0036       unsigned CurVersion, ThreadSafeModule &TSM)>;
0037
0038   /// ReOptimizeFunc will be called when ReOptimizeLayer reoptimization of a
0039   /// materialization unit was requested in order to reoptimize the IR module
0040   /// based on profile data. OldRT is the ResourceTracker that tracks the old
0041   /// function definitions. The OldRT must be kept alive until it can be
0042   /// guaranteed that every invocation of the old function definitions has been
0043   /// terminated.
0044   using ReOptimizeFunc = unique_function<Error(
0045       ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
0046       unsigned CurVersion, ResourceTrackerSP OldRT, ThreadSafeModule &TSM)>;
0047
0048   ReOptimizeLayer(ExecutionSession &ES, DataLayout &DL, IRLayer &BaseLayer,
0049                   RedirectableSymbolManager &RM)
0050       : IRLayer(ES, BaseLayer.getManglingOptions()), ES(ES), Mangle(ES, DL),
0051         BaseLayer(BaseLayer), RSManager(RM), ReOptFunc(identity),
0052         ProfilerFunc(reoptimizeIfCallFrequent) {}
0053
0054   void setReoptimizeFunc(ReOptimizeFunc ReOptFunc) {
0055     this->ReOptFunc = std::move(ReOptFunc);
0056   }
0057
0058   void setAddProfilerFunc(AddProfilerFunc ProfilerFunc) {
0059     this->ProfilerFunc = std::move(ProfilerFunc);
0060   }
0061
0062   /// Registers reoptimize runtime dispatch handlers to given PlatformJD. The
0063   /// reoptimization request will not be handled if dispatch handler is not
0064   /// registered by using this function.
0065   Error reigsterRuntimeFunctions(JITDylib &PlatformJD);
0066
0067   /// Emits the given module. This should not be called by clients: it will be
0068   /// called by the JIT when a definition added via the add method is requested.
0069   void emit(std::unique_ptr<MaterializationResponsibility> R,
0070             ThreadSafeModule TSM) override;
0071
0072   static const uint64_t CallCountThreshold = 10;
0073
0074   /// Basic AddProfilerFunc that reoptimizes the function when the call count
0075   /// exceeds CallCountThreshold.
0076   static Error reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
0077                                         ReOptMaterializationUnitID MUID,
0078                                         unsigned CurVersion,
0079                                         ThreadSafeModule &TSM);
0080
0081   static Error identity(ReOptimizeLayer &Parent,
0082                         ReOptMaterializationUnitID MUID, unsigned CurVersion,
0083                         ResourceTrackerSP OldRT, ThreadSafeModule &TSM) {
0084     return Error::success();
0085   }
0086
0087   // Create IR reoptimize request fucntion call.
0088   static void createReoptimizeCall(Module &M, Instruction &IP,
0089                                    GlobalVariable *ArgBuffer);
0090
0091   Error handleRemoveResources(JITDylib &JD, ResourceKey K) override;
0092   void handleTransferResources(JITDylib &JD, ResourceKey DstK,
0093                                ResourceKey SrcK) override;
0094
0095 private:
0096   class ReOptMaterializationUnitState {
0097   public:
0098     ReOptMaterializationUnitState() = default;
0099     ReOptMaterializationUnitState(ReOptMaterializationUnitID ID,
0100                                   ThreadSafeModule TSM)
0101         : ID(ID), TSM(std::move(TSM)) {}
0102     ReOptMaterializationUnitState(ReOptMaterializationUnitState &&Other)
0103         : ID(Other.ID), TSM(std::move(Other.TSM)), RT(std::move(Other.RT)),
0104           Reoptimizing(std::move(Other.Reoptimizing)),
0105           CurVersion(Other.CurVersion) {}
0106
0107     ReOptMaterializationUnitID getID() { return ID; }
0108
0109     const ThreadSafeModule &getThreadSafeModule() { return TSM; }
0110
0111     ResourceTrackerSP getResourceTracker() {
0112       std::unique_lock<std::mutex> Lock(Mutex);
0113       return RT;
0114     }
0115
0116     void setResourceTracker(ResourceTrackerSP RT) {
0117       std::unique_lock<std::mutex> Lock(Mutex);
0118       this->RT = RT;
0119     }
0120
0121     uint32_t getCurVersion() {
0122       std::unique_lock<std::mutex> Lock(Mutex);
0123       return CurVersion;
0124     }
0125
0126     bool tryStartReoptimize();
0127     void reoptimizeSucceeded();
0128     void reoptimizeFailed();
0129
0130   private:
0131     std::mutex Mutex;
0132     ReOptMaterializationUnitID ID;
0133     ThreadSafeModule TSM;
0134     ResourceTrackerSP RT;
0135     bool Reoptimizing = false;
0136     uint32_t CurVersion = 0;
0137   };
0138
0139   using SPSReoptimizeArgList =
0140       shared::SPSArgList<ReOptMaterializationUnitID, uint32_t>;
0141   using SendErrorFn = unique_function<void(Error)>;
0142
0143   Expected<SymbolMap> emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
0144                                         uint32_t Version, JITDylib &JD,
0145                                         ThreadSafeModule TSM);
0146
0147   void rt_reoptimize(SendErrorFn SendResult, ReOptMaterializationUnitID MUID,
0148                      uint32_t CurVersion);
0149
0150   static Expected<Constant *>
0151   createReoptimizeArgBuffer(Module &M, ReOptMaterializationUnitID MUID,
0152                             uint32_t CurVersion);
0153
0154   ReOptMaterializationUnitState &
0155   createMaterializationUnitState(const ThreadSafeModule &TSM);
0156
0157   void
0158   registerMaterializationUnitResource(ResourceKey Key,
0159                                       ReOptMaterializationUnitState &State);
0160
0161   ReOptMaterializationUnitState &
0162   getMaterializationUnitState(ReOptMaterializationUnitID MUID);
0163
0164   ExecutionSession &ES;
0165   MangleAndInterner Mangle;
0166   IRLayer &BaseLayer;
0167   RedirectableSymbolManager &RSManager;
0168
0169   ReOptimizeFunc ReOptFunc;
0170   AddProfilerFunc ProfilerFunc;
0171
0172   std::mutex Mutex;
0173   std::map<ReOptMaterializationUnitID, ReOptMaterializationUnitState> MUStates;
0174   DenseMap<ResourceKey, DenseSet<ReOptMaterializationUnitID>> MUResources;
0175   ReOptMaterializationUnitID NextID = 1;
0176 };
0177
0178 } // namespace orc
0179 } // namespace llvm
0180
0181 #endif // LLVM_EXECUTIONENGINE_ORC_REOPTIMIZELAYER_H