BreakCriticalEdges: do not split the critical edge from a CallBr indirect successor

Otherwise we'll fail the assertion in SplitBlockPredecessors() related
to splitting the edges from CallBr's.

Fixes: ClangBuiltLinux/linux#1161
Fixes: ClangBuiltLinux/linux#1252

Reviewed By: void, MaskRay, jyknight

Differential Revision: https://reviews.llvm.org/D88438

Author: nickdesaulniers

llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp

@@ -158,22 +158,21 @@ BasicBlock *llvm::SplitCriticalEdge(Instruction *TI, unsigned SuccNum,
   SmallVector<BasicBlock *, 4> LoopPreds;
   // Check if extra modifications will be required to preserve loop-simplify
   // form after splitting. If it would require splitting blocks with IndirectBr
-  // terminators, bail out if preserving loop-simplify form is requested.
+  // or CallBr terminators, bail out if preserving loop-simplify form is
+  // requested.
   if (LI) {
     if (Loop *TIL = LI->getLoopFor(TIBB)) {
 
-      // The only that we can break LoopSimplify form by splitting a critical
-      // edge is if after the split there exists some edge from TIL to DestBB
-      // *and* the only edge into DestBB from outside of TIL is that of
+      // The only way that we can break LoopSimplify form by splitting a
+      // critical edge is if after the split there exists some edge from TIL to
+      // DestBB *and* the only edge into DestBB from outside of TIL is that of
       // NewBB. If the first isn't true, then LoopSimplify still holds, NewBB
       // is the new exit block and it has no non-loop predecessors. If the
       // second isn't true, then DestBB was not in LoopSimplify form prior to
       // the split as it had a non-loop predecessor. In both of these cases,
       // the predecessor must be directly in TIL, not in a subloop, or again
       // LoopSimplify doesn't hold.
-      for (pred_iterator I = pred_begin(DestBB), E = pred_end(DestBB); I != E;
-           ++I) {
-        BasicBlock *P = *I;
+      for (BasicBlock *P : predecessors(DestBB)) {
         if (P == TIBB)
           continue; // The new block is known.
         if (LI->getLoopFor(P) != TIL) {
@@ -186,7 +185,10 @@ BasicBlock *llvm::SplitCriticalEdge(Instruction *TI, unsigned SuccNum,
       // Loop-simplify form can be preserved, if we can split all in-loop
       // predecessors.
       if (any_of(LoopPreds, [](BasicBlock *Pred) {
-            return isa<IndirectBrInst>(Pred->getTerminator());
+            const Instruction *T = Pred->getTerminator();
+            if (const auto *CBR = dyn_cast<CallBrInst>(T))
+              return CBR->getDefaultDest() != Pred;
+            return isa<IndirectBrInst>(T);
           })) {
         if (Options.PreserveLoopSimplify)
           return nullptr;

llvm/test/Transforms/LoopStrengthReduce/callbr-critical-edge-splitting.ll

@@ -0,0 +1,39 @@
+; RUN: opt -loop-reduce %s -o - -S | FileCheck %s
+; RUN: opt -passes='loop(loop-reduce)' %s -o - -S | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+define dso_local i32 @test1() local_unnamed_addr {
+entry:
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.inc, %entry
+; It's ok to modify this test in the future should be able to split critical
+; edges here, just noting that this is the critical edge that we care about.
+; CHECK: callbr void asm sideeffect "", "X,X,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@test1, %cond.true.i), i8* blockaddress(@test1, %for.end))
+; CHECK-NEXT: to label %asm.fallthrough.i.i [label %cond.true.i, label %for.end]
+  callbr void asm sideeffect "", "X,X,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@test1, %cond.true.i), i8* blockaddress(@test1, %for.end))
+          to label %asm.fallthrough.i.i [label %cond.true.i, label %for.end]
+
+asm.fallthrough.i.i:                              ; preds = %for.cond
+  unreachable
+
+cond.true.i:                                      ; preds = %for.cond
+  br label %do.body.i.i.do.body.i.i_crit_edge
+
+do.body.i.i.do.body.i.i_crit_edge:                ; preds = %do.body.i.i.do.body.i.i_crit_edge, %cond.true.i
+  %pgocount711 = phi i64 [ %0, %do.body.i.i.do.body.i.i_crit_edge ], [ 0, %cond.true.i ]
+  %0 = add nuw nsw i64 %pgocount711, 1
+  br i1 undef, label %do.body.i.i.rdrand_int.exit.i_crit_edge, label %do.body.i.i.do.body.i.i_crit_edge
+
+do.body.i.i.rdrand_int.exit.i_crit_edge:          ; preds = %do.body.i.i.do.body.i.i_crit_edge
+  %1 = add i64 %0, undef
+  br i1 undef, label %for.end, label %for.inc
+
+for.inc:                                          ; preds = %do.body.i.i.rdrand_int.exit.i_crit_edge
+  br label %for.cond
+
+for.end:                                          ; preds = %do.body.i.i.rdrand_int.exit.i_crit_edge, %for.cond
+  %pgocount.promoted24 = phi i64 [ undef, %for.cond ], [ %1, %do.body.i.i.rdrand_int.exit.i_crit_edge ]
+  ret i32 undef
+}

llvm/test/Transforms/LoopStrengthReduce/callbr-critical-edge-splitting2.ll

@@ -0,0 +1,61 @@
+; RUN: opt -passes='loop(loop-reduce)' %s -o - -S | FileCheck %s
+
+; Required metadata to trigger previously failing assertion.
+target datalayout = "e-m:e-i64:64-n32:64"
+
+@f = external dso_local local_unnamed_addr global i32, align 4
+
+declare i32 @a() local_unnamed_addr
+declare i32 @e(i32) local_unnamed_addr
+
+define dso_local i32 @b() {
+entry:
+  %call = tail call i32 @a()
+  %tobool.not = icmp eq i32 %call, 0
+  br i1 %tobool.not, label %cleanup.cont.critedge, label %if.then
+
+if.then:                                          ; preds = %entry
+; It's ok to modify this test in the future should be able to split critical
+; edges here, just noting that this is the critical edge that we care about.
+; CHECK: callbr void asm sideeffect "", "X"(i8* blockaddress(@b, %cleanup.cont.critedge))
+; CHECK-NEXT: to label %return [label %cleanup.cont.critedge]
+  callbr void asm sideeffect "", "X"(i8* blockaddress(@b, %cleanup.cont.critedge))
+          to label %return [label %cleanup.cont.critedge]
+
+cleanup.cont.critedge:                            ; preds = %entry, %if.then
+  br label %return
+
+return:                                           ; preds = %if.then, %cleanup.cont.critedge
+  %retval.0 = phi i32 [ 4, %cleanup.cont.critedge ], [ 0, %if.then ]
+  ret i32 %retval.0
+}
+
+define dso_local i32 @do_pages_move_nr_pages() local_unnamed_addr {
+entry:
+  br label %for.cond
+
+for.cond:                                         ; preds = %if.end3, %entry
+  %g.0 = phi i32 [ undef, %entry ], [ %inc, %if.end3 ]
+  %0 = load i32, i32* @f, align 4
+  %tobool.not = icmp eq i32 %0, 0
+  br i1 %tobool.not, label %if.end3, label %if.then
+
+if.then:                                          ; preds = %for.cond
+  %call.i = tail call i32 @a()
+  %tobool.not.i = icmp eq i32 %call.i, 0
+  br i1 %tobool.not.i, label %if.then2, label %if.then.i
+
+if.then.i:                                        ; preds = %if.then
+  callbr void asm sideeffect "", "X"(i8* blockaddress(@do_pages_move_nr_pages, %if.then2))
+          to label %if.end3 [label %if.then2]
+
+if.then2:                                         ; preds = %if.then, %if.then.i
+  %g.0.lcssa = phi i32 [ %g.0, %if.then ], [ %g.0, %if.then.i ]
+  %call4 = tail call i32 @e(i32 %g.0.lcssa)
+  ret i32 undef
+
+if.end3:                                          ; preds = %for.cond, %if.then.i
+  %inc = add nsw i32 %g.0, 1
+  br label %for.cond
+}
+