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

[via llvm-commits]

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

Fixes #64860.
When a select instruction comes in by PHINode, the phi's incoming block for it can flow indirectly past other BasicBlock into it. In this case, we cannot unfold select to the phi's BB.

>From f28c60ba63e10d693b22dec5983776eb395b6c3e 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

---
 llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp    |  6 ++++++
 .../DFAJumpThreading/dfa-unfold-select.ll          | 14 ++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
index f2efe60bdf886a2..711fe4c09d61a0c 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 = cast<PHINode>(SIUse);
+    if (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..d9df8aef60ab5a5 100644
--- a/llvm/test/Transforms/DFAJumpThreading/dfa-unfold-select.ll
+++ b/llvm/test/Transforms/DFAJumpThreading/dfa-unfold-select.ll
@@ -291,3 +291,17 @@ for.inc:
 for.end:
   ret i32 0
 }
+
+define void @select_coming_elsewhere(i1 %cond, i16 %a, i16 %b) {
+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
+}