[PartialInl] Enhance partial inliner to handle more complex conditions

[davidxl]

Summary

This handles case :

if (a || b)

return;

// Large_body of slow path code
do_something here.

[eraman]

I've a couple of trivial comments so far. Could you please generate a diff with a larger context?

lib/Transforms/IPO/PartialInlining.cpp

40

The comments doesn't convey what the function is doing.

47

nit: change -> chain

50

Why is this named Entries_ and not Entries?

[davidxl]

Addressed review comments .

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

40

Fixed.

50

Fixed.

[wmi]

lib/Transforms/IPO/PartialInlining.cpp

74

nit: ghain --> chain

207–216

Here we add a limitation that the function only has one return BB to be a partial inline candidate. Is it required? Can we walk forward from entry instead of backward from the return block so as to avoid the requirement of finding single return BB?

300–308

Here we allows a node on the chain with more than two successors. Will such case bring us problem later during extracting?

test/Transforms/CodeExtractor/MultiConditions.ll

5 ↗

(On Diff #95840)

we can use opt -instnamer to rename the temporaries.

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

207–216

the restriction no it is not required. Changed as suggested.

(Note if the extracted region has other return instruction, extra code will be generated (to handle the return status in caller), so the call to the outlined function is slightly higher).

300–308

Changed the code to limit 2 successors. More than one successors can cause problems or lead to very inefficient outline call (to handle 'multiple entries').

[davidxl]

Addressed review comments. Added more test including multi return case.

[wmi]

I am thinking of another two issues:

1. Is it possible that some of the BBs on the chain may be very big and we don't want to partial inline them?
2. The existing pattern handles if (a || b || c ...) case, but it may not be easy to extend for cases like (a && b && c) and ((a && b) || c). Basically, we want to find a collection of bbs with small sizes starting from entry. The bb collection only have two exits. one of them is ReturnBlock and the other is NonReturnBlock. The edges inside of the collection can be of any pattern.

Sorry I didn't come up with them at the first beginning and dive into details.

[davidxl]

Is it possible that some of the BBs on the chain may be very big and we don't want to partial inline them?

Yes, i was thinking adding a limit on the chain length as an option. Probably do this a a follow up -- which also needs to look at profile data. (Note that Partialinlininer is currently not enabled in any optimization pipeline yet).

The existing pattern handles if (a || b || c ...) case, but it may not be easy to extend for cases like (a && b && c) and ((a && b) || c). Basically, we want to find a collection of bbs with small sizes starting from entry. The bb collection only have two exits. one of them is ReturnBlock and the other is NonReturnBlock. The edges inside of the collection can be of any pattern.

a&b&c case is already partially handled -- the partial inliner can partially peel the first iteration which may be good enough. However I do like your suggestion of computing the entry predicate region idea.

[eraman]

Yes, i was thinking adding a limit on the chain length as an option. Probably do this a a follow up -- which also needs to look at profile data. (Note that Partialinlininer is currently not enabled in any optimization pipeline yet).

I think instruction count will be better limit than chain length.

lib/Transforms/IPO/PartialInlining.cpp

334

This will look cleaner if this is made to while (CommonSucc) and peel the else part outside of the while loop (since you always break from there). Perhaps you can even peel the first iteration out of the loop: first find the candidate non-return block, then walk the chain in a loop as long as one successor equals non-return block and the other succesor has a single predecessor. Some more comments will also increase readability.

[davidxl]

Address review feedback. The new version now can handle more general conditions (see test cases).

[davide]

lib/Transforms/IPO/PartialInlining.cpp

43

typo, partially inlined. BTW, do you have numbers for different values of MaxNumInlineBlocks ?

66–68

SmallVector<> maybe

118–129

These could probably live in BasicBlockUtils.

132–137

Also this one. Please remove else after return

160–164

Can you add comments about why we break in this case? (i.e. we don't outline blocks without exactly two successors, etc...)

205–207

braces unneeded.

330

explicit type maybe?

337

explicit type maybe?

346–348

return BB == NewReturnBlock || NewEntries.count(BB) , maybe?

382–385

unrelated change, also braces are unneeded.

408

unrelated?

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

43

Currently it is set to a high number for stress testing purpose. When partial inliner is enabled in the pipeline, the default value needs to be tuned depending on opt level.

118–129

will leave it here for now.

132–137

removed else.

[davidxl]

addressed review comments.

[wmi]

Some minor comments.

lib/Transforms/IPO/PartialInlining.cpp

68

I find NonReturnBlockPreds has been set but not used.

134

we can change the else if to if.

226–227

should we use an assertion here instead of an early exit, so we can catch mistake instead of giving up silently.

350–353

ToBeInlined(NewNonReturnBlock) is also false, so can we just do:

for (BasicBlock &BB : *DuplicateFunction)
  if (!ToBeInlined(&BB))
    ToExtract.push_back(&BB);

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

68

Removed

226–227

No. This is designed in way such that the candidate found may be illegal (to simplify the code), so a check here is necessary.

350–353

NewNonReturnBlock needs to be the head of the blocks to be cloned. It is already pushed in, so the filter is needed.

[davidxl]

Address Wei's review feedback. Also fixed a bug found in clang self build test.

[wmi]

LGTM.

[eraman]

lib/Transforms/IPO/PartialInlining.cpp

42

Expressing this as instructions will make this somewhat similar to inlining cost.

145

Cleaner to do return std::make_tuple(Succ1, Succ2) here and the following if

181

No need to initialize it here since GetCommonSucc would return a (nullptr, nullptr) tuple if there is no triangle shape.

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

42

This will be done in a follow up patch before the final patch that turn it on.

[davidxl]

Address review comments.

[eraman]

lib/Transforms/IPO/PartialInlining.cpp

160

Some comment here will be useful. You want to keep space for 1 return block and the CurrEntry and that's why you've the above. I initially thought the condition was wrong.

244

The HasNonEntryPred lambda checks for the presence of preds in Entries (the DenseSet above), but you're not updating the Entries in this loop when you insert Cand into OutliningInfo->Entries. So this loop will return after 1 iteration right?

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

160

Fixed. Added a member function.

244

good catch.Fixed.

[davidxl]

Addressed review feedback by eraman.

[eraman]

LGTM.

Please add a comment to explain the condition for exiting the loop for exceeding MaxNumInlineBlocks.

lib/Transforms/IPO/PartialInlining.cpp

160

You should still add a comment in the member function about the +1 (the return block) and here (we still have a pending CurrEntry block).

[gyiu]

llvm/trunk/lib/Transforms/IPO/PartialInlining.cpp

299 ↗

(On Diff #96932)

@davidxl Please correct me if I'm wrong, but in this loop, if the first instruction in the BB is not a PHINode we give up? If that's the case, then wouldn't 'Phi' be the first and only Phi we'll find?

[davidxl]

llvm/trunk/lib/Transforms/IPO/PartialInlining.cpp

299 ↗

(On Diff #96932)

The method returns the first phi of the block if it exists. It should add this:

if (isa<DbgInfoIntrinsic>(I))

continue;

lib/Transforms/IPO/PartialInlining.cpp

 #define DEBUG_TYPE "partial-inlining"
 
 STATISTIC(NumPartialInlined, "Number of functions partially inlined");
 
 // Command line option to disable partial-inlining. The default is false:
 static cl::opt<bool>
     DisablePartialInlining("disable-partial-inlining", cl::init(false),
                            cl::Hidden, cl::desc("Disable partial ininling"));
 
+static cl::opt<unsigned> MaxNumInlineBlocks(
+    "max-num-inline-blocks", cl::init(5), cl::Hidden,
+    cl::desc("Max Number of Blocks  To be Partially Inlined"));
+
 // Command line option to set the maximum number of partial inlining allowed
 // for the module. The default value of -1 means no limit.
 static cl::opt<int> MaxNumPartialInlining(
     "max-partial-inlining", cl::init(-1), cl::Hidden, cl::ZeroOrMore,
     cl::desc("Max number of partial inlining. The default is unlimited"));
 
 namespace {
+
+struct FunctionOutliningInfo {
+  FunctionOutliningInfo()
+      : Entries(), ReturnBlock(nullptr), NonReturnBlock(nullptr),
+        ReturnBlockPreds() {}
+
+  // A set of blocks including the function entry that guard
+  // the region to be outlined.
+  SmallVector<BasicBlock *, 4> Entries;
+  // The return block that is not included in the outlined region.
+  BasicBlock *ReturnBlock;
+  // The dominating block of the region ot be outlined.
+  BasicBlock *NonReturnBlock;
+  // The set of blocks in Entries that that are predecessors to ReturnBlock
+  SmallVector<BasicBlock *, 4> ReturnBlockPreds;
+};
+
 struct PartialInlinerImpl {
   PartialInlinerImpl(InlineFunctionInfo IFI) : IFI(std::move(IFI)) {}
   bool run(Module &M);
   Function *unswitchFunction(Function *F);
 
+  std::unique_ptr<FunctionOutliningInfo> computeOutliningInfo(Function *F);
+
 private:
   InlineFunctionInfo IFI;
   int NumPartialInlining = 0;
 
   bool IsLimitReached() {
     return (MaxNumPartialInlining != -1 &&
             NumPartialInlining >= MaxNumPartialInlining);
   }
 };
+
 struct PartialInlinerLegacyPass : public ModulePass {
   static char ID; // Pass identification, replacement for typeid
   PartialInlinerLegacyPass() : ModulePass(ID) {
     initializePartialInlinerLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
   }
   bool runOnModule(Module &M) override {
     if (skipModule(M))
       return false;
 
     AssumptionCacheTracker *ACT = &getAnalysis<AssumptionCacheTracker>();
     std::function<AssumptionCache &(Function &)> GetAssumptionCache =
         [&ACT](Function &F) -> AssumptionCache & {
       return ACT->getAssumptionCache(F);
     };
     InlineFunctionInfo IFI(nullptr, &GetAssumptionCache);
     return PartialInlinerImpl(IFI).run(M);
   }
 };
 }
 
-Function *PartialInlinerImpl::unswitchFunction(Function *F) {
-  // First, verify that this function is an unswitching candidate...
-  if (F->hasAddressTaken())
-    return nullptr;
-
+std::unique_ptr<FunctionOutliningInfo>
+PartialInlinerImpl::computeOutliningInfo(Function *F) {
   BasicBlock *EntryBlock = &F->front();
   BranchInst *BR = dyn_cast<BranchInst>(EntryBlock->getTerminator());
   if (!BR || BR->isUnconditional())
-    return nullptr;
+    return std::unique_ptr<FunctionOutliningInfo>();
+
+  // Returns true if Succ is BB's successor
+  auto IsSuccessor = [](BasicBlock *Succ, BasicBlock *BB) {
+    return is_contained(successors(BB), Succ);
+  };
+
+  auto SuccSize = [](BasicBlock *BB) {
+    return std::distance(succ_begin(BB), succ_end(BB));
+  };
+
+  auto IsReturnBlock = [](BasicBlock *BB) {
+    TerminatorInst *TI = BB->getTerminator();
+    return isa<ReturnInst>(TI);
+  };
+
+  auto GetReturnBlock = [=](BasicBlock *Succ1, BasicBlock *Succ2) {
+    if (IsReturnBlock(Succ1))
+      return std::make_tuple(Succ1, Succ2);
+    if (IsReturnBlock(Succ2))
+      return std::make_tuple(Succ2, Succ1);
+
+    return std::make_tuple<BasicBlock *, BasicBlock *>(nullptr, nullptr);
+  };
+
+  // Detect a triangular shape:
+  auto GetCommonSucc = [=](BasicBlock *Succ1, BasicBlock *Succ2) {
+    BasicBlock *CommSucc = nullptr;
+    BasicBlock *OtherSucc = nullptr;
+    if (IsSuccessor(Succ1, Succ2)) {
+      CommSucc = Succ1;
+      OtherSucc = Succ2;
+    } else if (IsSuccessor(Succ2, Succ1)) {
+      CommSucc = Succ2;
+      OtherSucc = Succ1;
+    }
+    return std::make_tuple(CommSucc, OtherSucc);
+  };
+
+  std::unique_ptr<FunctionOutliningInfo> OutliningInfo =
+      llvm::make_unique<FunctionOutliningInfo>();
+
+  BasicBlock *CurrEntry = EntryBlock;
+  bool CandidateFound = false;
+  do {
+    if (SuccSize(CurrEntry) != 2)
+      break;
+
+    // The number of blocks to be inlined is exceeding limit
+    if (OutliningInfo->Entries.size() + 1 >= MaxNumInlineBlocks)
+      break;
+
+    BasicBlock *Succ1 = *succ_begin(CurrEntry);
+    BasicBlock *Succ2 = *(succ_begin(CurrEntry) + 1);
+
+    BasicBlock *ReturnBlock, *NonReturnBlock;
+    std::tie(ReturnBlock, NonReturnBlock) = GetReturnBlock(Succ1, Succ2);
+
+    if (ReturnBlock) {
+      OutliningInfo->Entries.push_back(CurrEntry);
+      OutliningInfo->ReturnBlock = ReturnBlock;
+      OutliningInfo->NonReturnBlock = NonReturnBlock;
+      CandidateFound = true;
+      break;
+    }
+
+    BasicBlock *CommSucc = nullptr;
+    BasicBlock *OtherSucc = nullptr;
+    std::tie(CommSucc, OtherSucc) = GetCommonSucc(Succ1, Succ2);
+
+    if (!CommSucc)
+      break;
+
+    OutliningInfo->Entries.push_back(CurrEntry);
+    CurrEntry = OtherSucc;
+
+  } while (true);
+
+  if (!CandidateFound)
+    return std::unique_ptr<FunctionOutliningInfo>();
+
+  // Do sanity check of the entries: threre should not
+  // be any successors (not in the entry set) other than
+  // {ReturnBlock, NonReturnBlock}
+  assert(OutliningInfo->Entries[0] == &F->front());
+  DenseSet<BasicBlock *> Entries;
+  for (BasicBlock *E : OutliningInfo->Entries)
+    Entries.insert(E);
+
+  // Returns true of BB has Predecessor which is not
+  // in Entries set.
+  auto HasNonEntryPred = [Entries](BasicBlock *BB) {
+    for (auto Pred : predecessors(BB)) {
+      if (!Entries.count(Pred))
+        return true;
+    }
+    return false;
+  };
+  auto CheckAndNormalizeCandidate = [Entries, HasNonEntryPred](FunctionOutliningInfo *OutliningInfo) {
+    for (BasicBlock *E : OutliningInfo->Entries) {
+      for (auto Succ : successors(E)) {
+        if (Entries.count(Succ))
+          continue;
+        if (Succ == OutliningInfo->ReturnBlock)
+          OutliningInfo->ReturnBlockPreds.push_back(E);
+        else if (Succ != OutliningInfo->NonReturnBlock)
+          return false;
+      }
+      // There should not be any outside incoming edges either:
+      if (HasNonEntryPred(E))
+        return false;
+    }
+    return true;
+  };
+
+  if (!CheckAndNormalizeCandidate(OutliningInfo.get()))
+    return std::unique_ptr<FunctionOutliningInfo>();
+
+  // Now further growing the candidate's inlining region by
+  // peeling off dominating blocks from the outlining region:
+  while (OutliningInfo->Entries.size() + 1 < MaxNumInlineBlocks) {
+    BasicBlock *Cand = OutliningInfo->NonReturnBlock;
+    if (SuccSize(Cand) != 2)
+      break;
+
+    if (HasNonEntryPred(Cand))
+      break;
+
+    BasicBlock *Succ1 = *succ_begin(Cand);
+    BasicBlock *Succ2 = *(succ_begin(Cand) + 1);
+
+    BasicBlock *ReturnBlock, *NonReturnBlock;
+    std::tie(ReturnBlock, NonReturnBlock) = GetReturnBlock(Succ1, Succ2);
+    if (!ReturnBlock || ReturnBlock != OutliningInfo->ReturnBlock)
+      break;
+
+    if (NonReturnBlock->getSinglePredecessor() != Cand)
+      break;
+
+    // Now grow and update OutlininigInfo:
+    OutliningInfo->Entries.push_back(Cand);
+    OutliningInfo->NonReturnBlock = NonReturnBlock;
+    OutliningInfo->ReturnBlockPreds.push_back(Cand);
+  }
 
-  BasicBlock *ReturnBlock = nullptr;
-  BasicBlock *NonReturnBlock = nullptr;
-  unsigned ReturnCount = 0;
-  for (BasicBlock *BB : successors(EntryBlock)) {
-    if (isa<ReturnInst>(BB->getTerminator())) {
-      ReturnBlock = BB;
-      ReturnCount++;
-    } else
-      NonReturnBlock = BB;
-  }
+  return OutliningInfo;
 
-  if (ReturnCount != 1)
+Function *PartialInlinerImpl::unswitchFunction(Function *F) {
+
+  if (F->hasAddressTaken())
+    return nullptr;
+
+  std::unique_ptr<FunctionOutliningInfo> OutliningInfo =
+      computeOutliningInfo(F);
+
+  if (!OutliningInfo)
     return nullptr;
 
   // Clone the function, so that we can hack away on it.
   ValueToValueMapTy VMap;
   Function *DuplicateFunction = CloneFunction(F, VMap);
   DuplicateFunction->setLinkage(GlobalValue::InternalLinkage);
-  BasicBlock *NewEntryBlock = cast<BasicBlock>(VMap[EntryBlock]);
-  BasicBlock *NewReturnBlock = cast<BasicBlock>(VMap[ReturnBlock]);
-  BasicBlock *NewNonReturnBlock = cast<BasicBlock>(VMap[NonReturnBlock]);
+  BasicBlock *NewReturnBlock =
+      cast<BasicBlock>(VMap[OutliningInfo->ReturnBlock]);
+  BasicBlock *NewNonReturnBlock =
+      cast<BasicBlock>(VMap[OutliningInfo->NonReturnBlock]);
+  DenseSet<BasicBlock *> NewEntries;
+  for (BasicBlock *BB : OutliningInfo->Entries) {
+    NewEntries.insert(cast<BasicBlock>(VMap[BB]));
+  }
 
   // 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);
 
+  auto getFirstPHI = [](BasicBlock *BB) {
+    BasicBlock::iterator I = BB->begin();
+    PHINode *FirstPhi = nullptr;
+    while (I != BB->end()) {
+      PHINode *Phi = dyn_cast<PHINode>(I);
+      if (!Phi)
+        break;
+      if (!FirstPhi) {
+        FirstPhi = Phi;
+        break;
+      }
+    }
+    return FirstPhi;
+  };
   // 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;
+  // only split block when necessary:
+  PHINode *FirstPhi = getFirstPHI(PreReturn);
+  unsigned NumPredsFromEntries = OutliningInfo->ReturnBlockPreds.size();
+  if (FirstPhi && FirstPhi->getNumIncomingValues() > NumPredsFromEntries + 1) {
+
     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);
+      PHINode *RetPhi =
+          PHINode::Create(OldPhi->getType(), NumPredsFromEntries + 1, "", Ins);
       OldPhi->replaceAllUsesWith(RetPhi);
       Ins = NewReturnBlock->getFirstNonPHI();
 
       RetPhi->addIncoming(&*I, PreReturn);
-    RetPhi->addIncoming(OldPhi->getIncomingValueForBlock(NewEntryBlock),
-                        NewEntryBlock);
-    OldPhi->removeIncomingValue(NewEntryBlock);
-
+      for (BasicBlock *E : OutliningInfo->ReturnBlockPreds) {
+        BasicBlock *NewE = cast<BasicBlock>(VMap[E]);
+        RetPhi->addIncoming(OldPhi->getIncomingValueForBlock(NewE), NewE);
+        OldPhi->removeIncomingValue(NewE);
+      }
       ++I;
     }
-  NewEntryBlock->getTerminator()->replaceUsesOfWith(PreReturn, NewReturnBlock);
+    for (auto E : OutliningInfo->ReturnBlockPreds) {
+      BasicBlock *NewE = cast<BasicBlock>(VMap[E]);
+      NewE->getTerminator()->replaceUsesOfWith(PreReturn, NewReturnBlock);
+    }
+  }
 
+  // Returns true if the block is to be partial inlined into the caller
+  // (i.e. not to be extracted to the out of line function)
+  auto ToBeInlined = [=](BasicBlock *BB) {
+    return BB == NewReturnBlock || NewEntries.count(BB);
+  };
   // 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)
+    if (!ToBeInlined(&BB) && &BB != NewNonReturnBlock)
       ToExtract.push_back(&BB);
 
   // The CodeExtractor needs a dominator tree.
   DominatorTree DT;
   DT.recalculate(*DuplicateFunction);
 
   // Manually calculate a BlockFrequencyInfo and BranchProbabilityInfo.
   LoopInfo LI(DT);
   BranchProbabilityInfo BPI(*DuplicateFunction, LI);
for (User *User : Users) {
     CallSite CS;
     if (CallInst *CI = dyn_cast<CallInst>(User))
       CS = CallSite(CI);
     else if (InvokeInst *II = dyn_cast<InvokeInst>(User))
       CS = CallSite(II);
     else
       llvm_unreachable("All uses must be calls");
 
     if (IsLimitReached())
       continue;
+
     NumPartialInlining++;
 
     OptimizationRemarkEmitter ORE(CS.getCaller());
     DebugLoc DLoc = CS.getInstruction()->getDebugLoc();
     BasicBlock *Block = CS.getParent();
     ORE.emit(OptimizationRemark(DEBUG_TYPE, "PartiallyInlined", DLoc, Block)
              << ore::NV("Callee", F) << " partially inlined into "
              << ore::NV("Caller", CS.getCaller()));
 
     InlineFunction(CS, IFI);
   }
 
   // Ditch the duplicate, since we're done with it, and rewrite all remaining
   // users (function pointers, etc.) back to the original function.
   DuplicateFunction->replaceAllUsesWith(F);
   DuplicateFunction->eraseFromParent();
 
   ++NumPartialInlined;
 
   return ExtractedFunction;
 }
 
 bool PartialInlinerImpl::run(Module &M) {
   if (DisablePartialInlining)
     return false;
 
   std::vector<Function *> Worklist;
   Worklist.reserve(M.size());
   for (Function &F : M)
     if (!F.use_empty() && !F.isDeclaration())
       Worklist.push_back(&F);
 

test/Transforms/CodeExtractor/MultipleExitBranchProb.ll

-; RUN: opt < %s -partial-inliner -S | FileCheck %s
+; RUN: opt < %s -partial-inliner -max-num-inline-blocks=2 -S | FileCheck %s
 
 ; This test checks to make sure that CodeExtractor updates
 ;  the exit branch probabilities for multiple exit blocks.
 
 define i32 @inlinedFunc(i1 %cond) !prof !1 {
 entry:
   br i1 %cond, label %if.then, label %return, !prof !2
 if.then:
   br i1 %cond, label %return, label %return.2, !prof !3
 return.2:
   ret i32 10
 return:             ; preds = %entry
   ret i32 0
 }
 
 
 define internal i32 @dummyCaller(i1 %cond) !prof !1 {
 entry:
 %val = call i32 @inlinedFunc(i1 %cond)
 ret i32 %val
 
 ; CHECK-LABEL: @dummyCaller
 ; CHECK: call
-; CHECK-NEXT: br i1 {{.*}}!prof [[COUNT1:![0-9]+]]
+; CHECK-NEXT: br i1 {{.*}}return.i{{.*}}return.2{{.*}}!prof [[COUNT1:![0-9]+]]
 }
 
 !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}

test/Transforms/CodeExtractor/PartialInlineAnd.ll

+; RUN: opt < %s -partial-inliner -S | FileCheck %s
+; RUN: opt < %s -passes=partial-inliner -S | FileCheck %s
+; RUN: opt < %s -partial-inliner -max-num-inline-blocks=2 -S | FileCheck --check-prefix=LIMIT %s
+; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=2 -S | FileCheck  --check-prefix=LIMIT %s
+
+; Function Attrs: nounwind uwtable
+define i32 @bar(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = icmp slt i32 %arg, 0
+  br i1 %tmp, label %bb1, label %bb5
+
+bb1:                                              ; preds = %bb
+  %tmp2 = tail call i32 (...) @channels() #2
+  %tmp3 = icmp slt i32 %tmp2, %arg
+  br i1 %tmp3, label %bb4, label %bb5
+
+bb4:                                              ; preds = %bb1
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  br label %bb5
+
+bb5:                                              ; preds = %bb4, %bb1, %bb
+  %tmp6 = phi i32 [ 0, %bb4 ], [ 1, %bb1 ], [ 1, %bb ]
+  ret i32 %tmp6
+}
+
+declare i32 @channels(...) local_unnamed_addr #1
+
+declare void @foo(...) local_unnamed_addr #1
+
+; Function Attrs: nounwind uwtable
+define i32 @dummy_caller(i32 %arg) local_unnamed_addr #0 {
+bb:
+; CHECK-LABEL: @dummy_caller
+; CHECK: br i1
+; CHECK: br i1
+; CHECK: call void @bar.1_
+; LIMIT-LABEL: @dummy_caller
+; LIMIT: br i1
+; LIMIT-NOT: br
+; LIMIT: call void @bar.1_
+  %tmp = tail call i32 @bar(i32 %arg)
+  ret i32 %tmp
+}
+
+attributes #0 = { nounwind }
+attributes #1 = { nounwind }
+attributes #2 = { nounwind }
+

test/Transforms/CodeExtractor/PartialInlineAndOr.ll

+; RUN: opt < %s -partial-inliner -S | FileCheck %s
+; RUN: opt < %s -passes=partial-inliner -S | FileCheck %s
+; RUN: opt < %s -partial-inliner -max-num-inline-blocks=3 -S | FileCheck --check-prefix=LIMIT %s
+; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=3 -S | FileCheck  --check-prefix=LIMIT %s
+
+; Function Attrs: nounwind uwtable
+define i32 @bar(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = icmp slt i32 %arg, 0
+  br i1 %tmp, label %bb1, label %bb4
+
+bb1:                                              ; preds = %bb
+  %tmp2 = tail call i32 (...) @n() #2
+  %tmp3 = icmp slt i32 %tmp2, %arg
+  br i1 %tmp3, label %bb7, label %bb4
+
+bb4:                                              ; preds = %bb1, %bb
+  %tmp5 = tail call i32 (...) @m() #2
+  %tmp6 = icmp slt i32 %tmp5, %arg
+  br i1 %tmp6, label %bb7, label %bb8
+
+bb7:                                              ; preds = %bb4, %bb1
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  br label %bb8
+
+bb8:                                              ; preds = %bb7, %bb4
+  %tmp9 = phi i32 [ 0, %bb7 ], [ 1, %bb4 ]
+  ret i32 %tmp9
+}
+
+declare i32 @n(...) local_unnamed_addr #1
+
+declare i32 @m(...) local_unnamed_addr #1
+
+declare void @foo(...) local_unnamed_addr #1
+
+; Function Attrs: nounwind uwtable
+define i32 @dummy_caller(i32 %arg) local_unnamed_addr #0 {
+bb:
+; CHECK-LABEL: @dummy_caller
+; CHECK: br i1
+; CHECK: br i1
+; CHECK: br i1
+; CHECK: call void @bar.1_
+; LIMIT-LABEL: @dummy_caller
+; LIMIT-NOT: br i1
+; LIMIT: call i32 @bar
+  %tmp = tail call i32 @bar(i32 %arg)
+  ret i32 %tmp
+}
+
+attributes #0 = { nounwind } 
+attributes #1 = { nounwind }
+attributes #2 = { nounwind }
+

test/Transforms/CodeExtractor/PartialInlineOr.ll

+; RUN: opt < %s -partial-inliner -S | FileCheck %s
+; RUN: opt < %s -passes=partial-inliner -S | FileCheck %s
+; RUN: opt < %s -partial-inliner -max-num-inline-blocks=2 -S | FileCheck --check-prefix=LIMIT %s
+; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=2 -S | FileCheck  --check-prefix=LIMIT %s
+
+; Function Attrs: noinline nounwind uwtable
+define i32 @bar(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = icmp slt i32 %arg, 0
+  br i1 %tmp, label %bb4, label %bb1
+
+bb1:                                              ; preds = %bb
+  %tmp2 = tail call i32 (...) @channels() #1
+  %tmp3 = icmp slt i32 %tmp2, %arg
+  br i1 %tmp3, label %bb4, label %bb5
+
+bb4:                                              ; preds = %bb1, %bb
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  br label %bb5
+
+bb5:                                              ; preds = %bb4, %bb1
+  %.0 = phi i32 [ 0, %bb4 ], [ 1, %bb1 ]
+  ret i32 %.0
+}
+
+declare i32 @channels(...) local_unnamed_addr
+
+declare void @foo(...) local_unnamed_addr
+
+; Function Attrs: noinline nounwind uwtable
+define i32 @dummy_caller(i32 %arg) local_unnamed_addr #0 {
+bb:
+; CHECK-LABEL: @dummy_caller
+; CHECK: br i1
+; CHECK: br i1
+; CHECK: call void @bar.2_
+; LIMIT-LABEL: @dummy_caller
+; LIMIT-NOT: br
+; LIMIT: call i32 @bar(
+  %tmp = tail call i32 @bar(i32 %arg)
+  ret i32 %tmp
+}
+
+define i32 @bar_multi_ret(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = icmp slt i32 %arg, 0
+  br i1 %tmp, label %bb4, label %bb1
+
+bb1:                                              ; preds = %bb
+  %tmp2 = tail call i32 (...) @channels() #1
+  %tmp3 = icmp slt i32 %tmp2, %arg
+  br i1 %tmp3, label %bb4, label %bb5
+
+bb4:                                              ; preds = %bb1, %bb
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  tail call void (...) @foo() #1
+  %tmp4 = icmp slt i32 %arg, 10
+  br i1 %tmp4, label %bb6, label %bb5
+bb6:
+  tail call void (...) @foo() #1
+  %tmp5 = icmp slt i32 %arg, 3
+  br i1 %tmp5, label %bb7, label %bb5
+bb7:
+  tail call void (...) @foo() #1
+  br label %bb8
+bb8:
+  ret i32 0 
+
+bb5:                                              ; preds = %bb4, %bb1
+  %.0 = phi i32 [ 0, %bb4 ], [ 1, %bb1 ], [0, %bb6]
+  ret i32 %.0
+}
+
+define i32 @dummy_caller2(i32 %arg) local_unnamed_addr #0 {
+; CHECK: br i1
+; CHECK: br i1
+; CHECK: call {{.*}} @bar_multi_ret.1_
+  %tmp = tail call i32 @bar_multi_ret(i32 %arg)
+  ret i32 %tmp
+}
+
+attributes #0 = { noinline nounwind uwtable }
+attributes #1 = { nounwind }
+
+!llvm.ident = !{!0}
+
+!0 = !{!"clang version 5.0.0 (trunk 300576)"}

test/Transforms/CodeExtractor/PartialInlineOrAnd.ll

+; RUN: opt < %s -partial-inliner -S | FileCheck %s
+; RUN: opt < %s -passes=partial-inliner -S | FileCheck %s
+; RUN: opt < %s -partial-inliner -max-num-inline-blocks=3 -S | FileCheck --check-prefix=LIMIT3 %s
+; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=3 -S | FileCheck  --check-prefix=LIMIT3 %s
+; RUN: opt < %s -partial-inliner -max-num-inline-blocks=2 -S | FileCheck --check-prefix=LIMIT2 %s
+; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=2 -S | FileCheck  --check-prefix=LIMIT2 %s
+
+
+; Function Attrs: nounwind uwtable
+define i32 @bar(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = icmp slt i32 %arg, 0
+  br i1 %tmp, label %bb4, label %bb1
+
+bb1:                                              ; preds = %bb
+  %tmp2 = tail call i32 (...) @n() #2
+  %tmp3 = icmp slt i32 %tmp2, %arg
+  br i1 %tmp3, label %bb4, label %bb8
+
+bb4:                                              ; preds = %bb1, %bb
+  %tmp5 = tail call i32 (...) @m() #2
+  %tmp6 = icmp sgt i32 %tmp5, %arg
+  br i1 %tmp6, label %bb7, label %bb8
+
+bb7:                                              ; preds = %bb4
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  tail call void (...) @foo() #2
+  br label %bb8
+
+bb8:                                              ; preds = %bb7, %bb4, %bb1
+  %tmp9 = phi i32 [ 0, %bb7 ], [ 1, %bb4 ], [ 1, %bb1 ]
+  ret i32 %tmp9
+}
+
+declare i32 @n(...) local_unnamed_addr #1
+
+declare i32 @m(...) local_unnamed_addr #1
+
+declare void @foo(...) local_unnamed_addr #1
+
+; Function Attrs: nounwind uwtable
+define i32 @dummy_caller(i32 %arg) local_unnamed_addr #0 {
+bb:
+; CHECK-LABEL: @dummy_caller
+; CHECK: br i1
+; CHECK: br i1
+; CHECK: br i1
+; CHECK: call void @bar.1_
+; LIMIT3-LABEL: @dummy_caller
+; LIMIT3: br i1
+; LIMIT3: br i1
+; LIMIT3-NOT: br i1
+; LIMIT3: call void @bar.1_
+; LIMIT2-LABEL: @dummy_caller
+; LIMIT2-NOT: br i1
+; LIMIT2: call i32 @bar(
+  %tmp = tail call i32 @bar(i32 %arg)
+  ret i32 %tmp
+}
+
+attributes #0 = { nounwind } 
+attributes #1 = { nounwind }
+attributes #2 = { nounwind }
+

test/Transforms/CodeExtractor/SingleCondition.ll

+; RUN: opt < %s -partial-inliner -S  | FileCheck %s
+; RUN: opt < %s -passes=partial-inliner -S  | FileCheck %s
+
+define internal i32 @inlinedFunc(i1 %cond, i32* align 4 %align.val) {
+entry:
+  br i1 %cond, label %if.then, label %return
+if.then:
+  ; Dummy store to have more than 0 uses
+  store i32 10, i32* %align.val, align 4
+  br label %return
+return:             ; preds = %entry
+  ret i32 0
+}
+
+define internal i32 @dummyCaller(i1 %cond, i32* align 2 %align.val) {
+entry:
+; CHECK-LABEL: @dummyCaller
+; CHECK: br
+; CHECK: call void @inlinedFunc.1_ 
+  %val = call i32 @inlinedFunc(i1 %cond, i32* %align.val)
+  ret i32 %val
+}
+

test/Transforms/CodeExtractor/unreachable-block.ll

 ; RUN: opt -S -partial-inliner %s | FileCheck %s
 
 ; CHECK-LABEL: define void @dipsy(
 ; CHECK-NEXT:   call void @tinkywinky.1_ontrue()
 ; CHECK-NEXT:   call void @patatuccio()
 ; CHECK-NEXT:   ret void
 ; CHECK-NEXT: }
 
 ; CHECK-LABEL: define internal void @tinkywinky.1_ontrue() {
 ; CHECK-NEXT: newFuncRoot:
 ; CHECK-NEXT:   br label %ontrue
-; CHECK: .exitStub:
+; CHECK: onfalse{{.*}}:
 ; CHECK-NEXT:   ret void
 ; CHECK: ontrue:
 ; CHECK-NEXT:   call void @patatino()
-; CHECK-NEXT:   br label %onfalse
-; CHECK: onfalse:
-; CHECK-NEXT:   br label %.exitStub
+; CHECK-NEXT:   br label %onfalse{{.*}}
 ; CHECK-NEXT: }
 
 declare void @patatino()
 declare void @patatuccio()
 
 define fastcc void @tinkywinky() {
   br i1 true, label %ontrue, label %onfalse
 ontrue: