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

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

@@ -391,25 +391,24 @@ 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 getUncountableEdges().size(); }
 
   /// 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 &&
+    assert(getUncountableEdges().size() == 1 &&
            "Expected only a single uncountable exiting block");
-    return getUncountableExitingBlocks()[0];
+    return getUncountableEdges()[0].first;
   }
 
   /// Returns the destination of an uncountable early exiting block.
   BasicBlock *getUncountableEarlyExitBlock() const {
-    assert(getUncountableExitBlocks().size() == 1 &&
+    if (!hasUncountableEarlyExit())
+      return nullptr;
+    assert(getUncountableEdges().size() == 1 &&
            "Expected only a single uncountable exit block");
-    return getUncountableExitBlocks()[0];
+    return getUncountableEdges()[0].second;
   }
 
   /// Returns true if vector representation of the instruction \p I
@@ -463,14 +462,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 all the loop edges that have an uncountable exit.
+  const SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> &
+  getUncountableEdges() const {
+    return UncountableEdges;
   }
 
 private:
@@ -654,18 +649,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 all the loop edges with uncountable exits, where each entry
+  /// is a pair of (Exiting, Exit) blocks.
+  SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> UncountableEdges;
 };
 
 } // namespace llvm

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -1635,8 +1635,6 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
     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 +1651,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 +1662,7 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
   Predicates.clear();
 
   // We only support one uncountable early exit.
-  if (getUncountableExitingBlocks().size() != 1) {
+  if (getUncountableEdges().size() != 1) {
     reportVectorizationFailure(
         "Loop has too many uncountable exits",
         "Cannot vectorize early exit loop with more than one early exit",
@@ -1812,7 +1810,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 +1819,8 @@ bool LoopVectorizationLegality::canVectorize(bool UseVPlanNativePath) {
       else
         return false;
     } else {
-      HasUncountableEarlyExit = true;
       if (!isVectorizableEarlyExitLoop()) {
-        UncountableExitingBlocks.clear();
-        HasUncountableEarlyExit = false;
+        UncountableEdges.clear();
         if (DoExtraAnalysis)
           Result = false;
         else