[DFAJumpThreading] Only unfold select coming from directly where it is defined

[XChy]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: XChy (XChy)

Changes

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

[nikic]

There is a SIBB->getSingleSuccessor() != cast<Instruction>(SIUse)->getParent() check a few lines below -- can it be dropped now? I think it had the same purpose, but did not account for the loop entry case.

[XChy]

Thanks, I didn't notice it.

[nikic]

LGTM