From a0ffe2b21a5159f3f8eed8e98e488e723aa7cab3 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Thu, 10 Sep 2020 11:03:48 -0700
Subject: [PATCH] [PGO] Skip if an IndirectBrInst critical edge cannot be split

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.

Reviewed By: davidxl, lebedev.ri

Differential Revision: https://reviews.llvm.org/D87435
---
 .../Instrumentation/PGOInstrumentation.cpp         |  5 ++++-
 .../PGOProfile/split-indirectbr-critical-edges.ll  | 24 ++++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index be2e091..dd70c1f 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -807,8 +807,11 @@ BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
   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");
diff --git a/llvm/test/Transforms/PGOProfile/split-indirectbr-critical-edges.ll b/llvm/test/Transforms/PGOProfile/split-indirectbr-critical-edges.ll
index dc834b7..70daa54 100644
--- a/llvm/test/Transforms/PGOProfile/split-indirectbr-critical-edges.ll
+++ b/llvm/test/Transforms/PGOProfile/split-indirectbr-critical-edges.ll
@@ -37,3 +37,27 @@ if.end:                                           ; preds = %if.end.preheader, %
   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
+}
-- 
2.7.4