[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 cfcf9f1ad133e909af3674704bdbce678794e0f9 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    | 14 ++++++-----
 .../DFAJumpThreading/dfa-unfold-select.ll     | 25 +++++++++++++++++++
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
index f2efe60bdf886a2..8d681073ac6c9c0 100644
--- a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
@@ -100,10 +100,10 @@ static cl::opt<unsigned> MaxPathLength(
     cl::desc("Max number of blocks searched to find a threading path"),
     cl::Hidden, cl::init(20));
 
-static cl::opt<unsigned> MaxNumPaths(
-    "dfa-max-num-paths",
-    cl::desc("Max number of paths enumerated around a switch"),
-    cl::Hidden, cl::init(200));
+static cl::opt<unsigned>
+    MaxNumPaths("dfa-max-num-paths",
+                cl::desc("Max number of paths enumerated around a switch"),
+                cl::Hidden, cl::init(200));
 
 static cl::opt<unsigned>
     CostThreshold("dfa-cost-threshold",
@@ -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();
 }
 
@@ -466,8 +467,9 @@ struct MainSwitch {
     if (!SITerm || !SITerm->isUnconditional())
       return false;
 
-    if (isa<PHINode>(SIUse) &&
-        SIBB->getSingleSuccessor() != cast<Instruction>(SIUse)->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;
 
     // If select will not be sunk during unfolding, and it is in the same basic
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
+}