Split IndirectBr critical edges before PGO gen/use passes.

[hjyamauchi]

Summary

The PGO gen/use passes currently fail with an assert failure if there's a
critical edge whose source is an IndirectBr instruction and that edge
needs to be instrumented.

To avoid this in certain cases, split IndirectBr critical edges in the PGO
gen/use passes. This works for blocks with single indirectbr predecessors,
but not for those with multiple indirectbr predecessors (splitting an
IndirectBr critical edge isn't always possible.)

[efriedma]

The PGO gen/use passes currently fail with an assert failure if there's a critical edge whose source is an IndirectBr instruction.

findIBRPredecessor will fail for blocks with multiple multiple indirectbr predecessors. And then you're back at square one, with an assertion failure.

[davidxl]

1. As Eli pointed out, findIBRpredecessor may return null when there are multiple such predecessors, so perhaps create a new interface to return all such predecessors?

2. it is probably better to split out the refactoring part of the patch.

lib/Passes/PassBuilder.cpp

515

this is common to Use and Gen, so it is possible to move it into the previous SCC pipeline after InstCombinePass.

[hjyamauchi]

In D40699#941591, @efriedma wrote:

The PGO gen/use passes currently fail with an assert failure if there's a critical edge whose source is an IndirectBr instruction.

findIBRPredecessor will fail for blocks with multiple multiple indirectbr predecessors. And then you're back at square one, with an assertion failure.

True. This change (and SplitIndirectBrCriticalEdges) isn't intended to handle all possible cases (see the comment in BasicBlockUtils.h). It's an improvement in that it can handle single indirectbr predecessor cases.

[hjyamauchi]

lib/Passes/PassBuilder.cpp

515

Do you mean moving this to after line 498 above? If so, would it not work if isOptimizingForSize() is true?

[hjyamauchi]

Created https://reviews.llvm.org/D40750 for the refactoring part.

[hjyamauchi]

Rebased after D40750.

Changed the SplitIndirectBrCriticalEdges pass to a function pass.

Moved the SplitIndirectBrCriticalEdges pass into the previous SCC pipeline.

[efriedma]

It's an improvement in that it can handle single indirectbr predecessor cases.

It's not good to kick the can down the road like this. If an IR optimization pass is crashing, you should fix it by fixing that pass, not by changing the default pass pipeline to avoid generating IR which triggers the crash.

[davidxl]

This can be done in the PGO pass itself before the MST is formed. BFI has update interface such as setBlockFreq. The update can use simple heuristic to distribute the frequency.

[hjyamauchi]

Split indirectbr critical edges in the PGO gen/use passes.

[hjyamauchi]

Split only the IndirectBr critical edges that need to be split rather than
all. This avoids changing the indirectbr.ll test.

Minor cleanup.

[davidxl]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

567

Why not splitting the edges before MST is computed?

[hjyamauchi]

Detect indirectbr critical non-MST edges in CFGMST and if none, avoid iterating
over the edges in FuncPGOInstrumentation().

[hjyamauchi]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

567

This is trying to not split unless it's necessary (that is, if it's a non-MST indirectbr critical edges.) For that, we need to compute the MST once before making a decision to split. We could split all indirectbr critical edges unconditionally before computing the MST, but I thought it may be somewhat nicer that way as it would minimize CFG changes (eg, the indirectbr.ll test previously had to be changed even though it didn't suffer from this assert failure.)

[davidxl]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

572

It seems cleaner to move the IndirectBrTargets collection inside MST's buildEdges method.

576

There seems no guarantee that the non-split indirectBR critical edges will be be in MST when cfgmst is recomputed.

[hjyamauchi]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

576

Good point. I'll go back to splitting all indirectBR critical edges before computing the MST (the previous diff).

[hjyamauchi]

Going back to Diff 125662.

[hjyamauchi]

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

572

This comment doesn't apply any more.

[davidxl]

lgtm

[hjyamauchi]

Rebased.

[efriedma]

You still never answered the question of what happens when there's multiple indirectbrs whose edges can't be split.

[hjyamauchi]

I think it will be an assert failure like before. Like said above, this patch isn't intended to handle that case.

include/llvm/InitializePasses.h

 void initializeSimpleInlinerPass(PassRegistry&);
 void initializeSimpleLoopUnswitchLegacyPassPass(PassRegistry&);
 void initializeSingleLoopExtractorPass(PassRegistry&);
 void initializeSinkingLegacyPassPass(PassRegistry&);
 void initializeSjLjEHPreparePass(PassRegistry&);
 void initializeSlotIndexesPass(PassRegistry&);
 void initializeSpeculativeExecutionLegacyPassPass(PassRegistry&);
 void initializeSpillPlacementPass(PassRegistry&);
+void initializeSplitIndirectBrCriticalEdgesLegacyPassPass(PassRegistry&);
 void initializeStackColoringPass(PassRegistry&);
 void initializeStackMapLivenessPass(PassRegistry&);
 void initializeStackProtectorPass(PassRegistry&);
 void initializeStackSlotColoringPass(PassRegistry&);
 void initializeStraightLineStrengthReducePass(PassRegistry&);
 void initializeStripDeadDebugInfoPass(PassRegistry&);
 void initializeStripDeadPrototypesLegacyPassPass(PassRegistry&);
 void initializeStripDebugDeclarePass(PassRegistry&);

include/llvm/LinkAllPasses.h

ForcePassLinking() {
       (void) llvm::createDomPrinterPass();
       (void) llvm::createDomOnlyViewerPass();
       (void) llvm::createDomViewerPass();
       (void) llvm::createGCOVProfilerPass();
       (void) llvm::createPGOInstrumentationGenLegacyPass();
       (void) llvm::createPGOInstrumentationUseLegacyPass();
       (void) llvm::createPGOIndirectCallPromotionLegacyPass();
       (void) llvm::createPGOMemOPSizeOptLegacyPass();
+      (void) llvm::createSplitIndirectBrCriticalEdgesLegacyPass();
       (void) llvm::createInstrProfilingLegacyPass();
       (void) llvm::createFunctionImportPass();
       (void) llvm::createFunctionInliningPass();
       (void) llvm::createAlwaysInlinerLegacyPass();
       (void) llvm::createGlobalDCEPass();
       (void) llvm::createGlobalOptimizerPass();
       (void) llvm::createGlobalsAAWrapperPass();
       (void) llvm::createGuardWideningPass();

include/llvm/Transforms/Instrumentation.h

 // PGO Instrumention
 ModulePass *createPGOInstrumentationGenLegacyPass();
 ModulePass *
 createPGOInstrumentationUseLegacyPass(StringRef Filename = StringRef(""));
 ModulePass *createPGOIndirectCallPromotionLegacyPass(bool InLTO = false,
                                                      bool SamplePGO = false);
 FunctionPass *createPGOMemOPSizeOptLegacyPass();
 
+ModulePass *createSplitIndirectBrCriticalEdgesLegacyPass();
+
 // Helper function to check if it is legal to promote indirect call \p Inst
 // to a direct call of function \p F. Stores the reason in \p Reason.
 bool isLegalToPromote(Instruction *Inst, Function *F, const char **Reason);
 
 // Helper function that transforms Inst (either an indirect-call instruction, or
 // an invoke instruction , to a conditional call to F. This is like:
 //     if (Inst.CalledValue == F)
 //        F(...);

include/llvm/Transforms/PGOInstrumentation.h

public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
 void setProfMetadata(Module *M, Instruction *TI, ArrayRef<uint64_t> EdgeCounts,
                      uint64_t MaxCount);
 
 void setIrrLoopHeaderMetadata(Module *M, Instruction *TI, uint64_t Count);
 
+/// The split-indirectbr-critical-edges pass.
+class SplitIndirectBrCriticalEdgesPass :
+      public PassInfoMixin<SplitIndirectBrCriticalEdgesPass> {
+public:
+  SplitIndirectBrCriticalEdgesPass();
+
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_PGOINSTRUMENTATION_H

include/llvm/Transforms/Utils/BasicBlockUtils.h

 /// entered from if the condition is true, and the block that will be
 /// entered if the condition is false.
 ///
 /// This does no checking to see if the true/false blocks have large or unsavory
 /// instructions in them.
 Value *GetIfCondition(BasicBlock *BB, BasicBlock *&IfTrue,
                       BasicBlock *&IfFalse);
 
+// Split critical edges where the source of the edge is an indirectbr
+// instruction. This isn't always possible, but we can handle some easy cases.
+// This is useful because MI is unable to split such critical edges,
+// which means it will not be able to sink instructions along those edges.
+// This is especially painful for indirect branches with many successors, where
+// we end up having to prepare all outgoing values in the origin block.
+//
+// Our normal algorithm for splitting critical edges requires us to update
+// the outgoing edges of the edge origin block, but for an indirectbr this
+// is hard, since it would require finding and updating the block addresses
+// the indirect branch uses. But if a block only has a single indirectbr
+// predecessor, with the others being regular branches, we can do it in a
+// different way.
+// Say we have A -> D, B -> D, I -> D where only I -> D is an indirectbr.
+// We can split D into D0 and D1, where D0 contains only the PHIs from D,
+// and D1 is the D block body. We can then duplicate D0 as D0A and D0B, and
+// create the following structure:
+// A -> D0A, B -> D0A, I -> D0B, D0A -> D1, D0B -> D1
+bool SplitIndirectBrCriticalEdges(Function &F);
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_BASICBLOCKUTILS_H

lib/CodeGen/CodeGenPrepare.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopInfo.h"

 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/BypassSlowDivision.h"
-#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/SimplifyLibCalls.h"
-#include "llvm/Transforms/Utils/ValueMapper.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <iterator>
 #include <limits>
 #include <memory>
 #include <utility>
 #include <vector>
private:
     bool mergeSExts(Function &F);
     bool performAddressTypePromotion(
         Instruction *&Inst,
         bool AllowPromotionWithoutCommonHeader,
         bool HasPromoted, TypePromotionTransaction &TPT,
         SmallVectorImpl<Instruction *> &SpeculativelyMovedExts);
     bool splitBranchCondition(Function &F);
     bool simplifyOffsetableRelocate(Instruction &I);
-    bool splitIndirectCriticalEdges(Function &F);
   };
 
 } // end anonymous namespace
 
 char CodeGenPrepare::ID = 0;
 
 INITIALIZE_PASS_BEGIN(CodeGenPrepare, DEBUG_TYPE,
                       "Optimize for code generation", false, false)
bool CodeGenPrepare::runOnFunction(Function &F) {
   // find a node corresponding to the value.
   EverMadeChange |= placeDbgValues(F);
 
   if (!DisableBranchOpts)
     EverMadeChange |= splitBranchCondition(F);
 
   // Split some critical edges where one of the sources is an indirect branch,
   // to help generate sane code for PHIs involving such edges.
-  EverMadeChange |= splitIndirectCriticalEdges(F);
+  EverMadeChange |= SplitIndirectBrCriticalEdges(F);
 
   bool MadeChange = true;
   while (MadeChange) {
     MadeChange = false;
     SeenChainsForSExt.clear();
     ValToSExtendedUses.clear();
     RemovedInsts.clear();
     for (Function::iterator I = F.begin(); I != F.end(); ) {
if (DestBB == BB)
     return nullptr;
 
   if (!canMergeBlocks(BB, DestBB))
     DestBB = nullptr;
 
   return DestBB;
 }
 
-// Return the unique indirectbr predecessor of a block. This may return null
-// even if such a predecessor exists, if it's not useful for splitting.
-// If a predecessor is found, OtherPreds will contain all other (non-indirectbr)
-// predecessors of BB.
-static BasicBlock *
-findIBRPredecessor(BasicBlock *BB, SmallVectorImpl<BasicBlock *> &OtherPreds) {
-  // If the block doesn't have any PHIs, we don't care about it, since there's
-  // no point in splitting it.
-  PHINode *PN = dyn_cast<PHINode>(BB->begin());
-  if (!PN)
-    return nullptr;
-
-  // Verify we have exactly one IBR predecessor.
-  // Conservatively bail out if one of the other predecessors is not a "regular"
-  // terminator (that is, not a switch or a br).
-  BasicBlock *IBB = nullptr;
-  for (unsigned Pred = 0, E = PN->getNumIncomingValues(); Pred != E; ++Pred) {
-    BasicBlock *PredBB = PN->getIncomingBlock(Pred);
-    TerminatorInst *PredTerm = PredBB->getTerminator();
-    switch (PredTerm->getOpcode()) {
-    case Instruction::IndirectBr:
-      if (IBB)
-        return nullptr;
-      IBB = PredBB;
-      break;
-    case Instruction::Br:
-    case Instruction::Switch:
-      OtherPreds.push_back(PredBB);
-      continue;
-    default:
-      return nullptr;
-    }
-  }
-
-  return IBB;
-}
-
-// Split critical edges where the source of the edge is an indirectbr
-// instruction. This isn't always possible, but we can handle some easy cases.
-// This is useful because MI is unable to split such critical edges,
-// which means it will not be able to sink instructions along those edges.
-// This is especially painful for indirect branches with many successors, where
-// we end up having to prepare all outgoing values in the origin block.
-//
-// Our normal algorithm for splitting critical edges requires us to update
-// the outgoing edges of the edge origin block, but for an indirectbr this
-// is hard, since it would require finding and updating the block addresses
-// the indirect branch uses. But if a block only has a single indirectbr
-// predecessor, with the others being regular branches, we can do it in a
-// different way.
-// Say we have A -> D, B -> D, I -> D where only I -> D is an indirectbr.
-// We can split D into D0 and D1, where D0 contains only the PHIs from D,
-// and D1 is the D block body. We can then duplicate D0 as D0A and D0B, and
-// create the following structure:
-// A -> D0A, B -> D0A, I -> D0B, D0A -> D1, D0B -> D1
-bool CodeGenPrepare::splitIndirectCriticalEdges(Function &F) {
-  // Check whether the function has any indirectbrs, and collect which blocks
-  // they may jump to. Since most functions don't have indirect branches,
-  // this lowers the common case's overhead to O(Blocks) instead of O(Edges).
-  SmallSetVector<BasicBlock *, 16> Targets;
-  for (auto &BB : F) {
-    auto *IBI = dyn_cast<IndirectBrInst>(BB.getTerminator());
-    if (!IBI)
-      continue;
-
-    for (unsigned Succ = 0, E = IBI->getNumSuccessors(); Succ != E; ++Succ)
-      Targets.insert(IBI->getSuccessor(Succ));
-  }
-
-  if (Targets.empty())
-    return false;
-
-  bool Changed = false;
-  for (BasicBlock *Target : Targets) {
-    SmallVector<BasicBlock *, 16> OtherPreds;
-    BasicBlock *IBRPred = findIBRPredecessor(Target, OtherPreds);
-    // If we did not found an indirectbr, or the indirectbr is the only
-    // incoming edge, this isn't the kind of edge we're looking for.
-    if (!IBRPred || OtherPreds.empty())
-      continue;
-
-    // Don't even think about ehpads/landingpads.
-    Instruction *FirstNonPHI = Target->getFirstNonPHI();
-    if (FirstNonPHI->isEHPad() || Target->isLandingPad())
-      continue;
-
-    BasicBlock *BodyBlock = Target->splitBasicBlock(FirstNonPHI, ".split");
-    // It's possible Target was its own successor through an indirectbr.
-    // In this case, the indirectbr now comes from BodyBlock.
-    if (IBRPred == Target)
-      IBRPred = BodyBlock;
-
-    // At this point Target only has PHIs, and BodyBlock has the rest of the
-    // block's body. Create a copy of Target that will be used by the "direct"
-    // preds.
-    ValueToValueMapTy VMap;
-    BasicBlock *DirectSucc = CloneBasicBlock(Target, VMap, ".clone", &F);
-
-    for (BasicBlock *Pred : OtherPreds) {
-      // If the target is a loop to itself, then the terminator of the split
-      // block needs to be updated.
-      if (Pred == Target)
-        BodyBlock->getTerminator()->replaceUsesOfWith(Target, DirectSucc);
-      else
-        Pred->getTerminator()->replaceUsesOfWith(Target, DirectSucc);
-    }
-
-    // Ok, now fix up the PHIs. We know the two blocks only have PHIs, and that
-    // they are clones, so the number of PHIs are the same.
-    // (a) Remove the edge coming from IBRPred from the "Direct" PHI
-    // (b) Leave that as the only edge in the "Indirect" PHI.
-    // (c) Merge the two in the body block.
-    BasicBlock::iterator Indirect = Target->begin(),
-                         End = Target->getFirstNonPHI()->getIterator();
-    BasicBlock::iterator Direct = DirectSucc->begin();
-    BasicBlock::iterator MergeInsert = BodyBlock->getFirstInsertionPt();
-
-    assert(&*End == Target->getTerminator() &&
-           "Block was expected to only contain PHIs");
-
-    while (Indirect != End) {
-      PHINode *DirPHI = cast<PHINode>(Direct);
-      PHINode *IndPHI = cast<PHINode>(Indirect);
-
-      // Now, clean up - the direct block shouldn't get the indirect value,
-      // and vice versa.
-      DirPHI->removeIncomingValue(IBRPred);
-      Direct++;
-
-      // Advance the pointer here, to avoid invalidation issues when the old
-      // PHI is erased.
-      Indirect++;
-
-      PHINode *NewIndPHI = PHINode::Create(IndPHI->getType(), 1, "ind", IndPHI);
-      NewIndPHI->addIncoming(IndPHI->getIncomingValueForBlock(IBRPred),
-                             IBRPred);
-
-      // Create a PHI in the body block, to merge the direct and indirect
-      // predecessors.
-      PHINode *MergePHI =
-          PHINode::Create(IndPHI->getType(), 2, "merge", &*MergeInsert);
-      MergePHI->addIncoming(NewIndPHI, Target);
-      MergePHI->addIncoming(DirPHI, DirectSucc);
-
-      IndPHI->replaceAllUsesWith(MergePHI);
-      IndPHI->eraseFromParent();
-    }
-
-    Changed = true;
-  }
-
-  return Changed;
-}
-
 /// Eliminate blocks that contain only PHI nodes, debug info directives, and an
 /// unconditional branch. Passes before isel (e.g. LSR/loopsimplify) often split
 /// edges in ways that are non-optimal for isel. Start by eliminating these
 /// blocks so we can split them the way we want them.
 bool CodeGenPrepare::eliminateMostlyEmptyBlocks(Function &F) {
   SmallPtrSet<BasicBlock *, 16> Preheaders;
   SmallVector<Loop *, 16> LoopList(LI->begin(), LI->end());
   while (!LoopList.empty()) {

lib/Passes/PassBuilder.cpp

void PassBuilder::addPGOInstrPasses(ModulePassManager &MPM, bool DebugLogging,
   }
 
   // Delete anything that is now dead to make sure that we don't instrument
   // dead code. Instrumentation can end up keeping dead code around and
   // dramatically increase code size.
   MPM.addPass(GlobalDCEPass());
 
   if (RunProfileGen) {
+    // Split the indirectbr critical edges before the PGO instrumentation gen
+    // pass as that could change CFG enough to invalidate the BPI/BFI/MST if
+    // that's done during the pass.
+    MPM.addPass(SplitIndirectBrCriticalEdgesPass());
     MPM.addPass(PGOInstrumentationGen());
 
     FunctionPassManager FPM;
     FPM.addPass(createFunctionToLoopPassAdaptor(LoopRotatePass()));
     MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
 
     // Add the profile lowering pass.
     InstrProfOptions Options;
     if (!ProfileGenFile.empty())
       Options.InstrProfileOutput = ProfileGenFile;
     Options.DoCounterPromotion = true;
     MPM.addPass(InstrProfiling(Options));
   }
 
-  if (!ProfileUseFile.empty())
+  if (!ProfileUseFile.empty()) {
+    // Split the indirectbr critical edges before the PGO instrumentation use
+    // pass as that could change CFG enough to invalidate the BPI/BFI/MST if
+    // that's done during the pass.
+    MPM.addPass(SplitIndirectBrCriticalEdgesPass());
     MPM.addPass(PGOInstrumentationUse(ProfileUseFile));
   }
+}
 
 static InlineParams
 getInlineParamsFromOptLevel(PassBuilder::OptimizationLevel Level) {
   auto O3 = PassBuilder::O3;
   unsigned OptLevel = Level > O3 ? 2 : Level;
   unsigned SizeLevel = Level > O3 ? Level - O3 : 0;
   return getInlineParams(OptLevel, SizeLevel);
 }

lib/Passes/PassRegistry.def

 MODULE_PASS("ipsccp", IPSCCPPass())
 MODULE_PASS("lowertypetests", LowerTypeTestsPass())
 MODULE_PASS("name-anon-globals", NameAnonGlobalPass())
 MODULE_PASS("no-op-module", NoOpModulePass())
 MODULE_PASS("partial-inliner", PartialInlinerPass())
 MODULE_PASS("pgo-icall-prom", PGOIndirectCallPromotion())
 MODULE_PASS("pgo-instr-gen", PGOInstrumentationGen())
 MODULE_PASS("pgo-instr-use", PGOInstrumentationUse())
+MODULE_PASS("split-indirectbr-critical-edges",
+            SplitIndirectBrCriticalEdgesPass())
 MODULE_PASS("pre-isel-intrinsic-lowering", PreISelIntrinsicLoweringPass())
 MODULE_PASS("print-profile-summary", ProfileSummaryPrinterPass(dbgs()))
 MODULE_PASS("print-callgraph", CallGraphPrinterPass(dbgs()))
 MODULE_PASS("print", PrintModulePass(dbgs()))
 MODULE_PASS("print-lcg", LazyCallGraphPrinterPass(dbgs()))
 MODULE_PASS("print-lcg-dot", LazyCallGraphDOTPrinterPass(dbgs()))
 MODULE_PASS("rewrite-symbols", RewriteSymbolPass())
 MODULE_PASS("rpo-functionattrs", ReversePostOrderFunctionAttrsPass())

lib/Transforms/IPO/PassManagerBuilder.cpp

if (OptLevel > 0 && SizeLevel == 0 && !DisablePreInliner &&
     MPM.add(createFunctionInliningPass(IP));
     MPM.add(createSROAPass());
     MPM.add(createEarlyCSEPass());             // Catch trivial redundancies
     MPM.add(createCFGSimplificationPass());    // Merge & remove BBs
     MPM.add(createInstructionCombiningPass()); // Combine silly seq's
     addExtensionsToPM(EP_Peephole, MPM);
   }
   if (EnablePGOInstrGen) {
+    // Split the indirectbr critical edges before the PGO instrumentation gen
+    // pass as that could change CFG enough to invalidate the BPI/BFI/MST if
+    // that's done during the pass.
+    MPM.add(createSplitIndirectBrCriticalEdgesLegacyPass());
     MPM.add(createPGOInstrumentationGenLegacyPass());
     // Add the profile lowering pass.
     InstrProfOptions Options;
     if (!PGOInstrGen.empty())
       Options.InstrProfileOutput = PGOInstrGen;
     Options.DoCounterPromotion = true;
     MPM.add(createLoopRotatePass());
     MPM.add(createInstrProfilingLegacyPass(Options));
   }
-  if (!PGOInstrUse.empty())
+  if (!PGOInstrUse.empty()) {
+    // Split the indirectbr critical edges before the PGO instrumentation use
+    // pass as that could change CFG enough to invalidate the BPI/BFI/MST if
+    // that's done during the pass.
+    MPM.add(createSplitIndirectBrCriticalEdgesLegacyPass());
     MPM.add(createPGOInstrumentationUseLegacyPass(PGOInstrUse));
+  }
   // Indirect call promotion that promotes intra-module targets only.
   // For ThinLTO this is done earlier due to interactions with globalopt
   // for imported functions. We don't run this at -O0.
   if (OptLevel > 0)
     MPM.add(
         createPGOIndirectCallPromotionLegacyPass(false, !PGOSampleUse.empty()));
 }
 void PassManagerBuilder::addFunctionSimplificationPasses(

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

private:
 
   bool runOnModule(Module &M) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<BlockFrequencyInfoWrapperPass>();
   }
 };
 
+class SplitIndirectBrCriticalEdgesLegacyPass : public ModulePass {
+public:
+  static char ID;
+
+  SplitIndirectBrCriticalEdgesLegacyPass() : ModulePass(ID) {
+  }
+
+  StringRef getPassName() const override {
+    return "SplitIndirectBrCriticalEdgesPass";
+  }
+
+private:
+  bool runOnModule(Module &M) override;
+};
+
 } // end anonymous namespace
 
 char PGOInstrumentationGenLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(PGOInstrumentationGenLegacyPass, "pgo-instr-gen",
                       "PGO instrumentation.", false, false)
 INITIALIZE_PASS_DEPENDENCY(BlockFrequencyInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)

 INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)
 INITIALIZE_PASS_END(PGOInstrumentationUseLegacyPass, "pgo-instr-use",
                     "Read PGO instrumentation profile.", false, false)
 
 ModulePass *llvm::createPGOInstrumentationUseLegacyPass(StringRef Filename) {
   return new PGOInstrumentationUseLegacyPass(Filename.str());
 }
 
+char SplitIndirectBrCriticalEdgesLegacyPass::ID = 0;
+
+INITIALIZE_PASS_BEGIN(SplitIndirectBrCriticalEdgesLegacyPass,
+                      "split-indirectbr-critical-edges",
+                      "Split IndirectBr critical edges.", false, false)
+INITIALIZE_PASS_END(SplitIndirectBrCriticalEdgesLegacyPass,
+                    "split-indirectbr-critical-edges",
+                    "Split IndirectBr critical edges.", false, false)
+
+ModulePass *llvm::createSplitIndirectBrCriticalEdgesLegacyPass() {
+  return new SplitIndirectBrCriticalEdgesLegacyPass();
+}
+
 namespace {
 
 /// \brief An MST based instrumentation for PGO
 ///
 /// Implements a Minimum Spanning Tree (MST) based instrumentation for PGO
 /// in the function level.
 struct PGOEdge {
   // This class implements the CFG edges. Note the CFG can be a multi-graph.
FuncPGOInstrumentation(
       std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers,
       bool CreateGlobalVar = false, BranchProbabilityInfo *BPI = nullptr,
       BlockFrequencyInfo *BFI = nullptr)
       : F(Func), ComdatMembers(ComdatMembers), ValueSites(IPVK_Last + 1),
         SIVisitor(Func), MIVisitor(Func), MST(F, BPI, BFI) {
     // This should be done before CFG hash computation.
     SIVisitor.countSelects(Func);
     MIVisitor.countMemIntrinsics(Func);
     NumOfPGOSelectInsts += SIVisitor.getNumOfSelectInsts();
     NumOfPGOMemIntrinsics += MIVisitor.getNumOfMemIntrinsics();
     ValueSites[IPVK_IndirectCallTarget] = findIndirectCallSites(Func);
     ValueSites[IPVK_MemOPSize] = MIVisitor.findMemIntrinsics(Func);
 
     FuncName = getPGOFuncName(F);
     computeCFGHash();
     if (!ComdatMembers.empty())
       renameComdatFunction();
     DEBUG(dumpInfo("after CFGMST"));
 
     NumOfPGOBB += MST.BBInfos.size();
     for (auto &E : MST.AllEdges) {
       if (E->Removed)
         continue;
       NumOfPGOEdge++;
       if (!E->InMST)
         NumOfPGOInstrument++;
bool PGOInstrumentationUseLegacyPass::runOnModule(Module &M) {
   };
   auto LookupBFI = [this](Function &F) {
     return &this->getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI();
   };
 
   return annotateAllFunctions(M, ProfileFileName, LookupBPI, LookupBFI);
 }
 
+SplitIndirectBrCriticalEdgesPass::SplitIndirectBrCriticalEdgesPass() {
+}
+
+PreservedAnalyses
+SplitIndirectBrCriticalEdgesPass::run(Module &M,
+                                      ModuleAnalysisManager &AM) {
+  bool Changed = false;
+  for (auto &F : M) {
+    Changed |= SplitIndirectBrCriticalEdges(F);
+  }
+  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}
+
+bool SplitIndirectBrCriticalEdgesLegacyPass::runOnModule(Module &M) {
+  if (skipModule(M))
+    return false;
+
+  bool Changed = false;
+  for (auto &F : M) {
+    Changed |= SplitIndirectBrCriticalEdges(F);
+  }
+  return Changed;
+}
+
 static std::string getSimpleNodeName(const BasicBlock *Node) {
   if (!Node->getName().empty())
     return Node->getName();
 
   std::string SimpleNodeName;
   raw_string_ostream OS(SimpleNodeName);
   Node->printAsOperand(OS, false);
   return OS.str();

lib/Transforms/Utils/BreakCriticalEdges.cpp

 // inserting a dummy basic block.  This pass may be "required" by passes that
 // cannot deal with critical edges.  For this usage, the structure type is
 // forward declared.  This pass obviously invalidates the CFG, but can update
 // dominator trees.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/BreakCriticalEdges.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "break-crit-edges"
 
 STATISTIC(NumBroken, "Number of blocks inserted");
 
 namespace {
   struct BreakCriticalEdges : public FunctionPass {
if (Loop *TIL = LI->getLoopFor(TIBB)) {
             createPHIsForSplitLoopExit(LoopPreds, NewExitBB, DestBB);
         }
       }
     }
   }
 
   return NewBB;
 }
+
+// Return the unique indirectbr predecessor of a block. This may return null
+// even if such a predecessor exists, if it's not useful for splitting.
+// If a predecessor is found, OtherPreds will contain all other (non-indirectbr)
+// predecessors of BB.
+static BasicBlock *
+findIBRPredecessor(BasicBlock *BB, SmallVectorImpl<BasicBlock *> &OtherPreds) {
+  // If the block doesn't have any PHIs, we don't care about it, since there's
+  // no point in splitting it.
+  PHINode *PN = dyn_cast<PHINode>(BB->begin());
+  if (!PN)
+    return nullptr;
+
+  // Verify we have exactly one IBR predecessor.
+  // Conservatively bail out if one of the other predecessors is not a "regular"
+  // terminator (that is, not a switch or a br).
+  BasicBlock *IBB = nullptr;
+  for (unsigned Pred = 0, E = PN->getNumIncomingValues(); Pred != E; ++Pred) {
+    BasicBlock *PredBB = PN->getIncomingBlock(Pred);
+    TerminatorInst *PredTerm = PredBB->getTerminator();
+    switch (PredTerm->getOpcode()) {
+    case Instruction::IndirectBr:
+      if (IBB)
+        return nullptr;
+      IBB = PredBB;
+      break;
+    case Instruction::Br:
+    case Instruction::Switch:
+      OtherPreds.push_back(PredBB);
+      continue;
+    default:
+      return nullptr;
+    }
+  }
+
+  return IBB;
+}
+
+bool llvm::SplitIndirectBrCriticalEdges(Function &F) {
+  // Check whether the function has any indirectbrs, and collect which blocks
+  // they may jump to. Since most functions don't have indirect branches,
+  // this lowers the common case's overhead to O(Blocks) instead of O(Edges).
+  SmallSetVector<BasicBlock *, 16> Targets;
+  for (auto &BB : F) {
+    auto *IBI = dyn_cast<IndirectBrInst>(BB.getTerminator());
+    if (!IBI)
+      continue;
+
+    for (unsigned Succ = 0, E = IBI->getNumSuccessors(); Succ != E; ++Succ)
+      Targets.insert(IBI->getSuccessor(Succ));
+  }
+
+  if (Targets.empty())
+    return false;
+
+  bool Changed = false;
+  for (BasicBlock *Target : Targets) {
+    SmallVector<BasicBlock *, 16> OtherPreds;
+    BasicBlock *IBRPred = findIBRPredecessor(Target, OtherPreds);
+    // If we did not found an indirectbr, or the indirectbr is the only
+    // incoming edge, this isn't the kind of edge we're looking for.
+    if (!IBRPred || OtherPreds.empty())
+      continue;
+
+    // Don't even think about ehpads/landingpads.
+    Instruction *FirstNonPHI = Target->getFirstNonPHI();
+    if (FirstNonPHI->isEHPad() || Target->isLandingPad())
+      continue;
+
+    BasicBlock *BodyBlock = Target->splitBasicBlock(FirstNonPHI, ".split");
+    // It's possible Target was its own successor through an indirectbr.
+    // In this case, the indirectbr now comes from BodyBlock.
+    if (IBRPred == Target)
+      IBRPred = BodyBlock;
+
+    // At this point Target only has PHIs, and BodyBlock has the rest of the
+    // block's body. Create a copy of Target that will be used by the "direct"
+    // preds.
+    ValueToValueMapTy VMap;
+    BasicBlock *DirectSucc = CloneBasicBlock(Target, VMap, ".clone", &F);
+
+    for (BasicBlock *Pred : OtherPreds) {
+      // If the target is a loop to itself, then the terminator of the split
+      // block needs to be updated.
+      if (Pred == Target)
+        BodyBlock->getTerminator()->replaceUsesOfWith(Target, DirectSucc);
+      else
+        Pred->getTerminator()->replaceUsesOfWith(Target, DirectSucc);
+    }
+
+    // Ok, now fix up the PHIs. We know the two blocks only have PHIs, and that
+    // they are clones, so the number of PHIs are the same.
+    // (a) Remove the edge coming from IBRPred from the "Direct" PHI
+    // (b) Leave that as the only edge in the "Indirect" PHI.
+    // (c) Merge the two in the body block.
+    BasicBlock::iterator Indirect = Target->begin(),
+                         End = Target->getFirstNonPHI()->getIterator();
+    BasicBlock::iterator Direct = DirectSucc->begin();
+    BasicBlock::iterator MergeInsert = BodyBlock->getFirstInsertionPt();
+
+    assert(&*End == Target->getTerminator() &&
+           "Block was expected to only contain PHIs");
+
+    while (Indirect != End) {
+      PHINode *DirPHI = cast<PHINode>(Direct);
+      PHINode *IndPHI = cast<PHINode>(Indirect);
+
+      // Now, clean up - the direct block shouldn't get the indirect value,
+      // and vice versa.
+      DirPHI->removeIncomingValue(IBRPred);
+      Direct++;
+
+      // Advance the pointer here, to avoid invalidation issues when the old
+      // PHI is erased.
+      Indirect++;
+
+      PHINode *NewIndPHI = PHINode::Create(IndPHI->getType(), 1, "ind", IndPHI);
+      NewIndPHI->addIncoming(IndPHI->getIncomingValueForBlock(IBRPred),
+                             IBRPred);
+
+      // Create a PHI in the body block, to merge the direct and indirect
+      // predecessors.
+      PHINode *MergePHI =
+          PHINode::Create(IndPHI->getType(), 2, "merge", &*MergeInsert);
+      MergePHI->addIncoming(NewIndPHI, Target);
+      MergePHI->addIncoming(DirPHI, DirectSucc);
+
+      IndPHI->replaceAllUsesWith(MergePHI);
+      IndPHI->eraseFromParent();
+    }
+
+    Changed = true;
+  }
+
+  return Changed;
+}

test/Transforms/PGOProfile/split-indirectbr-critical-edges.ll

+; RUN: opt < %s -passes=split-indirectbr-critical-edges -S | FileCheck %s
+; RUN: opt < %s -passes=split-indirectbr-critical-edges,pgo-instr-gen -S | FileCheck %s
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @bar(i32 %v) local_unnamed_addr #0 {
+entry:
+  %mul = shl nsw i32 %v, 1
+  ret i32 %mul
+}
+
+; Function Attrs: norecurse nounwind readonly uwtable
+define i32 @foo(i8* nocapture readonly %p) #1 {
+entry:
+  %targets = alloca [256 x i8*], align 16
+  %arrayidx1 = getelementptr inbounds [256 x i8*], [256 x i8*]* %targets, i64 0, i64 93
+  store i8* blockaddress(@foo, %if.end), i8** %arrayidx1, align 8
+  br label %for.cond2
+
+for.cond2:                                        ; preds = %if.end, %for.cond2, %entry
+; CHECK: for.cond2:                                        ; preds = %.split1
+  %p.addr.0 = phi i8* [ %p, %entry ], [ %incdec.ptr5, %if.end ], [ %incdec.ptr, %for.cond2 ]
+  %incdec.ptr = getelementptr inbounds i8, i8* %p.addr.0, i64 1
+  %0 = load i8, i8* %p.addr.0, align 1
+  %cond = icmp eq i8 %0, 93
+  br i1 %cond, label %if.end.preheader, label %for.cond2
+
+if.end.preheader:                                 ; preds = %for.cond2
+  br label %if.end
+
+if.end:                                           ; preds = %if.end.preheader, %if.end
+; CHECK: if.end:                                           ; preds = %.split1
+  %p.addr.1 = phi i8* [ %incdec.ptr5, %if.end ], [ %incdec.ptr, %if.end.preheader ]
+  %incdec.ptr5 = getelementptr inbounds i8, i8* %p.addr.1, i64 1
+  %1 = load i8, i8* %p.addr.1, align 1
+  %idxprom6 = zext i8 %1 to i64
+  %arrayidx7 = getelementptr inbounds [256 x i8*], [256 x i8*]* %targets, i64 0, i64 %idxprom6
+  %2 = load i8*, i8** %arrayidx7, align 8
+  indirectbr i8* %2, [label %for.cond2, label %if.end]
+; CHECK: indirectbr i8* %2, [label %for.cond2, label %if.end]
+}