1//===- Transform/Utils/CodeExtractor.h - Code extraction util ---*- 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// A utility to support extracting code from one function into its own
10// stand-alone function.
11//
12//===----------------------------------------------------------------------===//
13
14#ifndef LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
15#define LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
16
17#include "llvm/ADT/ArrayRef.h"
18#include "llvm/ADT/DenseMap.h"
19#include "llvm/ADT/SetVector.h"
20#include <limits>
21
22namespace llvm {
23
24template <typename PtrType> class SmallPtrSetImpl;
25class AllocaInst;
26class BasicBlock;
27class BlockFrequency;
28class BlockFrequencyInfo;
29class BranchProbabilityInfo;
30class AssumptionCache;
31class CallInst;
32class DominatorTree;
33class Function;
34class Instruction;
35class Loop;
36class Module;
37class Type;
38class Value;
39
40/// A cache for the CodeExtractor analysis. The operation \ref
41/// CodeExtractor::extractCodeRegion is guaranteed not to invalidate this
42/// object. This object should conservatively be considered invalid if any
43/// other mutating operations on the IR occur.
44///
45/// Constructing this object is O(n) in the size of the function.
46class CodeExtractorAnalysisCache {
47  /// The allocas in the function.
48  SmallVector<AllocaInst *, 16> Allocas;
49
50  /// Base memory addresses of load/store instructions, grouped by block.
51  DenseMap<BasicBlock *, DenseSet<Value *>> BaseMemAddrs;
52
53  /// Blocks which contain instructions which may have unknown side-effects
54  /// on memory.
55  DenseSet<BasicBlock *> SideEffectingBlocks;
56
57  void findSideEffectInfoForBlock(BasicBlock &BB);
58
59public:
60  CodeExtractorAnalysisCache(Function &F);
61
62  /// Get the allocas in the function at the time the analysis was created.
63  /// Note that some of these allocas may no longer be present in the function,
64  /// due to \ref CodeExtractor::extractCodeRegion.
65  ArrayRef<AllocaInst *> getAllocas() const { return Allocas; }
66
67  /// Check whether \p BB contains an instruction thought to load from, store
68  /// to, or otherwise clobber the alloca \p Addr.
69  bool doesBlockContainClobberOfAddr(BasicBlock &BB, AllocaInst *Addr) const;
70};
71
72  /// Utility class for extracting code into a new function.
73  ///
74  /// This utility provides a simple interface for extracting some sequence of
75  /// code into its own function, replacing it with a call to that function. It
76  /// also provides various methods to query about the nature and result of
77  /// such a transformation.
78  ///
79  /// The rough algorithm used is:
80  /// 1) Find both the inputs and outputs for the extracted region.
81  /// 2) Pass the inputs as arguments, remapping them within the extracted
82  ///    function to arguments.
83  /// 3) Add allocas for any scalar outputs, adding all of the outputs' allocas
84  ///    as arguments, and inserting stores to the arguments for any scalars.
85  class CodeExtractor {
86    using ValueSet = SetVector<Value *>;
87
88    // Various bits of state computed on construction.
89    DominatorTree *const DT;
90    const bool AggregateArgs;
91    BlockFrequencyInfo *BFI;
92    BranchProbabilityInfo *BPI;
93    AssumptionCache *AC;
94
95    // A block outside of the extraction set where any intermediate
96    // allocations will be placed inside. If this is null, allocations
97    // will be placed in the entry block of the function.
98    BasicBlock *AllocationBlock;
99
100    // If true, varargs functions can be extracted.
101    bool AllowVarArgs;
102
103    // Bits of intermediate state computed at various phases of extraction.
104    SetVector<BasicBlock *> Blocks;
105    unsigned NumExitBlocks = std::numeric_limits<unsigned>::max();
106    Type *RetTy;
107
108    // Mapping from the original exit blocks, to the new blocks inside
109    // the function.
110    SmallVector<BasicBlock *, 4> OldTargets;
111
112    // Suffix to use when creating extracted function (appended to the original
113    // function name + "."). If empty, the default is to use the entry block
114    // label, if non-empty, otherwise "extracted".
115    std::string Suffix;
116
117    // If true, the outlined function has aggregate argument in zero address
118    // space.
119    bool ArgsInZeroAddressSpace;
120
121  public:
122    /// Create a code extractor for a sequence of blocks.
123    ///
124    /// Given a sequence of basic blocks where the first block in the sequence
125    /// dominates the rest, prepare a code extractor object for pulling this
126    /// sequence out into its new function. When a DominatorTree is also given,
127    /// extra checking and transformations are enabled. If AllowVarArgs is true,
128    /// vararg functions can be extracted. This is safe, if all vararg handling
129    /// code is extracted, including vastart. If AllowAlloca is true, then
130    /// extraction of blocks containing alloca instructions would be possible,
131    /// however code extractor won't validate whether extraction is legal.
132    /// Any new allocations will be placed in the AllocationBlock, unless
133    /// it is null, in which case it will be placed in the entry block of
134    /// the function from which the code is being extracted.
135    /// If ArgsInZeroAddressSpace param is set to true, then the aggregate
136    /// param pointer of the outlined function is declared in zero address
137    /// space.
138    CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
139                  bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
140                  BranchProbabilityInfo *BPI = nullptr,
141                  AssumptionCache *AC = nullptr, bool AllowVarArgs = false,
142                  bool AllowAlloca = false,
143                  BasicBlock *AllocationBlock = nullptr,
144                  std::string Suffix = "", bool ArgsInZeroAddressSpace = false);
145
146    /// Create a code extractor for a loop body.
147    ///
148    /// Behaves just like the generic code sequence constructor, but uses the
149    /// block sequence of the loop.
150    CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false,
151                  BlockFrequencyInfo *BFI = nullptr,
152                  BranchProbabilityInfo *BPI = nullptr,
153                  AssumptionCache *AC = nullptr,
154                  std::string Suffix = "");
155
156    /// Perform the extraction, returning the new function.
157    ///
158    /// Returns zero when called on a CodeExtractor instance where isEligible
159    /// returns false.
160    Function *extractCodeRegion(const CodeExtractorAnalysisCache &CEAC);
161
162    /// Perform the extraction, returning the new function and providing an
163    /// interface to see what was categorized as inputs and outputs.
164    ///
165    /// \param CEAC - Cache to speed up operations for the CodeExtractor when
166    /// hoisting, and extracting lifetime values and assumes.
167    /// \param Inputs [out] - filled with  values marked as inputs to the
168    /// newly outlined function.
169     /// \param Outputs [out] - filled with values marked as outputs to the
170    /// newly outlined function.
171    /// \returns zero when called on a CodeExtractor instance where isEligible
172    /// returns false.
173    Function *extractCodeRegion(const CodeExtractorAnalysisCache &CEAC,
174                                ValueSet &Inputs, ValueSet &Outputs);
175
176    /// Verify that assumption cache isn't stale after a region is extracted.
177    /// Returns true when verifier finds errors. AssumptionCache is passed as
178    /// parameter to make this function stateless.
179    static bool verifyAssumptionCache(const Function &OldFunc,
180                                      const Function &NewFunc,
181                                      AssumptionCache *AC);
182
183    /// Test whether this code extractor is eligible.
184    ///
185    /// Based on the blocks used when constructing the code extractor,
186    /// determine whether it is eligible for extraction.
187    ///
188    /// Checks that varargs handling (with vastart and vaend) is only done in
189    /// the outlined blocks.
190    bool isEligible() const;
191
192    /// Compute the set of input values and output values for the code.
193    ///
194    /// These can be used either when performing the extraction or to evaluate
195    /// the expected size of a call to the extracted function. Note that this
196    /// work cannot be cached between the two as once we decide to extract
197    /// a code sequence, that sequence is modified, including changing these
198    /// sets, before extraction occurs. These modifications won't have any
199    /// significant impact on the cost however.
200    void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
201                           const ValueSet &Allocas) const;
202
203    /// Check if life time marker nodes can be hoisted/sunk into the outline
204    /// region.
205    ///
206    /// Returns true if it is safe to do the code motion.
207    bool
208    isLegalToShrinkwrapLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
209                                       Instruction *AllocaAddr) const;
210
211    /// Find the set of allocas whose life ranges are contained within the
212    /// outlined region.
213    ///
214    /// Allocas which have life_time markers contained in the outlined region
215    /// should be pushed to the outlined function. The address bitcasts that
216    /// are used by the lifetime markers are also candidates for shrink-
217    /// wrapping. The instructions that need to be sunk are collected in
218    /// 'Allocas'.
219    void findAllocas(const CodeExtractorAnalysisCache &CEAC,
220                     ValueSet &SinkCands, ValueSet &HoistCands,
221                     BasicBlock *&ExitBlock) const;
222
223    /// Find or create a block within the outline region for placing hoisted
224    /// code.
225    ///
226    /// CommonExitBlock is block outside the outline region. It is the common
227    /// successor of blocks inside the region. If there exists a single block
228    /// inside the region that is the predecessor of CommonExitBlock, that block
229    /// will be returned. Otherwise CommonExitBlock will be split and the
230    /// original block will be added to the outline region.
231    BasicBlock *findOrCreateBlockForHoisting(BasicBlock *CommonExitBlock);
232
233    /// Exclude a value from aggregate argument passing when extracting a code
234    /// region, passing it instead as a scalar.
235    void excludeArgFromAggregate(Value *Arg);
236
237  private:
238    struct LifetimeMarkerInfo {
239      bool SinkLifeStart = false;
240      bool HoistLifeEnd = false;
241      Instruction *LifeStart = nullptr;
242      Instruction *LifeEnd = nullptr;
243    };
244
245    ValueSet ExcludeArgsFromAggregate;
246
247    LifetimeMarkerInfo
248    getLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
249                       Instruction *Addr, BasicBlock *ExitBlock) const;
250
251    void severSplitPHINodesOfEntry(BasicBlock *&Header);
252    void severSplitPHINodesOfExits(const SmallPtrSetImpl<BasicBlock *> &Exits);
253    void splitReturnBlocks();
254
255    Function *constructFunction(const ValueSet &inputs,
256                                const ValueSet &outputs,
257                                BasicBlock *header,
258                                BasicBlock *newRootNode, BasicBlock *newHeader,
259                                Function *oldFunction, Module *M);
260
261    void moveCodeToFunction(Function *newFunction);
262
263    void calculateNewCallTerminatorWeights(
264        BasicBlock *CodeReplacer,
265        DenseMap<BasicBlock *, BlockFrequency> &ExitWeights,
266        BranchProbabilityInfo *BPI);
267
268    CallInst *emitCallAndSwitchStatement(Function *newFunction,
269                                         BasicBlock *newHeader,
270                                         ValueSet &inputs, ValueSet &outputs);
271  };
272
273} // end namespace llvm
274
275#endif // LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
276