ExecutionEngine/Orc/IRPartitionLayer.h

0001 //===- IRPartitionLayer.h - Partition IR module on lookup -------*- 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 // JIT layer for breaking up modules into smaller submodules that only contains
0010 // looked up symbols.
0011 //
0012 //===----------------------------------------------------------------------===//
0013
0014 #ifndef LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
0015 #define LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
0016
0017 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
0018 #include "llvm/ExecutionEngine/Orc/Layer.h"
0019 #include "llvm/IR/Attributes.h"
0020 #include "llvm/IR/Constant.h"
0021 #include "llvm/IR/Constants.h"
0022 #include "llvm/IR/DataLayout.h"
0023 #include "llvm/IR/Function.h"
0024 #include "llvm/IR/GlobalAlias.h"
0025 #include "llvm/IR/GlobalValue.h"
0026 #include "llvm/IR/GlobalVariable.h"
0027 #include "llvm/IR/Instruction.h"
0028 #include "llvm/IR/Mangler.h"
0029 #include "llvm/IR/Module.h"
0030 #include "llvm/IR/Type.h"
0031
0032 namespace llvm {
0033 namespace orc {
0034
0035 /// A layer that breaks up IR modules into smaller submodules that only contains
0036 /// looked up symbols.
0037 class IRPartitionLayer : public IRLayer {
0038   friend class PartitioningIRMaterializationUnit;
0039
0040 public:
0041   using GlobalValueSet = std::set<const GlobalValue *>;
0042
0043   /// Partitioning function.
0044   using PartitionFunction =
0045       std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>;
0046
0047   /// Construct a IRPartitionLayer.
0048   IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer);
0049
0050   /// Off-the-shelf partitioning which compiles all requested symbols (usually
0051   /// a single function at a time).
0052   static std::optional<GlobalValueSet>
0053   compileRequested(GlobalValueSet Requested);
0054
0055   /// Off-the-shelf partitioning which compiles whole modules whenever any
0056   /// symbol in them is requested.
0057   static std::optional<GlobalValueSet>
0058   compileWholeModule(GlobalValueSet Requested);
0059
0060   /// Sets the partition function.
0061   void setPartitionFunction(PartitionFunction Partition);
0062
0063   /// Emits the given module. This should not be called by clients: it will be
0064   /// called by the JIT when a definition added via the add method is requested.
0065   void emit(std::unique_ptr<MaterializationResponsibility> R,
0066             ThreadSafeModule TSM) override;
0067
0068 private:
0069   void cleanUpModule(Module &M);
0070
0071   void expandPartition(GlobalValueSet &Partition);
0072
0073   void emitPartition(std::unique_ptr<MaterializationResponsibility> R,
0074                      ThreadSafeModule TSM,
0075                      IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
0076
0077   IRLayer &BaseLayer;
0078   PartitionFunction Partition = compileRequested;
0079   SymbolLinkagePromoter PromoteSymbols;
0080 };
0081
0082 } // namespace orc
0083 } // namespace llvm
0084
0085 #endif