xref: /llvm-project/llvm/include/llvm/Transforms/Utils/CodeExtractor.h (revision f6795e6b4f619cbecc59a92f7e5fad7ca90ece54)
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 
22 namespace llvm {
23 
24 template <typename PtrType> class SmallPtrSetImpl;
25 class AllocaInst;
26 class BasicBlock;
27 class BlockFrequency;
28 class BlockFrequencyInfo;
29 class BranchProbabilityInfo;
30 class AssumptionCache;
31 class CallInst;
32 class DominatorTree;
33 class Function;
34 class Instruction;
35 class Module;
36 class Type;
37 class Value;
38 class StructType;
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.
46 class 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 
59 public:
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 
106     /// Lists of blocks that are branched from the code region to be extracted,
107     /// also called the exit blocks. Each block is contained at most once. Its
108     /// order defines the return value of the extracted function.
109     ///
110     /// When there is just one (or no) exit block, the return value is
111     /// irrelevant.
112     ///
113     /// When there are exactly two exit blocks, the extracted function returns a
114     /// boolean. For ExtractedFuncRetVals[0], it returns 'true'. For
115     /// ExtractedFuncRetVals[1] it returns 'false'.
116     /// NOTE: Since a boolean is represented by i1, ExtractedFuncRetVals[0]
117     ///       returns 1 and ExtractedFuncRetVals[1] returns 0, which opposite
118     ///       of the regular pattern below.
119     ///
120     /// When there are 3 or more exit blocks, leaving the extracted function via
121     /// the first block it returns 0. When leaving via the second entry it
122     /// returns 1, etc.
123     SmallVector<BasicBlock *> ExtractedFuncRetVals;
124 
125     // Suffix to use when creating extracted function (appended to the original
126     // function name + "."). If empty, the default is to use the entry block
127     // label, if non-empty, otherwise "extracted".
128     std::string Suffix;
129 
130     // If true, the outlined function has aggregate argument in zero address
131     // space.
132     bool ArgsInZeroAddressSpace;
133 
134   public:
135     /// Create a code extractor for a sequence of blocks.
136     ///
137     /// Given a sequence of basic blocks where the first block in the sequence
138     /// dominates the rest, prepare a code extractor object for pulling this
139     /// sequence out into its new function. When a DominatorTree is also given,
140     /// extra checking and transformations are enabled. If AllowVarArgs is true,
141     /// vararg functions can be extracted. This is safe, if all vararg handling
142     /// code is extracted, including vastart. If AllowAlloca is true, then
143     /// extraction of blocks containing alloca instructions would be possible,
144     /// however code extractor won't validate whether extraction is legal.
145     /// Any new allocations will be placed in the AllocationBlock, unless
146     /// it is null, in which case it will be placed in the entry block of
147     /// the function from which the code is being extracted.
148     /// If ArgsInZeroAddressSpace param is set to true, then the aggregate
149     /// param pointer of the outlined function is declared in zero address
150     /// space.
151     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
152                   bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
153                   BranchProbabilityInfo *BPI = nullptr,
154                   AssumptionCache *AC = nullptr, bool AllowVarArgs = false,
155                   bool AllowAlloca = false,
156                   BasicBlock *AllocationBlock = nullptr,
157                   std::string Suffix = "", bool ArgsInZeroAddressSpace = false);
158 
159     /// Perform the extraction, returning the new function.
160     ///
161     /// Returns zero when called on a CodeExtractor instance where isEligible
162     /// returns false.
163     Function *extractCodeRegion(const CodeExtractorAnalysisCache &CEAC);
164 
165     /// Perform the extraction, returning the new function and providing an
166     /// interface to see what was categorized as inputs and outputs.
167     ///
168     /// \param CEAC - Cache to speed up operations for the CodeExtractor when
169     /// hoisting, and extracting lifetime values and assumes.
170     /// \param Inputs [out] - filled with  values marked as inputs to the
171     /// newly outlined function.
172      /// \param Outputs [out] - filled with values marked as outputs to the
173     /// newly outlined function.
174     /// \returns zero when called on a CodeExtractor instance where isEligible
175     /// returns false.
176     Function *extractCodeRegion(const CodeExtractorAnalysisCache &CEAC,
177                                 ValueSet &Inputs, ValueSet &Outputs);
178 
179     /// Verify that assumption cache isn't stale after a region is extracted.
180     /// Returns true when verifier finds errors. AssumptionCache is passed as
181     /// parameter to make this function stateless.
182     static bool verifyAssumptionCache(const Function &OldFunc,
183                                       const Function &NewFunc,
184                                       AssumptionCache *AC);
185 
186     /// Test whether this code extractor is eligible.
187     ///
188     /// Based on the blocks used when constructing the code extractor,
189     /// determine whether it is eligible for extraction.
190     ///
191     /// Checks that varargs handling (with vastart and vaend) is only done in
192     /// the outlined blocks.
193     bool isEligible() const;
194 
195     /// Compute the set of input values and output values for the code.
196     ///
197     /// These can be used either when performing the extraction or to evaluate
198     /// the expected size of a call to the extracted function. Note that this
199     /// work cannot be cached between the two as once we decide to extract
200     /// a code sequence, that sequence is modified, including changing these
201     /// sets, before extraction occurs. These modifications won't have any
202     /// significant impact on the cost however.
203     void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
204                            const ValueSet &Allocas,
205                            bool CollectGlobalInputs = false) const;
206 
207     /// Check if life time marker nodes can be hoisted/sunk into the outline
208     /// region.
209     ///
210     /// Returns true if it is safe to do the code motion.
211     bool
212     isLegalToShrinkwrapLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
213                                        Instruction *AllocaAddr) const;
214 
215     /// Find the set of allocas whose life ranges are contained within the
216     /// outlined region.
217     ///
218     /// Allocas which have life_time markers contained in the outlined region
219     /// should be pushed to the outlined function. The address bitcasts that
220     /// are used by the lifetime markers are also candidates for shrink-
221     /// wrapping. The instructions that need to be sunk are collected in
222     /// 'Allocas'.
223     void findAllocas(const CodeExtractorAnalysisCache &CEAC,
224                      ValueSet &SinkCands, ValueSet &HoistCands,
225                      BasicBlock *&ExitBlock) const;
226 
227     /// Find or create a block within the outline region for placing hoisted
228     /// code.
229     ///
230     /// CommonExitBlock is block outside the outline region. It is the common
231     /// successor of blocks inside the region. If there exists a single block
232     /// inside the region that is the predecessor of CommonExitBlock, that block
233     /// will be returned. Otherwise CommonExitBlock will be split and the
234     /// original block will be added to the outline region.
235     BasicBlock *findOrCreateBlockForHoisting(BasicBlock *CommonExitBlock);
236 
237     /// Exclude a value from aggregate argument passing when extracting a code
238     /// region, passing it instead as a scalar.
239     void excludeArgFromAggregate(Value *Arg);
240 
241   private:
242     struct LifetimeMarkerInfo {
243       bool SinkLifeStart = false;
244       bool HoistLifeEnd = false;
245       Instruction *LifeStart = nullptr;
246       Instruction *LifeEnd = nullptr;
247     };
248 
249     ValueSet ExcludeArgsFromAggregate;
250 
251     LifetimeMarkerInfo
252     getLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
253                        Instruction *Addr, BasicBlock *ExitBlock) const;
254 
255     /// Updates the list of SwitchCases (corresponding to exit blocks) after
256     /// changes of the control flow or the Blocks list.
257     void computeExtractedFuncRetVals();
258 
259     /// Return the type used for the return code of the extracted function to
260     /// indicate which exit block to jump to.
261     Type *getSwitchType();
262 
263     void severSplitPHINodesOfEntry(BasicBlock *&Header);
264     void severSplitPHINodesOfExits();
265     void splitReturnBlocks();
266 
267     void moveCodeToFunction(Function *newFunction);
268 
269     void calculateNewCallTerminatorWeights(
270         BasicBlock *CodeReplacer,
271         const DenseMap<BasicBlock *, BlockFrequency> &ExitWeights,
272         BranchProbabilityInfo *BPI);
273 
274     /// Normalizes the control flow of the extracted regions, such as ensuring
275     /// that the extracted region does not contain a return instruction.
276     void normalizeCFGForExtraction(BasicBlock *&header);
277 
278     /// Generates the function declaration for the function containing the
279     /// extracted code.
280     Function *constructFunctionDeclaration(const ValueSet &inputs,
281                                            const ValueSet &outputs,
282                                            BlockFrequency EntryFreq,
283                                            const Twine &Name,
284                                            ValueSet &StructValues,
285                                            StructType *&StructTy);
286 
287     /// Generates the code for the extracted function. That is: a prolog, the
288     /// moved or copied code from the original function, and epilogs for each
289     /// exit.
290     void emitFunctionBody(const ValueSet &inputs, const ValueSet &outputs,
291                           const ValueSet &StructValues, Function *newFunction,
292                           StructType *StructArgTy, BasicBlock *header,
293                           const ValueSet &SinkingCands);
294 
295     /// Generates a Basic Block that calls the extracted function.
296     CallInst *emitReplacerCall(const ValueSet &inputs, const ValueSet &outputs,
297                                const ValueSet &StructValues,
298                                Function *newFunction, StructType *StructArgTy,
299                                Function *oldFunction, BasicBlock *ReplIP,
300                                BlockFrequency EntryFreq,
301                                ArrayRef<Value *> LifetimesStart,
302                                std::vector<Value *> &Reloads);
303 
304     /// Connects the basic block containing the call to the extracted function
305     /// into the original function's control flow.
306     void insertReplacerCall(
307         Function *oldFunction, BasicBlock *header, BasicBlock *codeReplacer,
308         const ValueSet &outputs, ArrayRef<Value *> Reloads,
309         const DenseMap<BasicBlock *, BlockFrequency> &ExitWeights);
310   };
311 
312 } // end namespace llvm
313 
314 #endif // LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
315