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

[via llvm-commits]

Author: XChy
Date: 2023-11-02T21:25:54+08:00
New Revision: 7fa41d8a8f897219a6b22ab7a288f445e9d6119b

URL: https://github.com/llvm/llvm-project/commit/7fa41d8a8f897219a6b22ab7a288f445e9d6119b
DIFF: https://github.com/llvm/llvm-project/commit/7fa41d8a8f897219a6b22ab7a288f445e9d6119b.diff

LOG: [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.

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
    llvm/test/Transforms/DFAJumpThreading/dfa-unfold-select.ll

Removed: 
    


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