CodeExtractor : Move exit block splitting logic from PartialInliner to CodeExtractor

[rriddle]

Summary

Exit blocks (BBs that are branched to by blocks that are being extracted) that have phis with more than 1 incoming edge from the extracted region need to be split into two blocks, one in the extracted region and one on the outside. Move this logic from PartialInlining to CodeExtractor.

[silvas]

lib/Transforms/Utils/CodeExtractor.cpp

182

You may want to explicitly cast the result of 'Blocks.count(Pred)` to bool. Otherwise, I think the bool will promote to int which won't behave as intended.

852

We generally prefer DenseMap to std::unordered_map. Can you use DenseMap<BasicBlock *, unsigned>? If not, a comment would be useful.

http://llvm.org/docs/ProgrammersManual.html#llvm-adt-densemap-h

Also, this could use a comment explaining what the unsigned means. Maybe add to the comment just above ExitWeights?

test/Transforms/CodeExtractor/MultipleExitBranchProb.ll

34

The output of the compiler must be deterministic. You shouldn't be depending on the iteration order of a map or the numerical values of pointers (can be changed by ASLR etc.).

test/Transforms/CodeExtractor/SplitExitBlockPhiNodes.ll

1

You mentioned that this patch is just moving the logic from PartialInliner to CodeExtractor. So does this test case currently pass?

6

Explain which blocks get outlined and what split block we are supposed to be inserting.

27

You probably want a CHECK-LABEL for each function that you expect in the output. That way the CHECK for the phi node is restricted to being inside the function you want.

Also, you probably want to check a bit more, such as which BB's end up in which functions.

[davide]

River, are you still planning to work on this one?

include/llvm/Transforms/Utils/CodeExtractor.h

public:
     /// a code sequence, that sequence is modified, including changing these
     /// sets, before extraction occurs. These modifications won't have any
     /// significant impact on the cost however.
     void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs) const;
 
   private:
     void severSplitPHINodes(BasicBlock *&Header);
     void splitReturnBlocks();
-
+    BasicBlock *splitExitBlock(BasicBlock *Exit);
     Function *constructFunction(const ValueSet &inputs,
                                 const ValueSet &outputs,
                                 BasicBlock *header,
                                 BasicBlock *newRootNode, BasicBlock *newHeader,
                                 Function *oldFunction, Module *M);
 
     void moveCodeToFunction(Function *newFunction);
 

lib/Transforms/IPO/PartialInlining.cpp

Function *PartialInlinerImpl::unswitchFunction(Function *F) {
   BasicBlock *NewEntryBlock = cast<BasicBlock>(VMap[EntryBlock]);
   BasicBlock *NewReturnBlock = cast<BasicBlock>(VMap[ReturnBlock]);
   BasicBlock *NewNonReturnBlock = cast<BasicBlock>(VMap[NonReturnBlock]);
 
   // Go ahead and update all uses to the duplicate, so that we can just
   // use the inliner functionality when we're done hacking.
   F->replaceAllUsesWith(DuplicateFunction);
 
-  // Special hackery is needed with PHI nodes that have inputs from more than
-  // one extracted block.  For simplicity, just split the PHIs into a two-level
-  // sequence of PHIs, some of which will go in the extracted region, and some
-  // of which will go outside.
-  BasicBlock *PreReturn = NewReturnBlock;
-  NewReturnBlock = NewReturnBlock->splitBasicBlock(
-      NewReturnBlock->getFirstNonPHI()->getIterator());
-  BasicBlock::iterator I = PreReturn->begin();
-  Instruction *Ins = &NewReturnBlock->front();
-  while (I != PreReturn->end()) {
-    PHINode *OldPhi = dyn_cast<PHINode>(I);
-    if (!OldPhi)
-      break;
-
-    PHINode *RetPhi = PHINode::Create(OldPhi->getType(), 2, "", Ins);
-    OldPhi->replaceAllUsesWith(RetPhi);
-    Ins = NewReturnBlock->getFirstNonPHI();
-
-    RetPhi->addIncoming(&*I, PreReturn);
-    RetPhi->addIncoming(OldPhi->getIncomingValueForBlock(NewEntryBlock),
-                        NewEntryBlock);
-    OldPhi->removeIncomingValue(NewEntryBlock);
-
-    ++I;
-  }
-  NewEntryBlock->getTerminator()->replaceUsesOfWith(PreReturn, NewReturnBlock);
-
   // Gather up the blocks that we're going to extract.
   std::vector<BasicBlock *> ToExtract;
   ToExtract.push_back(NewNonReturnBlock);
   for (BasicBlock &BB : *DuplicateFunction)
     if (&BB != NewEntryBlock && &BB != NewReturnBlock &&
         &BB != NewNonReturnBlock)
       ToExtract.push_back(&BB);
 

lib/Transforms/Utils/CodeExtractor.cpp

 #include "llvm/Support/BlockFrequency.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <algorithm>
 #include <set>
+#include <unordered_map>
 using namespace llvm;
 
 #define DEBUG_TYPE "code-extractor"
 
 // Provide a command-line option to aggregate function arguments into a struct
 // for functions produced by the code extractor. This is useful when converting
 // extracted functions to pthread-based code, as only one argument (void*) can
 // be passed in to pthread_create().

 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;
 }
 
+/// movePredsTo - Move the predecessors of \p Block
+/// that are either within or outside of the provided \p Blocks depending
+/// on the \p IfWithinSet to \p New.
+static void movePredsTo(BasicBlock *Block, BasicBlock *New,
+                        const SetVector<BasicBlock *> Blocks,
+                        bool IfWithinSet) {
+  // Move the predecessors of Block to New.
+  SmallPtrSet<BasicBlock *, 4> Preds;
+  for (BasicBlock *Pred : predecessors(Block))
+    if (Blocks.count(Pred) == IfWithinSet)
+      Preds.insert(Pred);
+  for (BasicBlock *Pred : Preds)
+    Pred->getTerminator()->replaceUsesOfWith(Block, New);
+}
+
 void CodeExtractor::findInputsOutputs(ValueSet &Inputs,
                                       ValueSet &Outputs) const {
   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 (ReturnInst *RI = dyn_cast<ReturnInst>(Block->getTerminator())) {
         DomTreeNode *NewNode = DT->addNewBlock(New, Block);
 
         for (DomTreeNode *I : Children)
           DT->changeImmediateDominator(I, NewNode);
       }
     }
 }
 
+/// splitExitBlock - Split an exit block that has a phi with more than 1
+/// incoming
+/// block that is within the extracted region, returns the new exit block.
+///
+BasicBlock *CodeExtractor::splitExitBlock(BasicBlock *Exit) {
+  // Special hackery is needed with PHI nodes that have inputs from more than
+  // one extracted block.  For simplicity, just split the PHIs into a two-level
+  // sequence of PHIs, some of which will go in the extracted region, and some
+  // of which will go outside.
+  BasicBlock *ExtractedReturn = Exit;
+  Exit = SplitBlock(Exit, Exit->getFirstNonPHI(), DT);
+
+  // Move the branches of the non extracted blocks to the new exit block.
+  movePredsTo(ExtractedReturn, Exit, Blocks, /*IfWithinSet*/ false);
+
+  // Create new phi nodes in the exit block.
+  Instruction *InsPoint = &Exit->front();
+  for (Instruction &I : *ExtractedReturn) {
+    PHINode *Phi = dyn_cast<PHINode>(&I);
+    if (!Phi)
+      break;
+
+    // Create a new phi and redirect the extracted phi to it.
+    PHINode *NewPhi =
+        PHINode::Create(Phi->getType(), 0, "split.exit", InsPoint);
+    Phi->replaceAllUsesWith(NewPhi);
+    NewPhi->addIncoming(Phi, ExtractedReturn);
+
+    // Move non extracted incoming blocks to the new phi.
+    for (unsigned i = 0, e = Phi->getNumIncomingValues(); i < e; ++i) {
+      BasicBlock *IncomingBlock = Phi->getIncomingBlock(i);
+      if (Blocks.count(IncomingBlock))
+        continue;
+      NewPhi->addIncoming(Phi->getIncomingValue(i), IncomingBlock);
+      Phi->removeIncomingValue(i, /*DeletePHIIfEmpty*/ false);
+      --i;
+      --e;
+    }
+  }
+  Blocks.insert(ExtractedReturn);
+  return Exit;
+}
+
 /// constructFunction - make a function based on inputs and outputs, as follows:
 /// f(in0, ..., inN, out0, ..., outN)
 ///
 Function *CodeExtractor::constructFunction(const ValueSet &inputs,
                                            const ValueSet &outputs,
                                            BasicBlock *header,
                                            BasicBlock *newRootNode,
                                            BasicBlock *newHeader,
BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(),
                                                 header);
 
   // The new function needs a root node because other nodes can branch to the
   // 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);
-
   // Calculate the exit blocks for the extracted region and the total exit
   //  weights for each of those blocks.
   DenseMap<BasicBlock *, BlockFrequency> ExitWeights;
-  SmallPtrSet<BasicBlock *, 1> ExitBlocks;
+  std::unordered_map<BasicBlock *, unsigned> ExitBlocks;
   for (BasicBlock *Block : Blocks) {
     for (succ_iterator SI = succ_begin(Block), SE = succ_end(Block); SI != SE;
          ++SI) {
       if (!Blocks.count(*SI)) {
         // Update the branch weight for this successor.
         if (BFI) {
           BlockFrequency &BF = ExitWeights[*SI];
           BF += BFI->getBlockFreq(Block) * BPI->getEdgeProbability(Block, *SI);
         }
-        ExitBlocks.insert(*SI);
+        ++ExitBlocks[*SI];
       }
     }
   }
   NumExitBlocks = ExitBlocks.size();
 
+  // Split the phi nodes of each of the exit blocks if there is more than 1
+  // incoming block from the extracted region.
+  for (auto I = ExitBlocks.begin(), E = ExitBlocks.end(); I != E;) {
+    BasicBlock *Exit = I->first;
+    if (I->second < 2 || !isa<PHINode>(&Exit->front())) {
+      ++I;
+      continue;
+    }
+
+    // Split the block and swap the values with the old exit block before
+    // removing it.
+    BasicBlock *SplitBlock = splitExitBlock(Exit);
+    ExitWeights[SplitBlock] = ExitWeights[Exit];
+    ExitBlocks.emplace(SplitBlock, 0);
+    ExitWeights.erase(Exit);
+    ExitBlocks.erase(I++);
+  }
+
+  // Find inputs to, outputs from the code region.
+  findInputsOutputs(inputs, outputs);
+
   // Construct new function based on inputs/outputs & add allocas for all defs.
   Function *newFunction = constructFunction(inputs, outputs, header,
                                             newFuncRoot,
                                             codeReplacer, oldFunction,
                                             oldFunction->getParent());
 
   // Update the entry count of the function.
   if (BFI) {

test/Transforms/CodeExtractor/MultipleExitBranchProb.ll

 }
 
 !llvm.module.flags = !{!0}
 !0 = !{i32 1, !"MaxFunctionCount", i32 10000}
 !1 = !{!"function_entry_count", i64 10000}
 !2 = !{!"branch_weights", i32 5, i32 5}
 !3 = !{!"branch_weights", i32 4, i32 1}
 
-; CHECK: [[COUNT1]] = !{!"branch_weights", i32 8, i32 31}
+; CHECK: [[COUNT1]] = !{!"branch_weights", i32 {{(31, i32 8)|(8, i32 31)}}}

test/Transforms/CodeExtractor/SplitExitBlockPhiNodes.ll

+; RUN: opt < %s -partial-inliner -S | FileCheck %s
+; This testcase checks CodeExtractors ability to split
+; an exit block that contains a phi node with more than
+; 1 entry from an extracted block.
+
+define internal i32 @inlinedFunc(i1 %cond) {
+entry:
+  br i1 %cond, label %if.then, label %return
+if.then:
+  br i1 %cond, label %internal.if.then, label %internal.if.else
+internal.if.then:
+  br label %return
+internal.if.else:
+  br label %return
+return:             ; preds = %entry
+  %return.val = phi i32 [ 0, %entry ], [ 1, %internal.if.then ], [ 2, %internal.if.else ]
+  ret i32 %return.val
+}
+
+define internal i32 @dummyCaller(i1 %cond) {
+entry:
+  %val = call i32 @inlinedFunc(i1 %cond)
+  ret i32 %val
+}
+
+; CHECK-LABEL: @inlinedFunc.1_if.then
+; CHECK: phi i32