[JumpThreading] Unfold selects that depend on the same condition

Summary:
These are good candidates for jump threading. This enables later opts
(such as InstCombine) to combine instructions from the selects with
instructions out of the selects. SimplifyCFG will fold the select
again if unfolding wasn't worth it.

Patch by James Molloy and Pablo Barrio.

Reviewers: rengolin, haicheng, sebpop

Subscribers: jojo, jmolloy, llvm-commits

Differential Revision: https://reviews.llvm.org/D26391

llvm-svn: 286236

diff --git a/llvm/include/llvm/Transforms/Scalar/JumpThreading.h b/llvm/include/llvm/Transforms/Scalar/JumpThreading.h
index f96741c..306f0b6 100644 (file)
--- a/llvm/include/llvm/Transforms/Scalar/JumpThreading.h
+++ b/llvm/include/llvm/Transforms/Scalar/JumpThreading.h
@@ -129,6 +129,8 @@ private:
                                     BasicBlock *NewBB, BasicBlock *SuccBB);
   /// Check if the block has profile metadata for its outgoing edges.
   bool doesBlockHaveProfileData(BasicBlock *BB);
+  SelectInst *getSelectFedByPhi(PHINode *PN);
+  void expandSelect(SelectInst *SI);
 };
 
 } // end namespace llvm

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index ddad300..05ac11f 100644 (file)
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -1963,61 +1963,100 @@ bool JumpThreadingPass::TryToUnfoldSelect(CmpInst *CondCmp, BasicBlock *BB) {
   return false;
 }
 
-/// TryToUnfoldSelectInCurrBB - Look for PHI/Select in the same BB of the form
+/// GetSelectFedByPhi - Look for PHI/Select in the same BB of the form
 /// bb:
 ///   %p = phi [false, %bb1], [true, %bb2], [false, %bb3], [true, %bb4], ...
 ///   %s = select p, trueval, falseval
 ///
-/// And expand the select into a branch structure. This later enables
+/// And return the select. Unfolding it into a branch structure later enables
 /// jump-threading over bb in this pass.
 ///
-/// Using the similar approach of SimplifyCFG::FoldCondBranchOnPHI(), unfold
-/// select if the associated PHI has at least one constant.  If the unfolded
-/// select is not jump-threaded, it will be folded again in the later
-/// optimizations.
+/// Using the similar approach of SimplifyCFG::FoldCondBranchOnPHI(), return
+/// select if the associated PHI has at least one constant.
+SelectInst *JumpThreadingPass::getSelectFedByPhi(PHINode *PN) {
+
+  unsigned NumPHIValues = PN->getNumIncomingValues();
+  if (NumPHIValues == 0 || !PN->hasOneUse())
+    return nullptr;
+
+  SelectInst *SI = dyn_cast<SelectInst>(PN->user_back());
+  BasicBlock *BB = PN->getParent();
+  if (!SI || SI->getParent() != BB)
+    return nullptr;
+
+  Value *Cond = SI->getCondition();
+  if (!Cond || Cond != PN || !Cond->getType()->isIntegerTy(1))
+    return nullptr;
+
+  for (unsigned i = 0; i != NumPHIValues; ++i) {
+    if (PN->getIncomingBlock(i) == BB)
+      return nullptr;
+    if (isa<ConstantInt>(PN->getIncomingValue(i)))
+      return SI;
+  }
+
+  return nullptr;
+}
+
+/// ExpandSelect - Expand a select into an if-then-else construct.
+void JumpThreadingPass::expandSelect(SelectInst *SI) {
+
+  BasicBlock *BB = SI->getParent();
+  TerminatorInst *Term =
+      SplitBlockAndInsertIfThen(SI->getCondition(), SI, false);
+  PHINode *NewPN = PHINode::Create(SI->getType(), 2, "", SI);
+  NewPN->addIncoming(SI->getTrueValue(), Term->getParent());
+  NewPN->addIncoming(SI->getFalseValue(), BB);
+  SI->replaceAllUsesWith(NewPN);
+  SI->eraseFromParent();
+}
+
+/// TryToUnfoldSelectInCurrBB - Unfold selects that could be jump-threaded were
+/// they if-then-elses. If the unfolded selects are not jump-threaded, it will
+/// be folded again in the later optimizations.
 bool JumpThreadingPass::TryToUnfoldSelectInCurrBB(BasicBlock *BB) {
+
   // If threading this would thread across a loop header, don't thread the edge.
   // See the comments above FindLoopHeaders for justifications and caveats.
   if (LoopHeaders.count(BB))
     return false;
 
-  // Look for a Phi/Select pair in the same basic block.  The Phi feeds the
-  // condition of the Select and at least one of the incoming values is a
-  // constant.
-  for (BasicBlock::iterator BI = BB->begin();
-       PHINode *PN = dyn_cast<PHINode>(BI); ++BI) {
-    unsigned NumPHIValues = PN->getNumIncomingValues();
-    if (NumPHIValues == 0 || !PN->hasOneUse())
+  bool Changed = false;
+  for (auto &I : *BB) {
+
+    // Look for a Phi/Select pair in the same basic block.  The Phi feeds the
+    // condition of the Select and at least one of the incoming values is a
+    // constant.
+    PHINode *PN;
+    SelectInst *SI;
+    if ((PN = dyn_cast<PHINode>(&I)) && (SI = getSelectFedByPhi(PN))) {
+      expandSelect(SI);
+      Changed = true;
       continue;
+    }
 
-    SelectInst *SI = dyn_cast<SelectInst>(PN->user_back());
-    if (!SI || SI->getParent() != BB)
-      continue;
+    if (I.getType()->isIntegerTy(1)) {
 
-    Value *Cond = SI->getCondition();
-    if (!Cond || Cond != PN || !Cond->getType()->isIntegerTy(1))
-      continue;
+      SmallVector<SelectInst *, 4> Selects;
 
-    bool HasConst = false;
-    for (unsigned i = 0; i != NumPHIValues; ++i) {
-      if (PN->getIncomingBlock(i) == BB)
-        return false;
-      if (isa<ConstantInt>(PN->getIncomingValue(i)))
-        HasConst = true;
-    }
+      // Look for scalar booleans used in selects as conditions. If there are
+      // several selects that use the same boolean, they are candidates for jump
+      // threading and therefore we should unfold them.
+      for (Value *U : I.users())
+        if (auto *SI = dyn_cast<SelectInst>(U))
+          Selects.push_back(SI);
+      if (Selects.size() <= 1)
+        continue;
 
-    if (HasConst) {
-      // Expand the select.
-      TerminatorInst *Term =
-          SplitBlockAndInsertIfThen(SI->getCondition(), SI, false);
-      PHINode *NewPN = PHINode::Create(SI->getType(), 2, "", SI);
-      NewPN->addIncoming(SI->getTrueValue(), Term->getParent());
-      NewPN->addIncoming(SI->getFalseValue(), BB);
-      SI->replaceAllUsesWith(NewPN);
-      SI->eraseFromParent();
-      return true;
+      // Remove duplicates
+      std::sort(Selects.begin(), Selects.end());
+      auto NewEnd = std::unique(Selects.begin(), Selects.end());
+
+      Changed = true;
+      for (auto SI = Selects.begin(); SI != NewEnd; ++SI)
+        expandSelect(*SI);
     }
   }
-  
-  return false;
+
+  return Changed;
 }

diff --git a/llvm/test/Transforms/JumpThreading/unfold-selects-same-cond.ll b/llvm/test/Transforms/JumpThreading/unfold-selects-same-cond.ll
new file mode 100644 (file)
index 0000000..ef01b0d
--- /dev/null
+++ b/llvm/test/Transforms/JumpThreading/unfold-selects-same-cond.ll
@@ -0,0 +1,45 @@
+; RUN: opt < %s -jump-threading -instcombine -simplifycfg  -S | FileCheck %s
+
+; The three selects are jump-threaded so that instcombine can optimize, and
+; simplifycfg should turn the result into a single select.
+define i32 @f(i32 %a, i32 %b) {
+; CHECK: select
+; CHECK-NOT: select
+entry:
+  %0 = and i32 %a, 1
+  %1 = and i32 %b, 1
+  %xor = xor i32 %1, %a
+  %shr32 = lshr i32 %a, 1
+  %cmp10 = icmp eq i32 %xor, 1
+  %2 = xor i32 %b, 12345
+  %b.addr.1 = select i1 %cmp10, i32 %2, i32 %b
+  %shr1633 = lshr i32 %b.addr.1, 1
+  %3 = or i32 %shr1633, 54321
+  %b.addr.2 = select i1 %cmp10, i32 %3, i32 %shr1633
+  %shr1634 = lshr i32 %b.addr.2, 2
+  %4 = or i32 %shr1634, 54320
+  %b.addr.3 = select i1 %cmp10, i32 %4, i32 %shr1634
+  ret i32 %b.addr.3
+}
+
+; Case where the condition is not only used as condition but also as the
+; true or false value in at least one of the selects.
+define i1 @g(i32 %a, i32 %b) {
+; CHECK: select
+; CHECK-NOT: select
+entry:
+  %0 = and i32 %a, 1
+  %1 = and i32 %b, 1
+  %xor = xor i32 %1, %a
+  %shr32 = lshr i32 %a, 1
+  %cmp10 = icmp eq i32 %xor, 1
+  %2 = xor i32 %b, 12345
+  %b.addr.1 = select i1 %cmp10, i32 %2, i32 %b
+  %shr1633 = lshr i32 %b.addr.1, 1
+  %3 = or i32 %shr1633, 54321
+  %b.addr.2 = select i1 %cmp10, i32 %3, i32 %shr1633
+  %shr1634 = lshr i32 %b.addr.2, 2
+  %4 = icmp eq i32 %shr1634, 54320
+  %b.addr.3 = select i1 %cmp10, i1 %4, i1 %cmp10
+  ret i1 %b.addr.3
+}