Address review comments

* Combined changes from canVectorizeMemory into
isVectorizableEarlyExitLoop.
* Documented restrictions for isVectorizableEarlyExitLoop in
header file.
* Added reportVectorizationFailure call in LoopVectorize when
we bail out.
* General code refactor, clean-up.

Author: david-arm

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

@@ -470,7 +470,7 @@ class LoopVectorizationLegality {
   /// we read and write from memory. This method checks if it is
   /// legal to vectorize the code, considering only memory constrains.
   /// Returns true if the loop is vectorizable
-  bool canVectorizeMemory(bool IsEarlyExitLoop);
+  bool canVectorizeMemory();
 
   /// Return true if we can vectorize this loop using the IF-conversion
   /// transformation.
@@ -481,6 +481,21 @@ class LoopVectorizationLegality {
   bool canVectorizeOuterLoop();
 
   /// Returns true if this is an early exit loop that can be vectorized.
+  /// Currently, a loop with an uncountable early exit is considered
+  /// vectorizable if:
+  ///   1. There are no writes to memory in the loop.
+  ///   2. The loop has only one early uncountable exit
+  ///   3. The early exit block dominates the latch block.
+  ///   4. The latch block has an exact exit count.
+  ///   5. There are no loads after the early exit block.
+  ///   6. The loop does not contain reductions or recurrences.
+  ///   7. We can prove at compile-time that loops will not contain faulting
+  ///   loads.
+  ///   8. It is safe to speculatively execute instructions such as divide or
+  ///   call instructions.
+  /// The list above is not based on theoretical limitations of vectorization,
+  /// but simply a statement that more work is needed to support these
+  /// additional cases safely.
   bool isVectorizableEarlyExitLoop();
 
   /// Return true if all of the instructions in the block can be speculatively

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -1051,7 +1051,7 @@ bool LoopVectorizationLegality::canVectorizeInstrs() {
   return true;
 }
 
-bool LoopVectorizationLegality::canVectorizeMemory(bool IsEarlyExitLoop) {
+bool LoopVectorizationLegality::canVectorizeMemory() {
   LAI = &LAIs.getInfo(*TheLoop);
   const OptimizationRemarkAnalysis *LAR = LAI->getReport();
   if (LAR) {
@@ -1073,51 +1073,6 @@ bool LoopVectorizationLegality::canVectorizeMemory(bool IsEarlyExitLoop) {
     return false;
   }
 
-  // For loops with uncountable early exiting blocks that are not the latch
-  // it's necessary to perform extra checks, since the vectoriser is currently
-  // only capable of handling simple search loops.
-  if (IsEarlyExitLoop) {
-    // We don't support calls or any memory accesses that write to memory.
-    if (LAI->getNumStores()) {
-      reportVectorizationFailure(
-          "Writes to memory unsupported in early exit loops",
-          "Cannot vectorize early exit loop with writes to memory",
-          "WritesInEarlyExitLoop", ORE, TheLoop);
-      return false;
-    }
-
-    // The vectoriser cannot handle loads that occur after the early exit block.
-    BasicBlock *LatchBB = TheLoop->getLoopLatch();
-    assert(LatchBB->getUniquePredecessor() ==
-               getUncountableExitingBlocks()[0] &&
-           "Expected latch predecessor to be the early exiting block");
-
-    for (Instruction &I : *LatchBB) {
-      if (I.mayReadFromMemory()) {
-        reportVectorizationFailure(
-            "Loads not permitted after early exit",
-            "Cannot vectorize early exit loop with loads after early exit",
-            "LoadsAfterEarlyExit", ORE, TheLoop);
-        return false;
-      }
-      // Any other problematic instructions should have been caught earlier.
-      assert(!I.mayWriteToMemory() && !I.mayThrow() &&
-             !I.mayHaveSideEffects() &&
-             "Unexpected instructions in latch block of early exit loop");
-    }
-
-    // The vectoriser does not yet handle loops that may fault, but this will
-    // be improved in a follow-on patch.
-    // TODO: Handle loops that may fault.
-    if (!isDereferenceableReadOnlyLoop(TheLoop, PSE.getSE(), DT, AC)) {
-      reportVectorizationFailure(
-          "Loop may fault",
-          "Cannot vectorize potentially faulting early exit loop",
-          "PotentiallyFaultingEarlyExitLoop", ORE, TheLoop);
-      return false;
-    }
-  }
-
   // We can vectorize stores to invariant address when final reduction value is
   // guaranteed to be stored at the end of the loop. Also, if decision to
   // vectorize loop is made, runtime checks are added so as to make sure that
@@ -1491,7 +1446,6 @@ bool LoopVectorizationLegality::canVectorizeLoopNestCFG(
 }
 
 bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
-  // At least one of the exiting blocks must be the latch.
   BasicBlock *LatchBB = TheLoop->getLoopLatch();
   if (!LatchBB) {
     reportVectorizationFailure("Loop does not have a latch",
@@ -1553,33 +1507,38 @@ 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 || LatchPredBB != getUncountableExitingBlocks()[0]) {
+  if (LatchPredBB != getSpeculativeEarlyExitingBlock()) {
     reportVectorizationFailure("Early exit is not the latch predecessor",
                                "Cannot vectorize early exit loop",
                                "EarlyExitNotLatchPredecessor", ORE, TheLoop);
     return false;
   }
 
-  // Check all instructions in the loop to see if they could potentially
-  // generate exceptions or have side-effects.
+  // Check to see if there are instructions that could potentially generate
+  // exceptions or have side-effects.
   auto IsSafeOperation = [](Instruction *I) -> bool {
-    // Is this a divide?
     switch (I->getOpcode()) {
     case Instruction::Load:
     case Instruction::Store:
     case Instruction::PHI:
     case Instruction::Br:
-      // These are checked separately. For example, canVectorizeMemory will
-      // analyze the loads and stores in the loop.
+      // These are checked separately.
       return true;
     default:
       return isSafeToSpeculativelyExecute(I);
     }
   };
 
   for (auto *BB : TheLoop->blocks())
-    for (auto &I : *BB)
-      if (!IsSafeOperation(&I)) {
+    for (auto &I : *BB) {
+      if (I.mayWriteToMemory()) {
+        // We don't support writes to memory.
+        reportVectorizationFailure(
+            "Writes to memory unsupported in early exit loops",
+            "Cannot vectorize early exit loop with writes to memory",
+            "WritesInEarlyExitLoop", ORE, TheLoop);
+        return false;
+      } else if (!IsSafeOperation(&I)) {
         reportVectorizationFailure("Early exit loop contains operations that "
                                    "cannot be speculatively executed",
                                    "Early exit loop contains operations that "
@@ -1588,10 +1547,9 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
                                    TheLoop);
         return false;
       }
+    }
 
-  LLVM_DEBUG(
-      dbgs()
-      << "LV: Found an early exit. Retrying with speculative exit count.\n");
+  // At least one of the exiting blocks must be the latch.
   if (isa<SCEVCouldNotCompute>(
           PSE.getSE()->getPredicatedExitCount(TheLoop, LatchBB, &Predicates))) {
     reportVectorizationFailure(
@@ -1601,6 +1559,34 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
     return false;
   }
 
+  // The vectoriser cannot handle loads that occur after the early exit block.
+  assert(LatchBB->getUniquePredecessor() == getSpeculativeEarlyExitingBlock() &&
+         "Expected latch predecessor to be the early exiting block");
+  for (Instruction &I : *LatchBB) {
+    if (I.mayReadFromMemory()) {
+      reportVectorizationFailure(
+          "Loads not permitted after early exit",
+          "Cannot vectorize early exit loop with loads after early exit",
+          "LoadsAfterEarlyExit", ORE, TheLoop);
+      return false;
+    }
+    // Any other problematic instructions should have been caught earlier.
+    assert(!I.mayWriteToMemory() && !I.mayThrow() && !I.mayHaveSideEffects() &&
+           "Unexpected instructions in latch block of early exit loop");
+  }
+
+  // TODO: Handle loops that may fault.
+  if (!isDereferenceableReadOnlyLoop(TheLoop, PSE.getSE(), DT, AC)) {
+    reportVectorizationFailure(
+        "Loop may fault",
+        "Cannot vectorize potentially faulting early exit loop",
+        "PotentiallyFaultingEarlyExitLoop", ORE, TheLoop);
+    return false;
+  }
+
+  LLVM_DEBUG(
+      dbgs()
+      << "LV: Found an early exit. Retrying with speculative exit count.\n");
   const SCEV *SpecExitCount = PSE.getSymbolicMaxBackedgeTakenCount();
   assert(!isa<SCEVCouldNotCompute>(SpecExitCount) &&
          "Failed to get symbolic expression for backedge taken count");
@@ -1682,7 +1668,7 @@ bool LoopVectorizationLegality::canVectorize(bool UseVPlanNativePath) {
   }
 
   // Go over each instruction and look at memory deps.
-  if (!canVectorizeMemory(HasSpeculativeEarlyExit)) {
+  if (!canVectorizeMemory()) {
     LLVM_DEBUG(dbgs() << "LV: Can't vectorize due to memory conflicts\n");
     if (DoExtraAnalysis)
       Result = false;

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9816,8 +9816,10 @@ bool LoopVectorizePass::processLoop(Loop *L) {
   }
 
   if (LVL.hasSpeculativeEarlyExit()) {
-    LLVM_DEBUG(dbgs() << "LV: Not vectorizing: Auto-vectorization of early "
-                      << "exit loops is not yet supported.\n");
+    reportVectorizationFailure(
+        "Auto-vectorization of early exit loops is not yet supported.",
+        "Auto-vectorization of early exit loops is not yet supported.",
+        "EarlyExitLoopsUnsupported", ORE, L);
     return false;
   }
 

llvm/test/Transforms/LoopVectorize/AArch64/simple_early_exit.ll

@@ -1648,9 +1648,7 @@ loop.end:
 ; requesting the small constant max trip count.
 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-NEXT:  LV: Not vectorizing: Loop may fault.
+; DEBUG:       LV: Not vectorizing: Loop may fault.
 ; CHECK-LABEL: define i32 @diff_exit_block_needs_scev_check(
 ; CHECK-SAME: i32 [[END:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:  entry:
@@ -1962,7 +1960,7 @@ loop.end:
 
 
 
-declare i32 @foo(i32)
+declare i32 @foo(i32) readonly
 declare <vscale x 4 x i32> @foo_vec(<vscale x 4 x i32>)
 
 attributes #0 = { "vector-function-abi-variant"="_ZGVsNxv_foo(foo_vec)" }