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 ↗

(On Diff #125068)

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 ↗

(On Diff #125068)

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

539

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

539

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

544

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

548

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

548

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

544

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/Transforms/Utils/BasicBlockUtils.h

 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/InstrTypes.h"
 #include <cassert>
 
 namespace llvm {
 
+class BlockFrequencyInfo;
+class BranchProbabilityInfo;
 class DominatorTree;
 class Function;
 class Instruction;
 class LoopInfo;
 class MDNode;
 class MemoryDependenceResults;
 class ReturnInst;
 class TargetLibraryInfo;

 // 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);
+// If BPI and BFI aren't non-null, BPI/BFI will be updated accordingly.
+bool SplitIndirectBrCriticalEdges(Function &F,
+                                  BranchProbabilityInfo *BPI = nullptr,
+                                  BlockFrequencyInfo *BFI = nullptr);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_BASICBLOCKUTILS_H

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

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++;
BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
   return InstrBB;
 }
 
 // Visit all edge and instrument the edges not in MST, and do value profiling.
 // Critical edges will be split.
 static void instrumentOneFunc(
     Function &F, Module *M, BranchProbabilityInfo *BPI, BlockFrequencyInfo *BFI,
     std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers) {
+  // Split indirectbr critical edges here before computing the MST rather than
+  // later in getInstrBB() to avoid invalidating it.
+  SplitIndirectBrCriticalEdges(F, BPI, BFI);
   FuncPGOInstrumentation<PGOEdge, BBInfo> FuncInfo(F, ComdatMembers, true, BPI,
                                                    BFI);
   unsigned NumCounters = FuncInfo.getNumCounters();
 
   uint32_t I = 0;
   Type *I8PtrTy = Type::getInt8PtrTy(M->getContext());
   for (auto &E : FuncInfo.MST.AllEdges) {
     BasicBlock *InstrBB = FuncInfo.getInstrBB(E.get());
static bool annotateAllFunctions(
   collectComdatMembers(M, ComdatMembers);
   std::vector<Function *> HotFunctions;
   std::vector<Function *> ColdFunctions;
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
     auto *BPI = LookupBPI(F);
     auto *BFI = LookupBFI(F);
+    // Split indirectbr critical edges here before computing the MST rather than
+    // later in getInstrBB() to avoid invalidating it.
+    SplitIndirectBrCriticalEdges(F, BPI, BFI);
     PGOUseFunc Func(F, &M, ComdatMembers, BPI, BFI);
     if (!Func.readCounters(PGOReader.get()))
       continue;
     Func.populateCounters();
     Func.setBranchWeights();
     Func.annotateValueSites();
     Func.annotateIrrLoopHeaderWeights();
     PGOUseFunc::FuncFreqAttr FreqAttr = Func.getFuncFreqAttr();

lib/Transforms/Utils/BreakCriticalEdges.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #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/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BranchProbabilityInfo.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"
for (unsigned Pred = 0, E = PN->getNumIncomingValues(); Pred != E; ++Pred) {
     default:
       return nullptr;
     }
   }
 
   return IBB;
 }
 
-bool llvm::SplitIndirectBrCriticalEdges(Function &F) {
+bool llvm::SplitIndirectBrCriticalEdges(Function &F,
+                                        BranchProbabilityInfo *BPI,
+                                        BlockFrequencyInfo *BFI) {
   // 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 ShouldUpdateAnalysis = BPI && BFI;
   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");
+    if (ShouldUpdateAnalysis) {
+      // Copy the BFI/BPI from Target to BodyBlock.
+      for (unsigned I = 0, E = BodyBlock->getTerminator()->getNumSuccessors();
+           I < E; ++I)
+        BPI->setEdgeProbability(BodyBlock, I,
+                                BPI->getEdgeProbability(Target, I));
+      BFI->setBlockFreq(BodyBlock, BFI->getBlockFreq(Target).getFrequency());
+    }
     // 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);
 
+    BlockFrequency BlockFreqForDirectSucc;
     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);
+      // block (BodyBlock) needs to be updated.
+      BasicBlock *Src = Pred != Target ? Pred : BodyBlock;
+      Src->getTerminator()->replaceUsesOfWith(Target, DirectSucc);
+      if (ShouldUpdateAnalysis)
+        BlockFreqForDirectSucc += BFI->getBlockFreq(Src) *
+            BPI->getEdgeProbability(Src, DirectSucc);
+    }
+    if (ShouldUpdateAnalysis) {
+      BFI->setBlockFreq(DirectSucc, BlockFreqForDirectSucc.getFrequency());
+      BlockFrequency NewBlockFreqForTarget =
+          BFI->getBlockFreq(Target) - BlockFreqForDirectSucc;
+      BFI->setBlockFreq(Target, NewBlockFreqForTarget.getFrequency());
+      BPI->eraseBlock(Target);
     }
 
     // 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(),

test/Transforms/PGOProfile/Inputs/indirectbr.proftext

 # IR level Instrumentation Flag
 :ir
 foo
 # Func Hash:
-40197883220
[phone]
 # Num Counters:
 4
 # Counter Values:
-202
-88
+139
 20
 5
-
+63

test/Transforms/PGOProfile/indirectbr.ll

 return:
   %retval.0 = phi i32 [ 3, %label3 ], [ 2, %label2 ], [ 0, %entry ], [ 1, %if.then ]
   ret i32 %retval.0
 }
 
 ; BRANCHPROB: Printing analysis 'Branch Probability Analysis' for function 'foo':
 ; BRANCHPROB:---- Branch Probabilities ----
 ; BRANCHPROB:  edge entry -> if.then probability is 0x37c32b17 / 0x80000000 = 43.56%
-; BRANCHPROB:  edge entry -> return probability is 0x483cd4e9 / 0x80000000 = 56.44%
+; BRANCHPROB:  edge entry -> return.clone probability is 0x483cd4e9 / 0x80000000 = 56.44%
 ; BRANCHPROB:  edge if.then -> return probability is 0x5ba2e8ba / 0x80000000 = 71.59%
 ; BRANCHPROB:  edge if.then -> label2 probability is 0x1d1745d1 / 0x80000000 = 22.73%
 ; BRANCHPROB:  edge if.then -> label3 probability is 0x0745d174 / 0x80000000 = 5.68%
-; BRANCHPROB:  edge label2 -> return probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
-; BRANCHPROB:  edge label3 -> return probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+; BRANCHPROB:  edge label2 -> return.clone probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+; BRANCHPROB:  edge label3 -> return.clone probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+; BRANCHPROB:  edge return -> .split probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+; BRANCHPROB:  edge return.clone -> .split probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
 
 
 

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

+; RUN: opt < %s -passes=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]
+}