Address review comments

* Reverted LoopAccessAnalysis changes in favour of new code in
LoopVectorizationLegality.cpp that checks if instructions are
safe to speculatively execute.
* Updated some comments and done minor code refactoring.
* Added new tests for demonstrating under which circumstances
we can speculatively execute calls and divides.
* Added new test where the early exit is in a conditional block.
* Added new test containing store instructions.
* Added new test containing loads after early exit.

Author: david-arm

llvm/include/llvm/Analysis/LoopAccessAnalysis.h

@@ -661,8 +661,7 @@ class LoopAccessInfo {
   bool isInvariant(Value *V) const;
 
   unsigned getNumStores() const { return NumStores; }
-  unsigned getNumLoads() const { return NumLoads; }
-  unsigned getNumCalls() const { return NumCalls; }
+  unsigned getNumLoads() const { return NumLoads;}
 
   /// The diagnostics report generated for the analysis.  E.g. why we
   /// couldn't analyze the loop.
@@ -755,7 +754,6 @@ class LoopAccessInfo {
 
   unsigned NumLoads = 0;
   unsigned NumStores = 0;
-  unsigned NumCalls = 0;
 
   /// Cache the result of analyzeLoop.
   bool CanVecMem = false;

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

@@ -480,8 +480,7 @@ class LoopVectorizationLegality {
   /// specific checks for outer loop vectorization.
   bool canVectorizeOuterLoop();
 
-  /// Returns true if this is a supported early exit loop that we can
-  /// vectorize.
+  /// Returns true if this is an early exit loop that can be vectorized.
   bool isVectorizableEarlyExitLoop();
 
   /// Return true if all of the instructions in the block can be speculatively
@@ -594,10 +593,11 @@ class LoopVectorizationLegality {
   /// uncountable exiting block that is not the latch.
   bool HasSpeculativeEarlyExit = false;
 
-  /// Keeps track of all the exits with known or countable exit-not-taken
-  /// counts.
+  /// Keep track of all the loop exiting blocks.
   SmallVector<BasicBlock *, 4> CountableExitingBlocks;
   SmallVector<BasicBlock *, 4> UncountableExitingBlocks;
+
+  /// Keep track of the destinations of all uncountable exits.
   SmallVector<BasicBlock *, 4> UncountableExitBlocks;
 };
 

llvm/lib/Analysis/LoopAccessAnalysis.cpp

@@ -2445,11 +2445,8 @@ bool LoopAccessInfo::analyzeLoop(AAResults *AA, const LoopInfo *LI,
       // vectorize a loop if it contains known function calls that don't set
       // the flag. Therefore, it is safe to ignore this read from memory.
       auto *Call = dyn_cast<CallInst>(&I);
-      if (Call) {
-        NumCalls++;
-        if (getVectorIntrinsicIDForCall(Call, TLI))
-          continue;
-      }
+      if (Call && getVectorIntrinsicIDForCall(Call, TLI))
+        continue;
 
       // If this is a load, save it. If this instruction can read from memory
       // but is not a load, then we quit. Notice that we don't handle function

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -1086,16 +1086,12 @@ bool LoopVectorizationLegality::canVectorizeMemory(bool IsEarlyExitLoop) {
       return false;
     }
 
-    if (LAI->getNumCalls()) {
-      reportVectorizationFailure(
-          "Calls unsupported in early exit loops",
-          "Cannot vectorize early exit loop with function calls",
-          "CallsInEarlyExitLoop", 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(
@@ -1104,10 +1100,15 @@ bool LoopVectorizationLegality::canVectorizeMemory(bool IsEarlyExitLoop) {
             "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",
@@ -1499,32 +1500,42 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
     return false;
   }
 
+  if (Reductions.size() || FixedOrderRecurrences.size()) {
+    reportVectorizationFailure(
+        "Found reductions or recurrences in early-exit loop",
+        "Cannot vectorize early exit loop with reductions or recurrences",
+        "RecurrencesInEarlyExitLoop", ORE, TheLoop);
+    return false;
+  }
+
   SmallVector<BasicBlock *, 8> ExitingBlocks;
   TheLoop->getExitingBlocks(ExitingBlocks);
 
-  // Keep a record of all the exiting blocks with exact exit counts, as well as
-  // those with inexact counts.
+  // Keep a record of all the exiting blocks.
   SmallVector<const SCEVPredicate *, 4> Predicates;
   for (BasicBlock *BB1 : ExitingBlocks) {
     const SCEV *EC =
         PSE.getSE()->getPredicatedExitCount(TheLoop, BB1, &Predicates);
     if (isa<SCEVCouldNotCompute>(EC)) {
       UncountableExitingBlocks.push_back(BB1);
 
-      unsigned NumExitBlocks = 0;
-      for (BasicBlock *BB2 : successors(BB1)) {
-        if (!TheLoop->contains(BB2)) {
-          UncountableExitBlocks.push_back(BB2);
-          NumExitBlocks++;
-        }
-      }
-      if (NumExitBlocks > 1) {
+      SmallVector<BasicBlock *, 2> Succs(successors(BB1));
+      if (Succs.size() != 2) {
         reportVectorizationFailure(
-            "Early exiting block has more than one successor outside of loop",
-            "Too many successors from early exiting block",
+            "Early exiting block does not have exactly two successors",
+            "Incorrect number of successors from early exiting block",
             "EarlyExitTooManySuccessors", ORE, TheLoop);
         return false;
       }
+
+      BasicBlock *BB2;
+      if (!TheLoop->contains(Succs[0]))
+        BB2 = Succs[0];
+      else {
+        assert(!TheLoop->contains(Succs[1]));
+        BB2 = Succs[1];
+      }
+      UncountableExitBlocks.push_back(BB2);
     } else
       CountableExitingBlocks.push_back(BB1);
   }
@@ -1549,13 +1560,34 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
     return false;
   }
 
-  if (Reductions.size() || FixedOrderRecurrences.size()) {
-    reportVectorizationFailure(
-        "Found reductions or recurrences in early-exit loop",
-        "Cannot vectorize early exit loop with reductions or recurrences",
-        "RecurrencesInEarlyExitLoop", ORE, TheLoop);
-    return false;
-  }
+  // Check all instructions in the loop to see if they 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.
+      return true;
+    default:
+      return isSafeToSpeculativelyExecute(I);
+    }
+  };
+
+  for (auto *BB : TheLoop->blocks())
+    for (auto &I : *BB)
+      if (!IsSafeOperation(&I)) {
+        reportVectorizationFailure("Early exit loop contains operations that "
+                                   "cannot be speculatively executed",
+                                   "Early exit loop contains operations that "
+                                   "cannot be speculatively executed",
+                                   "UnsafeOperationsEarlyExitLoop", ORE,
+                                   TheLoop);
+        return false;
+      }
 
   LLVM_DEBUG(
       dbgs()