PGOInstrumentation, GCOVProfiling: Split indirectbr critical edges regardless of PHIs

[MatzeB]

Summary

The SplitIndirectBrCriticalEdges function was originally designed for CodeGenPrepare and skipped splitting of edges when the destination block didn't contain any PHI instructions. This only makes sense when reducing COPYs like CodeGenPrepare. In the case of PGOInstrumentation or GCOVProfiling it would result in missed counters and wrong result in functions with computed goto.

[modimo]

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

503–505

Please document this flag in the comment above

llvm/test/Transforms/GCOVProfiling/split-indirectbr-critical-edges.ll

37

nit:s/impossibel/impossible/g

nit2: If we're willing to change the labels the indirect gotos jump to we can technically split an arbitrary # of edges. Not that we should, of course..

llvm/test/Transforms/PGOProfile/irreducible.ll

142

I'm curious how the counters shifted around now that indirectgoto->sw.bb gets split. How does the changed counter values in the profiles above map to that?

llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp

475

s/bb3/bb2/g has no PHI

[MatzeB]

rebase, address review comments

llvm/test/Transforms/PGOProfile/irreducible.ll

142

This is hard to figure out given the size of the test and the fact that the profile files are checked-in as-is. I hope this is fine...

[MatzeB]

llvm/test/Transforms/GCOVProfiling/split-indirectbr-critical-edges.ll

37

Not completely sure what you mean in "nit2", but I changed the wording to make it clear that this is a trick to keep the edge from getting split so the test keeps working (I know more normalization etc. could "fix" this, but we're only running a single pass here)

[modimo]

LGTM

llvm/test/Transforms/GCOVProfiling/split-indirectbr-critical-edges.ll

37

You got it, "nit2" was that handling this case isn't technically impossible. Thanks for the wording change.

llvm/test/Transforms/PGOProfile/irreducible.ll

142

Took a look because I was curious how to get this information. Adding -debug-only=pgo-instrumentation dumps the edges and the instrumented edges line up with the counts in the profile. Comparing the two (irreducible.proftext) left is original right is with this diff:

The count of 100 gets inferred now and the new counts are for the split edges as expected. The change that causes the INDIRECTGOTO_IRR_LOOP to go from 400->399 is that the probes shifted so that edge 21 and edge 22 on RHS flipped compared to LHS. That's not a big deal though.

[xur]

LGTM.
Thanks for this fix. I kind of know this (I used to have an assert to after calling SplitIndirectBrCriticalEdges(), but got rid of it because some cases it does return false).
Even this fix, there are still chance that we cannot split the edge.
Maybe we should emit a warning in this case.

[MatzeB]

Even this fix, there are still chance that we cannot split the edge.
Maybe we should emit a warning in this case.

Yes I'm aware that it's not always possible to split critical edges; I think it should be possible in those case to just instrument the source and destination of the critical edge with a counter to get valid data?

Admittedly I did not try to fix that either as this fix here seems to be all we need for our software... I think the remaining cases should only effect exception handling code which is not perf critical anyway? For cases with two indirectbr predecessors I always saw LLVM normalizing control flow to a single one in earlier passes, so that was never a problem for instrumentation in my experiments.

[MatzeB]

llvm/test/Transforms/PGOProfile/irreducible.ll

142

Thanks!

[wenlei]

good catch, lgtm!

[MatzeB]

landed as 6a383369f9b800eac5de2456e49fa70577be8e33

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

 // 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
 // If BPI and BFI aren't non-null, BPI/BFI will be updated accordingly.
-bool SplitIndirectBrCriticalEdges(Function &F,
+// When `IgnoreBlocksWithoutPHI` is set to `true` critical edges leading to a
+// block without phi-instructions will not be split.
+bool SplitIndirectBrCriticalEdges(Function &F, bool IgnoreBlocksWithoutPHI,
                                   BranchProbabilityInfo *BPI = nullptr,
                                   BlockFrequencyInfo *BFI = nullptr);
 
 /// Given a set of incoming and outgoing blocks, create a "hub" such that every
 /// edge from an incoming block InBB to an outgoing block OutBB is now split
 /// into two edges, one from InBB to the hub and another from the hub to
 /// OutBB. The hub consists of a series of guard blocks, one for each outgoing
 /// block. Each guard block conditionally branches to the corresponding outgoing

llvm/lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::runOnFunction(Function &F) {
   EverMadeChange |= eliminateMostlyEmptyBlocks(F);
 
   bool ModifiedDT = false;
   if (!DisableBranchOpts)
     EverMadeChange |= splitBranchCondition(F, ModifiedDT);
 
   // Split some critical edges where one of the sources is an indirect branch,
   // to help generate sane code for PHIs involving such edges.
-  EverMadeChange |= SplitIndirectBrCriticalEdges(F);
+  EverMadeChange |=
+      SplitIndirectBrCriticalEdges(F, /*IgnoreBlocksWithoutPHI=*/true);
 
   bool MadeChange = true;
   while (MadeChange) {
     MadeChange = false;
     DT.reset();
     for (BasicBlock &BB : llvm::make_early_inc_range(F)) {
       bool ModifiedDTOnIteration = false;
       MadeChange |= optimizeBlock(BB, ModifiedDTOnIteration);

llvm/lib/Transforms/Instrumentation/GCOVProfiling.cpp

for (auto &F : M->functions()) {
       uint32_t Line = SP->getLine();
       auto Filename = getFilename(SP);
 
       BranchProbabilityInfo *BPI = GetBPI(F);
       BlockFrequencyInfo *BFI = GetBFI(F);
 
       // Split indirectbr critical edges here before computing the MST rather
       // than later in getInstrBB() to avoid invalidating it.
-      SplitIndirectBrCriticalEdges(F, BPI, BFI);
+      SplitIndirectBrCriticalEdges(F, /*IgnoreBlocksWithoutPHI=*/false, BPI,
+                                   BFI);
 
       CFGMST<Edge, BBInfo> MST(F, /*InstrumentFuncEntry_=*/false, BPI, BFI);
 
       // getInstrBB can split basic blocks and push elements to AllEdges.
       for (size_t I : llvm::seq<size_t>(0, MST.AllEdges.size())) {
         auto &E = *MST.AllEdges[I];
         // For now, disable spanning tree optimization when fork or exec* is
         // used.

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

 // Critical edges will be split.
 static void instrumentOneFunc(
     Function &F, Module *M, TargetLibraryInfo &TLI, BranchProbabilityInfo *BPI,
     BlockFrequencyInfo *BFI,
     std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers,
     bool IsCS) {
   // Split indirectbr critical edges here before computing the MST rather than
   // later in getInstrBB() to avoid invalidating it.
-  SplitIndirectBrCriticalEdges(F, BPI, BFI);
+  SplitIndirectBrCriticalEdges(F, /*IgnoreBlocksWithoutPHI=*/false, BPI, BFI);
 
   FuncPGOInstrumentation<PGOEdge, BBInfo> FuncInfo(
       F, TLI, ComdatMembers, true, BPI, BFI, IsCS, PGOInstrumentEntry);
 
   Type *I8PtrTy = Type::getInt8PtrTy(M->getContext());
   auto Name = ConstantExpr::getBitCast(FuncInfo.FuncNameVar, I8PtrTy);
   auto CFGHash = ConstantInt::get(Type::getInt64Ty(M->getContext()),
                                   FuncInfo.FunctionHash);
static bool annotateAllFunctions(
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
     auto &TLI = LookupTLI(F);
     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);
+    SplitIndirectBrCriticalEdges(F, /*IgnoreBlocksWithoutPHI=*/false, BPI, BFI);
     PGOUseFunc Func(F, &M, TLI, ComdatMembers, BPI, BFI, PSI, IsCS,
                     InstrumentFuncEntry);
     // When AllMinusOnes is true, it means the profile for the function
     // is unrepresentative and this function is actually hot. Set the
     // entry count of the function to be multiple times of hot threshold
     // and drop all its internal counters.
     bool AllMinusOnes = false;
     bool AllZeros = false;

llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp

 }
 
 // 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);
+  for (BasicBlock *PredBB : predecessors(BB)) {
     Instruction *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,
+                                        bool IgnoreBlocksWithoutPHI,
                                         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) {
+    if (IgnoreBlocksWithoutPHI && Target->phis().empty())
+      continue;
+
     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.

llvm/test/Transforms/GCOVProfiling/split-indirectbr-critical-edges.ll

indirect.preheader:                               ; preds = %for.cond
   %idxprom = sext i8 %1 to i64, !dbg !21
   %arrayidx4 = getelementptr inbounds <2 x i8*>, <2 x i8*>* %targets, i64 0, i64 %idxprom, !dbg !21
   %2 = load i8*, i8** %arrayidx4, align 8, !dbg !21
   br label %indirect
 
 indirect:                                         ; preds = %indirect.preheader, %indirect
   indirectbr i8* %2, [label %indirect, label %end]
 
+indirect2:
+  ; For this test we do not want critical edges split. Adding a 2nd `indirectbr`
+  ; does the trick.
+  indirectbr i8* %2, [label %indirect, label %end]
+
 end:                                              ; preds = %indirect
   ret i32 0, !dbg !22
 }
 
 attributes #0 = { norecurse nounwind readonly uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" "unsafe-fp-math"="false" "use-soft-float"="false" }
 
 !llvm.dbg.cu = !{!0}
 !llvm.module.flags = !{!3, !4, !5}

llvm/test/Transforms/PGOProfile/Inputs/irreducible.proftext

 :ir
 _Z11irreducibleii
 # Func Hash:
 287486624745028451
 # Num Counters:
 6
 # Counter Values:
 1000
 950
 100
 373
 1
 0
 
 _Z11irreduciblePh
 # Func Hash:
-331779889035882993
+52047014671956012
 # Num Counters:
 9
 # Counter Values:
-100
 300
 99
 300
 201
 1
 1
 0
 0
+0

llvm/test/Transforms/PGOProfile/Inputs/irreducible_entry.proftext

 1000
 950
 100
 373
 0
 
 _Z11irreduciblePh
 # Func Hash:
-331779889035882993
+52047014671956012
 # Num Counters:
 9
 # Counter Values:
 1
-100
 300
 99
 300
 201
 1
 0
 0
+0

llvm/test/Transforms/PGOProfile/irreducible.ll

 ; USE-SAME: !irr_loop ![[INDIRECTGOTO_IRR_LOOP:[0-9]+]]
 }
 
 ; USE: ![[FOR_COND2_IRR_LOOP]] = !{!"loop_header_weight", i64 1050}
 ; USE: ![[ENTRY8_IRR_LOOP]] = !{!"loop_header_weight", i64 373}
 ; USE: ![[IF_END9_IRR_LOOP]] = !{!"loop_header_weight", i64 1000}
 ; USE: ![[SW_BB6_IRR_LOOP]] = !{!"loop_header_weight", i64 501}
 ; USE: ![[SW_BB15_IRR_LOOP]] = !{!"loop_header_weight", i64 100}
-; USE: ![[INDIRECTGOTO_IRR_LOOP]] = !{!"loop_header_weight", i64 400}
+; USE: ![[INDIRECTGOTO_IRR_LOOP]] = !{!"loop_header_weight", i64 399}

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

 ; CHECK: indirectbr i8* %2, [label %for.cond2, label %if.end]
 }
 
 ;; If an indirectbr critical edge cannot be split, ignore it.
 ;; The edge will not be profiled.
 ; CHECK-LABEL: @cannot_split(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    call void @llvm.instrprof.increment
+; CHECK: indirect:
 ; CHECK-NOT:     call void @llvm.instrprof.increment
+; CHECK: indirect2:
+; CHECK-NEXT:    call void @llvm.instrprof.increment
 define i32 @cannot_split(i8* nocapture readonly %p) {
 entry:
   %targets = alloca <2 x i8*>, align 16
   store <2 x i8*> <i8* blockaddress(@cannot_split, %indirect), i8* blockaddress(@cannot_split, %end)>, <2 x i8*>* %targets, align 16
   %arrayidx2 = getelementptr inbounds i8, i8* %p, i64 1
   %0 = load i8, i8* %arrayidx2
   %idxprom = sext i8 %0 to i64
   %arrayidx3 = getelementptr inbounds <2 x i8*>, <2 x i8*>* %targets, i64 0, i64 %idxprom
   %1 = load i8*, i8** %arrayidx3, align 8
   br label %indirect
 
 indirect:                                         ; preds = %entry, %indirect
+  indirectbr i8* %1, [label %indirect, label %end, label %indirect2]
+
+indirect2:
+  ; For this test we do not want critical edges split. Adding a 2nd `indirectbr`
+  ; does the trick.
   indirectbr i8* %1, [label %indirect, label %end]
 
 end:                                              ; preds = %indirect
   ret i32 0
 }

llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp

)IR");
   PostDominatorTree PDT(*F);
 
   CriticalEdgeSplittingOptions CESO(&DT, nullptr, nullptr, &PDT);
   EXPECT_EQ(1u, SplitAllCriticalEdges(*F, CESO));
   EXPECT_TRUE(DT.verify());
   EXPECT_TRUE(PDT.verify());
 }
 
-TEST(BasicBlockUtils, SplitIndirectBrCriticalEdge) {
+TEST(BasicBlockUtils, SplitIndirectBrCriticalEdgesIgnorePHIs) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"IR(
-define void @crit_edge(i8* %cond0, i1 %cond1) {
+define void @crit_edge(i8* %tgt, i1 %cond0, i1 %cond1) {
 entry:
-  indirectbr i8* %cond0, [label %bb0, label %bb1]
+  indirectbr i8* %tgt, [label %bb0, label %bb1, label %bb2]
 bb0:
-  br label %bb1
+  br i1 %cond0, label %bb1, label %bb2
 bb1:
   %p = phi i32 [0, %bb0], [0, %entry]
-  br i1 %cond1, label %bb2, label %bb3
+  br i1 %cond1, label %bb3, label %bb4
 bb2:
   ret void
 bb3:
   ret void
+bb4:
+  ret void
 }
 )IR");
   Function *F = M->getFunction("crit_edge");
   DominatorTree DT(*F);
   LoopInfo LI(DT);
   BranchProbabilityInfo BPI(*F, LI);
   BlockFrequencyInfo BFI(*F, BPI, LI);
 
-  ASSERT_TRUE(SplitIndirectBrCriticalEdges(*F, &BPI, &BFI));
+  ASSERT_TRUE(SplitIndirectBrCriticalEdges(*F, /*IgnoreBlocksWithoutPHI=*/true,
+                                           &BPI, &BFI));
 
   // Check that successors of the split block get their probability correct.
   BasicBlock *BB1 = getBasicBlockByName(*F, "bb1");
   BasicBlock *SplitBB = BB1->getTerminator()->getSuccessor(0);
-  EXPECT_EQ(2u, SplitBB->getTerminator()->getNumSuccessors());
+  ASSERT_EQ(2u, SplitBB->getTerminator()->getNumSuccessors());
   EXPECT_EQ(BranchProbability(1, 2), BPI.getEdgeProbability(SplitBB, 0u));
   EXPECT_EQ(BranchProbability(1, 2), BPI.getEdgeProbability(SplitBB, 1u));
+
+  // bb2 has no PHI, so we shouldn't split bb0 -> bb2
+  BasicBlock *BB0 = getBasicBlockByName(*F, "bb0");
+  ASSERT_EQ(2u, BB0->getTerminator()->getNumSuccessors());
+  EXPECT_EQ(BB0->getTerminator()->getSuccessor(1),
+            getBasicBlockByName(*F, "bb2"));
+}
+
+TEST(BasicBlockUtils, SplitIndirectBrCriticalEdges) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+define void @crit_edge(i8* %tgt, i1 %cond0, i1 %cond1) {
+entry:
+  indirectbr i8* %tgt, [label %bb0, label %bb1, label %bb2]
+bb0:
+  br i1 %cond0, label %bb1, label %bb2
+bb1:
+  %p = phi i32 [0, %bb0], [0, %entry]
+  br i1 %cond1, label %bb3, label %bb4
+bb2:
+  ret void
+bb3:
+  ret void
+bb4:
+  ret void
+}
+)IR");
+  Function *F = M->getFunction("crit_edge");
+  DominatorTree DT(*F);
+  LoopInfo LI(DT);
+  BranchProbabilityInfo BPI(*F, LI);
+  BlockFrequencyInfo BFI(*F, BPI, LI);
+
+  ASSERT_TRUE(SplitIndirectBrCriticalEdges(*F, /*IgnoreBlocksWithoutPHI=*/false,
+                                           &BPI, &BFI));
+
+  // Check that successors of the split block get their probability correct.
+  BasicBlock *BB1 = getBasicBlockByName(*F, "bb1");
+  BasicBlock *SplitBB = BB1->getTerminator()->getSuccessor(0);
+  ASSERT_EQ(2u, SplitBB->getTerminator()->getNumSuccessors());
+  EXPECT_EQ(BranchProbability(1, 2), BPI.getEdgeProbability(SplitBB, 0u));
+  EXPECT_EQ(BranchProbability(1, 2), BPI.getEdgeProbability(SplitBB, 1u));
+
+  // Should split, resulting in:
+  //   bb0 -> bb2.clone; bb2 -> split1; bb2.clone -> split,
+  BasicBlock *BB0 = getBasicBlockByName(*F, "bb0");
+  ASSERT_EQ(2u, BB0->getTerminator()->getNumSuccessors());
+  BasicBlock *BB2Clone = BB0->getTerminator()->getSuccessor(1);
+  BasicBlock *BB2 = getBasicBlockByName(*F, "bb2");
+  EXPECT_NE(BB2Clone, BB2);
+  ASSERT_EQ(1u, BB2->getTerminator()->getNumSuccessors());
+  ASSERT_EQ(1u, BB2Clone->getTerminator()->getNumSuccessors());
+  EXPECT_EQ(BB2->getTerminator()->getSuccessor(0),
+            BB2Clone->getTerminator()->getSuccessor(0));
 }
 
 TEST(BasicBlockUtils, SetEdgeProbability) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"IR(
 define void @edge_probability(i32 %0) {
 entry:
 switch i32 %0, label %LD [