Address review comments + rebase test updates

Author: david-arm

llvm/include/llvm/Analysis/Loads.h

@@ -88,13 +88,6 @@ bool isDereferenceableAndAlignedInLoop(
     AssumptionCache *AC = nullptr,
     SmallVectorImpl<const SCEVPredicate *> *Predicates = nullptr);
 
-/// Return true if the loop \p L cannot fault on any iteration and only
-/// contains read-only memory accesses.
-bool isDereferenceableReadOnlyLoop(
-    Loop *L, ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
-    SmallVectorImpl<const SCEVPredicate *> *Predicates = nullptr,
-    SmallVectorImpl<LoadInst *> *NonDerefLoads = nullptr);
-
 /// Return true if we know that executing a load from this value cannot trap.
 ///
 /// If DT and ScanFrom are specified this method performs context-sensitive

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

@@ -382,11 +382,18 @@ class LoopVectorizationLegality {
   const LoopAccessInfo *getLAI() const { return LAI; }
 
   bool isSafeForAnyVectorWidth() const {
-    return LAI->getDepChecker().isSafeForAnyVectorWidth();
+    return LAI->getDepChecker().isSafeForAnyVectorWidth() &&
+           (!hasUncountableEarlyExit() || !getNumPotentiallyFaultingPointers());
   }
 
   uint64_t getMaxSafeVectorWidthInBits() const {
-    return LAI->getDepChecker().getMaxSafeVectorWidthInBits();
+    uint64_t MaxSafeVectorWidth =
+        LAI->getDepChecker().getMaxSafeVectorWidthInBits();
+    // The legalizer bails out if getMinPageSize does not return a value.
+    if (hasUncountableEarlyExit() && getNumPotentiallyFaultingPointers())
+      MaxSafeVectorWidth =
+          std::min(MaxSafeVectorWidth, uint64_t(*TTI->getMinPageSize()) * 8);
+    return MaxSafeVectorWidth;
   }
 
   /// Returns true if the loop has exactly one uncountable early exit, i.e. an
@@ -423,17 +430,17 @@ class LoopVectorizationLegality {
   unsigned getNumStores() const { return LAI->getNumStores(); }
   unsigned getNumLoads() const { return LAI->getNumLoads(); }
 
-  /// Return the number of loads in the loop that could potentially fault in a
-  /// loop with uncountable early exits.
-  unsigned getNumPotentiallyFaultingLoads() const {
-    return PotentiallyFaultingLoads.size();
+  /// Return the number of pointers in the loop that could potentially fault in
+  /// a loop with uncountable early exits.
+  unsigned getNumPotentiallyFaultingPointers() const {
+    return PotentiallyFaultingPtrs.size();
   }
 
-  /// Return a vector of all potentially faulting loads in a loop with
+  /// Return a vector of all potentially faulting pointers in a loop with
   /// uncountable early exits.
-  const SmallVectorImpl<std::pair<LoadInst *, const SCEV *>> *
-  getPotentiallyFaultingLoads() const {
-    return &PotentiallyFaultingLoads;
+  const SmallVectorImpl<std::pair<const SCEV *, Type *>> *
+  getPotentiallyFaultingPointers() const {
+    return &PotentiallyFaultingPtrs;
   }
 
   /// Returns a HistogramInfo* for the given instruction if it was determined
@@ -543,7 +550,7 @@ class LoopVectorizationLegality {
 
   /// Returns true if all loads in the loop contained in \p Loads can be
   /// analyzed as potentially faulting. Any loads that may fault are added to
-  /// the member variable PotentiallyFaultingLoads.
+  /// the member variable PotentiallyFaultingPtrs.
   bool analyzePotentiallyFaultingLoads(SmallVectorImpl<LoadInst *> *Loads);
 
   /// Return true if all of the instructions in the block can be speculatively
@@ -671,9 +678,9 @@ class LoopVectorizationLegality {
   /// of (Exiting, Exit) blocks, if there is exactly one early exit.
   std::optional<std::pair<BasicBlock *, BasicBlock *>> UncountableEdge;
 
-  /// Keep a record of all potentially faulting loads in loops with
+  /// Keep a record of all potentially faulting pointers in loops with
   /// uncountable early exits.
-  SmallVector<std::pair<LoadInst *, const SCEV *>, 4> PotentiallyFaultingLoads;
+  SmallVector<std::pair<const SCEV *, Type *>, 4> PotentiallyFaultingPtrs;
 };
 
 } // namespace llvm

llvm/lib/Analysis/Loads.cpp

@@ -813,29 +813,3 @@ bool llvm::canReplacePointersIfEqual(const Value *From, const Value *To,
 
   return isPointerAlwaysReplaceable(From, To, DL);
 }
-
-bool llvm::isDereferenceableReadOnlyLoop(
-    Loop *L, ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
-    SmallVectorImpl<const SCEVPredicate *> *Predicates,
-    SmallVectorImpl<LoadInst *> *NonDerefLoads) {
-  bool Result = true;
-  for (BasicBlock *BB : L->blocks()) {
-    for (Instruction &I : *BB) {
-      if (auto *LI = dyn_cast<LoadInst>(&I)) {
-        if (!isDereferenceableAndAlignedInLoop(LI, L, *SE, *DT, AC,
-                                               Predicates)) {
-          if (!NonDerefLoads)
-            return false;
-          NonDerefLoads->push_back(LI);
-          Result = false;
-        }
-      } else if (I.mayReadFromMemory() || I.mayWriteToMemory() ||
-                 I.mayThrow()) {
-        if (!NonDerefLoads)
-          return false;
-        Result = false;
-      }
-    }
-  }
-  return Result;
-}

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -1641,7 +1641,7 @@ bool LoopVectorizationLegality::analyzePotentiallyFaultingLoads(
       return false;
 
     LLVM_DEBUG(dbgs() << "LV: SCEV for Load Ptr: " << *Start << '\n');
-    PotentiallyFaultingLoads.push_back({LI, Start});
+    PotentiallyFaultingPtrs.push_back({Start, LI->getType()});
   }
   return true;
 }
@@ -1750,6 +1750,8 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
     }
   };
 
+  Predicates.clear();
+  SmallVector<LoadInst *, 4> NonDerefLoads;
   for (auto *BB : TheLoop->blocks())
     for (auto &I : *BB) {
       if (I.mayWriteToMemory()) {
@@ -1759,24 +1761,30 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
             "Cannot vectorize early exit loop with writes to memory",
             "WritesInEarlyExitLoop", ORE, TheLoop);
         return false;
-      } else if (!IsSafeOperation(&I)) {
+      } else if (I.mayThrow() || !IsSafeOperation(&I)) {
         reportVectorizationFailure("Early exit loop contains operations that "
                                    "cannot be speculatively executed",
                                    "UnsafeOperationsEarlyExitLoop", ORE,
                                    TheLoop);
         return false;
+      } else if (I.mayReadFromMemory()) {
+        auto *LI = dyn_cast<LoadInst>(&I);
+        if (!LI) {
+          reportVectorizationFailure(
+              "Loop may fault",
+              "Cannot vectorize potentially faulting early exit loop",
+              "PotentiallyFaultingEarlyExitLoop", ORE, TheLoop);
+          return false;
+        } else if (!isDereferenceableAndAlignedInLoop(LI, TheLoop, *PSE.getSE(),
+                                                      *DT, AC, &Predicates))
+          NonDerefLoads.push_back(LI);
       }
     }
 
-  // The vectoriser cannot handle loads that occur after the early exit block.
-  assert(LatchBB->getUniquePredecessor() == SingleUncountableEdge->first &&
-         "Expected latch predecessor to be the early exiting block");
-
-  Predicates.clear();
-  SmallVector<LoadInst *, 4> Loads;
-  if (!isDereferenceableReadOnlyLoop(TheLoop, PSE.getSE(), DT, AC, &Predicates,
-                                     &Loads)) {
-    if (!TTI->getMinPageSize() || !analyzePotentiallyFaultingLoads(&Loads)) {
+  if (!NonDerefLoads.empty()) {
+    if (!TTI->getMinPageSize() ||
+        !analyzePotentiallyFaultingLoads(&NonDerefLoads)) {
+      PotentiallyFaultingPtrs.clear();
       reportVectorizationFailure(
           "Loop may fault",
           "Cannot vectorize potentially faulting early exit loop",
@@ -1786,6 +1794,10 @@ bool LoopVectorizationLegality::isVectorizableEarlyExitLoop() {
     LLVM_DEBUG(dbgs() << "We can vectorize the loop with runtime checks.\n");
   }
 
+  // The vectoriser cannot handle loads that occur after the early exit block.
+  assert(LatchBB->getUniquePredecessor() == SingleUncountableEdge->first &&
+         "Expected latch predecessor to be the early exiting block");
+
   [[maybe_unused]] const SCEV *SymbolicMaxBTC =
       PSE.getSymbolicMaxBackedgeTakenCount();
   // Since we have an exact exit count for the latch and the early exit

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -402,7 +402,7 @@ static cl::opt<bool> EnableEarlyExitVectorization(
         "Enable vectorization of early exit loops with uncountable exits."));
 
 static cl::opt<unsigned> MaxNumPotentiallyFaultingPointers(
-    "max-num-faulting-pointers", cl::init(1), cl::Hidden,
+    "max-num-faulting-pointers", cl::init(0), cl::Hidden,
     cl::desc(
         "The maximum number of potentially faulting pointers we permit when "
         "vectorizing loops with uncountable exits."));
@@ -1621,22 +1621,6 @@ class LoopVectorizationCostModel {
                                        ElementCount MaxSafeVF,
                                        bool FoldTailByMasking);
 
-  bool isSafeForAnyVectorWidth() const {
-    return Legal->isSafeForAnyVectorWidth() &&
-           (!Legal->hasUncountableEarlyExit() ||
-            !Legal->getNumPotentiallyFaultingLoads());
-  }
-
-  uint64_t getMaxSafeVectorWidthInBits() const {
-    uint64_t MaxSafeVectorWidth = Legal->getMaxSafeVectorWidthInBits();
-    // The legalizer bails out if getMinPageSize does not return a value.
-    if (Legal->hasUncountableEarlyExit() &&
-        Legal->getNumPotentiallyFaultingLoads())
-      MaxSafeVectorWidth =
-          std::min(MaxSafeVectorWidth, uint64_t(*TTI.getMinPageSize()) * 8);
-    return MaxSafeVectorWidth;
-  }
-
   /// Checks if scalable vectorization is supported and enabled. Caches the
   /// result to avoid repeated debug dumps for repeated queries.
   bool isScalableVectorizationAllowed();
@@ -2185,38 +2169,24 @@ class GeneratedRTChecks {
 };
 } // namespace
 
-std::optional<unsigned> getMaxVScale(const Function &F,
-                                     const TargetTransformInfo &TTI) {
-  if (std::optional<unsigned> MaxVScale = TTI.getMaxVScale())
-    return MaxVScale;
-
-  if (F.hasFnAttribute(Attribute::VScaleRange))
-    return F.getFnAttribute(Attribute::VScaleRange).getVScaleRangeMax();
-
-  return std::nullopt;
-}
-
 static void addPointerAlignmentChecks(
-    const SmallVectorImpl<std::pair<LoadInst *, const SCEV *>> *Loads,
-    Function *F, PredicatedScalarEvolution &PSE, TargetTransformInfo *TTI,
-    ElementCount VF) {
+    const SmallVectorImpl<std::pair<const SCEV *, Type *>> *Ptrs, Function *F,
+    PredicatedScalarEvolution &PSE, TargetTransformInfo *TTI, ElementCount VF,
+    unsigned IC) {
   ScalarEvolution *SE = PSE.getSE();
   const DataLayout &DL = SE->getDataLayout();
-  Type *PtrIntType = DL.getIntPtrType(SE->getContext());
 
-  const SCEV *Zero = SE->getZero(PtrIntType);
-  const SCEV *ScevEC = SE->getElementCount(PtrIntType, VF);
-
-  for (auto Load : *Loads) {
-    APInt EltSize(
-        DL.getIndexTypeSizeInBits(Load.first->getPointerOperandType()),
-        DL.getTypeStoreSize(Load.first->getType()).getFixedValue());
-    const SCEV *Start = SE->getPtrToIntExpr(Load.second, PtrIntType);
+  for (auto Ptr : *Ptrs) {
+    Type *PtrIntType = DL.getIntPtrType(Ptr.first->getType());
+    APInt EltSize(PtrIntType->getScalarSizeInBits(),
+                  DL.getTypeStoreSize(Ptr.second).getFixedValue());
+    const SCEV *Start = SE->getPtrToIntExpr(Ptr.first, PtrIntType);
+    const SCEV *ScevEC = SE->getElementCount(PtrIntType, VF * IC);
     const SCEV *Align =
         SE->getMulExpr(ScevEC, SE->getConstant(EltSize),
                        (SCEV::NoWrapFlags)(SCEV::FlagNSW | SCEV::FlagNUW));
     const SCEV *Rem = SE->getURemExpr(Start, Align);
-    PSE.addPredicate(*(SE->getEqualPredicate(Rem, Zero)));
+    PSE.addPredicate(*(SE->getEqualPredicate(Rem, SE->getZero(PtrIntType))));
   }
 }
 
@@ -2389,6 +2359,17 @@ emitTransformedIndex(IRBuilderBase &B, Value *Index, Value *StartValue,
   llvm_unreachable("invalid enum");
 }
 
+std::optional<unsigned> getMaxVScale(const Function &F,
+                                     const TargetTransformInfo &TTI) {
+  if (std::optional<unsigned> MaxVScale = TTI.getMaxVScale())
+    return MaxVScale;
+
+  if (F.hasFnAttribute(Attribute::VScaleRange))
+    return F.getFnAttribute(Attribute::VScaleRange).getVScaleRangeMax();
+
+  return std::nullopt;
+}
+
 /// For the given VF and UF and maximum trip count computed for the loop, return
 /// whether the induction variable might overflow in the vectorized loop. If not,
 /// then we know a runtime overflow check always evaluates to false and can be
@@ -3881,15 +3862,15 @@ bool LoopVectorizationCostModel::isScalableVectorizationAllowed() {
     return false;
   }
 
-  if (!isSafeForAnyVectorWidth() && !getMaxVScale(*TheFunction, TTI)) {
+  if (!Legal->isSafeForAnyVectorWidth() && !getMaxVScale(*TheFunction, TTI)) {
     reportVectorizationInfo("The target does not provide maximum vscale value "
                             "for safe distance analysis.",
                             "ScalableVFUnfeasible", ORE, TheLoop);
     return false;
   }
 
   if (Legal->hasUncountableEarlyExit() &&
-      Legal->getNumPotentiallyFaultingLoads() &&
+      Legal->getNumPotentiallyFaultingPointers() &&
       !TTI.isVScaleKnownToBeAPowerOfTwo()) {
     reportVectorizationInfo("Cannot vectorize potentially faulting early exit "
                             "loop with scalable vectors.",
@@ -3908,7 +3889,7 @@ LoopVectorizationCostModel::getMaxLegalScalableVF(unsigned MaxSafeElements) {
 
   auto MaxScalableVF = ElementCount::getScalable(
       std::numeric_limits<ElementCount::ScalarTy>::max());
-  if (isSafeForAnyVectorWidth())
+  if (Legal->isSafeForAnyVectorWidth())
     return MaxScalableVF;
 
   std::optional<unsigned> MaxVScale = getMaxVScale(*TheFunction, TTI);
@@ -3935,11 +3916,11 @@ FixedScalableVFPair LoopVectorizationCostModel::computeFeasibleMaxVF(
   // the memory accesses that is most restrictive (involved in the smallest
   // dependence distance).
   unsigned MaxSafeElements =
-      llvm::bit_floor(getMaxSafeVectorWidthInBits() / WidestType);
+      llvm::bit_floor(Legal->getMaxSafeVectorWidthInBits() / WidestType);
 
   auto MaxSafeFixedVF = ElementCount::getFixed(MaxSafeElements);
   auto MaxSafeScalableVF = getMaxLegalScalableVF(MaxSafeElements);
-  if (!isSafeForAnyVectorWidth())
+  if (!Legal->isSafeForAnyVectorWidth())
     this->MaxSafeElements = MaxSafeElements;
 
   LLVM_DEBUG(dbgs() << "LV: The max safe fixed VF is: " << MaxSafeFixedVF
@@ -10492,7 +10473,7 @@ bool LoopVectorizePass::processLoop(Loop *L) {
     }
 
     unsigned NumPotentiallyFaultingPointers =
-        LVL.getNumPotentiallyFaultingLoads();
+        LVL.getNumPotentiallyFaultingPointers();
     if (NumPotentiallyFaultingPointers > MaxNumPotentiallyFaultingPointers) {
       reportVectorizationFailure("Not worth vectorizing loop with uncountable "
                                  "early exit, due to number of potentially "
@@ -10660,15 +10641,12 @@ bool LoopVectorizePass::processLoop(Loop *L) {
     //  Optimistically generate runtime checks if they are needed. Drop them if
     //  they turn out to not be profitable.
     if (VF.Width.isVector() || SelectedIC > 1) {
-      if (LVL.getNumPotentiallyFaultingLoads()) {
-        assert(SelectedIC == 1 &&
-               "Interleaving not supported for early exit loops and "
-               "potentially faulting loads");
+      if (LVL.getNumPotentiallyFaultingPointers()) {
         assert(!CM.foldTailWithEVL() &&
                "Explicit vector length unsupported for early exit loops and "
                "potentially faulting loads");
-        addPointerAlignmentChecks(LVL.getPotentiallyFaultingLoads(), F, PSE,
-                                  TTI, VF.Width);
+        addPointerAlignmentChecks(LVL.getPotentiallyFaultingPointers(), F, PSE,
+                                  TTI, VF.Width, SelectedIC);
       }
       Checks.create(L, *LVL.getLAI(), PSE.getPredicate(), VF.Width, SelectedIC);
     }