[PGO] Skip if an IndirectBrInst critical edge cannot be split

[MaskRay]

Summary

PGOInstrumentation runs SplitIndirectBrCriticalEdges but some IndirectBrInst
critical edge cannot be split. getInstrBB will crash when calling SplitCriticalEdge, e.g.

int foo(char *p) {
  void *targets[2];
  targets[0] = &&indirect;
  targets[1] = &&end;
for (;; p++)
  if (*p == 7) {
indirect:
    goto *targets[p[1]]; // the self loop is critical in -O
  }
end:
  return 0;
}

Skip such critical edges to prevent a crash.

[lebedev.ri]

I believe every other caller ensures that as a precondition.
Why is the existing approach incorrect, and PGOInstrumentation shouldn't be fixed itself?

[MaskRay]

In D87435#2265038, @lebedev.ri wrote:

I believe every other caller ensures that as a precondition.
Why is the existing approach incorrect, and PGOInstrumentation shouldn't be fixed itself?

How do other callers handle such a critical edge? SplitCriticalEdge is a generic API where some cases are already skipped (return nullptr), see isEHPad and other return nullptr in the function.

[lebedev.ri]

In D87435#2265077, @MaskRay wrote:

In D87435#2265038, @lebedev.ri wrote:

I believe every other caller ensures that as a precondition.
Why is the existing approach incorrect, and PGOInstrumentation shouldn't be fixed itself?

How do other callers handle such a critical edge?

By doing this check before calling the function?

SplitCriticalEdge is a generic API where some cases are already skipped (return nullptr), see isEHPad and other return nullptr in the function.

[lebedev.ri]

In D87435#2265101, @lebedev.ri wrote:

In D87435#2265077, @MaskRay wrote:

In D87435#2265038, @lebedev.ri wrote:

I believe every other caller ensures that as a precondition.
Why is the existing approach incorrect, and PGOInstrumentation shouldn't be fixed itself?

How do other callers handle such a critical edge?

By doing this check before calling the function?

SplitCriticalEdge is a generic API where some cases are already skipped (return nullptr), see isEHPad and other return nullptr in the function.

TBN i personally don't care whether all the callers do all these checks, or SplitCriticalEdge() itself does them (latter is likely slower),
but i would think we should be consistent.

[MaskRay]

Move the check from BreakCriticalEdges.cpp to PGOInstrumentation.cpp

[lebedev.ri]

SGTM, thanks

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

810

s/previous/following/

[davidxl]

lgtm

[MaskRay]

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

810

I think "previous" is correct. It is SplitIndirectBrCriticalEdges. The function calls SplitIndirectBrCriticalEdges in the very beginning.

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
   // Instrument the SrcBB if it has a single successor,
   // otherwise, the DestBB if this is not a critical edge.
   Instruction *TI = SrcBB->getTerminator();
   if (TI->getNumSuccessors() <= 1)
     return canInstrument(SrcBB);
   if (!E->IsCritical)
     return canInstrument(DestBB);
 
+  // Some IndirectBr critical edges cannot be split by the previous
+  // SplitIndirectBrCriticalEdges call. Bail out.
   unsigned SuccNum = GetSuccessorNumber(SrcBB, DestBB);
-  BasicBlock *InstrBB = SplitCriticalEdge(TI, SuccNum);
+  BasicBlock *InstrBB =
+      isa<IndirectBrInst>(TI) ? nullptr : SplitCriticalEdge(TI, SuccNum);
   if (!InstrBB) {
     LLVM_DEBUG(
         dbgs() << "Fail to split critical edge: not instrument this edge.\n");
     return nullptr;
   }
   // For a critical edge, we have to split. Instrument the newly
   // created BB.
   IsCS ? NumOfCSPGOSplit++ : NumOfPGOSplit++;

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

; CHECK: if.end:                                           ; preds = %.split1
   %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]
 }
+
+;; 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-NOT:     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]
+
+end:                                              ; preds = %indirect
+  ret i32 0
+}