[llvm] [JumpThreading] Fix bugs of simplifying BB with unreachable predecessor (PR #68576)

[via llvm-commits]

https://github.com/XChy created https://github.com/llvm/llvm-project/pull/68576

Resolve #68529.
JumpThreading calls TryToSimplifyUncondBranchFromEmptyBlock without checking/removing unreachable pred, which results in incorrectly redirecting value from unreachable predecessor.
This patch adds check for it.

>From 7e544cc01fe180b5bc2d02fc246e9afb1f2be60a Mon Sep 17 00:00:00 2001
From: XChy <xxs_chy at outlook.com>
Date: Mon, 9 Oct 2023 18:55:13 +0800
Subject: [PATCH 1/2] [JumpThreading][NFC] Precommit test for incorrectly
 simplifying BB

---
 .../bb-unreachable-from-entry.ll              | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/llvm/test/Transforms/JumpThreading/bb-unreachable-from-entry.ll b/llvm/test/Transforms/JumpThreading/bb-unreachable-from-entry.ll
index 2a9765233ec1043..882ea5d76747212 100644
--- a/llvm/test/Transforms/JumpThreading/bb-unreachable-from-entry.ll
+++ b/llvm/test/Transforms/JumpThreading/bb-unreachable-from-entry.ll
@@ -27,3 +27,36 @@ exit1:
 exit2:
   ret void
 }
+
+define i32 @bar(i32 noundef %a) {
+; CHECK-LABEL: @bar(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[DOTNOT:%.*]] = icmp eq i32 [[A:%.*]], 0
+; CHECK-NEXT:    br i1 [[DOTNOT]], label [[END:%.*]], label [[ELSE:%.*]]
+; CHECK:       else:
+; CHECK-NEXT:    [[B:%.*]] = and i32 undef, 65535
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       dead:
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    [[DOT0:%.*]] = phi i32 [ [[A]], [[ENTRY:%.*]] ], [ [[B]], [[DEAD:%.*]] ], [ [[B]], [[ELSE]] ]
+; CHECK-NEXT:    ret i32 [[DOT0]]
+;
+entry:
+  %.not = icmp eq i32 %a, 0
+  br i1 %.not, label %end, label %else
+
+  else:                                                ; preds = %1
+  %b = and i32 undef, 65535
+  br label %B2
+
+dead:                                                ; No predecessors!
+  br label %B2
+
+B2:                                                ; preds = %2, %4
+  br label %end
+
+end:                                                ; preds = %5, %1
+  %.0 = phi i32 [ %a, %entry ], [ %b, %B2 ]
+  ret i32 %.0
+}

>From 02b7933d1c0dfe8a0d8b62663e847a7a9ddb91e4 Mon Sep 17 00:00:00 2001
From: XChy <xxs_chy at outlook.com>
Date: Mon, 9 Oct 2023 18:48:30 +0800
Subject: [PATCH 2/2] [JumpThreading] Fix bugs of simplifying BB with
 unreachable predecessor

---
 llvm/lib/Transforms/Scalar/JumpThreading.cpp              | 8 ++++++++
 .../Transforms/JumpThreading/bb-unreachable-from-entry.ll | 6 ++++--
 llvm/test/Transforms/JumpThreading/unreachable-loops.ll   | 8 ++++++--
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index ed69f1938ec7921..fc10243cff19964 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -368,6 +368,13 @@ bool JumpThreadingPass::runImpl(Function &F_, FunctionAnalysisManager *FAM_,
         continue;
       }
 
+      auto PredsAreReachable = [&Unreachable](BasicBlock *BB) -> bool {
+        for (auto *Pred : predecessors(BB))
+          if (Unreachable.count(Pred))
+            return false;
+        return true;
+      };
+
       // processBlock doesn't thread BBs with unconditional TIs. However, if BB
       // is "almost empty", we attempt to merge BB with its sole successor.
       auto *BI = dyn_cast<BranchInst>(BB.getTerminator());
@@ -379,6 +386,7 @@ bool JumpThreadingPass::runImpl(Function &F_, FunctionAnalysisManager *FAM_,
             // Don't alter Loop headers and latches to ensure another pass can
             // detect and transform nested loops later.
             !LoopHeaders.count(&BB) && !LoopHeaders.count(Succ) &&
+            PredsAreReachable(&BB) &&
             TryToSimplifyUncondBranchFromEmptyBlock(&BB, DTU.get())) {
           RemoveRedundantDbgInstrs(Succ);
           // BB is valid for cleanup here because we passed in DTU. F remains
diff --git a/llvm/test/Transforms/JumpThreading/bb-unreachable-from-entry.ll b/llvm/test/Transforms/JumpThreading/bb-unreachable-from-entry.ll
index 882ea5d76747212..e8fa399ff8a6d99 100644
--- a/llvm/test/Transforms/JumpThreading/bb-unreachable-from-entry.ll
+++ b/llvm/test/Transforms/JumpThreading/bb-unreachable-from-entry.ll
@@ -35,11 +35,13 @@ define i32 @bar(i32 noundef %a) {
 ; CHECK-NEXT:    br i1 [[DOTNOT]], label [[END:%.*]], label [[ELSE:%.*]]
 ; CHECK:       else:
 ; CHECK-NEXT:    [[B:%.*]] = and i32 undef, 65535
-; CHECK-NEXT:    br label [[END]]
+; CHECK-NEXT:    br label [[B2:%.*]]
 ; CHECK:       dead:
+; CHECK-NEXT:    br label [[B2]]
+; CHECK:       B2:
 ; CHECK-NEXT:    br label [[END]]
 ; CHECK:       end:
-; CHECK-NEXT:    [[DOT0:%.*]] = phi i32 [ [[A]], [[ENTRY:%.*]] ], [ [[B]], [[DEAD:%.*]] ], [ [[B]], [[ELSE]] ]
+; CHECK-NEXT:    [[DOT0:%.*]] = phi i32 [ [[A]], [[ENTRY:%.*]] ], [ [[B]], [[B2]] ]
 ; CHECK-NEXT:    ret i32 [[DOT0]]
 ;
 entry:
diff --git a/llvm/test/Transforms/JumpThreading/unreachable-loops.ll b/llvm/test/Transforms/JumpThreading/unreachable-loops.ll
index 8d649761700a90c..7e1dcfb3e4c78a9 100644
--- a/llvm/test/Transforms/JumpThreading/unreachable-loops.ll
+++ b/llvm/test/Transforms/JumpThreading/unreachable-loops.ll
@@ -15,9 +15,11 @@ define void @unreachable_single_bb_loop() {
 ; CHECK:       bb2:
 ; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i32 [[TMP]], 1
 ; CHECK-NEXT:    switch i1 [[TMP4]], label [[BB2:%.*]] [
-; CHECK-NEXT:    i1 false, label [[BB8]]
+; CHECK-NEXT:    i1 false, label [[BB7:%.*]]
 ; CHECK-NEXT:    i1 true, label [[BB8]]
 ; CHECK-NEXT:    ]
+; CHECK:       bb7:
+; CHECK-NEXT:    br label [[BB8]]
 ; CHECK:       bb8:
 ; CHECK-NEXT:    ret void
 ;
@@ -57,9 +59,11 @@ define void @unreachable_multi_bbs_loop() {
 ; CHECK:       bb2:
 ; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i32 [[TMP]], 1
 ; CHECK-NEXT:    switch i1 [[TMP4]], label [[BB3:%.*]] [
-; CHECK-NEXT:    i1 false, label [[BB8]]
+; CHECK-NEXT:    i1 false, label [[BB7:%.*]]
 ; CHECK-NEXT:    i1 true, label [[BB8]]
 ; CHECK-NEXT:    ]
+; CHECK:       bb7:
+; CHECK-NEXT:    br label [[BB8]]
 ; CHECK:       bb8:
 ; CHECK-NEXT:    ret void
 ;