[LV][NFC] Refactor structures used to maintain uncountable exit info

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

@@ -391,25 +391,22 @@ class LoopVectorizationLegality {
 
   /// Returns true if the loop has an uncountable early exit, i.e. an
   /// uncountable exit that isn't the latch block.
-  bool hasUncountableEarlyExit() const { return HasUncountableEarlyExit; }
+  bool hasUncountableEarlyExit() const {
+    return getUncountableEdge().has_value();
+  }
 
   /// Returns the uncountable early exiting block.
   BasicBlock *getUncountableEarlyExitingBlock() const {
-    if (!HasUncountableEarlyExit) {
-      assert(getUncountableExitingBlocks().empty() &&
-             "Expected no uncountable exiting blocks");
+    if (!hasUncountableEarlyExit())
       return nullptr;
-    }
-    assert(getUncountableExitingBlocks().size() == 1 &&
-           "Expected only a single uncountable exiting block");
-    return getUncountableExitingBlocks()[0];
+    return (*getUncountableEdge()).first;
   }
 
   /// Returns the destination of an uncountable early exiting block.
   BasicBlock *getUncountableEarlyExitBlock() const {
-    assert(getUncountableExitBlocks().size() == 1 &&
-           "Expected only a single uncountable exit block");
-    return getUncountableExitBlocks()[0];
+    if (!hasUncountableEarlyExit())
+      return nullptr;
+    return (*getUncountableEdge()).second;
   }
 
   /// Returns true if vector representation of the instruction \p I
@@ -463,14 +460,10 @@ class LoopVectorizationLegality {
     return CountableExitingBlocks;
   }
 
-  /// Returns all the exiting blocks with an uncountable exit.
-  const SmallVector<BasicBlock *, 4> &getUncountableExitingBlocks() const {
-    return UncountableExitingBlocks;
-  }
-
-  /// Returns all the exit blocks from uncountable exiting blocks.
-  SmallVector<BasicBlock *, 4> getUncountableExitBlocks() const {
-    return UncountableExitBlocks;
+  /// Returns the loop edge with an uncountable exit.
+  std::optional<std::pair<BasicBlock *, BasicBlock *>>
+  getUncountableEdge() const {
+    return UncountableEdge;
   }
 
 private:
@@ -654,18 +647,13 @@ class LoopVectorizationLegality {
   /// supported.
   bool StructVecCallFound = false;
 
-  /// Indicates whether this loop has an uncountable early exit, i.e. an
-  /// uncountable exiting block that is not the latch.
-  bool HasUncountableEarlyExit = false;
-
   /// Keep track of all the countable and uncountable exiting blocks if
   /// the exact backedge taken count is not computable.
   SmallVector<BasicBlock *, 4> CountableExitingBlocks;
-  SmallVector<BasicBlock *, 4> UncountableExitingBlocks;
 
-  /// Keep track of the destinations of all uncountable exits if the
-  /// exact backedge taken count is not computable.
-  SmallVector<BasicBlock *, 4> UncountableExitBlocks;
+  /// Keep track of the loop edge with an uncountable exit, comprising a pair
+  /// of (Exiting, Exit) blocks.
+  std::optional<std::pair<BasicBlock *, BasicBlock *>> UncountableEdge;
 };
 
 } // namespace llvm

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -1631,12 +1631,11 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
 
   // Keep a record of all the exiting blocks.
   SmallVector<const SCEVPredicate *, 4> Predicates;
+  SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> UncountableEdges;
   for (BasicBlock *BB : ExitingBlocks) {
     const SCEV *EC =
         PSE.getSE()->getPredicatedExitCount(TheLoop, BB, &Predicates);
     if (isa<SCEVCouldNotCompute>(EC)) {
-      UncountableExitingBlocks.push_back(BB);
-
       SmallVector<BasicBlock *, 2> Succs(successors(BB));
       if (Succs.size() != 2) {
         reportVectorizationFailure(
@@ -1653,7 +1652,7 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
         assert(!TheLoop->contains(Succs[1]));
         ExitBlock = Succs[1];
       }
-      UncountableExitBlocks.push_back(ExitBlock);
+      UncountableEdges.push_back({BB, ExitBlock});
     } else
       CountableExitingBlocks.push_back(BB);
   }
@@ -1664,7 +1663,7 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
   Predicates.clear();
 
   // We only support one uncountable early exit.
-  if (getUncountableExitingBlocks().size() != 1) {
+  if (UncountableEdges.size() != 1) {
     reportVectorizationFailure(
         "Loop has too many uncountable exits",
         "Cannot vectorize early exit loop with more than one early exit",
@@ -1675,7 +1674,7 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
   // The only supported early exit loops so far are ones where the early
   // exiting block is a unique predecessor of the latch block.
   BasicBlock *LatchPredBB = LatchBB->getUniquePredecessor();
-  if (LatchPredBB != getUncountableEarlyExitingBlock()) {
+  if (LatchPredBB != UncountableEdges[0].first) {
     reportVectorizationFailure("Early exit is not the latch predecessor",
                                "Cannot vectorize early exit loop",
                                "EarlyExitNotLatchPredecessor", ORE, TheLoop);
@@ -1728,7 +1727,7 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
     }
 
   // The vectoriser cannot handle loads that occur after the early exit block.
-  assert(LatchBB->getUniquePredecessor() == getUncountableEarlyExitingBlock() &&
+  assert(LatchBB->getUniquePredecessor() == UncountableEdges[0].first &&
          "Expected latch predecessor to be the early exiting block");
 
   // TODO: Handle loops that may fault.
@@ -1751,6 +1750,7 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
   LLVM_DEBUG(dbgs() << "LV: Found an early exit loop with symbolic max "
                        "backedge taken count: "
                     << *SymbolicMaxBTC << '\n');
+  UncountableEdge = UncountableEdges[0];
   return true;
 }
 
@@ -1812,7 +1812,6 @@ bool LoopVectorizationLegality::canVectorize(bool UseVPlanNativePath) {
       return false;
   }
 
-  HasUncountableEarlyExit = false;
   if (isa<SCEVCouldNotCompute>(PSE.getBackedgeTakenCount())) {
     if (TheLoop->getExitingBlock()) {
       reportVectorizationFailure("Cannot vectorize uncountable loop",
@@ -1822,10 +1821,8 @@ bool LoopVectorizationLegality::canVectorize(bool UseVPlanNativePath) {
       else
         return false;
     } else {
-      HasUncountableEarlyExit = true;
       if (!isVectorizableEarlyExitLoop()) {
-        UncountableExitingBlocks.clear();
-        HasUncountableEarlyExit = false;
+        UncountableEdge = std::nullopt;
         if (DoExtraAnalysis)
           Result = false;
         else