[Orc] Refactor JITCompileCallbackManagerBase and CompileOnDemandLayer to support
[oota-llvm.git] / include / llvm / ExecutionEngine / Orc / IndirectionUtils.h
1 //===-- IndirectionUtils.h - Utilities for adding indirections --*- C++ -*-===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // Contains utilities for adding indirections and breaking up modules.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #ifndef LLVM_EXECUTIONENGINE_ORC_INDIRECTIONUTILS_H
15 #define LLVM_EXECUTIONENGINE_ORC_INDIRECTIONUTILS_H
16
17 #include "JITSymbol.h"
18 #include "llvm/ADT/DenseSet.h"
19 #include "llvm/IR/IRBuilder.h"
20 #include "llvm/IR/Mangler.h"
21 #include "llvm/IR/Module.h"
22 #include <sstream>
23
24 namespace llvm {
25 namespace orc {
26
27 /// @brief Base class for JITLayer independent aspects of
28 ///        JITCompileCallbackManager.
29 class JITCompileCallbackManagerBase {
30 public:
31
32   typedef std::function<TargetAddress()> CompileFtor;
33   typedef std::function<void(TargetAddress)> UpdateFtor;
34
35   /// @brief Handle to a newly created compile callback. Can be used to get an
36   ///        IR constant representing the address of the trampoline, and to set
37   ///        the compile and update actions for the callback.
38   class CompileCallbackInfo {
39   public:
40     CompileCallbackInfo(Constant *Addr, CompileFtor &Compile,
41                         UpdateFtor &Update)
42       : Addr(Addr), Compile(Compile), Update(Update) {}
43
44     Constant* getAddress() const { return Addr; }
45     void setCompileAction(CompileFtor Compile) {
46       this->Compile = std::move(Compile);
47     }
48     void setUpdateAction(UpdateFtor Update) {
49       this->Update = std::move(Update);
50     }
51   private:
52     Constant *Addr;
53     CompileFtor &Compile;
54     UpdateFtor &Update;
55   };
56
57   /// @brief Construct a JITCompileCallbackManagerBase.
58   /// @param ErrorHandlerAddress The address of an error handler in the target
59   ///                            process to be used if a compile callback fails.
60   /// @param NumTrampolinesPerBlock Number of trampolines to emit if there is no
61   ///                             available trampoline when getCompileCallback is
62   ///                             called.
63   JITCompileCallbackManagerBase(TargetAddress ErrorHandlerAddress,
64                                 unsigned NumTrampolinesPerBlock)
65     : ErrorHandlerAddress(ErrorHandlerAddress),
66       NumTrampolinesPerBlock(NumTrampolinesPerBlock) {}
67
68   virtual ~JITCompileCallbackManagerBase() {}
69
70   /// @brief Execute the callback for the given trampoline id. Called by the JIT
71   ///        to compile functions on demand.
72   TargetAddress executeCompileCallback(TargetAddress TrampolineID) {
73     typename TrampolineMapT::iterator I = ActiveTrampolines.find(TrampolineID);
74     // FIXME: Also raise an error in the Orc error-handler when we finally have
75     //        one.
76     if (I == ActiveTrampolines.end())
77       return ErrorHandlerAddress;
78
79     // Found a callback handler. Yank this trampoline out of the active list and
80     // put it back in the available trampolines list, then try to run the
81     // handler's compile and update actions.
82     // Moving the trampoline ID back to the available list first means there's at
83     // least one available trampoline if the compile action triggers a request for
84     // a new one.
85     AvailableTrampolines.push_back(I->first);
86     auto CallbackHandler = std::move(I->second);
87     ActiveTrampolines.erase(I);
88
89     if (auto Addr = CallbackHandler.Compile()) {
90       CallbackHandler.Update(Addr);
91       return Addr;
92     }
93     return ErrorHandlerAddress;
94   }
95
96   /// @brief Get/create a compile callback with the given signature.
97   virtual CompileCallbackInfo getCompileCallback(FunctionType &FT) = 0;
98
99 protected:
100
101   struct CallbackHandler {
102     CompileFtor Compile;
103     UpdateFtor Update;
104   };
105
106   TargetAddress ErrorHandlerAddress;
107   unsigned NumTrampolinesPerBlock;
108
109   typedef std::map<TargetAddress, CallbackHandler> TrampolineMapT;
110   TrampolineMapT ActiveTrampolines;
111   std::vector<TargetAddress> AvailableTrampolines;
112 };
113
114 /// @brief Manage compile callbacks.
115 template <typename JITLayerT, typename TargetT>
116 class JITCompileCallbackManager : public JITCompileCallbackManagerBase {
117 public:
118
119   /// @brief Construct a JITCompileCallbackManager.
120   /// @param JIT JIT layer to emit callback trampolines, etc. into.
121   /// @param Context LLVMContext to use for trampoline & resolve block modules.
122   /// @param ErrorHandlerAddress The address of an error handler in the target
123   ///                            process to be used if a compile callback fails.
124   /// @param NumTrampolinesPerBlock Number of trampolines to allocate whenever
125   ///                               there is no existing callback trampoline.
126   ///                               (Trampolines are allocated in blocks for
127   ///                               efficiency.)
128   JITCompileCallbackManager(JITLayerT &JIT, LLVMContext &Context,
129                             TargetAddress ErrorHandlerAddress,
130                             unsigned NumTrampolinesPerBlock)
131     : JITCompileCallbackManagerBase(ErrorHandlerAddress,
132                                     NumTrampolinesPerBlock),
133       JIT(JIT) {
134     emitResolverBlock(Context);
135   }
136
137   /// @brief Get/create a compile callback with the given signature.
138   CompileCallbackInfo getCompileCallback(FunctionType &FT) final {
139     TargetAddress TrampolineAddr = getAvailableTrampolineAddr(FT.getContext());
140     auto &CallbackHandler =
141       this->ActiveTrampolines[TrampolineAddr];
142     Constant *AddrIntVal =
143       ConstantInt::get(Type::getInt64Ty(FT.getContext()), TrampolineAddr);
144     Constant *AddrPtrVal =
145       ConstantExpr::getCast(Instruction::IntToPtr, AddrIntVal,
146                             PointerType::get(&FT, 0));
147
148     return CompileCallbackInfo(AddrPtrVal, CallbackHandler.Compile,
149                                CallbackHandler.Update);
150   }
151
152 private:
153
154   std::vector<std::unique_ptr<Module>>
155   SingletonSet(std::unique_ptr<Module> M) {
156     std::vector<std::unique_ptr<Module>> Ms;
157     Ms.push_back(std::move(M));
158     return Ms;
159   }
160
161   void emitResolverBlock(LLVMContext &Context) {
162     std::unique_ptr<Module> M(new Module("resolver_block_module",
163                                          Context));
164     TargetT::insertResolverBlock(*M, *this);
165     auto H = JIT.addModuleSet(SingletonSet(std::move(M)), nullptr);
166     JIT.emitAndFinalize(H);
167     auto ResolverBlockSymbol =
168       JIT.findSymbolIn(H, TargetT::ResolverBlockName, false);
169     assert(ResolverBlockSymbol && "Failed to insert resolver block");
170     ResolverBlockAddr = ResolverBlockSymbol.getAddress();
171   }
172
173   TargetAddress getAvailableTrampolineAddr(LLVMContext &Context) {
174     if (this->AvailableTrampolines.empty())
175       grow(Context);
176     assert(!this->AvailableTrampolines.empty() &&
177            "Failed to grow available trampolines.");
178     TargetAddress TrampolineAddr = this->AvailableTrampolines.back();
179     this->AvailableTrampolines.pop_back();
180     return TrampolineAddr;
181   }
182
183   void grow(LLVMContext &Context) {
184     assert(this->AvailableTrampolines.empty() && "Growing prematurely?");
185     std::unique_ptr<Module> M(new Module("trampoline_block", Context));
186     auto GetLabelName =
187       TargetT::insertCompileCallbackTrampolines(*M, ResolverBlockAddr,
188                                                 this->NumTrampolinesPerBlock,
189                                                 this->ActiveTrampolines.size());
190     auto H = JIT.addModuleSet(SingletonSet(std::move(M)), nullptr);
191     JIT.emitAndFinalize(H);
192     for (unsigned I = 0; I < this->NumTrampolinesPerBlock; ++I) {
193       std::string Name = GetLabelName(I);
194       auto TrampolineSymbol = JIT.findSymbolIn(H, Name, false);
195       assert(TrampolineSymbol && "Failed to emit trampoline.");
196       this->AvailableTrampolines.push_back(TrampolineSymbol.getAddress());
197     }
198   }
199
200   JITLayerT &JIT;
201   TargetAddress ResolverBlockAddr;
202 };
203
204 /// @brief Get an update functor for updating the value of a named function
205 ///        pointer.
206 template <typename JITLayerT>
207 JITCompileCallbackManagerBase::UpdateFtor
208 getLocalFPUpdater(JITLayerT &JIT, typename JITLayerT::ModuleSetHandleT H,
209                   std::string Name) {
210     // FIXME: Move-capture Name once we can use C++14.
211     return [=,&JIT](TargetAddress Addr) {
212       auto FPSym = JIT.findSymbolIn(H, Name, true);
213       assert(FPSym && "Cannot find function pointer to update.");
214       void *FPAddr = reinterpret_cast<void*>(
215                        static_cast<uintptr_t>(FPSym.getAddress()));
216       memcpy(FPAddr, &Addr, sizeof(uintptr_t));
217     };
218   }
219
220 GlobalVariable* createImplPointer(Function &F, const Twine &Name,
221                                   Constant *Initializer);
222
223 void makeStub(Function &F, GlobalVariable &ImplPointer);
224
225 typedef std::map<Module*, DenseSet<const GlobalValue*>> ModulePartitionMap;
226
227 void partition(Module &M, const ModulePartitionMap &PMap);
228
229 /// @brief Struct for trivial "complete" partitioning of a module.
230 class FullyPartitionedModule {
231 public:
232   std::unique_ptr<Module> GlobalVars;
233   std::unique_ptr<Module> Commons;
234   std::vector<std::unique_ptr<Module>> Functions;
235
236   FullyPartitionedModule() = default;
237   FullyPartitionedModule(FullyPartitionedModule &&S)
238       : GlobalVars(std::move(S.GlobalVars)), Commons(std::move(S.Commons)),
239         Functions(std::move(S.Functions)) {}
240 };
241
242 FullyPartitionedModule fullyPartition(Module &M);
243
244 } // End namespace orc.
245 } // End namespace llvm.
246
247 #endif // LLVM_EXECUTIONENGINE_ORC_INDIRECTIONUTILS_H