[llvm] [DFAJumpThreading] Only unfold select coming from directly where it is defined (PR #70966)

[via llvm-commits]

https://github.com/XChy updated https://github.com/llvm/llvm-project/pull/70966

>From fcb35cfa5fdb3f791c5d7bedad0907c9cbb7fc5f Mon Sep 17 00:00:00 2001
From: XChy <xxs_chy at outlook.com>
Date: Thu, 2 Nov 2023 01:55:25 +0800
Subject: [PATCH] [DFAJumpThreading] Only unfold select coming from directly
 where it is defined

---
 .../Transforms/Scalar/DFAJumpThreading.cpp    |  6 +++++
 .../DFAJumpThreading/dfa-unfold-select.ll     | 25 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
index f2efe60bdf886a2..a72a8d903970531 100644
--- a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
@@ -297,6 +297,7 @@ void unfold(DomTreeUpdater *DTU, SelectInstToUnfold SIToUnfold,
                      {DominatorTree::Insert, StartBlock, FT}});
 
   // The select is now dead.
+  assert(SI->use_empty() && "Select must be dead now");
   SI->eraseFromParent();
 }
 
@@ -460,6 +461,11 @@ struct MainSwitch {
 
     BasicBlock *SIBB = SI->getParent();
 
+    // Only fold the select coming from directly where it is defined.
+    PHINode *PHIUser = dyn_cast<PHINode>(SIUse);
+    if (PHIUser && PHIUser->getIncomingBlock(*SI->use_begin()) != SIBB)
+      return false;
+
     // Currently, we can only expand select instructions in basic blocks with
     // one successor.
     BranchInst *SITerm = dyn_cast<BranchInst>(SIBB->getTerminator());
diff --git a/llvm/test/Transforms/DFAJumpThreading/dfa-unfold-select.ll b/llvm/test/Transforms/DFAJumpThreading/dfa-unfold-select.ll
index 746f2037d8e8688..c7bce505b717892 100644
--- a/llvm/test/Transforms/DFAJumpThreading/dfa-unfold-select.ll
+++ b/llvm/test/Transforms/DFAJumpThreading/dfa-unfold-select.ll
@@ -291,3 +291,28 @@ for.inc:
 for.end:
   ret i32 0
 }
+
+define void @select_coming_elsewhere(i1 %cond, i16 %a, i16 %b) {
+; CHECK-LABEL: @select_coming_elsewhere(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[DIV:%.*]] = select i1 [[COND:%.*]], i16 [[A:%.*]], i16 [[B:%.*]]
+; CHECK-NEXT:    br label [[FOR_COND:%.*]]
+; CHECK:       for.cond:
+; CHECK-NEXT:    [[E_ADDR_0:%.*]] = phi i16 [ 0, [[ENTRY:%.*]] ], [ [[DIV]], [[LOR_END:%.*]] ]
+; CHECK-NEXT:    switch i16 [[E_ADDR_0]], label [[LOR_END]] [
+; CHECK-NEXT:    ]
+; CHECK:       lor.end:
+; CHECK-NEXT:    br label [[FOR_COND]]
+;
+entry:
+  %div = select i1 %cond, i16 %a, i16 %b
+  br label %for.cond
+
+for.cond:                                         ; preds = %lor.end, %entry
+  %e.addr.0 = phi i16 [ 0, %entry ], [ %div, %lor.end ]
+  switch i16 %e.addr.0, label %lor.end [
+  ]
+
+lor.end:                                          ; preds = %for.cond
+  br label %for.cond
+}