[NFC][GuardUtils] Add util to extract widenable conditions

This is the next preparation patch to support widenable conditions
widening instead of branches widening.

We've added parseWidenableGuard util which parses guard condition and
collects all checks existing in the expression tree: D157276

Here we are adding util which walks similar way through the expression
tree but looks up for widenable condition without collecting the checks.
Therefore llvm::extractWidenableCondition could parse widenable branches
with arbitrary position of widenable condition in the expression tree.

llvm::parseWidenableBranch which is we are going to get rid of is being
replaced by llvm::extractWidenableCondition where it's possible.

Reviewed By: anna

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

Author: aleks-tmb

llvm/include/llvm/Analysis/GuardUtils.h

@@ -28,8 +28,8 @@ bool isGuard(const User *U);
 /// call
 bool isWidenableCondition(const Value *V);
 
-/// Returns true iff \p U is a widenable branch (that is, parseWidenableBranch
-/// returns true).
+/// Returns true iff \p U is a widenable branch (that is,
+/// extractWidenableCondition returns widenable condition).
 bool isWidenableBranch(const User *U);
 
 /// Returns true iff \p U has semantics of a guard expressed in a form of a
@@ -60,6 +60,10 @@ bool parseWidenableBranch(User *U, Use *&Cond, Use *&WC, BasicBlock *&IfTrueBB,
 //   cond1 && cond2 && cond3 && widenable_contidion ...
 // Method collects the list of checks, but skips widenable_condition.
 void parseWidenableGuard(const User *U, llvm::SmallVectorImpl<Value *> &Checks);
+
+// Returns widenable_condition if it exists in the expression tree rooting from
+// \p U and has only one use.
+Value *extractWidenableCondition(const User *U);
 } // llvm
 
 #endif // LLVM_ANALYSIS_GUARDUTILS_H

llvm/lib/Analysis/GuardUtils.cpp

@@ -24,18 +24,13 @@ bool llvm::isWidenableCondition(const Value *V) {
 }
 
 bool llvm::isWidenableBranch(const User *U) {
-  Value *Condition, *WidenableCondition;
-  BasicBlock *GuardedBB, *DeoptBB;
-  return parseWidenableBranch(U, Condition, WidenableCondition, GuardedBB,
-                              DeoptBB);
+  return extractWidenableCondition(U) != nullptr;
 }
 
 bool llvm::isGuardAsWidenableBranch(const User *U) {
-  Value *Condition, *WidenableCondition;
-  BasicBlock *GuardedBB, *DeoptBB;
-  if (!parseWidenableBranch(U, Condition, WidenableCondition, GuardedBB,
-                            DeoptBB))
+  if (!isWidenableBranch(U))
     return false;
+  BasicBlock *DeoptBB = cast<BranchInst>(U)->getSuccessor(1);
   SmallPtrSet<const BasicBlock *, 2> Visited;
   Visited.insert(DeoptBB);
   do {
@@ -117,7 +112,8 @@ bool llvm::parseWidenableBranch(User *U, Use *&C,Use *&WC,
 }
 
 template <typename CallbackType>
-static void parseCondition(Value *Condition, CallbackType Callback) {
+static void parseCondition(Value *Condition,
+                           CallbackType RecordCheckOrWidenableCond) {
   SmallVector<Value *, 4> Worklist(1, Condition);
   SmallPtrSet<Value *, 4> Visited;
   Visited.insert(Condition);
@@ -131,7 +127,8 @@ static void parseCondition(Value *Condition, CallbackType Callback) {
         Worklist.push_back(RHS);
       continue;
     }
-    Callback(Check);
+    if (!RecordCheckOrWidenableCond(Check))
+      break;
   } while (!Worklist.empty());
 }
 
@@ -144,5 +141,28 @@ void llvm::parseWidenableGuard(const User *U,
   parseCondition(Condition, [&](Value *Check) {
     if (!isWidenableCondition(Check))
       Checks.push_back(Check);
+    return true;
+  });
+}
+
+Value *llvm::extractWidenableCondition(const User *U) {
+  auto *BI = dyn_cast<BranchInst>(U);
+  if (!BI || !BI->isConditional())
+    return nullptr;
+
+  auto Condition = BI->getCondition();
+  if (!Condition->hasOneUse())
+    return nullptr;
+
+  Value *WidenableCondition = nullptr;
+  parseCondition(Condition, [&](Value *Check) {
+    // We require widenable_condition has only one use, otherwise we don't
+    // consider appropriate branch as widenable.
+    if (isWidenableCondition(Check) && Check->hasOneUse()) {
+      WidenableCondition = Check;
+      return false;
+    }
+    return true;
   });
+  return WidenableCondition;
 }

llvm/lib/Transforms/Scalar/GuardWidening.cpp

@@ -124,9 +124,7 @@ static void eliminateGuard(Instruction *GuardInst, MemorySSAUpdater *MSSAU) {
 /// the one described at https://github.com/llvm/llvm-project/issues/60234. The
 /// safest way to do it is to expand the new condition at WC's block.
 static Instruction *findInsertionPointForWideCondition(Instruction *Guard) {
-  Value *Condition, *WC;
-  BasicBlock *IfTrue, *IfFalse;
-  if (parseWidenableBranch(Guard, Condition, WC, IfTrue, IfFalse))
+  if (auto WC = extractWidenableCondition(Guard))
     return cast<Instruction>(WC);
   return Guard;
 }

llvm/lib/Transforms/Scalar/LoopPredication.cpp

@@ -802,18 +802,13 @@ bool LoopPredication::widenWidenableBranchGuardConditions(
   LLVM_DEBUG(dbgs() << "Processing guard:\n");
   LLVM_DEBUG(BI->dump());
 
-  Value *Cond, *WC;
-  BasicBlock *IfTrueBB, *IfFalseBB;
-  bool Parsed = parseWidenableBranch(BI, Cond, WC, IfTrueBB, IfFalseBB);
-  assert(Parsed && "Must be able to parse widenable branch");
-  (void)Parsed;
-
   TotalConsidered++;
   SmallVector<Value *, 4> Checks;
   SmallVector<Value *> WidenedChecks;
   parseWidenableGuard(BI, Checks);
   // At the moment, our matching logic for wideable conditions implicitly
   // assumes we preserve the form: (br (and Cond, WC())).  FIXME
+  auto WC = extractWidenableCondition(BI);
   Checks.push_back(WC);
   widenChecks(Checks, WidenedChecks, Expander, BI);
   if (WidenedChecks.empty())
@@ -827,6 +822,7 @@ bool LoopPredication::widenWidenableBranchGuardConditions(
   auto *OldCond = BI->getCondition();
   BI->setCondition(AllChecks);
   if (InsertAssumesOfPredicatedGuardsConditions) {
+    BasicBlock *IfTrueBB = BI->getSuccessor(0);
     Builder.SetInsertPoint(IfTrueBB, IfTrueBB->getFirstInsertionPt());
     // If this block has other predecessors, we might not be able to use Cond.
     // In this case, create a Phi where every other input is `true` and input
@@ -1033,13 +1029,9 @@ static BranchInst *FindWidenableTerminatorAboveLoop(Loop *L, LoopInfo &LI) {
   } while (true);
 
   if (BasicBlock *Pred = BB->getSinglePredecessor()) {
-    auto *Term = Pred->getTerminator();
-
-    Value *Cond, *WC;
-    BasicBlock *IfTrueBB, *IfFalseBB;
-    if (parseWidenableBranch(Term, Cond, WC, IfTrueBB, IfFalseBB) &&
-        IfTrueBB == BB)
-      return cast<BranchInst>(Term);
+    if (auto *BI = dyn_cast<BranchInst>(Pred->getTerminator()))
+      if (BI->getSuccessor(0) == BB && isWidenableBranch(BI))
+        return BI;
   }
   return nullptr;
 }
@@ -1127,13 +1119,13 @@ bool LoopPredication::predicateLoopExits(Loop *L, SCEVExpander &Rewriter) {
     if (!BI)
       continue;
 
-    Use *Cond, *WC;
-    BasicBlock *IfTrueBB, *IfFalseBB;
-    if (parseWidenableBranch(BI, Cond, WC, IfTrueBB, IfFalseBB) &&
-        L->contains(IfTrueBB)) {
-      WC->set(ConstantInt::getTrue(IfTrueBB->getContext()));
-      ChangedLoop = true;
-    }
+    if (auto WC = extractWidenableCondition(BI))
+      if (L->contains(BI->getSuccessor(0))) {
+        assert(WC->hasOneUse() && "Not appropriate widenable branch!");
+        WC->user_back()->replaceUsesOfWith(
+            WC, ConstantInt::getTrue(BI->getContext()));
+        ChangedLoop = true;
+      }
   }
   if (ChangedLoop)
     SE->forgetLoop(L);

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -4149,10 +4149,10 @@ static bool tryWidenCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
   // 2) We can sink side effecting instructions into BI's fallthrough
   //    successor provided they doesn't contribute to computation of
   //    BI's condition.
-  Value *CondWB, *WC;
-  BasicBlock *IfTrueBB, *IfFalseBB;
-  if (!parseWidenableBranch(PBI, CondWB, WC, IfTrueBB, IfFalseBB) ||
-      IfTrueBB != BI->getParent() || !BI->getParent()->getSinglePredecessor())
+  BasicBlock *IfTrueBB = PBI->getSuccessor(0);
+  BasicBlock *IfFalseBB = PBI->getSuccessor(1);
+  if (!isWidenableBranch(PBI) || IfTrueBB != BI->getParent() ||
+      !BI->getParent()->getSinglePredecessor())
     return false;
   if (!IfFalseBB->phis().empty())
     return false; // TODO