[JumpThreading] In updatePredecessorProfileMetadata, stop searching predecessor when the current bb is an unreachable single bb loop

[wmi]

Summary

updatePredecessorProfileMetadata in jumpthreading tries to find the first dominating predecessor block for a PHI value by searching upwards the predecessor block chain.

But jumpthreading may see some temporary IR state which contains unreachable bb not being cleaned up. If an unreachable single bb loop happens to be on the predecessor block chain, keeping following the predecessor block will run into an infinite loop.

The patch fixes it.

[efriedma]

Is there any reason to expect that unreachable loops have exactly one block?

[wmi]

In D65310#1601929, @efriedma wrote:

Is there any reason to expect that unreachable loops have exactly one block?

I was wrong. I thought getSinglePredecessor already handle the case but I didn't think it carefully enough. Thanks for catching it.

I update the patch to stop searching whenever unreachable code is seen, and update the test with unreachable multi-bbs loop case.

[wmi]

Address Eli's comment.

[efriedma]

Please rebase against trunk.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

1214 ↗

(On Diff #211871)

Calling getDomTree() here is expensive; it flushes all updates in the DomTreeUpdater.

[wmi]

Address Eli's comment.

[efriedma]

LGTM

llvm/trunk/lib/Transforms/Scalar/JumpThreading.cpp

static void updatePredecessorProfileMetadata(PHINode *PN, BasicBlock *BB) {
 
   // Returns the outgoing edge of the dominating predecessor block
   // that leads to the PhiNode's incoming block:
   auto GetPredOutEdge =
       [](BasicBlock *IncomingBB,
          BasicBlock *PhiBB) -> std::pair<BasicBlock *, BasicBlock *> {
     auto *PredBB = IncomingBB;
     auto *SuccBB = PhiBB;
+    SmallPtrSet<BasicBlock *, 16> Visited;
     while (true) {
       BranchInst *PredBr = dyn_cast<BranchInst>(PredBB->getTerminator());
       if (PredBr && PredBr->isConditional())
         return {PredBB, SuccBB};
+      Visited.insert(PredBB);
       auto *SinglePredBB = PredBB->getSinglePredecessor();
       if (!SinglePredBB)
         return {nullptr, nullptr};
+
+      // Stop searching when SinglePredBB has been visited. It means we see
+      // an unreachable loop.
+      if (Visited.count(SinglePredBB))
+        return {nullptr, nullptr};
+
       SuccBB = PredBB;
       PredBB = SinglePredBB;
     }
   };
 
   for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i) {
     Value *PhiOpnd = PN->getIncomingValue(i);
     ConstantInt *CI = dyn_cast<ConstantInt>(PhiOpnd);
 
     if (!CI || !CI->getType()->isIntegerTy(1))
       continue;
 
     BP = (CI->isOne() ? BranchProbability::getBranchProbability(
                             TrueWeight, TrueWeight + FalseWeight)
                       : BranchProbability::getBranchProbability(
                             FalseWeight, TrueWeight + FalseWeight));
 
     auto PredOutEdge = GetPredOutEdge(PN->getIncomingBlock(i), BB);
     if (!PredOutEdge.first)
       return;
 
     BasicBlock *PredBB = PredOutEdge.first;
-    BranchInst *PredBr = cast<BranchInst>(PredBB->getTerminator());
+    BranchInst *PredBr = dyn_cast<BranchInst>(PredBB->getTerminator());
+    if (!PredBr)
+      return;
 
     uint64_t PredTrueWeight, PredFalseWeight;
     // FIXME: We currently only set the profile data when it is missing.
     // With PGO, this can be used to refine even existing profile data with
     // context information. This needs to be done after more performance
     // testing.
     if (PredBr->extractProfMetadata(PredTrueWeight, PredFalseWeight))
       continue;

llvm/trunk/test/Transforms/JumpThreading/unreachable-loops.ll

+; RUN: opt -jump-threading -S < %s | FileCheck %s
+; RUN: opt -passes=jump-threading -S < %s | FileCheck %s
+; Check the unreachable loop won't cause infinite loop
+; in jump-threading when it tries to update the predecessors'
+; profile metadata from a phi node.
+
+define void @unreachable_single_bb_loop() {
+; CHECK-LABEL: @unreachable_single_bb_loop()
+bb:
+  %tmp = call i32 @a()
+  %tmp1 = icmp eq i32 %tmp, 1
+  br i1 %tmp1, label %bb5, label %bb8
+
+; unreachable single bb loop.
+bb2:                                              ; preds = %bb2
+  %tmp4 = icmp ne i32 %tmp, 1
+  switch i1 %tmp4, label %bb2 [
+    i1 0, label %bb5
+    i1 1, label %bb8
+  ]
+
+bb5:                                              ; preds = %bb2, %bb
+  %tmp6 = phi i1 [ %tmp1, %bb ], [ false, %bb2 ]
+  br i1 %tmp6, label %bb8, label %bb7, !prof !0
+
+bb7:                                              ; preds = %bb5
+  br label %bb8
+
+bb8:                                              ; preds = %bb8, %bb7, %bb5, %bb2
+  ret void
+}
+
+define void @unreachable_multi_bbs_loop() {
+; CHECK-LABEL: @unreachable_multi_bbs_loop()
+bb:
+  %tmp = call i32 @a()
+  %tmp1 = icmp eq i32 %tmp, 1
+  br i1 %tmp1, label %bb5, label %bb8
+
+; unreachable two bbs loop.
+bb3:                                              ; preds = %bb2
+  br label %bb2
+
+bb2:                                              ; preds = %bb3
+  %tmp4 = icmp ne i32 %tmp, 1
+  switch i1 %tmp4, label %bb3 [
+    i1 0, label %bb5
+    i1 1, label %bb8
+  ]
+
+bb5:                                              ; preds = %bb2, %bb
+  %tmp6 = phi i1 [ %tmp1, %bb ], [ false, %bb2 ]
+  br i1 %tmp6, label %bb8, label %bb7, !prof !0
+
+bb7:                                              ; preds = %bb5
+  br label %bb8
+
+bb8:                                              ; preds = %bb8, %bb7, %bb5, %bb2
+  ret void
+}
+declare i32 @a()
+
+!0 = !{!"branch_weights", i32 2146410443, i32 1073205}