Replace CodeExtractor callbacks with subclasses and simplify their creation based on OutlineInfo structures

Author: skatrak

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -31,6 +31,7 @@
 
 namespace llvm {
 class CanonicalLoopInfo;
+class CodeExtractor;
 class ScanInfo;
 struct TargetRegionEntryInfo;
 class OffloadEntriesInfoManager;
@@ -2260,27 +2261,31 @@ class OpenMPIRBuilder {
   /// during finalization.
   struct OutlineInfo {
     using PostOutlineCBTy = std::function<void(Function &)>;
-    using CustomArgAllocatorCBTy = std::function<Instruction *(
-        BasicBlock *, BasicBlock::iterator, Type *, const Twine &)>;
-    using CustomArgDeallocatorCBTy = std::function<Instruction *(
-        BasicBlock *, BasicBlock::iterator, Value *, Type *)>;
     PostOutlineCBTy PostOutlineCB;
-    CustomArgAllocatorCBTy CustomArgAllocatorCB;
-    CustomArgDeallocatorCBTy CustomArgDeallocatorCB;
     BasicBlock *EntryBB, *ExitBB, *OuterAllocaBB;
     SmallVector<Value *, 2> ExcludeArgsFromAggregate;
 
+    LLVM_ABI virtual ~OutlineInfo() = default;
+
     /// Collect all blocks in between EntryBB and ExitBB in both the given
     /// vector and set.
     LLVM_ABI void collectBlocks(SmallPtrSetImpl<BasicBlock *> &BlockSet,
                                 SmallVectorImpl<BasicBlock *> &BlockVector);
 
+    /// Create a CodeExtractor instance based on the information stored in this
+    /// structure, the list of collected blocks from a previous call to
+    /// \c collectBlocks and a flag stating whether arguments must be passed in
+    /// address space 0.
+    LLVM_ABI virtual std::unique_ptr<CodeExtractor>
+    createCodeExtractor(ArrayRef<BasicBlock *> Blocks,
+                        bool ArgsInZeroAddressSpace, Twine Suffix = Twine(""));
+
     /// Return the function that contains the region to be outlined.
     Function *getFunction() const { return EntryBB->getParent(); }
   };
 
   /// Collection of regions that need to be outlined during finalization.
-  SmallVector<OutlineInfo, 16> OutlineInfos;
+  SmallVector<std::unique_ptr<OutlineInfo>, 16> OutlineInfos;
 
   /// A collection of candidate target functions that's constant allocas will
   /// attempt to be raised on a call of finalize after all currently enqueued
@@ -2295,7 +2300,9 @@ class OpenMPIRBuilder {
   std::forward_list<ScanInfo> ScanInfos;
 
   /// Add a new region that will be outlined later.
-  void addOutlineInfo(OutlineInfo &&OI) { OutlineInfos.emplace_back(OI); }
+  void addOutlineInfo(std::unique_ptr<OutlineInfo> &&OI) {
+    OutlineInfos.emplace_back(std::move(OI));
+  }
 
   /// An ordered map of auto-generated variables to their unique names.
   /// It stores variables with the following names: 1) ".gomp_critical_user_" +

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

@@ -24,6 +24,7 @@
 namespace llvm {
 
 template <typename PtrType> class SmallPtrSetImpl;
+class AddrSpaceCastInst;
 class AllocaInst;
 class BlockFrequency;
 class BlockFrequencyInfo;
@@ -85,10 +86,6 @@ class CodeExtractorAnalysisCache {
   /// 3) Add allocas for any scalar outputs, adding all of the outputs' allocas
   ///    as arguments, and inserting stores to the arguments for any scalars.
   class CodeExtractor {
-    using CustomArgAllocatorCBTy = std::function<Instruction *(
-        BasicBlock *, BasicBlock::iterator, Type *, const Twine &)>;
-    using CustomArgDeallocatorCBTy = std::function<Instruction *(
-        BasicBlock *, BasicBlock::iterator, Value *, Type *)>;
     using ValueSet = SetVector<Value *>;
 
     // Various bits of state computed on construction.
@@ -103,6 +100,14 @@ class CodeExtractorAnalysisCache {
     /// will be placed in the entry block of the function.
     BasicBlock *AllocationBlock;
 
+    /// A block outside of the extraction set where deallocations for
+    /// intermediate allocations can be placed inside. Not used for
+    /// automatically deallocated memory (e.g. `alloca`), which is the default.
+    ///
+    /// If it is null and needed, the end of the replacement basic block will be
+    /// used to place deallocations.
+    BasicBlock *DeallocationBlock;
+
     /// If true, varargs functions can be extracted.
     bool AllowVarArgs;
 
@@ -137,25 +142,6 @@ class CodeExtractorAnalysisCache {
     /// space.
     bool ArgsInZeroAddressSpace;
 
-    /// If set, this callback will be used to allocate the arguments in the
-    /// caller before passing it to the outlined function holding the extracted
-    /// piece of code.
-    CustomArgAllocatorCBTy *CustomArgAllocatorCB;
-
-    /// A block outside of the extraction set where previously introduced
-    /// intermediate allocations can be deallocated. This is only used when a
-    /// custom deallocator is specified.
-    BasicBlock *DeallocationBlock;
-
-    /// If set, this callback will be used to deallocate the arguments in the
-    /// caller after running the outlined function holding the extracted piece
-    /// of code. It will not be called if a custom allocator isn't also present.
-    ///
-    /// By default, this will be done at the end of the basic block containing
-    /// the call to the outlined function, except if a deallocation block is
-    /// specified. In that case, that will take precedence.
-    CustomArgDeallocatorCBTy *CustomArgDeallocatorCB;
-
   public:
     /// Create a code extractor for a sequence of blocks.
     ///
@@ -169,23 +155,23 @@ class CodeExtractorAnalysisCache {
     /// however code extractor won't validate whether extraction is legal.
     /// Any new allocations will be placed in the AllocationBlock, unless
     /// it is null, in which case it will be placed in the entry block of
-    /// the function from which the code is being extracted.
+    /// the function from which the code is being extracted. Explicit
+    /// deallocations for the aforementioned allocations will be placed in the
+    /// DeallocationBlock or the end of the replacement block, if needed.
     /// If ArgsInZeroAddressSpace param is set to true, then the aggregate
     /// param pointer of the outlined function is declared in zero address
-    /// space. If a CustomArgAllocatorCB callback is specified, it will be used
-    /// to allocate any structures or variable copies needed to pass arguments
-    /// to the outlined function, rather than using regular allocas.
+    /// space.
     LLVM_ABI
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
                   bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
                   BranchProbabilityInfo *BPI = nullptr,
                   AssumptionCache *AC = nullptr, bool AllowVarArgs = false,
                   bool AllowAlloca = false,
                   BasicBlock *AllocationBlock = nullptr,
-                  std::string Suffix = "", bool ArgsInZeroAddressSpace = false,
-                  CustomArgAllocatorCBTy *CustomArgAllocatorCB = nullptr,
                   BasicBlock *DeallocationBlock = nullptr,
-                  CustomArgDeallocatorCBTy *CustomArgDeallocatorCB = nullptr);
+                  std::string Suffix = "", bool ArgsInZeroAddressSpace = false);
+
+    LLVM_ABI virtual ~CodeExtractor() = default;
 
     /// Perform the extraction, returning the new function.
     ///
@@ -271,6 +257,19 @@ class CodeExtractorAnalysisCache {
     /// region, passing it instead as a scalar.
     LLVM_ABI void excludeArgFromAggregate(Value *Arg);
 
+  protected:
+    /// Allocate an intermediate variable at the specified point.
+    LLVM_ABI virtual Instruction *
+    allocateVar(BasicBlock *BB, BasicBlock::iterator AllocIP, Type *VarType,
+                const Twine &Name = Twine(""),
+                AddrSpaceCastInst **CastedAlloc = nullptr);
+
+    /// Deallocate a previously-allocated intermediate variable at the specified
+    /// point.
+    LLVM_ABI virtual Instruction *deallocateVar(BasicBlock *BB,
+                                                BasicBlock::iterator DeallocIP,
+                                                Value *Var, Type *VarType);
+
   private:
     struct LifetimeMarkerInfo {
       bool SinkLifeStart = false;