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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: XChy (XChy)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/70966.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp (+6) 
- (modified) llvm/test/Transforms/DFAJumpThreading/dfa-unfold-select.ll (+14) 


``````````diff
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
+}

``````````

</details>


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