[DFAJumpThreading] Only unfold select coming from directly where it is defined (#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.

Author: XChy

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

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