1//===- CompileOnDemandLayer.h - Compile each function on demand -*- 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// JIT layer for breaking up modules and inserting callbacks to allow
11// individual functions to be compiled on demand.
12//
13//===----------------------------------------------------------------------===//
14
15#ifndef LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
16#define LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
17
18#include "llvm/ADT/APInt.h"
19#include "llvm/ADT/STLExtras.h"
20#include "llvm/ADT/StringRef.h"
21#include "llvm/ADT/Twine.h"
22#include "llvm/ExecutionEngine/JITSymbol.h"
23#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
24#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
25#include "llvm/ExecutionEngine/RuntimeDyld.h"
26#include "llvm/IR/Attributes.h"
27#include "llvm/IR/Constants.h"
28#include "llvm/IR/DataLayout.h"
29#include "llvm/IR/Function.h"
30#include "llvm/IR/GlobalAlias.h"
31#include "llvm/IR/GlobalValue.h"
32#include "llvm/IR/GlobalVariable.h"
33#include "llvm/IR/Mangler.h"
34#include "llvm/IR/Module.h"
35#include "llvm/Support/Casting.h"
36#include "llvm/Support/raw_ostream.h"
37#include <algorithm>
38#include <cassert>
39#include <functional>
40#include <iterator>
41#include <list>
42#include <memory>
43#include <set>
44#include <string>
45#include <utility>
46#include <vector>
47
48namespace llvm {
49namespace orc {
50
51/// @brief Compile-on-demand layer.
52///
53///   When a module is added to this layer a stub is created for each of its
54/// function definitions. The stubs and other global values are immediately
55/// added to the layer below. When a stub is called it triggers the extraction
56/// of the function body from the original module. The extracted body is then
57/// compiled and executed.
58template <typename BaseLayerT,
59          typename CompileCallbackMgrT = JITCompileCallbackManager,
60          typename IndirectStubsMgrT = IndirectStubsManager>
61class CompileOnDemandLayer {
62private:
63  template <typename MaterializerFtor>
64  class LambdaMaterializer final : public ValueMaterializer {
65  public:
66    LambdaMaterializer(MaterializerFtor M) : M(std::move(M)) {}
67
68    Value *materialize(Value *V) final { return M(V); }
69
70  private:
71    MaterializerFtor M;
72  };
73
74  template <typename MaterializerFtor>
75  LambdaMaterializer<MaterializerFtor>
76  createLambdaMaterializer(MaterializerFtor M) {
77    return LambdaMaterializer<MaterializerFtor>(std::move(M));
78  }
79
80  typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
81
82  // Provide type-erasure for the Modules and MemoryManagers.
83  template <typename ResourceT>
84  class ResourceOwner {
85  public:
86    ResourceOwner() = default;
87    ResourceOwner(const ResourceOwner&) = delete;
88    ResourceOwner& operator=(const ResourceOwner&) = delete;
89    virtual ~ResourceOwner() = default;
90
91    virtual ResourceT& getResource() const = 0;
92  };
93
94  template <typename ResourceT, typename ResourcePtrT>
95  class ResourceOwnerImpl : public ResourceOwner<ResourceT> {
96  public:
97    ResourceOwnerImpl(ResourcePtrT ResourcePtr)
98      : ResourcePtr(std::move(ResourcePtr)) {}
99
100    ResourceT& getResource() const override { return *ResourcePtr; }
101
102  private:
103    ResourcePtrT ResourcePtr;
104  };
105
106  template <typename ResourceT, typename ResourcePtrT>
107  std::unique_ptr<ResourceOwner<ResourceT>>
108  wrapOwnership(ResourcePtrT ResourcePtr) {
109    typedef ResourceOwnerImpl<ResourceT, ResourcePtrT> RO;
110    return llvm::make_unique<RO>(std::move(ResourcePtr));
111  }
112
113  class StaticGlobalRenamer {
114  public:
115    StaticGlobalRenamer() = default;
116    StaticGlobalRenamer(StaticGlobalRenamer &&) = default;
117    StaticGlobalRenamer &operator=(StaticGlobalRenamer &&) = default;
118
119    void rename(Module &M) {
120      for (auto &F : M)
121        if (F.hasLocalLinkage())
122          F.setName("$static." + Twine(NextId++));
123      for (auto &G : M.globals())
124        if (G.hasLocalLinkage())
125          G.setName("$static." + Twine(NextId++));
126    }
127
128  private:
129    unsigned NextId = 0;
130  };
131
132  struct LogicalDylib {
133    typedef std::function<JITSymbol(const std::string&)> SymbolResolverFtor;
134
135    typedef std::function<typename BaseLayerT::ModuleSetHandleT(
136                            BaseLayerT&,
137                            std::unique_ptr<Module>,
138                            std::unique_ptr<JITSymbolResolver>)>
139      ModuleAdderFtor;
140
141    struct SourceModuleEntry {
142      std::unique_ptr<ResourceOwner<Module>> SourceMod;
143      std::set<Function*> StubsToClone;
144    };
145
146    typedef std::vector<SourceModuleEntry> SourceModulesList;
147    typedef typename SourceModulesList::size_type SourceModuleHandle;
148
149    SourceModuleHandle
150    addSourceModule(std::unique_ptr<ResourceOwner<Module>> M) {
151      SourceModuleHandle H = SourceModules.size();
152      SourceModules.push_back(SourceModuleEntry());
153      SourceModules.back().SourceMod = std::move(M);
154      return H;
155    }
156
157    Module& getSourceModule(SourceModuleHandle H) {
158      return SourceModules[H].SourceMod->getResource();
159    }
160
161    std::set<Function*>& getStubsToClone(SourceModuleHandle H) {
162      return SourceModules[H].StubsToClone;
163    }
164
165    JITSymbol findSymbol(BaseLayerT &BaseLayer, const std::string &Name,
166                         bool ExportedSymbolsOnly) {
167      if (auto Sym = StubsMgr->findStub(Name, ExportedSymbolsOnly))
168        return Sym;
169      for (auto BLH : BaseLayerHandles)
170        if (auto Sym = BaseLayer.findSymbolIn(BLH, Name, ExportedSymbolsOnly))
171          return Sym;
172      return nullptr;
173    }
174
175    void removeModulesFromBaseLayer(BaseLayerT &BaseLayer) {
176      for (auto &BLH : BaseLayerHandles)
177        BaseLayer.removeModuleSet(BLH);
178    }
179
180    std::unique_ptr<JITSymbolResolver> ExternalSymbolResolver;
181    std::unique_ptr<ResourceOwner<RuntimeDyld::MemoryManager>> MemMgr;
182    std::unique_ptr<IndirectStubsMgrT> StubsMgr;
183    StaticGlobalRenamer StaticRenamer;
184    ModuleAdderFtor ModuleAdder;
185    SourceModulesList SourceModules;
186    std::vector<BaseLayerModuleSetHandleT> BaseLayerHandles;
187  };
188
189  typedef std::list<LogicalDylib> LogicalDylibList;
190
191public:
192  /// @brief Handle to a set of loaded modules.
193  typedef typename LogicalDylibList::iterator ModuleSetHandleT;
194
195  /// @brief Module partitioning functor.
196  typedef std::function<std::set<Function*>(Function&)> PartitioningFtor;
197
198  /// @brief Builder for IndirectStubsManagers.
199  typedef std::function<std::unique_ptr<IndirectStubsMgrT>()>
200    IndirectStubsManagerBuilderT;
201
202  /// @brief Construct a compile-on-demand layer instance.
203  CompileOnDemandLayer(BaseLayerT &BaseLayer, PartitioningFtor Partition,
204                       CompileCallbackMgrT &CallbackMgr,
205                       IndirectStubsManagerBuilderT CreateIndirectStubsManager,
206                       bool CloneStubsIntoPartitions = true)
207      : BaseLayer(BaseLayer), Partition(std::move(Partition)),
208        CompileCallbackMgr(CallbackMgr),
209        CreateIndirectStubsManager(std::move(CreateIndirectStubsManager)),
210        CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
211
212  ~CompileOnDemandLayer() {
213    while (!LogicalDylibs.empty())
214      removeModuleSet(LogicalDylibs.begin());
215  }
216
217  /// @brief Add a module to the compile-on-demand layer.
218  template <typename ModuleSetT, typename MemoryManagerPtrT,
219            typename SymbolResolverPtrT>
220  ModuleSetHandleT addModuleSet(ModuleSetT Ms,
221                                MemoryManagerPtrT MemMgr,
222                                SymbolResolverPtrT Resolver) {
223
224    LogicalDylibs.push_back(LogicalDylib());
225    auto &LD = LogicalDylibs.back();
226    LD.ExternalSymbolResolver = std::move(Resolver);
227    LD.StubsMgr = CreateIndirectStubsManager();
228
229    auto &MemMgrRef = *MemMgr;
230    LD.MemMgr = wrapOwnership<RuntimeDyld::MemoryManager>(std::move(MemMgr));
231
232    LD.ModuleAdder =
233      [&MemMgrRef](BaseLayerT &B, std::unique_ptr<Module> M,
234                   std::unique_ptr<JITSymbolResolver> R) {
235        std::vector<std::unique_ptr<Module>> Ms;
236        Ms.push_back(std::move(M));
237        return B.addModuleSet(std::move(Ms), &MemMgrRef, std::move(R));
238      };
239
240    // Process each of the modules in this module set.
241    for (auto &M : Ms)
242      addLogicalModule(LogicalDylibs.back(), std::move(M));
243
244    return std::prev(LogicalDylibs.end());
245  }
246
247  /// @brief Remove the module represented by the given handle.
248  ///
249  ///   This will remove all modules in the layers below that were derived from
250  /// the module represented by H.
251  void removeModuleSet(ModuleSetHandleT H) {
252    H->removeModulesFromBaseLayer(BaseLayer);
253    LogicalDylibs.erase(H);
254  }
255
256  /// @brief Search for the given named symbol.
257  /// @param Name The name of the symbol to search for.
258  /// @param ExportedSymbolsOnly If true, search only for exported symbols.
259  /// @return A handle for the given named symbol, if it exists.
260  JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
261    for (auto LDI = LogicalDylibs.begin(), LDE = LogicalDylibs.end();
262         LDI != LDE; ++LDI) {
263      if (auto Sym = LDI->StubsMgr->findStub(Name, ExportedSymbolsOnly))
264        return Sym;
265      if (auto Sym = findSymbolIn(LDI, Name, ExportedSymbolsOnly))
266        return Sym;
267    }
268    return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
269  }
270
271  /// @brief Get the address of a symbol provided by this layer, or some layer
272  ///        below this one.
273  JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
274                         bool ExportedSymbolsOnly) {
275    return H->findSymbol(BaseLayer, Name, ExportedSymbolsOnly);
276  }
277
278  /// @brief Update the stub for the given function to point at FnBodyAddr.
279  /// This can be used to support re-optimization.
280  /// @return true if the function exists and the stub is updated, false
281  ///         otherwise.
282  //
283  // FIXME: We should track and free associated resources (unused compile
284  //        callbacks, uncompiled IR, and no-longer-needed/reachable function
285  //        implementations).
286  // FIXME: Return Error once the JIT APIs are Errorized.
287  bool updatePointer(std::string FuncName, JITTargetAddress FnBodyAddr) {
288    //Find out which logical dylib contains our symbol
289    auto LDI = LogicalDylibs.begin();
290    for (auto LDE = LogicalDylibs.end(); LDI != LDE; ++LDI) {
291      if (auto LMResources = LDI->getLogicalModuleResourcesForSymbol(FuncName, false)) {
292        Module &SrcM = LMResources->SourceModule->getResource();
293        std::string CalledFnName = mangle(FuncName, SrcM.getDataLayout());
294        if (auto EC = LMResources->StubsMgr->updatePointer(CalledFnName, FnBodyAddr))
295          return false;
296        else
297          return true;
298      }
299    }
300    return false;
301  }
302
303private:
304  template <typename ModulePtrT>
305  void addLogicalModule(LogicalDylib &LD, ModulePtrT SrcMPtr) {
306
307    // Rename all static functions / globals to $static.X :
308    // This will unique the names across all modules in the logical dylib,
309    // simplifying symbol lookup.
310    LD.StaticRenamer.rename(*SrcMPtr);
311
312    // Bump the linkage and rename any anonymous/privote members in SrcM to
313    // ensure that everything will resolve properly after we partition SrcM.
314    makeAllSymbolsExternallyAccessible(*SrcMPtr);
315
316    // Create a logical module handle for SrcM within the logical dylib.
317    Module &SrcM = *SrcMPtr;
318    auto LMId = LD.addSourceModule(wrapOwnership<Module>(std::move(SrcMPtr)));
319
320    // Create stub functions.
321    const DataLayout &DL = SrcM.getDataLayout();
322    {
323      typename IndirectStubsMgrT::StubInitsMap StubInits;
324      for (auto &F : SrcM) {
325        // Skip declarations.
326        if (F.isDeclaration())
327          continue;
328
329        // Skip weak functions for which we already have definitions.
330        auto MangledName = mangle(F.getName(), DL);
331        if (F.hasWeakLinkage() || F.hasLinkOnceLinkage())
332          if (auto Sym = LD.findSymbol(BaseLayer, MangledName, false))
333            continue;
334
335        // Record all functions defined by this module.
336        if (CloneStubsIntoPartitions)
337          LD.getStubsToClone(LMId).insert(&F);
338
339        // Create a callback, associate it with the stub for the function,
340        // and set the compile action to compile the partition containing the
341        // function.
342        auto CCInfo = CompileCallbackMgr.getCompileCallback();
343        StubInits[MangledName] =
344          std::make_pair(CCInfo.getAddress(),
345                         JITSymbolFlags::fromGlobalValue(F));
346        CCInfo.setCompileAction([this, &LD, LMId, &F]() {
347          return this->extractAndCompile(LD, LMId, F);
348        });
349      }
350
351      auto EC = LD.StubsMgr->createStubs(StubInits);
352      (void)EC;
353      // FIXME: This should be propagated back to the user. Stub creation may
354      //        fail for remote JITs.
355      assert(!EC && "Error generating stubs");
356    }
357
358    // If this module doesn't contain any globals, aliases, or module flags then
359    // we can bail out early and avoid the overhead of creating and managing an
360    // empty globals module.
361    if (SrcM.global_empty() && SrcM.alias_empty() &&
362        !SrcM.getModuleFlagsMetadata())
363      return;
364
365    // Create the GlobalValues module.
366    auto GVsM = llvm::make_unique<Module>((SrcM.getName() + ".globals").str(),
367                                          SrcM.getContext());
368    GVsM->setDataLayout(DL);
369
370    ValueToValueMapTy VMap;
371
372    // Clone global variable decls.
373    for (auto &GV : SrcM.globals())
374      if (!GV.isDeclaration() && !VMap.count(&GV))
375        cloneGlobalVariableDecl(*GVsM, GV, &VMap);
376
377    // And the aliases.
378    for (auto &A : SrcM.aliases())
379      if (!VMap.count(&A))
380        cloneGlobalAliasDecl(*GVsM, A, VMap);
381
382    // Clone the module flags.
383    cloneModuleFlagsMetadata(*GVsM, SrcM, VMap);
384
385    // Now we need to clone the GV and alias initializers.
386
387    // Initializers may refer to functions declared (but not defined) in this
388    // module. Build a materializer to clone decls on demand.
389    auto Materializer = createLambdaMaterializer(
390      [&LD, &GVsM](Value *V) -> Value* {
391        if (auto *F = dyn_cast<Function>(V)) {
392          // Decls in the original module just get cloned.
393          if (F->isDeclaration())
394            return cloneFunctionDecl(*GVsM, *F);
395
396          // Definitions in the original module (which we have emitted stubs
397          // for at this point) get turned into a constant alias to the stub
398          // instead.
399          const DataLayout &DL = GVsM->getDataLayout();
400          std::string FName = mangle(F->getName(), DL);
401          auto StubSym = LD.StubsMgr->findStub(FName, false);
402          unsigned PtrBitWidth = DL.getPointerTypeSizeInBits(F->getType());
403          ConstantInt *StubAddr =
404            ConstantInt::get(GVsM->getContext(),
405                             APInt(PtrBitWidth, StubSym.getAddress()));
406          Constant *Init = ConstantExpr::getCast(Instruction::IntToPtr,
407                                                 StubAddr, F->getType());
408          return GlobalAlias::create(F->getFunctionType(),
409                                     F->getType()->getAddressSpace(),
410                                     F->getLinkage(), F->getName(),
411                                     Init, GVsM.get());
412        }
413        // else....
414        return nullptr;
415      });
416
417    // Clone the global variable initializers.
418    for (auto &GV : SrcM.globals())
419      if (!GV.isDeclaration())
420        moveGlobalVariableInitializer(GV, VMap, &Materializer);
421
422    // Clone the global alias initializers.
423    for (auto &A : SrcM.aliases()) {
424      auto *NewA = cast<GlobalAlias>(VMap[&A]);
425      assert(NewA && "Alias not cloned?");
426      Value *Init = MapValue(A.getAliasee(), VMap, RF_None, nullptr,
427                             &Materializer);
428      NewA->setAliasee(cast<Constant>(Init));
429    }
430
431    // Build a resolver for the globals module and add it to the base layer.
432    auto GVsResolver = createLambdaResolver(
433        [this, &LD](const std::string &Name) {
434          if (auto Sym = LD.StubsMgr->findStub(Name, false))
435            return Sym;
436          if (auto Sym = LD.findSymbol(BaseLayer, Name, false))
437            return Sym;
438          return LD.ExternalSymbolResolver->findSymbolInLogicalDylib(Name);
439        },
440        [&LD](const std::string &Name) {
441          return LD.ExternalSymbolResolver->findSymbol(Name);
442        });
443
444    auto GVsH = LD.ModuleAdder(BaseLayer, std::move(GVsM),
445                               std::move(GVsResolver));
446    LD.BaseLayerHandles.push_back(GVsH);
447  }
448
449  static std::string mangle(StringRef Name, const DataLayout &DL) {
450    std::string MangledName;
451    {
452      raw_string_ostream MangledNameStream(MangledName);
453      Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
454    }
455    return MangledName;
456  }
457
458  JITTargetAddress
459  extractAndCompile(LogicalDylib &LD,
460                    typename LogicalDylib::SourceModuleHandle LMId,
461                    Function &F) {
462    Module &SrcM = LD.getSourceModule(LMId);
463
464    // If F is a declaration we must already have compiled it.
465    if (F.isDeclaration())
466      return 0;
467
468    // Grab the name of the function being called here.
469    std::string CalledFnName = mangle(F.getName(), SrcM.getDataLayout());
470
471    auto Part = Partition(F);
472    auto PartH = emitPartition(LD, LMId, Part);
473
474    JITTargetAddress CalledAddr = 0;
475    for (auto *SubF : Part) {
476      std::string FnName = mangle(SubF->getName(), SrcM.getDataLayout());
477      auto FnBodySym = BaseLayer.findSymbolIn(PartH, FnName, false);
478      assert(FnBodySym && "Couldn't find function body.");
479
480      JITTargetAddress FnBodyAddr = FnBodySym.getAddress();
481
482      // If this is the function we're calling record the address so we can
483      // return it from this function.
484      if (SubF == &F)
485        CalledAddr = FnBodyAddr;
486
487      // Update the function body pointer for the stub.
488      if (auto EC = LD.StubsMgr->updatePointer(FnName, FnBodyAddr))
489        return 0;
490    }
491
492    LD.BaseLayerHandles.push_back(PartH);
493
494    return CalledAddr;
495  }
496
497  template <typename PartitionT>
498  BaseLayerModuleSetHandleT
499  emitPartition(LogicalDylib &LD,
500                typename LogicalDylib::SourceModuleHandle LMId,
501                const PartitionT &Part) {
502    Module &SrcM = LD.getSourceModule(LMId);
503
504    // Create the module.
505    std::string NewName = SrcM.getName();
506    for (auto *F : Part) {
507      NewName += ".";
508      NewName += F->getName();
509    }
510
511    auto M = llvm::make_unique<Module>(NewName, SrcM.getContext());
512    M->setDataLayout(SrcM.getDataLayout());
513    ValueToValueMapTy VMap;
514
515    auto Materializer = createLambdaMaterializer([&LD, &LMId,
516                                                  &M](Value *V) -> Value * {
517      if (auto *GV = dyn_cast<GlobalVariable>(V))
518        return cloneGlobalVariableDecl(*M, *GV);
519
520      if (auto *F = dyn_cast<Function>(V)) {
521        // Check whether we want to clone an available_externally definition.
522        if (!LD.getStubsToClone(LMId).count(F))
523          return cloneFunctionDecl(*M, *F);
524
525        // Ok - we want an inlinable stub. For that to work we need a decl
526        // for the stub pointer.
527        auto *StubPtr = createImplPointer(*F->getType(), *M,
528                                          F->getName() + "$stub_ptr", nullptr);
529        auto *ClonedF = cloneFunctionDecl(*M, *F);
530        makeStub(*ClonedF, *StubPtr);
531        ClonedF->setLinkage(GlobalValue::AvailableExternallyLinkage);
532        ClonedF->addFnAttr(Attribute::AlwaysInline);
533        return ClonedF;
534      }
535
536      if (auto *A = dyn_cast<GlobalAlias>(V)) {
537        auto *Ty = A->getValueType();
538        if (Ty->isFunctionTy())
539          return Function::Create(cast<FunctionType>(Ty),
540                                  GlobalValue::ExternalLinkage, A->getName(),
541                                  M.get());
542
543        return new GlobalVariable(*M, Ty, false, GlobalValue::ExternalLinkage,
544                                  nullptr, A->getName(), nullptr,
545                                  GlobalValue::NotThreadLocal,
546                                  A->getType()->getAddressSpace());
547      }
548
549      return nullptr;
550    });
551
552    // Create decls in the new module.
553    for (auto *F : Part)
554      cloneFunctionDecl(*M, *F, &VMap);
555
556    // Move the function bodies.
557    for (auto *F : Part)
558      moveFunctionBody(*F, VMap, &Materializer);
559
560    // Create memory manager and symbol resolver.
561    auto Resolver = createLambdaResolver(
562        [this, &LD](const std::string &Name) {
563          if (auto Sym = LD.findSymbol(BaseLayer, Name, false))
564            return Sym;
565          return LD.ExternalSymbolResolver->findSymbolInLogicalDylib(Name);
566        },
567        [&LD](const std::string &Name) {
568          return LD.ExternalSymbolResolver->findSymbol(Name);
569        });
570
571    return LD.ModuleAdder(BaseLayer, std::move(M), std::move(Resolver));
572  }
573
574  BaseLayerT &BaseLayer;
575  PartitioningFtor Partition;
576  CompileCallbackMgrT &CompileCallbackMgr;
577  IndirectStubsManagerBuilderT CreateIndirectStubsManager;
578
579  LogicalDylibList LogicalDylibs;
580  bool CloneStubsIntoPartitions;
581};
582
583} // end namespace orc
584} // end namespace llvm
585
586#endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
587