Add option to CodeExtractor for hoisting static allocas with uses contained within extracted region

[rriddle]

Summary

Give CodeExtractor the ability to hoist any static allocas that have all of their uses inside of the extracted region to be hoisted into the new function. Also changed the internal ValueSet of CodeExtractor to be a SmallVector to allow for the inputs and outputs to be modified after collection. This doesn't really change anything given that the inputs and outputs ValueSet are never used as a set after collection.

[rriddle]

Formatting

[silvas]

A couple nits. Will look closer tomorrow.

We don't use the date convention on test cases anymore. You can just call the test case ExtractDominatedStaticAllocas.ll

/Users/rriddle/Desktop/llvm/llvm/lib/Transforms/Utils/CodeExtractor.cpp

728 ↗

(On Diff #65254)

It's a bit unusual to have this extra space here on all the lines except the first (I noticed this in one of your other patches too).

730 ↗

(On Diff #65254)

This can just be regular int or unsigned for the loop variable. Also, for whatever reason, the standard integer variables are i and e in LLVM even though otherwise our convention is for CamelCase variable names. I and E would be used for iterators however.

Also, even better, this can become a range-for loop.

[silvas]

Also, I'm not sure if you're using the terminology "static alloca" correct, as you seem to be using it to refer to simply constant size. An alloca must be in the entry block as well to be considered a static alloca:

/// isStaticAlloca - Return true if this alloca is in the entry block of the
/// function and is a constant size.  If so, the code generator will fold it
/// into the prolog/epilog code, so it is basically free.
bool isStaticAlloca() const;

[silvas]

Another quick comment.

/Users/rriddle/Desktop/llvm/llvm/lib/Transforms/Utils/CodeExtractor.cpp

732 ↗

(On Diff #65254)

There seem to be very few uses of this pattern of const_cast combined with dyn_cast, and it smells a bit fishy to me. Is there some other preferred pattern in use in the codebase? Or is there another way to avoid the const_cast?

[rriddle]

Updated to reflect Sean's comments

[davide]

Some comments online. Also please clang-format this patch.

lib/Transforms/IPO/PartialInlining.cpp

138–139

is this clang-formatted?

lib/Transforms/Utils/CodeExtractor.cpp

176–177

this needs a comment.

180–187

Can you add a test exercising this behaviour?

189–193

same.

[davide]

s/online/inline/

include/llvm/Transforms/Utils/CodeExtractor.h

template <typename T> class ArrayRef;
   ///
   /// The rough algorithm used is:
   /// 1) Find both the inputs and outputs for the extracted region.
   /// 2) Pass the inputs as arguments, remapping them within the extracted
   ///    function to arguments.
   /// 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 {
-    typedef SetVector<Value *> ValueSet;
+    typedef SmallVector<Value *, 4> ValueSet;
 
     // Various bits of state computed on construction.
     DominatorTree *const DT;
     const bool AggregateArgs;
+    const bool ExtractContainedStaticAllocas;
 
     // Bits of intermediate state computed at various phases of extraction.
     SetVector<BasicBlock *> Blocks;
     unsigned NumExitBlocks;
     Type *RetTy;
 
   public:
 
     /// \brief Check to see if a block is valid for extraction.
     ///
     /// Blocks containing EHPads, allocas, invokes, or vastarts are not valid.
     static bool isBlockValidForExtraction(const BasicBlock &BB);
 
     /// \brief Create a code extractor for a single basic block.
     ///
     /// In this formation, we don't require a dominator tree. The given basic
     /// block is set up for extraction.
-    CodeExtractor(BasicBlock *BB, bool AggregateArgs = false);
+    CodeExtractor(BasicBlock *BB, bool AggregateArgs = false,
+                  bool ExtractContainedStaticAllocas = false);
 
     /// \brief Create a code extractor for a sequence of blocks.
     ///
     /// Given a sequence of basic blocks where the first block in the sequence
     /// dominates the rest, prepare a code extractor object for pulling this
     /// sequence out into its new function. When a DominatorTree is also given,
     /// extra checking and transformations are enabled.
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
-                  bool AggregateArgs = false);
+                  bool AggregateArgs = false,
+                  bool ExtractContainedStaticAllocas = false);
 
     /// \brief Create a code extractor for a loop body.
     ///
     /// Behaves just like the generic code sequence constructor, but uses the
     /// block sequence of the loop.
-    CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false);
+    CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false,
+                  bool ExtractContainedStaticAllocas = false);
 
     /// \brief Create a code extractor for a region node.
     ///
     /// Behaves just like the generic code sequence constructor, but uses the
     /// block sequence of the region node passed in.
     CodeExtractor(DominatorTree &DT, const RegionNode &RN,
-                  bool AggregateArgs = false);
+                  bool AggregateArgs = false,
+                  bool ExtractContainedStaticAllocas = false);
 
     /// \brief Perform the extraction, returning the new function.
     ///
     /// Returns zero when called on a CodeExtractor instance where isEligible
     /// returns false.
     Function *extractCodeRegion();
 
     /// \brief Test whether this code extractor is eligible.

lib/Transforms/IPO/PartialInlining.cpp

if (&BB != NewEntryBlock && &BB != NewReturnBlock &&
       ToExtract.push_back(&BB);
 
   // The CodeExtractor needs a dominator tree.
   DominatorTree DT;
   DT.recalculate(*DuplicateFunction);
 
   // Extract the body of the if.
   Function *ExtractedFunction =
-      CodeExtractor(ToExtract, &DT).extractCodeRegion();
+      CodeExtractor(ToExtract, &DT, /*AggregateArgs*/false,
+                    /*ExtractContainedStaticAllocas*/true).extractCodeRegion();
 
   // Inline the top-level if test into all callers.
   std::vector<User *> Users(DuplicateFunction->user_begin(),
                             DuplicateFunction->user_end());
   for (User *User : Users)
     if (CallInst *CI = dyn_cast<CallInst>(User))
       InlineFunction(CI, IFI);
     else if (InvokeInst *II = dyn_cast<InvokeInst>(User))

lib/Transforms/Utils/CodeExtractor.cpp

if (!RN.isSubRegion())
     // Just a single BasicBlock.
     return buildExtractionBlockSet(RN.getNodeAs<BasicBlock>());
 
   const Region &R = *RN.getNodeAs<Region>();
 
   return buildExtractionBlockSet(R.block_begin(), R.block_end());
 }
 
-CodeExtractor::CodeExtractor(BasicBlock *BB, bool AggregateArgs)
+CodeExtractor::CodeExtractor(BasicBlock *BB, bool AggregateArgs,
+                             bool ExtractContainedStaticAllocas)
   : DT(nullptr), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+    ExtractContainedStaticAllocas(ExtractContainedStaticAllocas),
     Blocks(buildExtractionBlockSet(BB)), NumExitBlocks(~0U) {}
 
 CodeExtractor::CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
-                             bool AggregateArgs)
+                             bool AggregateArgs,
+                             bool ExtractContainedStaticAllocas)
   : DT(DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+    ExtractContainedStaticAllocas(ExtractContainedStaticAllocas),
     Blocks(buildExtractionBlockSet(BBs)), NumExitBlocks(~0U) {}
 
-CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs)
+CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs,
+                             bool ExtractContainedStaticAllocas)
   : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+    ExtractContainedStaticAllocas(ExtractContainedStaticAllocas),
     Blocks(buildExtractionBlockSet(L.getBlocks())), NumExitBlocks(~0U) {}
 
 CodeExtractor::CodeExtractor(DominatorTree &DT, const RegionNode &RN,
-                             bool AggregateArgs)
+                             bool AggregateArgs,
+                             bool ExtractContainedStaticAllocas)
   : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+    ExtractContainedStaticAllocas(ExtractContainedStaticAllocas),
     Blocks(buildExtractionBlockSet(RN)), NumExitBlocks(~0U) {}
 
 /// definedInRegion - Return true if the specified value is defined in the
 /// extracted region.
 static bool definedInRegion(const SetVector<BasicBlock *> &Blocks, Value *V) {
   if (Instruction *I = dyn_cast<Instruction>(V))
     if (Blocks.count(I->getParent()))
       return true;
   return false;
 }
 
 /// definedInCaller - Return true if the specified value is defined in the
 /// function being code extracted, but not in the region being extracted.
 /// These values must be passed in as live-ins to the function.
 static bool definedInCaller(const SetVector<BasicBlock *> &Blocks, Value *V) {
   if (isa<Argument>(V)) return true;
   if (Instruction *I = dyn_cast<Instruction>(V))
     if (!Blocks.count(I->getParent()))
       return true;
   return false;
 }
 
+/// hasOutsideUses - Return true if the specified instruction is used outside of
+/// the region that is being extracted.
+static bool hasOutsideUses(const SetVector<BasicBlock *> &Blocks,
+                           Instruction *I) {
+  for (Use &U : I->uses()) {
+    Instruction *User = dyn_cast<Instruction>(U.getUser());
+    if (!User)
+      continue;
+    if (!Blocks.count(User->getParent()))
+      return true;
+    if (!I->getType()->isPointerTy())
+      continue;
+    // Check to see if there is a gep of this instruction that geps this
+    // instruction.
+    if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(User)) {
+      if (GEP->getPointerOperand() != I)
+        continue;
+      if (hasOutsideUses(Blocks, GEP))
+        return true;
+    }
+    // Do the same for bit cast instructions.
+    else if (BitCastInst *BC = dyn_cast<BitCastInst>(User)) {
+      if (hasOutsideUses(Blocks, BC))
+        return true;
+    }
+  }
+  return false;
+}
+
 void CodeExtractor::findInputsOutputs(ValueSet &Inputs,
                                       ValueSet &Outputs) const {
+  SmallPtrSet<Value *, 4> InputSet, OutputSet;
   for (BasicBlock *BB : Blocks) {
     // If a used value is defined outside the region, it's an input.  If an
     // instruction is used outside the region, it's an output.
     for (Instruction &II : *BB) {
       for (User::op_iterator OI = II.op_begin(), OE = II.op_end(); OI != OE;
            ++OI)
         if (definedInCaller(Blocks, *OI))
-          Inputs.insert(*OI);
+          if (InputSet.insert(*OI).second)
+            Inputs.push_back(*OI);
 
       for (User *U : II.users())
         if (!definedInRegion(Blocks, U)) {
-          Outputs.insert(&II);
+          if (OutputSet.insert(&II).second)
+            Outputs.push_back(&II);
           break;
         }
     }
   }
 }
 
 /// severSplitPHINodes - If a PHI node has multiple inputs from outside of the
 /// region, we need to split the entry block of the region so that the PHI node

 /// emitCallAndSwitchStatement - This method sets up the caller side by adding
 /// the call instruction, splitting any PHI nodes in the header block as
 /// necessary.
 void CodeExtractor::
 emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
                            ValueSet &inputs, ValueSet &outputs) {
   // Emit a call to the new function, passing in: *pointer to struct (if
   // aggregating parameters), or plan inputs and allocated memory for outputs
-  std::vector<Value*> params, StructValues, ReloadOutputs, Reloads;
-  
+  SmallVector<Value *, 4> params, StructValues, ReloadOutputs, Reloads;
   LLVMContext &Context = newFunction->getContext();
 
   // Add inputs as params, or to be filled into the struct
   for (Value *input : inputs)
     if (AggregateArgs)
       StructValues.push_back(input);
     else
       params.push_back(input);
Function *CodeExtractor::extractCodeRegion() {
   // head of the region, but the entry node of a function cannot have preds.
   BasicBlock *newFuncRoot = BasicBlock::Create(header->getContext(), 
                                                "newFuncRoot");
   newFuncRoot->getInstList().push_back(BranchInst::Create(header));
 
   // Find inputs to, outputs from the code region.
   findInputsOutputs(inputs, outputs);
 
+  // Loop through all of the inputs looking for static allocas that do not have
+  //  any uses outside of the extracted region.
+  if (ExtractContainedStaticAllocas)
+    for (unsigned i = 0, e = inputs.size(); i < e;) {
+      AllocaInst *AI = dyn_cast<AllocaInst>(inputs[i]);
+      if (AI && AI->isStaticAlloca() && !hasOutsideUses(Blocks, AI)) {
+        AI->moveBefore(&newFuncRoot->front());
+        Value *Back = inputs.pop_back_val();
+
+        // If this is the last input value then just break from the loop.
+        if (AI == Back)
+          break;
+        inputs[i] = Back;
+        --e;
+        continue;
+      }
+      ++i;
+    }
+
   SmallPtrSet<BasicBlock *, 1> ExitBlocks;
   for (BasicBlock *Block : Blocks)
     for (succ_iterator SI = succ_begin(Block), SE = succ_end(Block); SI != SE;
          ++SI)
       if (!Blocks.count(*SI))
         ExitBlocks.insert(*SI);
   NumExitBlocks = ExitBlocks.size();
 

test/Transforms/CodeExtractor/ExtractContainedStaticAllocas.ll

+; RUN: opt < %s -partial-inliner -S | FileCheck %s
+; This testcase hoists a static alloca that has all
+;  of it uses within the extracted region.
+
+define internal i32 @inlinedFunc(i1 %cond) {
+entry:
+  %dominated.alloca = alloca i32, align 4
+  br i1 %cond, label %if.then, label %return
+if.then:
+; Dummy store to represent the only use
+  store i32 10, i32* %dominated.alloca, align 4
+  br label %return
+return:             ; preds = %entry
+  ret i32 0
+}
+
+define internal i32 @dummyCaller(i1 %cond) {
+entry:
+  %val = call i32 @inlinedFunc(i1 %cond)
+  ret i32 %val
+}
+
+; CHECK-LABEL: @inlinedFunc.1_if.then
+; CHECK: alloca i32