[LoopVectorize] Address comments on PR #107004 left post-commit (#109300)

david-arm · web-flow · commit f4eeae1244c8 · 2024-09-23T13:36:25.000+01:00
* Rename Speculative -&gt; Uncountable and update tests.
* Add comments explaining why it's safe to ignore the predicates when
building up a list of exiting blocks.
* Reshuffle some code to do (hopefully) cheaper checks first.
diff --git a/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h b/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h
@@ -377,19 +377,19 @@ class LoopVectorizationLegality {
     return LAI->getDepChecker().getMaxSafeVectorWidthInBits();
   }
 
-  /// Returns true if the loop has a speculative early exit, i.e. an
+  /// Returns true if the loop has an uncountable early exit, i.e. an
   /// uncountable exit that isn't the latch block.
-  bool hasSpeculativeEarlyExit() const { return HasSpeculativeEarlyExit; }
+  bool hasUncountableEarlyExit() const { return HasUncountableEarlyExit; }
 
-  /// Returns the speculative early exiting block.
-  BasicBlock *getSpeculativeEarlyExitingBlock() const {
+  /// Returns the uncountable early exiting block.
+  BasicBlock *getUncountableEarlyExitingBlock() const {
     assert(getUncountableExitingBlocks().size() == 1 &&
            "Expected only a single uncountable exiting block");
     return getUncountableExitingBlocks()[0];
   }
 
-  /// Returns the destination of a speculative early exiting block.
-  BasicBlock *getSpeculativeEarlyExitBlock() const {
+  /// 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];
@@ -603,15 +603,17 @@ class LoopVectorizationLegality {
   /// the use of those function variants.
   bool VecCallVariantsFound = false;
 
-  /// Indicates whether this loop has a speculative early exit, i.e. an
+  /// Indicates whether this loop has an uncountable early exit, i.e. an
   /// uncountable exiting block that is not the latch.
-  bool HasSpeculativeEarlyExit = false;
+  bool HasUncountableEarlyExit = false;
 
-  /// Keep track of all the loop exiting blocks.
+  /// 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.
+  /// Keep track of the destinations of all uncountable exits if the
+  /// exact backedge taken count is not computable.
   SmallVector<BasicBlock *, 4> UncountableExitBlocks;
 };
 
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
@@ -1473,13 +1473,13 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
 
   // Keep a record of all the exiting blocks.
   SmallVector<const SCEVPredicate *, 4> Predicates;
-  for (BasicBlock *BB1 : ExitingBlocks) {
+  for (BasicBlock *BB : ExitingBlocks) {
     const SCEV *EC =
-        PSE.getSE()->getPredicatedExitCount(TheLoop, BB1, &Predicates);
+        PSE.getSE()->getPredicatedExitCount(TheLoop, BB, &Predicates);
     if (isa<SCEVCouldNotCompute>(EC)) {
-      UncountableExitingBlocks.push_back(BB1);
+      UncountableExitingBlocks.push_back(BB);
 
-      SmallVector<BasicBlock *, 2> Succs(successors(BB1));
+      SmallVector<BasicBlock *, 2> Succs(successors(BB));
       if (Succs.size() != 2) {
         reportVectorizationFailure(
             "Early exiting block does not have exactly two successors",
@@ -1488,17 +1488,21 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
         return false;
       }
 
-      BasicBlock *BB2;
+      BasicBlock *ExitBlock;
       if (!TheLoop->contains(Succs[0]))
-        BB2 = Succs[0];
+        ExitBlock = Succs[0];
       else {
         assert(!TheLoop->contains(Succs[1]));
-        BB2 = Succs[1];
+        ExitBlock = Succs[1];
       }
-      UncountableExitBlocks.push_back(BB2);
+      UncountableExitBlocks.push_back(ExitBlock);
     } else
-      CountableExitingBlocks.push_back(BB1);
+      CountableExitingBlocks.push_back(BB);
   }
+  // We can safely ignore the predicates here because when vectorizing the loop
+  // the PredicatatedScalarEvolution class will keep track of all predicates
+  // for each exiting block anyway. This happens when calling
+  // PSE.getSymbolicMaxBackedgeTakenCount() below.
   Predicates.clear();
 
   // We only support one uncountable early exit.
@@ -1513,13 +1517,25 @@ 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 != getSpeculativeEarlyExitingBlock()) {
+  if (LatchPredBB != getUncountableEarlyExitingBlock()) {
     reportVectorizationFailure("Early exit is not the latch predecessor",
                                "Cannot vectorize early exit loop",
                                "EarlyExitNotLatchPredecessor", ORE, TheLoop);
     return false;
   }
 
+  // The latch block must have a countable exit.
+  if (isa<SCEVCouldNotCompute>(
+          PSE.getSE()->getPredicatedExitCount(TheLoop, LatchBB, &Predicates))) {
+    reportVectorizationFailure(
+        "Cannot determine exact exit count for latch block",
+        "Cannot vectorize early exit loop",
+        "UnknownLatchExitCountEarlyExitLoop", ORE, TheLoop);
+    return false;
+  }
+  assert(llvm::is_contained(CountableExitingBlocks, LatchBB) &&
+         "Latch block not found in list of countable exits!");
+
   // Check to see if there are instructions that could potentially generate
   // exceptions or have side-effects.
   auto IsSafeOperation = [](Instruction *I) -> bool {
@@ -1555,18 +1571,8 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
       }
     }
 
-  // The latch block must have a countable exit.
-  if (isa<SCEVCouldNotCompute>(
-          PSE.getSE()->getPredicatedExitCount(TheLoop, LatchBB, &Predicates))) {
-    reportVectorizationFailure(
-        "Cannot determine exact exit count for latch block",
-        "Cannot vectorize early exit loop",
-        "UnknownLatchExitCountEarlyExitLoop", ORE, TheLoop);
-    return false;
-  }
-
   // The vectoriser cannot handle loads that occur after the early exit block.
-  assert(LatchBB->getUniquePredecessor() == getSpeculativeEarlyExitingBlock() &&
+  assert(LatchBB->getUniquePredecessor() == getUncountableEarlyExitingBlock() &&
          "Expected latch predecessor to be the early exiting block");
 
   // TODO: Handle loops that may fault.
@@ -1580,16 +1586,15 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
     return false;
   }
 
-  LLVM_DEBUG(
-      dbgs()
-      << "LV: Found an early exit. Retrying with speculative exit count.\n");
-  [[maybe_unused]] const SCEV *SpecExitCount =
+  [[maybe_unused]] const SCEV *SymbolicMaxBTC =
       PSE.getSymbolicMaxBackedgeTakenCount();
-  assert(!isa<SCEVCouldNotCompute>(SpecExitCount) &&
+  // Since we have an exact exit count for the latch and the early exit
+  // dominates the latch, then this should guarantee a computed SCEV value.
+  assert(!isa<SCEVCouldNotCompute>(SymbolicMaxBTC) &&
          "Failed to get symbolic expression for backedge taken count");
-
-  LLVM_DEBUG(dbgs() << "LV: Found speculative backedge taken count: "
-                    << *SpecExitCount << '\n');
+  LLVM_DEBUG(dbgs() << "LV: Found an early exit loop with symbolic max "
+                       "backedge taken count: "
+                    << *SymbolicMaxBTC << '\n');
   return true;
 }
 
@@ -1653,15 +1658,15 @@ bool LoopVectorizationLegality::canVectorize(bool UseVPlanNativePath) {
       return false;
   }
 
-  HasSpeculativeEarlyExit = false;
+  HasUncountableEarlyExit = false;
   if (isa<SCEVCouldNotCompute>(PSE.getBackedgeTakenCount())) {
     if (!isVectorizableEarlyExitLoop()) {
       if (DoExtraAnalysis)
         Result = false;
       else
         return false;
     } else
-      HasSpeculativeEarlyExit = true;
+      HasUncountableEarlyExit = true;
   }
 
   // Go over each instruction and look at memory deps.
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -9794,11 +9794,12 @@ bool LoopVectorizePass::processLoop(Loop *L) {
     return false;
   }
 
-  if (LVL.hasSpeculativeEarlyExit()) {
-    reportVectorizationFailure(
-        "Auto-vectorization of early exit loops is not yet supported.",
-        "Auto-vectorization of early exit loops is not yet supported.",
-        "EarlyExitLoopsUnsupported", ORE, L);
+  if (LVL.hasUncountableEarlyExit()) {
+    reportVectorizationFailure("Auto-vectorization of loops with uncountable "
+                               "early exit is not yet supported",
+                               "Auto-vectorization of loops with uncountable "
+                               "early exit is not yet supported",
+                               "UncountableEarlyExitLoopsUnsupported", ORE, L);
     return false;
   }
 
diff --git a/llvm/test/Transforms/LoopVectorize/simple_early_exit.ll b/llvm/test/Transforms/LoopVectorize/simple_early_exit.ll
@@ -7,10 +7,9 @@ declare void @init_mem(ptr, i64);
 
 define i64 @same_exit_block_pre_inc_use1() {
 ; DEBUG-LABEL: LV: Checking a loop in 'same_exit_block_pre_inc_use1'
-; DEBUG:       LV: Found an early exit. Retrying with speculative exit count.
-; DEBUG-NEXT:  LV: Found speculative backedge taken count: 63
+; DEBUG:       LV: Found an early exit loop with symbolic max backedge taken count: 63
 ; DEBUG-NEXT:  LV: We can vectorize this loop!
-; DEBUG-NEXT:  LV: Not vectorizing: Auto-vectorization of early exit loops is not yet supported.
+; DEBUG-NEXT:  LV: Not vectorizing: Auto-vectorization of loops with uncountable early exit is not yet supported.
 ; CHECK-LABEL: define i64 @same_exit_block_pre_inc_use1() {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[P1:%.*]] = alloca [1024 x i8], align 1
@@ -1089,8 +1088,7 @@ loop.end:
 
 define i64 @loop_contains_safe_call() {
 ; DEBUG-LABEL: LV: Checking a loop in 'loop_contains_safe_call'
-; DEBUG:       LV: Found an early exit. Retrying with speculative exit count.
-; DEBUG-NEXT:  LV: Found speculative backedge taken count: 63
+; DEBUG:       LV: Found an early exit loop with symbolic max backedge taken count: 63
 ; DEBUG-NEXT:  LV: We can vectorize this loop!
 ; CHECK-LABEL: define i64 @loop_contains_safe_call() {
 ; CHECK-NEXT:  entry:
@@ -1193,8 +1191,7 @@ loop.end:
 
 define i64 @loop_contains_safe_div() {
 ; DEBUG-LABEL: LV: Checking a loop in 'loop_contains_safe_div'
-; DEBUG:       LV: Found an early exit. Retrying with speculative exit count.
-; DEBUG-NEXT:  LV: Found speculative backedge taken count: 63
+; DEBUG:       LV: Found an early exit loop with symbolic max backedge taken count: 63
 ; DEBUG-NEXT:  LV: We can vectorize this loop!
 ; CHECK-LABEL: define i64 @loop_contains_safe_div() {
 ; CHECK-NEXT:  entry:
@@ -1347,10 +1344,9 @@ loop.end:
 
 define i64 @loop_contains_load_after_early_exit(ptr dereferenceable(1024) align(8) %p2) {
 ; DEBUG-LABEL: LV: Checking a loop in 'loop_contains_load_after_early_exit'
-; DEBUG:       LV: Found an early exit. Retrying with speculative exit count.
-; DEBUG-NEXT:  LV: Found speculative backedge taken count: 63
+; DEBUG:       LV: Found an early exit loop with symbolic max backedge taken count: 63
 ; DEBUG-NEXT:  LV: We can vectorize this loop!
-; DEBUG-NEXT:  LV: Not vectorizing: Auto-vectorization of early exit loops is not yet supported.
+; DEBUG-NEXT:  LV: Not vectorizing: Auto-vectorization of loops with uncountable early exit is not yet supported.
 ; CHECK-LABEL: define i64 @loop_contains_load_after_early_exit(
 ; CHECK-SAME: ptr align 8 dereferenceable(1024) [[P2:%.*]]) {
 ; CHECK-NEXT:  entry:
@@ -1623,10 +1619,9 @@ loop.end:
 ; The form of the induction variables requires SCEV predicates.
 define i32 @diff_exit_block_needs_scev_check(i32 %end) {
 ; DEBUG-LABEL: LV: Checking a loop in 'diff_exit_block_needs_scev_check'
-; DEBUG:       LV: Found an early exit. Retrying with speculative exit count.
-; DEBUG-NEXT:  LV: Found speculative backedge taken count: (-1 + (1 umax (zext i10 (trunc i32 %end to i10) to i32)))<nsw>
+; DEBUG:       Found an early exit loop with symbolic max backedge taken count: (-1 + (1 umax (zext i10 (trunc i32 %end to i10) to i32)))<nsw>
 ; DEBUG-NEXT:  LV: We can vectorize this loop!
-; DEBUG-NEXT:  LV: Not vectorizing: Auto-vectorization of early exit loops is not yet supported.
+; DEBUG-NEXT:  LV: Not vectorizing: Auto-vectorization of loops with uncountable early exit is not yet supported.
 ; CHECK-LABEL: define i32 @diff_exit_block_needs_scev_check(
 ; CHECK-SAME: i32 [[END:%.*]]) {
 ; CHECK-NEXT:  entry:
@@ -1695,9 +1690,8 @@ declare void @abort()
 ; early is loop invariant.
 define i32 @diff_blocks_invariant_early_exit_cond(ptr %s) {
 ; DEBUG-LABEL: LV: Checking a loop in 'diff_blocks_invariant_early_exit_cond'
-; DEBUG:       LV: Found an early exit. Retrying with speculative exit count.
-; DEBUG-NEXT:  LV: Found speculative backedge taken count: 275
-; DEBUG:       LV: Not vectorizing: Auto-vectorization of early exit loops is not yet supported.
+; DEBUG:       LV: Found an early exit loop with symbolic max backedge taken count: 275
+; DEBUG:       LV: Not vectorizing: Auto-vectorization of loops with uncountable early exit is not yet supported.
 ; CHECK-LABEL: define i32 @diff_blocks_invariant_early_exit_cond(
 ; CHECK-SAME: ptr [[S:%.*]]) {
 ; CHECK-NEXT:  entry:
