[JumpThreading] Unfold selects that depend on the same condition

[pbarrio]

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.

[pbarrio]

Hi @sebpop, @rengolin,

Could you review this patch to JumpThreading? This is an update to D25477, which was reverted two weeks ago as it was breaking LNT and some bootstrap buildbots.

The problem was that an instruction that is used by selects as the condition can be also used as the true or false value. This made the old code insert the "user" select instruction twice in the list of candidates for expansion. When we traversed that list, the instruction was therefore processed twice, but the second time it didn't exist anymore because it had been expanded already into an if-then-else.

The fix is to remove duplicates in the vector of select instructions with sort+unique. A second test has also been added.

[sebpop]

LGTM.

lib/Transforms/Scalar/JumpThreading.cpp

2053

Note to other reviewers: the above two lines "sort | uniq" are the only change from the previous patch D25477. The explanation of why this is needed is in the first comment from Pablo.

[pbarrio]

@sebpop , thanks for the quick response!

[efriedma]

llvm/trunk/lib/Transforms/Scalar/JumpThreading.cpp

2053 ↗

(On Diff #77189)

This needs a large comment explaining why sorting by pointer addresses doesn't introduce non-determinism.

[efriedma]

llvm/trunk/lib/Transforms/Scalar/JumpThreading.cpp

2053 ↗

(On Diff #77189)

Actually, thinking about it a bit more, I'm 99% sure this does introduce non-determinism, at least in the ordering of use-lists. Could you use a SetVector here instead?

[pbarrio]

Thanks for the comment @efriedma. I committed before I saw it, but I agree with you that this change will affect determinism of the use-lists when two of the selects on the list depend on each other. I will create a separate revision for the change and add you as reviewer.

include/llvm/Transforms/Scalar/JumpThreading.h

 
 private:
   BasicBlock *SplitBlockPreds(BasicBlock *BB, ArrayRef<BasicBlock *> Preds,
                               const char *Suffix);
   void UpdateBlockFreqAndEdgeWeight(BasicBlock *PredBB, BasicBlock *BB,
                                     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
 
 #endif

lib/Transforms/Scalar/JumpThreading.cpp

if ((LHSFolds != LazyValueInfo::Unknown ||
         if (Phi != CondLHS)
           Phi->addIncoming(Phi->getIncomingValueForBlock(Pred), NewBB);
       return true;
     }
   }
   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.
-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;
+/// Using the similar approach of SimplifyCFG::FoldCondBranchOnPHI(), return
+/// select if the associated PHI has at least one constant.
+SelectInst *JumpThreadingPass::getSelectFedByPhi(PHINode *PN) {
 
-  // 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())
-      continue;
+    return nullptr;
 
-    SelectInst *SI = dyn_cast<SelectInst>(PN->user_back());
-    if (!SI || SI->getParent() != BB)
-      continue;
+  BasicBlock *BB = PN->getParent();
+    return nullptr;
 
-    Value *Cond = SI->getCondition();
-    if (!Cond || Cond != PN || !Cond->getType()->isIntegerTy(1))
-      continue;
+    return nullptr;
 
-    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;
-    }
+      return nullptr;
+      return SI;
 
-    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;
-    }
-  }
-  
+  return nullptr;
+}
+
+/// ExpandSelect - Expand a select into an if-then-else construct.
+void JumpThreadingPass::expandSelect(SelectInst *SI) {
+
+  BasicBlock *BB = SI->getParent();
+/// 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;
+
+  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;
+    }
+
+    if (I.getType()->isIntegerTy(1)) {
+
+      SmallVector<SelectInst *, 4> Selects;
+
+      // 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;
+
+      // 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 Changed;
 }

test/Transforms/JumpThreading/unfold-selects-same-cond.ll

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