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

[aleksandr.popov]

Summary

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.

[anna]

llvm/lib/Analysis/GuardUtils.cpp

130–131

Can you pls rename Callback to recordCheckOrWidenableCond ?

159

Why do you need the oneUse check here? the extractWidenableCondition API and following patches only talk about extracting widenable condition. Pls move this additional check to the patch where this context applies.

[anna]

LGTM w/ comments.

llvm/lib/Analysis/GuardUtils.cpp

160

llvm/lib/Transforms/Scalar/LoopPredication.cpp

1124

[aleksandr.popov]

Rebased

llvm/include/llvm/Analysis/GuardUtils.h

 /// Returns true iff \p U has semantics of a guard expressed in a form of call
 /// of llvm.experimental.guard intrinsic.
 bool isGuard(const User *U);
 
 /// Returns true iff \p V has semantics of llvm.experimental.widenable.condition
 /// 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
 /// widenable conditional branch to deopt block.
 bool isGuardAsWidenableBranch(const User *U);
 
 /// If U is widenable branch looking like:
 ///   %cond = ...
bool parseWidenableBranch(User *U, Use *&Cond, Use *&WC, BasicBlock *&IfTrueBB,
                           BasicBlock *&IfFalseBB);
 
 // The guard condition is expected to be in form of:
 //   cond1 && cond2 && cond3 ...
 // or in case of widenable branch:
 //   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

bool llvm::isGuard(const User *U) {
   return match(U, m_Intrinsic<Intrinsic::experimental_guard>());
 }
 
 bool llvm::isWidenableCondition(const Value *V) {
   return match(V, m_Intrinsic<Intrinsic::experimental_widenable_condition>());
 }
 
 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 {
     for (auto &Insn : *DeoptBB) {
       if (match(&Insn, m_Intrinsic<Intrinsic::experimental_deoptimize>()))
         return true;
       if (Insn.mayHaveSideEffects())
         return false;
if (match(B, m_Intrinsic<Intrinsic::experimental_widenable_condition>()) &&
     WC = &And->getOperandUse(1);
     C = &And->getOperandUse(0);
     return true;
   }
   return false;
 }
 
 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);
   do {
     Value *Check = Worklist.pop_back_val();
     Value *LHS, *RHS;
     if (match(Check, m_And(m_Value(LHS), m_Value(RHS)))) {
       if (Visited.insert(LHS).second)
         Worklist.push_back(LHS);
       if (Visited.insert(RHS).second)
         Worklist.push_back(RHS);
       continue;
     }
-    Callback(Check);
+    if (!RecordCheckOrWidenableCond(Check))
+      break;
   } while (!Worklist.empty());
 }
 
 void llvm::parseWidenableGuard(const User *U,
                                llvm::SmallVectorImpl<Value *> &Checks) {
   assert((isGuard(U) || isWidenableBranch(U)) && "Should be");
   Value *Condition = isGuard(U) ? cast<IntrinsicInst>(U)->getArgOperand(0)
                                 : cast<BranchInst>(U)->getCondition();
 
   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

 /// instruction. However, when we are dealing with widenable branch, we must
 /// account for the following situation: widening should not turn a
 /// loop-invariant condition into a loop-variant. It means that if
 /// widenable.condition() call is invariant (w.r.t. any loop), the new wide
 /// condition should stay invariant. Otherwise there can be a miscompile, like
 /// 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;
 }
 
 class GuardWideningImpl {
   DominatorTree &DT;
   PostDominatorTree *PDT;
   LoopInfo &LI;

llvm/lib/Transforms/Scalar/LoopPredication.cpp

 }
 
 bool LoopPredication::widenWidenableBranchGuardConditions(
     BranchInst *BI, SCEVExpander &Expander) {
   assert(isGuardAsWidenableBranch(BI) && "Must be!");
   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())
     return false;
 
   TotalWidened += WidenedChecks.size();
 
   // Emit the new guard condition
   IRBuilder<> Builder(findInsertPt(BI, Checks));
   Value *AllChecks = Builder.CreateAnd(Checks);
   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
     // from guard block is Cond.
     Value *AssumeCond = Builder.CreateAnd(WidenedChecks);
     if (!IfTrueBB->getUniquePredecessor()) {
       auto *GuardBB = BI->getParent();
       auto *PN = Builder.CreatePHI(AssumeCond->getType(), pred_size(IfTrueBB),
if (BasicBlock *Pred = BB->getSinglePredecessor())
       if (BB == Pred->getSingleSuccessor()) {
         BB = Pred;
         continue;
       }
     break;
   } 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;
 }
 
 /// Return the minimum of all analyzeable exit counts.  This is an upper bound
 /// on the actual exit count.  If there are not at least two analyzeable exits,
 /// returns SCEVCouldNotCompute.
 static const SCEV *getMinAnalyzeableBackedgeTakenCount(ScalarEvolution &SE,
bool LoopPredication::predicateLoopExits(Loop *L, SCEVExpander &Rewriter) {
   for (auto *ExitingBB : ExitingBlocks) {
     if (LI->getLoopFor(ExitingBB) != L)
       continue;
 
     auto *BI = dyn_cast<BranchInst>(ExitingBB->getTerminator());
     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()));
+    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);
 
   // The insertion point for the widening should be at the widenably call, not
   // at the WidenableBR. If we do this at the widenableBR, we can incorrectly
   // change a loop-invariant condition to a loop-varying one.
   auto *IP = cast<Instruction>(WidenableBR->getCondition());

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

 static bool tryWidenCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
                                            DomTreeUpdater *DTU) {
   // TODO: This can be generalized in two important ways:
   // 1) We can allow phi nodes in IfFalseBB and simply reuse all the input
   //    values from the PBI edge.
   // 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
   // This helps avoid infinite loop with SimplifyCondBranchToCondBranch which
   // may undo the transform done here.
   // TODO: There might be a more fine-grained solution to this.
   if (!llvm::succ_empty(IfFalseBB))
     return false;