[BasicBlockUtils] Allow splitting predecessors with callbr terminators

[nikic]

Summary

SplitBlockPredecessors currently asserts if one of the predecessor terminators is a callbr. This limitation was originally necessary, because just like with indirectbr, it was not possible to replace successors of a callbr. However, this is no longer the case since D67252. As the requirement nowadays is that callbr must reference all blockaddrs directly in the call arguments, and these get automatically updated when setSuccessor() is called, we no longer need this limitation.

The only thing we need to do here is use replaceSuccessorWith() instead of replaceUsesOfWith(), because only the former does the necessary blockaddr updating magic.

I believe there's other similar limitations that can be removed, e.g. related to critical edge splitting.

[nikic]

Remove one more check.

[nickdesaulniers]

instead of replaceUsesOfWith(), because only the former does the necessary blockaddr updating magic.

Indeed, see: https://reviews.llvm.org/D74947.

Thanks for the patch, will do review and more testing later. Just getting back online after a long holiday weekend.

[nickdesaulniers]

llvm/lib/Transforms/Scalar/LICM.cpp

1510–1512

I think I attempted this in https://reviews.llvm.org/D94833, see the attached bug report from @nathanchance .

[nikic]

llvm/lib/Transforms/Scalar/LICM.cpp

1510–1512

I think that patch was missing the change to SplitBlockPredecessors() done here.

[efriedma]

If we don't support using arbitrary blockaddresses as destinations, can we stop using blockaddress to represent the destinations? The presence of blockaddress blocks optimizations even if we aren't imposing the full control flow restrictions.

[nikic]

In D129205#3633528, @efriedma wrote:

If we don't support using arbitrary blockaddresses as destinations, can we stop using blockaddress to represent the destinations? The presence of blockaddress blocks optimizations even if we aren't imposing the full control flow restrictions.

Yeah, dropping the blockaddress arguments sounds reasonable and should largely remove the need for special handling of callbr in transforms. I think a possible way to do this would be to say that we implicitly pass the blockaddresses for all callbr targets at the end of the argument list. (Not sure I want to volunteer for that change though...)

[nickdesaulniers]

In D129205#3633528, @efriedma wrote:

If we don't support using arbitrary blockaddresses as destinations, can we stop using blockaddress to represent the destinations? The presence of blockaddress blocks optimizations even if we aren't imposing the full control flow restrictions.

That's what I had in mind in https://github.com/ClangBuiltLinux/linux/issues/972. Particularly because of https://reviews.llvm.org/D74947. On my TODO list, but seemingly never high enough priority.

[nickdesaulniers]

Thanks for the patch and feedback on my linked issue. Testing x86_64 and arm64 defconfig, defconfig+thinLTO Linux kernel builds + boots.

[nathanchance]

I did a set of builds with this patch on mainline and saw no issues.

[jyknight]

In D129205#3633528, @efriedma wrote:

If we don't support using arbitrary blockaddresses as destinations, can we stop using blockaddress to represent the destinations? The presence of blockaddress blocks optimizations even if we aren't imposing the full control flow restrictions.

Just to be clear: the original intent of the LLVM IR design here was that it would be valid to store a blockaddress into some memory, and then call that stored address inside a callbr asm. Thus the callbr asm might or might not even have _any_ blockaddress inputs itself -- that a blockaddress was an argument to the callbr had no special significance vs the blockaddress anywhere else. All possible destinations the callbr could jump to must be listed as indirect labels, however.

I do think that was a reasonable design choice to make, but we've already moved halfway away from it in a sort of ad-hoc fashion...

[nikic]

In D129205#3635611, @jyknight wrote:

In D129205#3633528, @efriedma wrote:

If we don't support using arbitrary blockaddresses as destinations, can we stop using blockaddress to represent the destinations? The presence of blockaddress blocks optimizations even if we aren't imposing the full control flow restrictions.

Just to be clear: the original intent of the LLVM IR design here was that it would be valid to store a blockaddress into some memory, and then call that stored address inside a callbr asm. Thus the callbr asm might or might not even have _any_ blockaddress inputs itself -- that a blockaddress was an argument to the callbr had no special significance vs the blockaddress anywhere else. All possible destinations the callbr could jump to must be listed as indirect labels, however.

I do think that was a reasonable design choice to make, but we've already moved halfway away from it in a sort of ad-hoc fashion...

Yes, that was the original design, which essentially made callbr the same as indirectbr. I assume that this design was later discarded because this generality is not necessary (or particularly useful) to support asm goto. At the same time, giving callbr the same semantics as indirectbr prevents many CFG optimizations and increases the chances of miscompiles and assertion failures for any CFG transform that does not have explicit handling for callbr. It's a reasonable design in principle, but not probably not a good tradeoff for its intended use-case.

I've put up https://reviews.llvm.org/D129288 to decouple callbr representation from blockaddresses entirely (on the IR level), which should allow us to remove the remaining callbr-related special cases, and also make it clearer that callbr cannot be used with arbitrary blockaddresses.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

if (PredsScanned.size() == AvailablePreds.size()+1 &&
     SmallPtrSet<BasicBlock*, 8> AvailablePredSet;
 
     for (const auto &AvailablePred : AvailablePreds)
       AvailablePredSet.insert(AvailablePred.first);
 
     // Add all the unavailable predecessors to the PredsToSplit list.
     for (BasicBlock *P : predecessors(LoadBB)) {
       // If the predecessor is an indirect goto, we can't split the edge.
-      // Same for CallBr.
-      if (isa<IndirectBrInst>(P->getTerminator()) ||
-          isa<CallBrInst>(P->getTerminator()))
+      if (isa<IndirectBrInst>(P->getTerminator()))
         return false;
 
       if (!AvailablePredSet.count(P))
         PredsToSplit.push_back(P);
     }
 
     // Split them out to their own block.
     UnavailablePred = splitBlockPreds(LoadBB, PredsToSplit, "thread-pre-split");
if (PredToDestList.empty()) {
         OnlyDest = MultipleDestSentinel;
       // It possible we have same destination, but different value, e.g. default
       // case in switchinst.
       if (Val != OnlyVal)
         OnlyVal = MultipleVal;
     }
 
     // If the predecessor ends with an indirect goto, we can't change its
-    // destination. Same for CallBr.
-    if (isa<IndirectBrInst>(Pred->getTerminator()) ||
-        isa<CallBrInst>(Pred->getTerminator()))
+    // destination.
+    if (isa<IndirectBrInst>(Pred->getTerminator()))
       continue;
 
     PredToDestList.emplace_back(Pred, DestBB);
   }
 
   // If all edges were unthreadable, we fail.
   if (PredToDestList.empty())
     return false;
if (!SplitVal) {
       // If all preds provide 1, set the computed value to 1.
       BO->setOperand(!isLHS, SplitVal);
     }
 
     return true;
   }
 
   // If any of predecessors end with an indirect goto, we can't change its
-  // destination. Same for CallBr.
+  // destination.
   if (any_of(BlocksToFoldInto, [](BasicBlock *Pred) {
-        return isa<IndirectBrInst>(Pred->getTerminator()) ||
-               isa<CallBrInst>(Pred->getTerminator());
+        return isa<IndirectBrInst>(Pred->getTerminator());
       }))
     return false;
 
   // Try to duplicate BB into PredBB.
   return duplicateCondBranchOnPHIIntoPred(BB, BlocksToFoldInto);
 }
 
 /// addPHINodeEntriesForMappedBlock - We're adding 'NewPred' as a new

llvm/lib/Transforms/Scalar/LICM.cpp

static bool canSplitPredecessors(PHINode *PN, LoopSafetyInfo *SafetyInfo) {
   if (!BB->canSplitPredecessors())
     return false;
   // It's not impossible to split EHPad blocks, but if BlockColors already exist
   // it require updating BlockColors for all offspring blocks accordingly. By
   // skipping such corner case, we can make updating BlockColors after splitting
   // predecessor fairly simple.
   if (!SafetyInfo->getBlockColors().empty() && BB->getFirstNonPHI()->isEHPad())
     return false;
   for (BasicBlock *BBPred : predecessors(BB)) {
-    if (isa<IndirectBrInst>(BBPred->getTerminator()) ||
-        isa<CallBrInst>(BBPred->getTerminator()))
+    if (isa<IndirectBrInst>(BBPred->getTerminator()))
       return false;
   }
   return true;
 }
 
 static void splitPredecessorsOfLoopExit(PHINode *PN, DominatorTree *DT,
                                         LoopInfo *LI, const Loop *CurLoop,
                                         LoopSafetyInfo *SafetyInfo,

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

SplitBlockPredecessorsImpl(BasicBlock *BB, ArrayRef<BasicBlock *> Preds,
 
   // Move the edges from Preds to point to NewBB instead of BB.
   for (unsigned i = 0, e = Preds.size(); i != e; ++i) {
     // This is slightly more strict than necessary; the minimum requirement
     // is that there be no more than one indirectbr branching to BB. And
     // all BlockAddress uses would need to be updated.
     assert(!isa<IndirectBrInst>(Preds[i]->getTerminator()) &&
            "Cannot split an edge from an IndirectBrInst");
-    assert(!isa<CallBrInst>(Preds[i]->getTerminator()) &&
-           "Cannot split an edge from a CallBrInst");
-    Preds[i]->getTerminator()->replaceUsesOfWith(BB, NewBB);
+    Preds[i]->getTerminator()->replaceSuccessorWith(BB, NewBB);
   }
 
   // Insert a new PHI node into NewBB for every PHI node in BB and that new PHI
   // node becomes an incoming value for BB's phi node.  However, if the Preds
   // list is empty, we need to insert dummy entries into the PHI nodes in BB to
   // account for the newly created predecessor.
   if (Preds.empty()) {
     // Insert dummy values as the incoming value.

llvm/lib/Transforms/Utils/LoopUtils.cpp

auto RewriteExit = [&](BasicBlock *BB) {
     // See if there are any non-loop predecessors of this exit block and
     // keep track of the in-loop predecessors.
     bool IsDedicatedExit = true;
     for (auto *PredBB : predecessors(BB))
       if (L->contains(PredBB)) {
         if (isa<IndirectBrInst>(PredBB->getTerminator()))
           // We cannot rewrite exiting edges from an indirectbr.
           return false;
-        if (isa<CallBrInst>(PredBB->getTerminator()))
-          // We cannot rewrite exiting edges from a callbr.
-          return false;
 
         InLoopPredecessors.push_back(PredBB);
       } else {
         IsDedicatedExit = false;
       }
 
     assert(!InLoopPredecessors.empty() && "Must have *some* loop predecessor!");
 

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

for (const auto &Pair : KnownValues) {
     ArrayRef<BasicBlock *> PredBBs = Pair.second.getArrayRef();
     BasicBlock *RealDest = BI->getSuccessor(!CB->getZExtValue());
 
     if (RealDest == BB)
       continue; // Skip self loops.
 
     // Skip if the predecessor's terminator is an indirect branch.
     if (any_of(PredBBs, [](BasicBlock *PredBB) {
-          return isa<IndirectBrInst>(PredBB->getTerminator()) ||
-                 isa<CallBrInst>(PredBB->getTerminator());
+          return isa<IndirectBrInst>(PredBB->getTerminator());
         }))
       continue;
 
     LLVM_DEBUG({
       dbgs() << "Condition " << *Cond << " in " << BB->getName()
              << " has value " << *Pair.first << " in predecessors:\n";
       for (const BasicBlock *PredBB : Pair.second)
         dbgs() << "  " << PredBB->getName() << "\n";

llvm/test/Transforms/JumpThreading/callbr-edge-split.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -S -jump-threading | FileCheck %s
 
 ; This test used to cause jump threading to try to split an edge of a callbr.
 
 @a = global i32 0
 
 define i32 @c() {
 ; CHECK-LABEL: @c(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* @a
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* @a, align 4
 ; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i32 [[TMP0]], 0
-; CHECK-NEXT:    br i1 [[TOBOOL]], label [[IF_ELSE:%.*]], label [[IF_THEN:%.*]]
-; CHECK:       if.then:
-; CHECK-NEXT:    [[CALL:%.*]] = call i32 @b()
-; CHECK-NEXT:    [[PHITMP:%.*]] = icmp ne i32 [[CALL]], 0
-; CHECK-NEXT:    br i1 [[PHITMP]], label [[IF_THEN2:%.*]], label [[IF_END4:%.*]]
+; CHECK-NEXT:    br i1 [[TOBOOL]], label [[IF_ELSE:%.*]], label [[IF_END:%.*]]
 ; CHECK:       if.else:
-; CHECK-NEXT:    callbr void asm sideeffect "", "i"(i8* blockaddress(@c, [[IF_THEN2]]))
-; CHECK-NEXT:    to label [[IF_END_THREAD:%.*]] [label %if.then2]
-; CHECK:       if.end.thread:
+; CHECK-NEXT:    callbr void asm sideeffect "", "i"(i8* blockaddress(@c, [[IF_THEN2:%.*]]))
+; CHECK-NEXT:    to label [[NORMAL:%.*]] [label %if.then2]
+; CHECK:       normal:
 ; CHECK-NEXT:    br label [[IF_THEN2]]
+; CHECK:       if.end:
+; CHECK-NEXT:    [[CALL:%.*]] = call i32 @b()
+; CHECK-NEXT:    [[PHITMP:%.*]] = icmp ne i32 [[CALL]], 0
+; CHECK-NEXT:    br i1 [[PHITMP]], label [[IF_THEN2]], label [[IF_END4:%.*]]
 ; CHECK:       if.then2:
 ; CHECK-NEXT:    [[CALL3:%.*]] = call i32 @b()
 ; CHECK-NEXT:    br label [[IF_END4]]
 ; CHECK:       if.end4:
 ; CHECK-NEXT:    ret i32 undef
 ;
 entry:
   %0 = load i32, i32* @a

llvm/test/Transforms/JumpThreading/pr46857-callbr.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -jump-threading -S | FileCheck %s
 
 ; CHECK-ALL-LABEL: @func(
 
 define i1 @func(i1 %arg, i32 %arg1, i1 %arg2) {
 ; CHECK-LABEL: @func(
 ; CHECK-NEXT:  bb:
-; CHECK-NEXT:    br i1 [[ARG:%.*]], label [[BB3:%.*]], label [[BB4:%.*]]
-; CHECK:       bb3:
-; CHECK-NEXT:    [[I:%.*]] = icmp eq i32 [[ARG1:%.*]], 0
-; CHECK-NEXT:    br label [[BB7:%.*]]
+; CHECK-NEXT:    br i1 [[ARG:%.*]], label [[BB7:%.*]], label [[BB4:%.*]]
 ; CHECK:       bb4:
-; CHECK-NEXT:    callbr void asm sideeffect "", "i"(i8* blockaddress(@func, [[BB7]]))
-; CHECK-NEXT:    to label [[BB5:%.*]] [label %bb7]
+; CHECK-NEXT:    callbr void asm sideeffect "", "i"(i8* blockaddress(@func, [[BB7_THR_COMM:%.*]]))
+; CHECK-NEXT:    to label [[BB5:%.*]] [label %bb7.thr_comm]
 ; CHECK:       bb5:
-; CHECK-NEXT:    br label [[BB7]]
+; CHECK-NEXT:    br label [[BB7_THR_COMM]]
+; CHECK:       bb7.thr_comm:
+; CHECK-NEXT:    [[I91:%.*]] = xor i1 [[ARG2:%.*]], [[ARG]]
+; CHECK-NEXT:    br i1 [[I91]], label [[BB11:%.*]], label [[BB11]]
 ; CHECK:       bb7:
-; CHECK-NEXT:    [[I8:%.*]] = phi i1 [ [[I]], [[BB3]] ], [ [[ARG2:%.*]], [[BB5]] ], [ [[ARG2]], [[BB4]] ]
-; CHECK-NEXT:    [[I9:%.*]] = xor i1 [[I8]], [[ARG]]
-; CHECK-NEXT:    br i1 [[I9]], label [[BB11:%.*]], label [[BB11]]
+; CHECK-NEXT:    [[I:%.*]] = icmp eq i32 [[ARG1:%.*]], 0
+; CHECK-NEXT:    [[I9:%.*]] = xor i1 [[I]], [[ARG]]
+; CHECK-NEXT:    br i1 [[I9]], label [[BB11]], label [[BB11]]
 ; CHECK:       bb11:
-; CHECK-NEXT:    ret i1 [[I9]]
+; CHECK-NEXT:    [[I93:%.*]] = phi i1 [ [[I91]], [[BB7_THR_COMM]] ], [ [[I9]], [[BB7]] ], [ [[I91]], [[BB7_THR_COMM]] ], [ [[I9]], [[BB7]] ]
+; CHECK-NEXT:    ret i1 [[I93]]
 ;
 bb:
   br i1 %arg, label %bb3, label %bb4
 
 bb3:
   %i = icmp eq i32 %arg1, 0
   br label %bb7
 

llvm/test/Transforms/LICM/callbr-crash.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -licm < %s | FileCheck %s
 
 define i32 @j() {
 ; CHECK-LABEL: @j(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_COND:%.*]]
 ; CHECK:       for.cond:
-; CHECK-NEXT:    callbr void asm sideeffect "", "i,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@j, [[FOR_END:%.*]]))
-; CHECK-NEXT:    to label [[COND_TRUE_I:%.*]] [label %for.end]
+; CHECK-NEXT:    callbr void asm sideeffect "", "i,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@j, [[FOR_END_SPLIT_LOOP_EXIT1:%.*]]))
+; CHECK-NEXT:    to label [[COND_TRUE_I:%.*]] [label %for.end.split.loop.exit1]
 ; CHECK:       cond.true.i:
-; CHECK-NEXT:    br i1 true, label [[FOR_END]], label [[FOR_COND]]
+; CHECK-NEXT:    br i1 true, label [[FOR_END_SPLIT_LOOP_EXIT:%.*]], label [[FOR_COND]]
+; CHECK:       for.end.split.loop.exit:
+; CHECK-NEXT:    [[ASMRESULT1_I_I_LE:%.*]] = extractvalue { i8, i32 } zeroinitializer, 1
+; CHECK-NEXT:    br label [[FOR_END:%.*]]
+; CHECK:       for.end.split.loop.exit1:
+; CHECK-NEXT:    [[PHI_PH2:%.*]] = phi i32 [ undef, [[FOR_COND]] ]
+; CHECK-NEXT:    br label [[FOR_END]]
 ; CHECK:       for.end:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ 0, [[COND_TRUE_I]] ], [ undef, [[FOR_COND]] ]
+; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[ASMRESULT1_I_I_LE]], [[FOR_END_SPLIT_LOOP_EXIT]] ], [ [[PHI_PH2]], [[FOR_END_SPLIT_LOOP_EXIT1]] ]
 ; CHECK-NEXT:    ret i32 [[PHI]]
 ;
 entry:
   br label %for.cond
 
 for.cond:                                         ; preds = %cond.true.i, %entry
   callbr void asm sideeffect "", "i,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@j, %for.end))
   to label %cond.true.i [label %for.end]

llvm/test/Transforms/SimplifyCFG/jump-threading.ll

 exit:
   ret void
 }
 
 define void @callbr() {
 ; CHECK-LABEL: @callbr(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    callbr void asm sideeffect "", "i,~{dirflag},~{fpsr},~{flags}"(ptr blockaddress(@callbr, [[TARGET:%.*]]))
-; CHECK-NEXT:    to label [[JOIN:%.*]] [label %target]
+; CHECK-NEXT:    to label [[IF_END:%.*]] [label %target]
 ; CHECK:       target:
-; CHECK-NEXT:    br label [[JOIN]]
-; CHECK:       join:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i1 [ false, [[TARGET]] ], [ false, [[ENTRY:%.*]] ]
-; CHECK-NEXT:    br i1 [[PHI]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
-; CHECK:       if.then:
-; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       if.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   callbr void asm sideeffect "", "i,~{dirflag},~{fpsr},~{flags}"(ptr blockaddress(@callbr, %target))
   to label %join [label %target]