10b57cec5SDimitry Andric //===- Transform/Utils/CodeExtractor.h - Code extraction util ---*- C++ -*-===//
20b57cec5SDimitry Andric //
30b57cec5SDimitry Andric // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
40b57cec5SDimitry Andric // See https://llvm.org/LICENSE.txt for license information.
50b57cec5SDimitry Andric // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
60b57cec5SDimitry Andric //
70b57cec5SDimitry Andric //===----------------------------------------------------------------------===//
80b57cec5SDimitry Andric //
90b57cec5SDimitry Andric // A utility to support extracting code from one function into its own
100b57cec5SDimitry Andric // stand-alone function.
110b57cec5SDimitry Andric //
120b57cec5SDimitry Andric //===----------------------------------------------------------------------===//
130b57cec5SDimitry Andric 
140b57cec5SDimitry Andric #ifndef LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
150b57cec5SDimitry Andric #define LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
160b57cec5SDimitry Andric 
170b57cec5SDimitry Andric #include "llvm/ADT/ArrayRef.h"
180b57cec5SDimitry Andric #include "llvm/ADT/DenseMap.h"
190b57cec5SDimitry Andric #include "llvm/ADT/SetVector.h"
200b57cec5SDimitry Andric #include <limits>
210b57cec5SDimitry Andric 
220b57cec5SDimitry Andric namespace llvm {
230b57cec5SDimitry Andric 
2481ad6265SDimitry Andric template <typename PtrType> class SmallPtrSetImpl;
258bcb0991SDimitry Andric class AllocaInst;
260b57cec5SDimitry Andric class BasicBlock;
270b57cec5SDimitry Andric class BlockFrequency;
280b57cec5SDimitry Andric class BlockFrequencyInfo;
290b57cec5SDimitry Andric class BranchProbabilityInfo;
300b57cec5SDimitry Andric class AssumptionCache;
310b57cec5SDimitry Andric class CallInst;
320b57cec5SDimitry Andric class DominatorTree;
330b57cec5SDimitry Andric class Function;
340b57cec5SDimitry Andric class Instruction;
350b57cec5SDimitry Andric class Loop;
360b57cec5SDimitry Andric class Module;
370b57cec5SDimitry Andric class Type;
380b57cec5SDimitry Andric class Value;
390b57cec5SDimitry Andric 
408bcb0991SDimitry Andric /// A cache for the CodeExtractor analysis. The operation \ref
418bcb0991SDimitry Andric /// CodeExtractor::extractCodeRegion is guaranteed not to invalidate this
428bcb0991SDimitry Andric /// object. This object should conservatively be considered invalid if any
438bcb0991SDimitry Andric /// other mutating operations on the IR occur.
448bcb0991SDimitry Andric ///
458bcb0991SDimitry Andric /// Constructing this object is O(n) in the size of the function.
468bcb0991SDimitry Andric class CodeExtractorAnalysisCache {
478bcb0991SDimitry Andric   /// The allocas in the function.
488bcb0991SDimitry Andric   SmallVector<AllocaInst *, 16> Allocas;
498bcb0991SDimitry Andric 
508bcb0991SDimitry Andric   /// Base memory addresses of load/store instructions, grouped by block.
518bcb0991SDimitry Andric   DenseMap<BasicBlock *, DenseSet<Value *>> BaseMemAddrs;
528bcb0991SDimitry Andric 
538bcb0991SDimitry Andric   /// Blocks which contain instructions which may have unknown side-effects
548bcb0991SDimitry Andric   /// on memory.
558bcb0991SDimitry Andric   DenseSet<BasicBlock *> SideEffectingBlocks;
568bcb0991SDimitry Andric 
578bcb0991SDimitry Andric   void findSideEffectInfoForBlock(BasicBlock &BB);
588bcb0991SDimitry Andric 
598bcb0991SDimitry Andric public:
608bcb0991SDimitry Andric   CodeExtractorAnalysisCache(Function &F);
618bcb0991SDimitry Andric 
628bcb0991SDimitry Andric   /// Get the allocas in the function at the time the analysis was created.
638bcb0991SDimitry Andric   /// Note that some of these allocas may no longer be present in the function,
648bcb0991SDimitry Andric   /// due to \ref CodeExtractor::extractCodeRegion.
getAllocas()658bcb0991SDimitry Andric   ArrayRef<AllocaInst *> getAllocas() const { return Allocas; }
668bcb0991SDimitry Andric 
678bcb0991SDimitry Andric   /// Check whether \p BB contains an instruction thought to load from, store
688bcb0991SDimitry Andric   /// to, or otherwise clobber the alloca \p Addr.
698bcb0991SDimitry Andric   bool doesBlockContainClobberOfAddr(BasicBlock &BB, AllocaInst *Addr) const;
708bcb0991SDimitry Andric };
718bcb0991SDimitry Andric 
720b57cec5SDimitry Andric   /// Utility class for extracting code into a new function.
730b57cec5SDimitry Andric   ///
740b57cec5SDimitry Andric   /// This utility provides a simple interface for extracting some sequence of
750b57cec5SDimitry Andric   /// code into its own function, replacing it with a call to that function. It
760b57cec5SDimitry Andric   /// also provides various methods to query about the nature and result of
770b57cec5SDimitry Andric   /// such a transformation.
780b57cec5SDimitry Andric   ///
790b57cec5SDimitry Andric   /// The rough algorithm used is:
800b57cec5SDimitry Andric   /// 1) Find both the inputs and outputs for the extracted region.
810b57cec5SDimitry Andric   /// 2) Pass the inputs as arguments, remapping them within the extracted
820b57cec5SDimitry Andric   ///    function to arguments.
830b57cec5SDimitry Andric   /// 3) Add allocas for any scalar outputs, adding all of the outputs' allocas
840b57cec5SDimitry Andric   ///    as arguments, and inserting stores to the arguments for any scalars.
850b57cec5SDimitry Andric   class CodeExtractor {
860b57cec5SDimitry Andric     using ValueSet = SetVector<Value *>;
870b57cec5SDimitry Andric 
880b57cec5SDimitry Andric     // Various bits of state computed on construction.
890b57cec5SDimitry Andric     DominatorTree *const DT;
900b57cec5SDimitry Andric     const bool AggregateArgs;
910b57cec5SDimitry Andric     BlockFrequencyInfo *BFI;
920b57cec5SDimitry Andric     BranchProbabilityInfo *BPI;
930b57cec5SDimitry Andric     AssumptionCache *AC;
940b57cec5SDimitry Andric 
9581ad6265SDimitry Andric     // A block outside of the extraction set where any intermediate
9681ad6265SDimitry Andric     // allocations will be placed inside. If this is null, allocations
9781ad6265SDimitry Andric     // will be placed in the entry block of the function.
9881ad6265SDimitry Andric     BasicBlock *AllocationBlock;
9981ad6265SDimitry Andric 
1000b57cec5SDimitry Andric     // If true, varargs functions can be extracted.
1010b57cec5SDimitry Andric     bool AllowVarArgs;
1020b57cec5SDimitry Andric 
1030b57cec5SDimitry Andric     // Bits of intermediate state computed at various phases of extraction.
1040b57cec5SDimitry Andric     SetVector<BasicBlock *> Blocks;
1050b57cec5SDimitry Andric     unsigned NumExitBlocks = std::numeric_limits<unsigned>::max();
1060b57cec5SDimitry Andric     Type *RetTy;
1070b57cec5SDimitry Andric 
108349cc55cSDimitry Andric     // Mapping from the original exit blocks, to the new blocks inside
109349cc55cSDimitry Andric     // the function.
110349cc55cSDimitry Andric     SmallVector<BasicBlock *, 4> OldTargets;
111349cc55cSDimitry Andric 
1120b57cec5SDimitry Andric     // Suffix to use when creating extracted function (appended to the original
1130b57cec5SDimitry Andric     // function name + "."). If empty, the default is to use the entry block
1140b57cec5SDimitry Andric     // label, if non-empty, otherwise "extracted".
1150b57cec5SDimitry Andric     std::string Suffix;
1160b57cec5SDimitry Andric 
1175f757f3fSDimitry Andric     // If true, the outlined function has aggregate argument in zero address
1185f757f3fSDimitry Andric     // space.
1195f757f3fSDimitry Andric     bool ArgsInZeroAddressSpace;
1205f757f3fSDimitry Andric 
1210b57cec5SDimitry Andric   public:
1220b57cec5SDimitry Andric     /// Create a code extractor for a sequence of blocks.
1230b57cec5SDimitry Andric     ///
1240b57cec5SDimitry Andric     /// Given a sequence of basic blocks where the first block in the sequence
1250b57cec5SDimitry Andric     /// dominates the rest, prepare a code extractor object for pulling this
1260b57cec5SDimitry Andric     /// sequence out into its new function. When a DominatorTree is also given,
1270b57cec5SDimitry Andric     /// extra checking and transformations are enabled. If AllowVarArgs is true,
1280b57cec5SDimitry Andric     /// vararg functions can be extracted. This is safe, if all vararg handling
1290b57cec5SDimitry Andric     /// code is extracted, including vastart. If AllowAlloca is true, then
1300b57cec5SDimitry Andric     /// extraction of blocks containing alloca instructions would be possible,
1310b57cec5SDimitry Andric     /// however code extractor won't validate whether extraction is legal.
13281ad6265SDimitry Andric     /// Any new allocations will be placed in the AllocationBlock, unless
13381ad6265SDimitry Andric     /// it is null, in which case it will be placed in the entry block of
13481ad6265SDimitry Andric     /// the function from which the code is being extracted.
1355f757f3fSDimitry Andric     /// If ArgsInZeroAddressSpace param is set to true, then the aggregate
1365f757f3fSDimitry Andric     /// param pointer of the outlined function is declared in zero address
1375f757f3fSDimitry Andric     /// space.
1380b57cec5SDimitry Andric     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
1390b57cec5SDimitry Andric                   bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
1400b57cec5SDimitry Andric                   BranchProbabilityInfo *BPI = nullptr,
14181ad6265SDimitry Andric                   AssumptionCache *AC = nullptr, bool AllowVarArgs = false,
14281ad6265SDimitry Andric                   bool AllowAlloca = false,
14381ad6265SDimitry Andric                   BasicBlock *AllocationBlock = nullptr,
1445f757f3fSDimitry Andric                   std::string Suffix = "", bool ArgsInZeroAddressSpace = false);
1450b57cec5SDimitry Andric 
1460b57cec5SDimitry Andric     /// Create a code extractor for a loop body.
1470b57cec5SDimitry Andric     ///
1480b57cec5SDimitry Andric     /// Behaves just like the generic code sequence constructor, but uses the
1490b57cec5SDimitry Andric     /// block sequence of the loop.
1500b57cec5SDimitry Andric     CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false,
1510b57cec5SDimitry Andric                   BlockFrequencyInfo *BFI = nullptr,
1520b57cec5SDimitry Andric                   BranchProbabilityInfo *BPI = nullptr,
1530b57cec5SDimitry Andric                   AssumptionCache *AC = nullptr,
1540b57cec5SDimitry Andric                   std::string Suffix = "");
1550b57cec5SDimitry Andric 
1560b57cec5SDimitry Andric     /// Perform the extraction, returning the new function.
1570b57cec5SDimitry Andric     ///
1580b57cec5SDimitry Andric     /// Returns zero when called on a CodeExtractor instance where isEligible
1590b57cec5SDimitry Andric     /// returns false.
1608bcb0991SDimitry Andric     Function *extractCodeRegion(const CodeExtractorAnalysisCache &CEAC);
1618bcb0991SDimitry Andric 
162349cc55cSDimitry Andric     /// Perform the extraction, returning the new function and providing an
163349cc55cSDimitry Andric     /// interface to see what was categorized as inputs and outputs.
164349cc55cSDimitry Andric     ///
165349cc55cSDimitry Andric     /// \param CEAC - Cache to speed up operations for the CodeExtractor when
166349cc55cSDimitry Andric     /// hoisting, and extracting lifetime values and assumes.
167349cc55cSDimitry Andric     /// \param Inputs [out] - filled with  values marked as inputs to the
168349cc55cSDimitry Andric     /// newly outlined function.
169349cc55cSDimitry Andric      /// \param Outputs [out] - filled with values marked as outputs to the
170349cc55cSDimitry Andric     /// newly outlined function.
171349cc55cSDimitry Andric     /// \returns zero when called on a CodeExtractor instance where isEligible
172349cc55cSDimitry Andric     /// returns false.
173349cc55cSDimitry Andric     Function *extractCodeRegion(const CodeExtractorAnalysisCache &CEAC,
174349cc55cSDimitry Andric                                 ValueSet &Inputs, ValueSet &Outputs);
175349cc55cSDimitry Andric 
1768bcb0991SDimitry Andric     /// Verify that assumption cache isn't stale after a region is extracted.
1775ffd83dbSDimitry Andric     /// Returns true when verifier finds errors. AssumptionCache is passed as
1788bcb0991SDimitry Andric     /// parameter to make this function stateless.
1795ffd83dbSDimitry Andric     static bool verifyAssumptionCache(const Function &OldFunc,
1805ffd83dbSDimitry Andric                                       const Function &NewFunc,
1815ffd83dbSDimitry Andric                                       AssumptionCache *AC);
1820b57cec5SDimitry Andric 
1830b57cec5SDimitry Andric     /// Test whether this code extractor is eligible.
1840b57cec5SDimitry Andric     ///
1850b57cec5SDimitry Andric     /// Based on the blocks used when constructing the code extractor,
1860b57cec5SDimitry Andric     /// determine whether it is eligible for extraction.
1878bcb0991SDimitry Andric     ///
1888bcb0991SDimitry Andric     /// Checks that varargs handling (with vastart and vaend) is only done in
1898bcb0991SDimitry Andric     /// the outlined blocks.
1908bcb0991SDimitry Andric     bool isEligible() const;
1910b57cec5SDimitry Andric 
1920b57cec5SDimitry Andric     /// Compute the set of input values and output values for the code.
1930b57cec5SDimitry Andric     ///
1940b57cec5SDimitry Andric     /// These can be used either when performing the extraction or to evaluate
1950b57cec5SDimitry Andric     /// the expected size of a call to the extracted function. Note that this
1960b57cec5SDimitry Andric     /// work cannot be cached between the two as once we decide to extract
1970b57cec5SDimitry Andric     /// a code sequence, that sequence is modified, including changing these
1980b57cec5SDimitry Andric     /// sets, before extraction occurs. These modifications won't have any
1990b57cec5SDimitry Andric     /// significant impact on the cost however.
2000b57cec5SDimitry Andric     void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
2010b57cec5SDimitry Andric                            const ValueSet &Allocas) const;
2020b57cec5SDimitry Andric 
2030b57cec5SDimitry Andric     /// Check if life time marker nodes can be hoisted/sunk into the outline
2040b57cec5SDimitry Andric     /// region.
2050b57cec5SDimitry Andric     ///
2060b57cec5SDimitry Andric     /// Returns true if it is safe to do the code motion.
2078bcb0991SDimitry Andric     bool
2088bcb0991SDimitry Andric     isLegalToShrinkwrapLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
2098bcb0991SDimitry Andric                                        Instruction *AllocaAddr) const;
2100b57cec5SDimitry Andric 
2110b57cec5SDimitry Andric     /// Find the set of allocas whose life ranges are contained within the
2120b57cec5SDimitry Andric     /// outlined region.
2130b57cec5SDimitry Andric     ///
2140b57cec5SDimitry Andric     /// Allocas which have life_time markers contained in the outlined region
2150b57cec5SDimitry Andric     /// should be pushed to the outlined function. The address bitcasts that
2160b57cec5SDimitry Andric     /// are used by the lifetime markers are also candidates for shrink-
2170b57cec5SDimitry Andric     /// wrapping. The instructions that need to be sunk are collected in
2180b57cec5SDimitry Andric     /// 'Allocas'.
2198bcb0991SDimitry Andric     void findAllocas(const CodeExtractorAnalysisCache &CEAC,
2208bcb0991SDimitry Andric                      ValueSet &SinkCands, ValueSet &HoistCands,
2210b57cec5SDimitry Andric                      BasicBlock *&ExitBlock) const;
2220b57cec5SDimitry Andric 
2230b57cec5SDimitry Andric     /// Find or create a block within the outline region for placing hoisted
2240b57cec5SDimitry Andric     /// code.
2250b57cec5SDimitry Andric     ///
2260b57cec5SDimitry Andric     /// CommonExitBlock is block outside the outline region. It is the common
2270b57cec5SDimitry Andric     /// successor of blocks inside the region. If there exists a single block
2280b57cec5SDimitry Andric     /// inside the region that is the predecessor of CommonExitBlock, that block
2290b57cec5SDimitry Andric     /// will be returned. Otherwise CommonExitBlock will be split and the
2300b57cec5SDimitry Andric     /// original block will be added to the outline region.
2310b57cec5SDimitry Andric     BasicBlock *findOrCreateBlockForHoisting(BasicBlock *CommonExitBlock);
2320b57cec5SDimitry Andric 
23304eeddc0SDimitry Andric     /// Exclude a value from aggregate argument passing when extracting a code
23404eeddc0SDimitry Andric     /// region, passing it instead as a scalar.
23504eeddc0SDimitry Andric     void excludeArgFromAggregate(Value *Arg);
23604eeddc0SDimitry Andric 
2370b57cec5SDimitry Andric   private:
2380b57cec5SDimitry Andric     struct LifetimeMarkerInfo {
2390b57cec5SDimitry Andric       bool SinkLifeStart = false;
2400b57cec5SDimitry Andric       bool HoistLifeEnd = false;
2410b57cec5SDimitry Andric       Instruction *LifeStart = nullptr;
2420b57cec5SDimitry Andric       Instruction *LifeEnd = nullptr;
2430b57cec5SDimitry Andric     };
2440b57cec5SDimitry Andric 
24504eeddc0SDimitry Andric     ValueSet ExcludeArgsFromAggregate;
24604eeddc0SDimitry Andric 
2478bcb0991SDimitry Andric     LifetimeMarkerInfo
2488bcb0991SDimitry Andric     getLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
2498bcb0991SDimitry Andric                        Instruction *Addr, BasicBlock *ExitBlock) const;
2500b57cec5SDimitry Andric 
2510b57cec5SDimitry Andric     void severSplitPHINodesOfEntry(BasicBlock *&Header);
2520b57cec5SDimitry Andric     void severSplitPHINodesOfExits(const SmallPtrSetImpl<BasicBlock *> &Exits);
2530b57cec5SDimitry Andric     void splitReturnBlocks();
2540b57cec5SDimitry Andric 
2550b57cec5SDimitry Andric     Function *constructFunction(const ValueSet &inputs,
2560b57cec5SDimitry Andric                                 const ValueSet &outputs,
2570b57cec5SDimitry Andric                                 BasicBlock *header,
2580b57cec5SDimitry Andric                                 BasicBlock *newRootNode, BasicBlock *newHeader,
2590b57cec5SDimitry Andric                                 Function *oldFunction, Module *M);
2600b57cec5SDimitry Andric 
2610b57cec5SDimitry Andric     void moveCodeToFunction(Function *newFunction);
2620b57cec5SDimitry Andric 
2630b57cec5SDimitry Andric     void calculateNewCallTerminatorWeights(
2640b57cec5SDimitry Andric         BasicBlock *CodeReplacer,
2650b57cec5SDimitry Andric         DenseMap<BasicBlock *, BlockFrequency> &ExitWeights,
2660b57cec5SDimitry Andric         BranchProbabilityInfo *BPI);
2670b57cec5SDimitry Andric 
2680b57cec5SDimitry Andric     CallInst *emitCallAndSwitchStatement(Function *newFunction,
2690b57cec5SDimitry Andric                                          BasicBlock *newHeader,
2700b57cec5SDimitry Andric                                          ValueSet &inputs, ValueSet &outputs);
2710b57cec5SDimitry Andric   };
2720b57cec5SDimitry Andric 
2730b57cec5SDimitry Andric } // end namespace llvm
2740b57cec5SDimitry Andric 
2750b57cec5SDimitry Andric #endif // LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
276